mito2/manifest/storage/
staging.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 std::sync::Arc;
16
17use common_datasource::compression::CompressionType;
18use common_telemetry::debug;
19use object_store::{Lister, ObjectStore, util};
20use snafu::ResultExt;
21use store_api::ManifestVersion;
22
23use crate::error::{OpenDalSnafu, Result};
24use crate::manifest::storage::delta::DeltaStorage;
25use crate::manifest::storage::size_tracker::NoopTracker;
26use crate::manifest::storage::utils::{sort_manifests, write_and_put_cache};
27use crate::manifest::storage::{file_version, is_delta_file};
28
29/// A simple blob storage for arbitrary binary data in the staging directory.
30///
31/// This is primarily used during repartition operations to store generated
32/// manifests that will be consumed by other regions via [`ApplyStagingManifestRequest`](store_api::region_request::ApplyStagingManifestRequest).
33/// The blobs are stored in `{region_dir}/staging/blob/` directory.
34#[derive(Debug, Clone)]
35pub(crate) struct StagingBlobStorage {
36    object_store: ObjectStore,
37    path: String,
38}
39
40/// Returns the staging path from the blob path.
41///
42/// # Example
43/// - Input: `"data/table/region_0001/manifest/"`
44/// - Output: `"data/table/region_0001/staging/blob/"`
45pub fn staging_blob_path(manifest_path: &str) -> String {
46    let parent_dir = manifest_path
47        .trim_end_matches("manifest/")
48        .trim_end_matches('/');
49    util::normalize_dir(&format!("{}/staging/blob", parent_dir))
50}
51
52impl StagingBlobStorage {
53    pub fn new(path: String, object_store: ObjectStore) -> Self {
54        let path = util::normalize_dir(&path);
55        common_telemetry::debug!(
56            "Staging blob storage path: {}, root: {}",
57            path,
58            object_store.info().root()
59        );
60        Self { object_store, path }
61    }
62
63    /// Returns the path.
64    pub fn path(&self) -> &str {
65        &self.path
66    }
67
68    /// Returns the object store.
69    pub fn object_store(&self) -> &ObjectStore {
70        &self.object_store
71    }
72
73    /// Put the bytes to the blob storage.
74    pub async fn put(&self, path: &str, bytes: Vec<u8>) -> Result<()> {
75        let path = format!("{}{}", self.path, path);
76        common_telemetry::debug!(
77            "Putting blob to staging blob storage, path: {}, root: {}, bytes: {}",
78            path,
79            self.object_store.info().root(),
80            bytes.len()
81        );
82        write_and_put_cache(&self.object_store, None, &path, bytes).await?;
83        Ok(())
84    }
85
86    /// Get the bytes from the blob storage.
87    pub async fn get(&self, path: &str) -> Result<Vec<u8>> {
88        let path = format!("{}{}", self.path, path);
89        common_telemetry::debug!(
90            "Reading blob from staging blob storage, path: {}, root: {}",
91            path,
92            self.object_store.info().root()
93        );
94        let bytes = self.object_store.read(&path).await.context(OpenDalSnafu)?;
95
96        Ok(bytes.to_vec())
97    }
98}
99
100/// Storage for staging manifest files and blobs used during repartition operations.
101///
102/// Fields:
103/// - `delta_storage`: Manages incremental manifest delta files specific to the staging region.
104/// - `blob_storage`: Manages arbitrary blobs, such as generated manifests for regions.
105///
106/// Directory structure:
107/// - `{region_dir}/staging/manifest/` — for incremental manifest delta files for the staging region.
108/// - `{region_dir}/staging/blob/` — for arbitrary blobs (e.g., generated region manifests).
109#[derive(Debug, Clone)]
110pub(crate) struct StagingStorage {
111    delta_storage: DeltaStorage<NoopTracker>,
112    blob_storage: StagingBlobStorage,
113}
114
115/// Returns the staging path from the manifest path.
116///
117/// # Example
118/// - Input: `"data/table/region_0001/manifest/"`
119/// - Output: `"data/table/region_0001/staging/manifest/"`
120pub fn staging_manifest_path(manifest_path: &str) -> String {
121    let parent_dir = manifest_path
122        .trim_end_matches("manifest/")
123        .trim_end_matches('/');
124    util::normalize_dir(&format!("{}/staging/manifest", parent_dir))
125}
126
127impl StagingStorage {
128    pub fn new(path: String, object_store: ObjectStore, compress_type: CompressionType) -> Self {
129        let staging_blob_path = staging_blob_path(&path);
130        let blob_storage = StagingBlobStorage::new(staging_blob_path, object_store.clone());
131        let staging_manifest_path = staging_manifest_path(&path);
132        let delta_storage = DeltaStorage::new(
133            staging_manifest_path.clone(),
134            object_store.clone(),
135            compress_type,
136            // StagingStorage does not use a manifest cache; set to None.
137            None,
138            // StagingStorage does not track file sizes, since all staging files are
139            // deleted after exiting staging mode.
140            Arc::new(NoopTracker),
141        );
142
143        Self {
144            delta_storage,
145            blob_storage,
146        }
147    }
148
149    /// Returns the blob storage.
150    pub(crate) fn blob_storage(&self) -> &StagingBlobStorage {
151        &self.blob_storage
152    }
153
154    /// Returns an iterator of manifests from staging directory.
155    pub(crate) async fn manifest_lister(&self) -> Result<Option<Lister>> {
156        self.delta_storage.manifest_lister(None).await
157    }
158
159    /// Fetch all staging manifest files and return them as (version, action_list) pairs.
160    pub(crate) async fn fetch_manifests(&self) -> Result<Vec<(ManifestVersion, Vec<u8>)>> {
161        let manifest_entries = self
162            .delta_storage
163            .get_paths(None, |entry| {
164                let file_name = entry.name();
165                if is_delta_file(file_name) {
166                    let version = file_version(file_name);
167                    Some((version, entry))
168                } else {
169                    None
170                }
171            })
172            .await?;
173
174        let mut sorted_entries = manifest_entries;
175        sort_manifests(&mut sorted_entries);
176
177        self.delta_storage
178            .fetch_manifests_from_entries(sorted_entries)
179            .await
180    }
181
182    /// Save the delta manifest file.
183    pub(crate) async fn save(&mut self, version: ManifestVersion, bytes: &[u8]) -> Result<()> {
184        self.delta_storage.save(version, bytes).await
185    }
186
187    /// Clean all staging manifest files.
188    pub(crate) async fn clear(&self) -> Result<()> {
189        self.delta_storage
190            .object_store()
191            .delete_with(self.delta_storage.path())
192            .recursive(true)
193            .await
194            .context(OpenDalSnafu)?;
195
196        self.blob_storage
197            .object_store()
198            .delete_with(self.blob_storage.path())
199            .recursive(true)
200            .await
201            .context(OpenDalSnafu)?;
202
203        debug!(
204            "Cleared all staging manifest files from {}, blob directory: {}",
205            self.delta_storage.path(),
206            self.blob_storage.path()
207        );
208
209        Ok(())
210    }
211}
212
213#[cfg(test)]
214impl StagingStorage {
215    pub fn set_compress_type(&mut self, compress_type: CompressionType) {
216        self.delta_storage.set_compress_type(compress_type);
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use std::sync::{Arc, Mutex};
223
224    use common_datasource::compression::CompressionType;
225    use common_test_util::temp_dir::create_temp_dir;
226    use futures::TryStreamExt;
227    use object_store::layers::mock::{MockLayerBuilder, ObservedWriter, WriterObservations, oio};
228    use object_store::services::Fs;
229    use object_store::{ErrorKind, ObjectStore};
230
231    use super::{StagingBlobStorage, StagingStorage, staging_blob_path, staging_manifest_path};
232
233    // A test-only per-write limit, not a storage provider's default limit.
234    const WRITE_LIMIT: usize = 64 * 1024 * 1024;
235
236    struct LimitedWriter {
237        inner: oio::Writer,
238    }
239
240    impl oio::Write for LimitedWriter {
241        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
242            if bytes.len() > WRITE_LIMIT {
243                return Err(object_store::Error::new(
244                    ErrorKind::Unexpected,
245                    "test write size limit exceeded",
246                ));
247            }
248            self.inner.write(bytes).await
249        }
250
251        async fn close(&mut self) -> object_store::Result<object_store::layers::mock::Metadata> {
252            self.inner.close().await
253        }
254
255        async fn abort(&mut self) -> object_store::Result<()> {
256            self.inner.abort().await
257        }
258    }
259
260    #[tokio::test]
261    async fn test_staging_blob_put_chunked() {
262        let tmp_dir = create_temp_dir("test_staging_blob_put_chunked");
263        let observations = Arc::new(Mutex::new(WriterObservations::default()));
264        let captured = observations.clone();
265        let layer = MockLayerBuilder::default()
266            .writer_factory(Arc::new(move |_, _, inner| {
267                Box::new(ObservedWriter::new(
268                    Box::new(LimitedWriter { inner }),
269                    captured.clone(),
270                ))
271            }))
272            .build()
273            .unwrap();
274        let store = ObjectStore::new(Fs::default().root(&tmp_dir.path().to_string_lossy()))
275            .unwrap()
276            .layer(layer);
277        let storage = StagingBlobStorage::new("staging/blob/".to_string(), store);
278
279        for size in [17, WRITE_LIMIT, WRITE_LIMIT + 17] {
280            let bytes = (0..size).map(|i| (i % 251) as u8).collect::<Vec<_>>();
281            *observations.lock().unwrap() = WriterObservations::default();
282            storage.put("manifest", bytes.clone()).await.unwrap();
283            assert_eq!(storage.get("manifest").await.unwrap(), bytes);
284            let observed = observations.lock().unwrap();
285            assert_eq!(observed.write_sizes.iter().sum::<usize>(), size);
286            assert!(observed.write_sizes.iter().all(|size| *size <= WRITE_LIMIT));
287            assert_eq!(observed.write_sizes.len(), size.div_ceil(WRITE_LIMIT));
288            assert_eq!(observed.close_calls, 1);
289        }
290    }
291
292    #[test]
293    fn test_staging_path() {
294        let path = "/data/table/region_0001/manifest/";
295        let expected = "/data/table/region_0001/staging/manifest/";
296        assert_eq!(staging_manifest_path(path), expected);
297    }
298
299    #[test]
300    fn test_staging_blob_path() {
301        let path = "/data/table/region_0001/manifest/";
302        let expected = "/data/table/region_0001/staging/blob/";
303        assert_eq!(staging_blob_path(path), expected);
304    }
305
306    async fn count_entries(object_store: &ObjectStore, path: &str) -> usize {
307        match object_store.lister_with(path).await {
308            Ok(lister) => lister.try_collect::<Vec<_>>().await.unwrap().len(),
309            Err(err) if err.kind() == ErrorKind::NotFound => 0,
310            Err(err) => panic!("failed to list {path}: {err}"),
311        }
312    }
313
314    #[tokio::test]
315    async fn test_staging_clear_removes_blobs() {
316        common_telemetry::init_default_ut_logging();
317
318        let tmp_dir = create_temp_dir("test_staging_storage_clear");
319        let builder = Fs::default().root(&tmp_dir.path().to_string_lossy());
320        let object_store = ObjectStore::new(builder).unwrap();
321        let manifest_path = "/data/table/region_0001/manifest/";
322        let mut storage = StagingStorage::new(
323            manifest_path.to_string(),
324            object_store.clone(),
325            CompressionType::Uncompressed,
326        );
327        let manifest_dir = staging_manifest_path(manifest_path);
328        let blob_dir = staging_blob_path(manifest_path);
329
330        storage.save(0, b"manifest").await.unwrap();
331        storage
332            .blob_storage()
333            .put("region_0002_manifest", b"blob".to_vec())
334            .await
335            .unwrap();
336
337        assert!(count_entries(&object_store, &manifest_dir).await > 0);
338        assert!(count_entries(&object_store, &blob_dir).await > 0);
339
340        storage.clear().await.unwrap();
341
342        assert_eq!(count_entries(&object_store, &manifest_dir).await, 0);
343        assert_eq!(count_entries(&object_store, &blob_dir).await, 0);
344    }
345}
mito2/manifest/storage/staging.rs

mito2/manifest/storage/
staging.rs
