mito2/sst/index/
puffin_manager.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::path::Path;
16use std::sync::Arc;
17use std::time::Duration;
18
19use async_trait::async_trait;
20use common_error::ext::BoxedError;
21use object_store::{FuturesAsyncWriter, ObjectStore};
22use puffin::error::{self as puffin_error, Result as PuffinResult};
23use puffin::puffin_manager::file_accessor::PuffinFileAccessor;
24use puffin::puffin_manager::fs_puffin_manager::FsPuffinManager;
25use puffin::puffin_manager::stager::{BoundedStager, Stager};
26use puffin::puffin_manager::{BlobGuard, PuffinManager, PuffinReader};
27use snafu::ResultExt;
28
29use crate::access_layer::FilePathProvider;
30use crate::error::{PuffinInitStagerSnafu, PuffinPurgeStagerSnafu, Result};
31use crate::metrics::{
32    INDEX_PUFFIN_FLUSH_OP_TOTAL, INDEX_PUFFIN_READ_BYTES_TOTAL, INDEX_PUFFIN_READ_OP_TOTAL,
33    INDEX_PUFFIN_WRITE_BYTES_TOTAL, INDEX_PUFFIN_WRITE_OP_TOTAL, StagerMetrics,
34};
35use crate::sst::file::RegionIndexId;
36use crate::sst::index::store::{self, InstrumentedStore};
37
38type InstrumentedRangeReader = store::InstrumentedRangeReader<'static>;
39type InstrumentedAsyncWrite = store::InstrumentedAsyncWrite<'static, FuturesAsyncWriter>;
40
41pub(crate) type SstPuffinManager =
42    FsPuffinManager<Arc<BoundedStager<RegionIndexId>>, ObjectStorePuffinFileAccessor>;
43pub(crate) type SstPuffinReader = <SstPuffinManager as PuffinManager>::Reader;
44pub(crate) type SstPuffinWriter = <SstPuffinManager as PuffinManager>::Writer;
45pub(crate) type SstPuffinBlob = <SstPuffinReader as PuffinReader>::Blob;
46pub(crate) type SstPuffinDir = <SstPuffinReader as PuffinReader>::Dir;
47pub(crate) type BlobReader = <SstPuffinBlob as BlobGuard>::Reader;
48
49const STAGING_DIR: &str = "staging";
50
51/// A factory for creating `SstPuffinManager` instances.
52#[derive(Clone)]
53pub struct PuffinManagerFactory {
54    /// The stager used by the puffin manager.
55    stager: Arc<BoundedStager<RegionIndexId>>,
56
57    /// The size of the write buffer used to create object store.
58    write_buffer_size: Option<usize>,
59}
60
61impl PuffinManagerFactory {
62    /// Creates a new `PuffinManagerFactory` instance.
63    pub async fn new(
64        aux_path: impl AsRef<Path>,
65        staging_capacity: u64,
66        write_buffer_size: Option<usize>,
67        staging_ttl: Option<Duration>,
68    ) -> Result<Self> {
69        let staging_dir = aux_path.as_ref().join(STAGING_DIR);
70        let stager = BoundedStager::new(
71            staging_dir,
72            staging_capacity,
73            Some(Arc::new(StagerMetrics::default())),
74            staging_ttl,
75        )
76        .await
77        .context(PuffinInitStagerSnafu)?;
78        Ok(Self {
79            stager: Arc::new(stager),
80            write_buffer_size,
81        })
82    }
83
84    /// Overrides only this task's upload buffers, sharing the existing stager.
85    pub(crate) fn with_write_buffer_size(mut self, bytes: usize) -> Self {
86        self.write_buffer_size = Some(bytes);
87        self
88    }
89
90    pub(crate) fn build(
91        &self,
92        store: ObjectStore,
93        path_provider: impl FilePathProvider + 'static,
94    ) -> SstPuffinManager {
95        let store = InstrumentedStore::new(store).with_write_buffer_size(self.write_buffer_size);
96        let puffin_file_accessor =
97            ObjectStorePuffinFileAccessor::new(store, Arc::new(path_provider));
98        SstPuffinManager::new(self.stager.clone(), puffin_file_accessor)
99    }
100
101    pub(crate) async fn purge_stager(&self, file_id: RegionIndexId) -> Result<()> {
102        self.stager
103            .purge(&file_id)
104            .await
105            .context(PuffinPurgeStagerSnafu)
106    }
107}
108
109#[cfg(test)]
110impl PuffinManagerFactory {
111    pub(crate) async fn new_for_test_async(
112        prefix: &str,
113    ) -> (common_test_util::temp_dir::TempDir, Self) {
114        let tempdir = common_test_util::temp_dir::create_temp_dir(prefix);
115        let factory = Self::new(tempdir.path().to_path_buf(), 1024, None, None)
116            .await
117            .unwrap();
118        (tempdir, factory)
119    }
120
121    pub(crate) fn new_for_test_block(prefix: &str) -> (common_test_util::temp_dir::TempDir, Self) {
122        let tempdir = common_test_util::temp_dir::create_temp_dir(prefix);
123
124        let f = Self::new(tempdir.path().to_path_buf(), 1024, None, None);
125        let factory = common_runtime::block_on_global(f).unwrap();
126
127        (tempdir, factory)
128    }
129}
130
131/// A `PuffinFileAccessor` implementation that uses an object store as the underlying storage.
132#[derive(Clone)]
133pub(crate) struct ObjectStorePuffinFileAccessor {
134    object_store: InstrumentedStore,
135    path_provider: Arc<dyn FilePathProvider>,
136}
137
138impl ObjectStorePuffinFileAccessor {
139    pub fn new(object_store: InstrumentedStore, path_provider: Arc<dyn FilePathProvider>) -> Self {
140        Self {
141            object_store,
142            path_provider,
143        }
144    }
145
146    pub fn store(&self) -> &InstrumentedStore {
147        &self.object_store
148    }
149}
150
151#[async_trait]
152impl PuffinFileAccessor for ObjectStorePuffinFileAccessor {
153    type Reader = InstrumentedRangeReader;
154    type Writer = InstrumentedAsyncWrite;
155    type FileHandle = RegionIndexId;
156
157    async fn reader(&self, handle: &RegionIndexId) -> PuffinResult<Self::Reader> {
158        let file_path = self
159            .path_provider
160            .build_index_file_path_with_version(*handle);
161        self.object_store
162            .range_reader(
163                &file_path,
164                &INDEX_PUFFIN_READ_BYTES_TOTAL,
165                &INDEX_PUFFIN_READ_OP_TOTAL,
166            )
167            .await
168            .map_err(BoxedError::new)
169            .context(puffin_error::ExternalSnafu)
170    }
171
172    async fn writer(&self, handle: &RegionIndexId) -> PuffinResult<Self::Writer> {
173        let file_path = self
174            .path_provider
175            .build_index_file_path_with_version(*handle);
176        self.object_store
177            .writer(
178                &file_path,
179                &INDEX_PUFFIN_WRITE_BYTES_TOTAL,
180                &INDEX_PUFFIN_WRITE_OP_TOTAL,
181                &INDEX_PUFFIN_FLUSH_OP_TOTAL,
182            )
183            .await
184            .map_err(BoxedError::new)
185            .context(puffin_error::ExternalSnafu)
186    }
187}
188
189#[cfg(test)]
190mod tests {
191
192    use common_base::range_read::RangeReader;
193    use common_test_util::temp_dir::create_temp_dir;
194    use futures::io::Cursor;
195    use object_store::services::Memory;
196    use puffin::blob_metadata::CompressionCodec;
197    use puffin::puffin_manager::{PuffinManager, PuffinReader, PuffinWriter, PutOptions};
198    use store_api::storage::FileId;
199
200    use super::*;
201    use crate::sst::file::{RegionFileId, RegionIndexId};
202
203    struct TestFilePathProvider;
204
205    impl FilePathProvider for TestFilePathProvider {
206        fn build_index_file_path(&self, file_id: RegionFileId) -> String {
207            file_id.file_id().to_string()
208        }
209
210        fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String {
211            index_id.file_id.file_id().to_string()
212        }
213
214        fn build_sst_file_path(&self, file_id: RegionFileId) -> String {
215            file_id.file_id().to_string()
216        }
217    }
218
219    #[tokio::test]
220    async fn test_puffin_manager_factory() {
221        let (_dir, factory) =
222            PuffinManagerFactory::new_for_test_async("test_puffin_manager_factory_").await;
223
224        let object_store = ObjectStore::new(Memory::default()).unwrap();
225        let manager = factory.build(object_store, TestFilePathProvider);
226
227        let file_id = RegionIndexId::new(RegionFileId::new(0.into(), FileId::random()), 0);
228        let blob_key = "blob-key";
229        let dir_key = "dir-key";
230        let raw_data = b"hello world!";
231
232        let mut writer = manager.writer(&file_id).await.unwrap();
233        writer
234            .put_blob(
235                blob_key,
236                Cursor::new(raw_data),
237                PutOptions::default(),
238                Default::default(),
239            )
240            .await
241            .unwrap();
242        let dir_data = create_temp_dir("test_puffin_manager_factory_dir_data_");
243        tokio::fs::write(dir_data.path().join("hello"), raw_data)
244            .await
245            .unwrap();
246        writer
247            .put_dir(
248                dir_key,
249                dir_data.path().into(),
250                PutOptions {
251                    compression: Some(CompressionCodec::Zstd),
252                },
253                Default::default(),
254            )
255            .await
256            .unwrap();
257        writer.finish().await.unwrap();
258
259        let reader = manager.reader(&file_id).await.unwrap();
260        let blob_guard = reader.blob(blob_key).await.unwrap();
261        let blob_reader = blob_guard.reader().await.unwrap();
262        let meta = blob_reader.metadata().await.unwrap();
263        let bs = blob_reader.read(0..meta.content_length).await.unwrap();
264        assert_eq!(&*bs, raw_data);
265
266        let (dir_guard, _metrics) = reader.dir(dir_key).await.unwrap();
267        let file = dir_guard.path().join("hello");
268        let data = tokio::fs::read(file).await.unwrap();
269        assert_eq!(data, raw_data);
270    }
271}
mito2/sst/index/puffin_manager.rs

mito2/sst/index/
puffin_manager.rs
