Skip to main content

risingwave_backup/
storage.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::sync::Arc;
17
18use itertools::Itertools;
19use risingwave_common::config::ObjectStoreConfig;
20use risingwave_object_store::object::object_metrics::ObjectStoreMetrics;
21use risingwave_object_store::object::{
22    InMemObjectStore, MonitoredObjectStore, MonitoredStreamingReader, ObjectError, ObjectStoreImpl,
23    ObjectStoreRef,
24};
25use tokio::sync::RwLock;
26
27use crate::meta_snapshot::{MetaSnapshot, Metadata};
28use crate::{
29    BackupError, BackupResult, MetaSnapshotId, MetaSnapshotManifest, MetaSnapshotMetadata,
30};
31
32pub type MetaSnapshotStorageRef = Arc<ObjectStoreMetaSnapshotStorage>;
33
34#[async_trait::async_trait]
35pub trait MetaSnapshotStorage: 'static + Sync + Send {
36    /// Creates a snapshot.
37    async fn create<S: Metadata>(
38        &self,
39        snapshot: &MetaSnapshot<S>,
40        remarks: Option<String>,
41    ) -> BackupResult<()>;
42
43    /// Gets a snapshot by id.
44    async fn get<S: Metadata>(&self, id: MetaSnapshotId) -> BackupResult<MetaSnapshot<S>>;
45
46    /// Gets encoded snapshot bytes stream by id.
47    async fn get_bytes_stream(&self, id: MetaSnapshotId) -> BackupResult<MonitoredStreamingReader>;
48
49    /// Gets local snapshot manifest.
50    async fn manifest(&self) -> Arc<MetaSnapshotManifest>;
51
52    /// Refreshes local snapshot manifest.
53    async fn refresh_manifest(&self) -> BackupResult<()>;
54
55    /// Deletes snapshots by ids.
56    async fn delete(&self, ids: &[MetaSnapshotId]) -> BackupResult<()>;
57}
58
59#[derive(Clone)]
60pub struct ObjectStoreMetaSnapshotStorage {
61    path: String,
62    store: ObjectStoreRef,
63    manifest: Arc<RwLock<Arc<MetaSnapshotManifest>>>,
64}
65
66// TODO #6482: purge stale snapshots that is not in manifest.
67impl ObjectStoreMetaSnapshotStorage {
68    pub async fn new(path: &str, store: ObjectStoreRef) -> BackupResult<Self> {
69        let instance = Self {
70            path: path.to_owned(),
71            store,
72            manifest: Default::default(),
73        };
74        instance.refresh_manifest().await?;
75        Ok(instance)
76    }
77
78    async fn update_manifest(
79        &self,
80        update: impl FnOnce(MetaSnapshotManifest) -> MetaSnapshotManifest,
81    ) -> BackupResult<()> {
82        let mut guard = self.manifest.write().await;
83        let new_manifest = update((**guard).clone());
84        let bytes =
85            serde_json::to_vec(&new_manifest).map_err(|e| BackupError::Encoding(e.into()))?;
86        self.store
87            .upload(&self.get_manifest_path(), bytes.into())
88            .await?;
89        *guard = Arc::new(new_manifest);
90        Ok(())
91    }
92
93    async fn get_manifest(&self) -> BackupResult<Option<MetaSnapshotManifest>> {
94        let manifest_path = self.get_manifest_path();
95        let bytes = match self.store.read(&manifest_path, ..).await {
96            Ok(bytes) => bytes,
97            Err(e) => {
98                if e.is_object_not_found_error() {
99                    return Ok(None);
100                }
101                return Err(e.into());
102            }
103        };
104        let manifest: MetaSnapshotManifest =
105            serde_json::from_slice(&bytes).map_err(|e| BackupError::Encoding(e.into()))?;
106        Ok(Some(manifest))
107    }
108
109    fn get_manifest_path(&self) -> String {
110        format!("{}/manifest.json", self.path)
111    }
112
113    fn get_snapshot_path(&self, id: MetaSnapshotId) -> String {
114        format!("{}/{}.snapshot", self.path, id)
115    }
116
117    #[expect(dead_code)]
118    fn get_snapshot_id_from_path(path: &str) -> MetaSnapshotId {
119        let split = path.split(&['/', '.']).collect_vec();
120        debug_assert!(split.len() > 2);
121        debug_assert!(split[split.len() - 1] == "snapshot");
122        split[split.len() - 2]
123            .parse::<MetaSnapshotId>()
124            .expect("valid meta snapshot id")
125    }
126}
127
128#[async_trait::async_trait]
129impl MetaSnapshotStorage for ObjectStoreMetaSnapshotStorage {
130    async fn create<S: Metadata>(
131        &self,
132        snapshot: &MetaSnapshot<S>,
133        remarks: Option<String>,
134    ) -> BackupResult<()> {
135        let path = self.get_snapshot_path(snapshot.id);
136        let uploader = self.store.streaming_upload(&path).await?;
137        snapshot.encode_to_uploader(uploader).await?;
138        self.update_manifest(|mut manifest: MetaSnapshotManifest| {
139            manifest.manifest_id += 1;
140            manifest.snapshot_metadata.push(MetaSnapshotMetadata::new(
141                snapshot.id,
142                snapshot.metadata.hummock_version_ref(),
143                snapshot.format_version,
144                remarks,
145                snapshot.metadata.table_change_log_object_ids().into_iter(),
146            ));
147            manifest
148        })
149        .await?;
150        Ok(())
151    }
152
153    async fn get<S: Metadata>(&self, id: MetaSnapshotId) -> BackupResult<MetaSnapshot<S>> {
154        let reader = self.get_bytes_stream(id).await?;
155        MetaSnapshot::decode_from_stream(reader).await
156    }
157
158    async fn get_bytes_stream(&self, id: MetaSnapshotId) -> BackupResult<MonitoredStreamingReader> {
159        let path = self.get_snapshot_path(id);
160        Ok(self.store.streaming_read(&path, ..).await?)
161    }
162
163    async fn manifest(&self) -> Arc<MetaSnapshotManifest> {
164        self.manifest.read().await.clone()
165    }
166
167    async fn refresh_manifest(&self) -> BackupResult<()> {
168        if let Some(manifest) = self.get_manifest().await? {
169            let mut guard = self.manifest.write().await;
170            if manifest.manifest_id > guard.manifest_id {
171                *guard = Arc::new(manifest);
172            }
173        }
174        Ok(())
175    }
176
177    async fn delete(&self, ids: &[MetaSnapshotId]) -> BackupResult<()> {
178        let to_delete: HashSet<MetaSnapshotId> = HashSet::from_iter(ids.iter().cloned());
179        self.update_manifest(|mut manifest: MetaSnapshotManifest| {
180            manifest.manifest_id += 1;
181            manifest
182                .snapshot_metadata
183                .retain(|m| !to_delete.contains(&m.id));
184            manifest
185        })
186        .await?;
187        let paths = ids
188            .iter()
189            .map(|id| self.get_snapshot_path(*id))
190            .collect_vec();
191        self.store.delete_objects(&paths).await?;
192        Ok(())
193    }
194}
195
196impl From<ObjectError> for BackupError {
197    fn from(e: ObjectError) -> Self {
198        BackupError::BackupStorage(e.into())
199    }
200}
201
202// #[cfg(test)]
203pub async fn unused() -> ObjectStoreMetaSnapshotStorage {
204    ObjectStoreMetaSnapshotStorage::new(
205        "",
206        Arc::new(ObjectStoreImpl::InMem(MonitoredObjectStore::new(
207            InMemObjectStore::for_test(),
208            Arc::new(ObjectStoreMetrics::unused()),
209            Arc::new(ObjectStoreConfig::default()),
210        ))),
211    )
212    .await
213    .unwrap()
214}
215
216#[cfg(test)]
217mod tests {
218    use risingwave_hummock_sdk::HummockVersionId;
219    use risingwave_meta_model::hummock_sequence;
220
221    use super::{MetaSnapshotStorage, unused};
222    use crate::meta_snapshot::MetaSnapshot;
223    use crate::meta_snapshot_v2::{MetadataV2, decode_hummock_sequences_from_stream};
224
225    #[tokio::test]
226    async fn test_create_v2_snapshot_with_streaming_upload() {
227        let storage = unused().await;
228        let mut metadata = MetadataV2::default();
229        metadata.hummock_version.id = HummockVersionId::new(321);
230        let snapshot = MetaSnapshot {
231            format_version: 2,
232            id: 123,
233            metadata,
234        };
235
236        storage.create(&snapshot, None).await.unwrap();
237        let decoded: MetaSnapshot<MetadataV2> = storage.get(snapshot.id).await.unwrap();
238
239        assert_eq!(snapshot.format_version, decoded.format_version);
240        assert_eq!(snapshot.id, decoded.id);
241        assert_eq!(
242            snapshot.metadata.hummock_version.id,
243            decoded.metadata.hummock_version.id
244        );
245    }
246
247    #[tokio::test]
248    async fn test_decode_hummock_sequences_with_streaming_read() {
249        let storage = unused().await;
250        let snapshot = MetaSnapshot {
251            format_version: 2,
252            id: 456,
253            metadata: MetadataV2 {
254                hummock_sequences: vec![
255                    hummock_sequence::Model {
256                        name: "meta_backup".to_owned(),
257                        seq: 42,
258                    },
259                    hummock_sequence::Model {
260                        name: "sstable_object".to_owned(),
261                        seq: 100,
262                    },
263                ],
264                ..Default::default()
265            },
266        };
267
268        storage.create(&snapshot, None).await.unwrap();
269        let decoded = decode_hummock_sequences_from_stream(
270            storage.get_bytes_stream(snapshot.id).await.unwrap(),
271        )
272        .await
273        .unwrap();
274
275        assert_eq!(decoded, snapshot.metadata.hummock_sequences);
276    }
277}