mito2/
wal.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Write ahead log of the engine.
16
17pub mod encoder;
18pub(crate) mod entry_distributor;
19pub(crate) mod entry_reader;
20pub(crate) mod raw_entry_reader;
21
22use std::collections::HashMap;
23use std::mem;
24use std::sync::Arc;
25
26use api::v1::WalEntry;
27use common_error::ext::BoxedError;
28use common_telemetry::debug;
29use encoder::WalEntryEncoder;
30use entry_reader::NoopEntryReader;
31use futures::future::BoxFuture;
32use futures::stream::BoxStream;
33use snafu::ResultExt;
34use store_api::logstore::entry::Entry;
35use store_api::logstore::provider::Provider;
36use store_api::logstore::{AppendBatchResponse, LogStore};
37use store_api::storage::RegionId;
38
39use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, WriteWalSnafu};
40use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader};
41use crate::wal::raw_entry_reader::{LogStoreRawEntryReader, RegionRawEntryReader};
42
43/// WAL entry id.
44pub type EntryId = store_api::logstore::entry::Id;
45/// A stream that yields tuple of WAL entry id and corresponding entry.
46pub type WalEntryStream<'a> = BoxStream<'a, Result<(EntryId, WalEntry)>>;
47
48/// Waits until the WAL of one region is durable through an entry id.
49///
50/// A flush calls it before it records an entry id as flushed in the manifest:
51/// a log store that acknowledges an append before its entries are durable
52/// hands out entry ids that a crash can lose, and a manifest watermark that
53/// names such an id would skip the entries assigned again after a restart.
54#[derive(Clone)]
55pub(crate) struct DurabilityBarrier(
56    Arc<dyn Fn(EntryId) -> BoxFuture<'static, Result<()>> + Send + Sync>,
57);
58
59impl DurabilityBarrier {
60    /// Returns once the region is durable through `entry_id`.
61    pub(crate) async fn wait(&self, entry_id: EntryId) -> Result<()> {
62        (self.0)(entry_id).await
63    }
64
65    /// A barrier that never waits, for tests without a log store.
66    #[cfg(test)]
67    pub(crate) fn noop() -> Self {
68        Self(Arc::new(|_| Box::pin(async { Ok(()) })))
69    }
70}
71
72impl std::fmt::Debug for DurabilityBarrier {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.write_str("DurabilityBarrier")
75    }
76}
77
78/// Write ahead log.
79///
80/// All regions in the engine shares the same WAL instance.
81#[derive(Debug)]
82pub struct Wal<S> {
83    /// The underlying log store.
84    store: Arc<S>,
85}
86
87impl<S> Wal<S> {
88    /// Creates a new [Wal] from the log store.
89    pub fn new(store: Arc<S>) -> Self {
90        Self { store }
91    }
92
93    pub fn store(&self) -> &Arc<S> {
94        &self.store
95    }
96}
97
98impl<S> Clone for Wal<S> {
99    fn clone(&self) -> Self {
100        Self {
101            store: Arc::clone(&self.store),
102        }
103    }
104}
105
106impl<S: LogStore> Wal<S> {
107    /// Returns a writer to write to the WAL.
108    pub fn writer(&self) -> WalWriter<S> {
109        WalWriter {
110            store: self.store.clone(),
111            entries: Vec::new(),
112            providers: HashMap::new(),
113            encoder: WalEntryEncoder::new(),
114        }
115    }
116
117    /// Returns a [OnRegionOpened] function.
118    pub(crate) fn on_region_opened(
119        &self,
120    ) -> impl FnOnce(RegionId, EntryId, &Provider) -> BoxFuture<Result<()>> {
121        let store = self.store.clone();
122        move |region_id, last_entry_id, provider| -> BoxFuture<'_, Result<()>> {
123            if let Provider::Noop = provider {
124                debug!("Skip obsolete for region: {}", region_id);
125                return Box::pin(async move { Ok(()) });
126            }
127            Box::pin(async move {
128                store
129                    .obsolete(provider, region_id, last_entry_id)
130                    .await
131                    .map_err(BoxedError::new)
132                    .context(DeleteWalSnafu { region_id })
133            })
134        }
135    }
136
137    /// Returns the [DurabilityBarrier] of the region written through `provider`.
138    pub(crate) fn durability_barrier(
139        &self,
140        region_id: RegionId,
141        provider: &Provider,
142    ) -> DurabilityBarrier {
143        let store = self.store.clone();
144        let provider = provider.clone();
145        DurabilityBarrier(Arc::new(move |entry_id| {
146            let store = store.clone();
147            let provider = provider.clone();
148            Box::pin(async move {
149                if let Provider::Noop = provider {
150                    return Ok(());
151                }
152                store
153                    .wait_durable(&provider, entry_id)
154                    .await
155                    .map_err(BoxedError::new)
156                    .context(WaitWalDurableSnafu { region_id })
157            })
158        }))
159    }
160
161    /// Returns a [WalEntryReader]
162    pub(crate) fn wal_entry_reader(
163        &self,
164        provider: &Provider,
165        region_id: RegionId,
166    ) -> Box<dyn WalEntryReader> {
167        match provider {
168            Provider::RaftEngine(_) => Box::new(LogStoreEntryReader::new(
169                LogStoreRawEntryReader::new(self.store.clone()),
170            )),
171            // Entries of many regions share one namespace, so the reader filters by region.
172            Provider::Kafka(_) | Provider::ObjectStore(_) => {
173                Box::new(LogStoreEntryReader::new(RegionRawEntryReader::new(
174                    LogStoreRawEntryReader::new(self.store.clone()),
175                    region_id,
176                )))
177            }
178            Provider::Noop => Box::new(NoopEntryReader),
179        }
180    }
181
182    /// Scan entries of specific region starting from `start_id` (inclusive).
183    /// Currently only used in tests.
184    pub fn scan<'a>(
185        &'a self,
186        region_id: RegionId,
187        start_id: EntryId,
188        provider: &'a Provider,
189    ) -> Result<WalEntryStream<'a>> {
190        let mut reader = self.wal_entry_reader(provider, region_id);
191        reader.read(provider, start_id)
192    }
193
194    /// Mark entries whose ids `<= last_id` as deleted.
195    pub async fn obsolete(
196        &self,
197        region_id: RegionId,
198        last_id: EntryId,
199        provider: &Provider,
200    ) -> Result<()> {
201        if let Provider::Noop = provider {
202            return Ok(());
203        }
204        self.store
205            .obsolete(provider, region_id, last_id)
206            .await
207            .map_err(BoxedError::new)
208            .context(DeleteWalSnafu { region_id })
209    }
210
211    /// Deletes all WAL entries in the namespace represented by `provider`.
212    pub async fn delete_namespace(&self, region_id: RegionId, provider: &Provider) -> Result<()> {
213        if let Provider::Noop = provider {
214            return Ok(());
215        }
216        self.store
217            .delete_namespace(provider)
218            .await
219            .map_err(BoxedError::new)
220            .context(DeleteWalSnafu { region_id })
221    }
222
223    /// Releases every WAL entry of a dropped region, see [LogStore::release_region].
224    pub async fn release_region(&self, region_id: RegionId, provider: &Provider) -> Result<()> {
225        if let Provider::Noop = provider {
226            return Ok(());
227        }
228        self.store
229            .release_region(provider, region_id)
230            .await
231            .map_err(BoxedError::new)
232            .context(DeleteWalSnafu { region_id })
233    }
234
235    /// Marks all WAL entries of a region as obsolete and removes its dedicated namespace when
236    /// supported by the backend.
237    pub async fn obsolete_all(&self, region_id: RegionId, provider: &Provider) -> Result<()> {
238        self.store
239            .obsolete_all(provider, region_id)
240            .await
241            .map_err(BoxedError::new)
242            .context(DeleteWalSnafu { region_id })
243    }
244}
245
246/// WAL batch writer.
247pub struct WalWriter<S: LogStore> {
248    /// Log store of the WAL.
249    store: Arc<S>,
250    /// Entries to write.
251    entries: Vec<Entry>,
252    /// Providers of regions being written into.
253    providers: HashMap<RegionId, Provider>,
254    /// Cached-size single-pass encoder, reused across entries in this batch.
255    encoder: WalEntryEncoder,
256}
257
258impl<S: LogStore> WalWriter<S> {
259    /// Add a wal entry for specific region to the writer's buffer.
260    pub fn add_entry(
261        &mut self,
262        region_id: RegionId,
263        entry_id: EntryId,
264        wal_entry: &WalEntry,
265        provider: &Provider,
266    ) -> Result<()> {
267        // Gets or inserts with a newly built provider.
268        let provider = self
269            .providers
270            .entry(region_id)
271            .or_insert_with(|| provider.clone());
272
273        let data = self.encoder.encode_to_vec(wal_entry);
274        let entry = self
275            .store
276            .entry(data, entry_id, region_id, provider)
277            .map_err(BoxedError::new)
278            .context(BuildEntrySnafu { region_id })?;
279
280        self.entries.push(entry);
281
282        Ok(())
283    }
284
285    /// Write all buffered entries to the WAL.
286    pub async fn write_to_wal(&mut self) -> Result<AppendBatchResponse> {
287        // TODO(yingwen): metrics.
288
289        let entries = mem::take(&mut self.entries);
290        self.store
291            .append_batch(entries)
292            .await
293            .map_err(BoxedError::new)
294            .context(WriteWalSnafu)
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use api::v1::helper::{tag_column_schema, time_index_column_schema};
301    use api::v1::{
302        ArrowIpc, BulkWalEntry, ColumnDataType, Mutation, OpType, Row, Rows, Value, bulk_wal_entry,
303        value,
304    };
305    use common_recordbatch::DfRecordBatch;
306    use common_test_util::flight::encode_to_flight_data;
307    use common_test_util::temp_dir::{TempDir, create_temp_dir};
308    use datatypes::arrow;
309    use datatypes::arrow::array::{ArrayRef, TimestampMillisecondArray};
310    use datatypes::arrow::datatypes::Field;
311    use datatypes::arrow_array::StringArray;
312    use futures::TryStreamExt;
313    use log_store::raft_engine::log_store::RaftEngineLogStore;
314    use log_store::test_util::log_store_util;
315    use store_api::storage::SequenceNumber;
316
317    use super::*;
318
319    struct WalEnv {
320        _wal_dir: TempDir,
321        log_store: Option<Arc<RaftEngineLogStore>>,
322    }
323
324    impl WalEnv {
325        async fn new() -> WalEnv {
326            let wal_dir = create_temp_dir("");
327            let log_store =
328                log_store_util::create_tmp_local_file_log_store(wal_dir.path().to_str().unwrap())
329                    .await;
330            WalEnv {
331                _wal_dir: wal_dir,
332                log_store: Some(Arc::new(log_store)),
333            }
334        }
335
336        fn new_wal(&self) -> Wal<RaftEngineLogStore> {
337            let log_store = self.log_store.clone().unwrap();
338            Wal::new(log_store)
339        }
340    }
341
342    /// Create a new mutation from rows.
343    ///
344    /// The row format is (string, i64).
345    fn new_mutation(op_type: OpType, sequence: SequenceNumber, rows: &[(&str, i64)]) -> Mutation {
346        let rows = rows
347            .iter()
348            .map(|(str_col, int_col)| {
349                let values = vec![
350                    Value {
351                        value_data: Some(value::ValueData::StringValue(str_col.to_string())),
352                    },
353                    Value {
354                        value_data: Some(value::ValueData::TimestampMillisecondValue(*int_col)),
355                    },
356                ];
357                Row { values }
358            })
359            .collect();
360        let schema = vec![
361            tag_column_schema("tag", ColumnDataType::String),
362            time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
363        ];
364
365        Mutation {
366            op_type: op_type as i32,
367            sequence,
368            rows: Some(Rows { schema, rows }),
369            write_hint: None,
370        }
371    }
372
373    #[tokio::test]
374    async fn test_write_wal() {
375        let env = WalEnv::new().await;
376        let wal = env.new_wal();
377
378        let entry = WalEntry {
379            mutations: vec![
380                new_mutation(OpType::Put, 1, &[("k1", 1), ("k2", 2)]),
381                new_mutation(OpType::Put, 2, &[("k3", 3), ("k4", 4)]),
382            ],
383            bulk_entries: vec![],
384        };
385        let mut writer = wal.writer();
386        // Region 1 entry 1.
387        let region_id = RegionId::new(1, 1);
388        writer
389            .add_entry(
390                region_id,
391                1,
392                &entry,
393                &Provider::raft_engine_provider(region_id.as_u64()),
394            )
395            .unwrap();
396        // Region 2 entry 1.
397        let region_id = RegionId::new(1, 2);
398        writer
399            .add_entry(
400                region_id,
401                1,
402                &entry,
403                &Provider::raft_engine_provider(region_id.as_u64()),
404            )
405            .unwrap();
406        // Region 1 entry 2.
407        let region_id = RegionId::new(1, 2);
408        writer
409            .add_entry(
410                region_id,
411                2,
412                &entry,
413                &Provider::raft_engine_provider(region_id.as_u64()),
414            )
415            .unwrap();
416
417        // Test writing multiple region to wal.
418        writer.write_to_wal().await.unwrap();
419    }
420
421    fn build_record_batch(rows: &[(&str, i64)]) -> DfRecordBatch {
422        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
423            Field::new("tag", arrow::datatypes::DataType::Utf8, false),
424            Field::new(
425                "ts",
426                arrow::datatypes::DataType::Timestamp(
427                    arrow::datatypes::TimeUnit::Millisecond,
428                    None,
429                ),
430                false,
431            ),
432        ]));
433
434        let tag = Arc::new(StringArray::from_iter_values(
435            rows.iter().map(|r| r.0.to_string()),
436        )) as ArrayRef;
437        let ts = Arc::new(TimestampMillisecondArray::from_iter_values(
438            rows.iter().map(|r| r.1),
439        )) as ArrayRef;
440        DfRecordBatch::try_new(schema, vec![tag, ts]).unwrap()
441    }
442
443    fn build_bulk_wal_entry(sequence_number: SequenceNumber, rows: &[(&str, i64)]) -> BulkWalEntry {
444        let rb = build_record_batch(rows);
445        let (schema, rb) = encode_to_flight_data(rb);
446        let max_ts = rows.iter().map(|r| r.1).max().unwrap();
447        let min_ts = rows.iter().map(|r| r.1).min().unwrap();
448        BulkWalEntry {
449            sequence: sequence_number,
450            max_ts,
451            min_ts,
452            timestamp_index: 1,
453            body: Some(bulk_wal_entry::Body::ArrowIpc(ArrowIpc {
454                schema: schema.data_header,
455                data_header: rb.data_header,
456                payload: rb.data_body,
457            })),
458        }
459    }
460
461    fn sample_entries() -> Vec<WalEntry> {
462        vec![
463            WalEntry {
464                mutations: vec![
465                    new_mutation(OpType::Put, 1, &[("k1", 1), ("k2", 2)]),
466                    new_mutation(OpType::Put, 2, &[("k3", 3), ("k4", 4)]),
467                ],
468                bulk_entries: vec![],
469            },
470            WalEntry {
471                mutations: vec![new_mutation(OpType::Put, 3, &[("k1", 1), ("k2", 2)])],
472                bulk_entries: vec![],
473            },
474            WalEntry {
475                mutations: vec![
476                    new_mutation(OpType::Put, 4, &[("k1", 1), ("k2", 2)]),
477                    new_mutation(OpType::Put, 5, &[("k3", 3), ("k4", 4)]),
478                ],
479                bulk_entries: vec![],
480            },
481            WalEntry {
482                mutations: vec![new_mutation(OpType::Put, 6, &[("k1", 1), ("k2", 2)])],
483                bulk_entries: vec![build_bulk_wal_entry(7, &[("k1", 8), ("k2", 9)])],
484            },
485        ]
486    }
487
488    fn check_entries(
489        expect: &[WalEntry],
490        expect_start_id: EntryId,
491        actual: &[(EntryId, WalEntry)],
492    ) {
493        for (idx, (expect_entry, (actual_id, actual_entry))) in
494            expect.iter().zip(actual.iter()).enumerate()
495        {
496            let expect_id_entry = (expect_start_id + idx as u64, expect_entry);
497            assert_eq!(expect_id_entry, (*actual_id, actual_entry));
498        }
499        assert_eq!(expect.len(), actual.len());
500    }
501
502    #[tokio::test]
503    async fn test_scan_wal() {
504        let env = WalEnv::new().await;
505        let wal = env.new_wal();
506
507        let entries = sample_entries();
508        let (id1, id2) = (RegionId::new(1, 1), RegionId::new(1, 2));
509        let ns1 = Provider::raft_engine_provider(id1.as_u64());
510        let ns2 = Provider::raft_engine_provider(id2.as_u64());
511        let mut writer = wal.writer();
512        writer.add_entry(id1, 1, &entries[0], &ns1).unwrap();
513        // Insert one entry into region2. Scan should not return this entry.
514        writer.add_entry(id2, 1, &entries[0], &ns2).unwrap();
515        writer.add_entry(id1, 2, &entries[1], &ns1).unwrap();
516        writer.add_entry(id1, 3, &entries[2], &ns1).unwrap();
517        writer.add_entry(id1, 4, &entries[3], &ns1).unwrap();
518
519        writer.write_to_wal().await.unwrap();
520
521        // Scan all contents region1
522        let stream = wal.scan(id1, 1, &ns1).unwrap();
523        let actual: Vec<_> = stream.try_collect().await.unwrap();
524        check_entries(&entries, 1, &actual);
525
526        // Scan parts of contents
527        let stream = wal.scan(id1, 2, &ns1).unwrap();
528        let actual: Vec<_> = stream.try_collect().await.unwrap();
529        check_entries(&entries[1..], 2, &actual);
530
531        // Scan out of range
532        let stream = wal.scan(id1, 5, &ns1).unwrap();
533        let actual: Vec<_> = stream.try_collect().await.unwrap();
534        assert!(actual.is_empty());
535    }
536
537    #[tokio::test]
538    async fn test_obsolete_wal() {
539        let env = WalEnv::new().await;
540        let wal = env.new_wal();
541
542        let entries = sample_entries();
543        let mut writer = wal.writer();
544        let region_id = RegionId::new(1, 1);
545        let ns = Provider::raft_engine_provider(region_id.as_u64());
546        writer.add_entry(region_id, 1, &entries[0], &ns).unwrap();
547        writer.add_entry(region_id, 2, &entries[1], &ns).unwrap();
548        writer.add_entry(region_id, 3, &entries[2], &ns).unwrap();
549
550        writer.write_to_wal().await.unwrap();
551
552        // Delete 1, 2.
553        wal.obsolete(region_id, 2, &ns).await.unwrap();
554
555        // Put 4.
556        let mut writer = wal.writer();
557        writer.add_entry(region_id, 4, &entries[3], &ns).unwrap();
558        writer.write_to_wal().await.unwrap();
559
560        // Scan all
561        let stream = wal.scan(region_id, 1, &ns).unwrap();
562        let actual: Vec<_> = stream.try_collect().await.unwrap();
563        check_entries(&entries[2..], 3, &actual);
564    }
565}
mito2/wal.rs

mito2/
wal.rs
