mito2/wal/
entry_reader.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
15use api::v1::WalEntry;
16use async_stream::stream;
17use common_telemetry::tracing::warn;
18use futures::StreamExt;
19use object_store::Buffer;
20use prost::Message;
21use snafu::ResultExt;
22use store_api::logstore::entry::Entry;
23use store_api::logstore::provider::Provider;
24
25use crate::error::{DecodeWalSnafu, Result};
26use crate::wal::raw_entry_reader::RawEntryReader;
27use crate::wal::{EntryId, WalEntryStream};
28
29/// Decodes the [Entry] into [WalEntry].
30///
31/// The caller must ensure the [Entry] is complete.
32pub(crate) fn decode_raw_entry(raw_entry: Entry) -> Result<(EntryId, WalEntry)> {
33    let entry_id = raw_entry.entry_id();
34    let region_id = raw_entry.region_id();
35    debug_assert!(raw_entry.is_complete());
36    let buffer = into_buffer(raw_entry);
37    let wal_entry = WalEntry::decode(buffer).context(DecodeWalSnafu { region_id })?;
38    Ok((entry_id, wal_entry))
39}
40
41fn into_buffer(raw_entry: Entry) -> Buffer {
42    match raw_entry {
43        Entry::Naive(entry) => Buffer::from(entry.data),
44        Entry::MultiplePart(entry) => {
45            Buffer::from_iter(entry.parts.into_iter().map(bytes::Bytes::from))
46        }
47    }
48}
49
50/// [WalEntryReader] provides the ability to read and decode entries from the underlying store.
51///
52/// Notes: It will consume the inner stream and only allow invoking the `read` at once.
53pub(crate) trait WalEntryReader: Send + Sync {
54    fn read(&mut self, ns: &'_ Provider, start_id: EntryId) -> Result<WalEntryStream<'static>>;
55}
56
57pub(crate) struct NoopEntryReader;
58
59impl WalEntryReader for NoopEntryReader {
60    fn read(&mut self, _ns: &'_ Provider, _start_id: EntryId) -> Result<WalEntryStream<'static>> {
61        Ok(futures::stream::empty().boxed())
62    }
63}
64
65/// A Reader reads the [Entry] from [RawEntryReader] and decodes [Entry] into [WalEntry].
66pub struct LogStoreEntryReader<R> {
67    reader: R,
68}
69
70impl<R> LogStoreEntryReader<R> {
71    pub fn new(reader: R) -> Self {
72        Self { reader }
73    }
74}
75
76impl<R: RawEntryReader> WalEntryReader for LogStoreEntryReader<R> {
77    fn read(&mut self, ns: &'_ Provider, start_id: EntryId) -> Result<WalEntryStream<'static>> {
78        let LogStoreEntryReader { reader } = self;
79        let mut stream = reader.read(ns, start_id)?;
80
81        let stream = stream! {
82            while let Some(next_entry) = stream.next().await {
83                let entry = next_entry?;
84                if entry.is_complete() {
85                    yield decode_raw_entry(entry);
86                } else {
87                    warn!("Ignoring incomplete entry: {}", entry);
88                }
89            }
90        };
91
92        Ok(Box::pin(stream))
93    }
94}
95
96#[cfg(test)]
97mod tests {
98
99    use std::time::Duration;
100
101    use api::v1::{Mutation, OpType, WalEntry};
102    use futures::{StreamExt, TryStreamExt, stream};
103    use prost::Message;
104    use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader};
105    use store_api::logstore::provider::Provider;
106    use store_api::storage::RegionId;
107
108    use crate::error::Result;
109    use crate::test_util::wal_util::MockRawEntryStream;
110    use crate::wal::EntryId;
111    use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader};
112    use crate::wal::raw_entry_reader::{EntryStream, RawEntryReader};
113
114    struct LiveRawEntryReader {
115        entry: Entry,
116    }
117
118    impl RawEntryReader for LiveRawEntryReader {
119        fn read(&self, _ns: &Provider, _start_id: EntryId) -> Result<EntryStream<'static>> {
120            let stream = stream::iter([Ok(self.entry.clone())]).chain(stream::pending());
121            Ok(stream.boxed())
122        }
123    }
124
125    #[tokio::test]
126    async fn test_delivers_complete_entry_while_stream_is_alive() {
127        let provider = Provider::kafka_provider("my_topic".to_string());
128        let wal_entry = WalEntry::default();
129        let raw_entry_reader = LiveRawEntryReader {
130            entry: Entry::Naive(store_api::logstore::entry::NaiveEntry {
131                provider: provider.clone(),
132                region_id: RegionId::new(1, 1),
133                entry_id: 1,
134                data: wal_entry.encode_to_vec(),
135            }),
136        };
137        let mut reader = LogStoreEntryReader::new(raw_entry_reader);
138        let mut entries = reader.read(&provider, 0).unwrap();
139
140        let entry = tokio::time::timeout(Duration::from_secs(1), entries.next())
141            .await
142            .unwrap()
143            .unwrap()
144            .unwrap();
145        assert_eq!(entry, (1, wal_entry));
146    }
147
148    #[tokio::test]
149    async fn test_tail_corrupted_stream() {
150        common_telemetry::init_default_ut_logging();
151        let provider = Provider::kafka_provider("my_topic".to_string());
152        let wal_entry = WalEntry {
153            mutations: vec![Mutation {
154                op_type: OpType::Put as i32,
155                sequence: 1u64,
156                rows: None,
157                write_hint: None,
158            }],
159            bulk_entries: vec![],
160        };
161        let encoded_entry = wal_entry.encode_to_vec();
162        let parts = encoded_entry
163            .chunks(encoded_entry.len() / 2)
164            .map(Into::into)
165            .collect::<Vec<_>>();
166        let raw_entry_stream = MockRawEntryStream {
167            entries: vec![
168                Entry::MultiplePart(MultiplePartEntry {
169                    provider: provider.clone(),
170                    region_id: RegionId::new(1, 1),
171                    entry_id: 2,
172                    headers: vec![MultiplePartHeader::First, MultiplePartHeader::Last],
173                    parts,
174                }),
175                // The tail incomplete entry.
176                Entry::MultiplePart(MultiplePartEntry {
177                    provider: provider.clone(),
178                    region_id: RegionId::new(1, 1),
179                    entry_id: 1,
180                    headers: vec![MultiplePartHeader::Last],
181                    parts: vec![vec![1; 100]],
182                }),
183            ],
184        };
185
186        let mut reader = LogStoreEntryReader::new(raw_entry_stream);
187        let entries = reader
188            .read(&provider, 0)
189            .unwrap()
190            .try_collect::<Vec<_>>()
191            .await
192            .unwrap()
193            .into_iter()
194            .map(|(_, entry)| entry)
195            .collect::<Vec<_>>();
196
197        assert_eq!(entries, vec![wal_entry]);
198    }
199
200    #[tokio::test]
201    async fn test_corrupted_stream() {
202        let provider = Provider::kafka_provider("my_topic".to_string());
203        let raw_entry_stream = MockRawEntryStream {
204            entries: vec![
205                // The incomplete entry.
206                Entry::MultiplePart(MultiplePartEntry {
207                    provider: provider.clone(),
208                    region_id: RegionId::new(1, 1),
209                    entry_id: 1,
210                    headers: vec![MultiplePartHeader::Last],
211                    parts: vec![vec![1; 100]],
212                }),
213                Entry::MultiplePart(MultiplePartEntry {
214                    provider: provider.clone(),
215                    region_id: RegionId::new(1, 1),
216                    entry_id: 2,
217                    headers: vec![MultiplePartHeader::First],
218                    parts: vec![vec![1; 100]],
219                }),
220            ],
221        };
222
223        let mut reader = LogStoreEntryReader::new(raw_entry_stream);
224        let entries = reader
225            .read(&provider, 0)
226            .unwrap()
227            .try_collect::<Vec<_>>()
228            .await
229            .unwrap();
230        assert!(entries.is_empty());
231    }
232}
mito2/wal/entry_reader.rs

mito2/wal/
entry_reader.rs
