risingwave_backup/
meta_snapshot.rs

1// Copyright 2022 RisingWave Labs
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 std::collections::HashSet;
16use std::fmt::{Display, Formatter};
17use std::future::Future;
18use std::hash::Hasher;
19use std::mem::size_of;
20
21use bytes::{Buf, BufMut, Bytes, BytesMut};
22use futures::StreamExt;
23use risingwave_hummock_sdk::HummockRawObjectId;
24use risingwave_hummock_sdk::version::HummockVersion;
25use risingwave_object_store::object::{
26    MonitoredStreamingReader, ObjectDataStream, ObjectDataStreamReader, ObjectStreamingUploader,
27};
28use tokio::io::AsyncReadExt;
29use twox_hash::XxHash64;
30
31use crate::MetaSnapshotId;
32use crate::error::{BackupError, BackupResult};
33
34pub trait Metadata: Display + Send + Sync {
35    fn encode_to_writer(
36        &self,
37        writer: SnapshotPayloadWriter,
38    ) -> impl Future<Output = BackupResult<()>> + Send;
39
40    fn decode_from_reader(
41        reader: SnapshotPayloadReader,
42    ) -> impl Future<Output = BackupResult<Self>> + Send
43    where
44        Self: Sized;
45
46    fn hummock_version_ref(&self) -> &HummockVersion;
47
48    fn hummock_version(self) -> HummockVersion;
49
50    fn storage_url(&self) -> BackupResult<String>;
51
52    fn storage_directory(&self) -> BackupResult<String>;
53
54    fn table_change_log_object_ids(&self) -> HashSet<HummockRawObjectId>;
55}
56
57#[derive(Debug, Default, Clone, PartialEq)]
58pub struct MetaSnapshot<T: Metadata> {
59    pub format_version: u32,
60    pub id: MetaSnapshotId,
61    /// Snapshot of meta store.
62    pub metadata: T,
63}
64
65impl<T: Metadata> MetaSnapshot<T> {
66    pub async fn encode_to_uploader(&self, uploader: ObjectStreamingUploader) -> BackupResult<()> {
67        let mut writer = SnapshotPayloadWriter::new(uploader);
68
69        let mut header = BytesMut::with_capacity(size_of::<u32>() + size_of::<u64>());
70        header.put_u32_le(self.format_version);
71        header.put_u64_le(self.id);
72        writer.write_snapshot_bytes(header.freeze()).await?;
73
74        self.metadata.encode_to_writer(writer).await
75    }
76
77    pub async fn decode_from_stream(reader: MonitoredStreamingReader) -> BackupResult<Self> {
78        let mut reader =
79            SnapshotPayloadReader::new(ObjectDataStreamReader::new(reader.into_stream()));
80        let (format_version, id) = reader.read_snapshot_header().await?;
81        let metadata = T::decode_from_reader(reader).await?;
82        Ok(Self {
83            format_version,
84            id,
85            metadata,
86        })
87    }
88}
89
90pub struct SnapshotPayloadWriter {
91    uploader: ObjectStreamingUploader,
92    hasher: XxHash64,
93    size: usize,
94}
95
96impl SnapshotPayloadWriter {
97    pub fn new(uploader: ObjectStreamingUploader) -> Self {
98        Self {
99            uploader,
100            hasher: XxHash64::with_seed(0),
101            size: 0,
102        }
103    }
104
105    pub async fn write_snapshot_bytes(&mut self, data: Bytes) -> BackupResult<()> {
106        self.hasher.write(data.as_ref());
107        self.size += data.len();
108        self.uploader.write_bytes(data).await?;
109        Ok(())
110    }
111
112    pub async fn finish(mut self) -> BackupResult<usize> {
113        let mut checksum = BytesMut::with_capacity(size_of::<u64>());
114        checksum.put_u64_le(self.hasher.finish());
115        let checksum = checksum.freeze();
116        self.size += checksum.len();
117        self.uploader.write_bytes(checksum).await?;
118        self.uploader.finish().await?;
119        Ok(self.size)
120    }
121}
122
123type SnapshotPayloadStreamReader = ObjectDataStreamReader<ObjectDataStream>;
124
125pub struct SnapshotPayloadReader {
126    reader: SnapshotPayloadStreamReader,
127    hasher: XxHash64,
128}
129
130impl SnapshotPayloadReader {
131    pub fn new(reader: SnapshotPayloadStreamReader) -> Self {
132        Self {
133            reader,
134            hasher: XxHash64::with_seed(0),
135        }
136    }
137
138    pub async fn read_snapshot_header(&mut self) -> BackupResult<(u32, MetaSnapshotId)> {
139        let format_version = self.read_u32_le().await?;
140        let id = self.read_u64_le().await?;
141        Ok((format_version, id))
142    }
143
144    pub async fn read_u32_le(&mut self) -> BackupResult<u32> {
145        let bytes = self.read_exact(size_of::<u32>()).await?;
146        Ok(u32::from_le_bytes(
147            bytes[..].try_into().expect("u32 length"),
148        ))
149    }
150
151    pub async fn read_u64_le(&mut self) -> BackupResult<u64> {
152        let bytes = self.read_exact(size_of::<u64>()).await?;
153        Ok(u64::from_le_bytes(
154            bytes[..].try_into().expect("u64 length"),
155        ))
156    }
157
158    pub async fn read_exact(&mut self, len: usize) -> BackupResult<BytesMut> {
159        let mut bytes = BytesMut::with_capacity(len);
160        bytes.resize(len, 0);
161        self.reader.read_exact(&mut bytes).await?;
162        self.hasher.write(&bytes);
163        Ok(bytes)
164    }
165
166    pub async fn skip_exact(&mut self, mut len: usize) -> BackupResult<()> {
167        const BUFFER_SIZE: usize = 1024;
168        let mut buf = [0; BUFFER_SIZE];
169        while len > 0 {
170            let read_len = len.min(BUFFER_SIZE);
171            self.reader.read_exact(&mut buf[..read_len]).await?;
172            self.hasher.write(&buf[..read_len]);
173            len -= read_len;
174        }
175        Ok(())
176    }
177
178    pub async fn read_to_end(mut self) -> BackupResult<Vec<u8>> {
179        let mut data = vec![];
180        self.reader.read_to_end(&mut data).await?;
181        if data.len() < size_of::<u64>() {
182            return Err(BackupError::Decoding(
183                anyhow::anyhow!("meta snapshot is missing checksum").into(),
184            ));
185        }
186
187        let checksum = data.split_off(data.len() - size_of::<u64>());
188        self.hasher.write(&data);
189        self.verify_checksum(&checksum)?;
190        Ok(data)
191    }
192
193    pub async fn finish_after_skipping_to_end(self) -> BackupResult<()> {
194        let Self { reader, mut hasher } = self;
195        let mut bytes_stream = reader.into_bytes_stream();
196        let mut tail = BytesMut::with_capacity(size_of::<u64>());
197
198        while let Some(bytes) = bytes_stream.next().await.transpose()? {
199            let payload_len = tail.len() + bytes.len();
200            if payload_len <= size_of::<u64>() {
201                tail.extend_from_slice(&bytes);
202                continue;
203            }
204
205            let payload_len = payload_len - size_of::<u64>();
206            if payload_len <= tail.len() {
207                hasher.write(&tail[..payload_len]);
208                tail.advance(payload_len);
209                tail.extend_from_slice(&bytes);
210            } else {
211                let bytes_payload_len = payload_len - tail.len();
212                hasher.write(&tail);
213                hasher.write(&bytes[..bytes_payload_len]);
214                tail.clear();
215                tail.extend_from_slice(&bytes[bytes_payload_len..]);
216            }
217        }
218
219        if tail.len() != size_of::<u64>() {
220            return Err(BackupError::Decoding(
221                anyhow::anyhow!("meta snapshot is missing checksum").into(),
222            ));
223        }
224        Self::verify_checksum_with_hasher(&hasher, &tail)
225    }
226
227    pub async fn finish(mut self) -> BackupResult<()> {
228        let mut checksum = [0; size_of::<u64>()];
229        self.reader.read_exact(&mut checksum).await?;
230        self.verify_checksum(&checksum)?;
231
232        let mut trailing = [0; 1];
233        let n = self.reader.read(&mut trailing).await?;
234        if n != 0 {
235            return Err(BackupError::Decoding(
236                anyhow::anyhow!("unexpected bytes after meta snapshot checksum").into(),
237            ));
238        }
239        Ok(())
240    }
241
242    fn verify_checksum(&self, checksum: &[u8]) -> BackupResult<()> {
243        Self::verify_checksum_with_hasher(&self.hasher, checksum)
244    }
245
246    fn verify_checksum_with_hasher(hasher: &XxHash64, checksum: &[u8]) -> BackupResult<()> {
247        let expected = u64::from_le_bytes(checksum.try_into().expect("u64 length"));
248        let found = hasher.finish();
249        if expected != found {
250            return Err(BackupError::ChecksumMismatch { expected, found });
251        }
252        Ok(())
253    }
254}
255
256pub(crate) fn read_u32_le(buf: &mut &[u8]) -> BackupResult<u32> {
257    if buf.remaining() < 4 {
258        return Err(BackupError::Other(anyhow::anyhow!(
259            "metadata snapshot is truncated while reading u32"
260        )));
261    }
262    Ok(buf.get_u32_le())
263}
264
265impl<T: Metadata> Display for MetaSnapshot<T> {
266    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
267        writeln!(f, "format_version: {}", self.format_version)?;
268        writeln!(f, "id: {}", self.id)?;
269        writeln!(f, "{}", self.metadata)?;
270        Ok(())
271    }
272}
risingwave_backup/meta_snapshot.rs

risingwave_backup/
meta_snapshot.rs
