risingwave_backup/
storage.rs1use 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 async fn create<S: Metadata>(
38 &self,
39 snapshot: &MetaSnapshot<S>,
40 remarks: Option<String>,
41 ) -> BackupResult<()>;
42
43 async fn get<S: Metadata>(&self, id: MetaSnapshotId) -> BackupResult<MetaSnapshot<S>>;
45
46 async fn get_bytes_stream(&self, id: MetaSnapshotId) -> BackupResult<MonitoredStreamingReader>;
48
49 async fn manifest(&self) -> Arc<MetaSnapshotManifest>;
51
52 async fn refresh_manifest(&self) -> BackupResult<()>;
54
55 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
66impl 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
202pub 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}