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

mito2/
wal.rs
