mito2/sst/index/
intermediate.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::hash::BuildHasher;
16use std::io::ErrorKind;
17use std::path::PathBuf;
18
19use async_trait::async_trait;
20use common_base::hash::FixedRandomState;
21use common_error::ext::BoxedError;
22use common_telemetry::{debug, warn};
23use futures::{AsyncRead, AsyncWrite};
24use index::error as index_error;
25use index::error::Result as IndexResult;
26use index::external_provider::ExternalTempFileProvider;
27use object_store::util::{self, normalize_dir};
28use snafu::ResultExt;
29use store_api::storage::{ColumnId, FileId, RegionId};
30use uuid::Uuid;
31
32use crate::access_layer::new_fs_cache_store;
33use crate::error::Result;
34use crate::metrics::{
35    INDEX_INTERMEDIATE_FLUSH_OP_TOTAL, INDEX_INTERMEDIATE_READ_BYTES_TOTAL,
36    INDEX_INTERMEDIATE_READ_OP_TOTAL, INDEX_INTERMEDIATE_SEEK_OP_TOTAL,
37    INDEX_INTERMEDIATE_WRITE_BYTES_TOTAL, INDEX_INTERMEDIATE_WRITE_OP_TOTAL,
38};
39use crate::sst::index::store::InstrumentedStore;
40
41const INTERMEDIATE_DIR: &str = "__intm";
42
43/// `IntermediateManager` provides store to access to intermediate files.
44#[derive(Clone)]
45pub struct IntermediateManager {
46    base_dir: PathBuf,
47    store: InstrumentedStore,
48}
49
50impl IntermediateManager {
51    /// Create a new `IntermediateManager` with the given root path.
52    /// It will clean up all garbage intermediate files from previous runs.
53    pub async fn init_fs(aux_path: impl AsRef<str>) -> Result<Self> {
54        common_telemetry::info!(
55            "Initializing intermediate manager, aux_path: {}",
56            aux_path.as_ref()
57        );
58
59        // Remove the intermediate directory on bankground
60        let aux_pb = PathBuf::from(aux_path.as_ref());
61        let intm_dir = aux_pb.join(INTERMEDIATE_DIR);
62        let deleted_dir = intm_dir.with_extension(format!("deleted-{}", Uuid::new_v4()));
63        match tokio::fs::rename(&intm_dir, &deleted_dir).await {
64            Ok(_) => {
65                tokio::spawn(async move {
66                    if let Err(err) = tokio::fs::remove_dir_all(deleted_dir).await {
67                        warn!(err; "Failed to remove intermediate directory");
68                    }
69                });
70            }
71            Err(err) => {
72                if err.kind() != ErrorKind::NotFound {
73                    warn!(err; "Failed to rename intermediate directory");
74                }
75            }
76        }
77
78        let store = new_fs_cache_store(&normalize_dir(aux_path.as_ref())).await?;
79        let store = InstrumentedStore::new(store);
80
81        Ok(Self {
82            base_dir: PathBuf::from(aux_path.as_ref()),
83            store,
84        })
85    }
86
87    /// Set the write buffer size for the store.
88    pub fn with_buffer_size(mut self, write_buffer_size: Option<usize>) -> Self {
89        self.store = self.store.with_write_buffer_size(write_buffer_size);
90        self
91    }
92
93    /// Returns the store to access to intermediate files.
94    pub(crate) fn store(&self) -> &InstrumentedStore {
95        &self.store
96    }
97
98    /// Returns the intermediate directory path for building fulltext index.
99    /// The format is `{aux_path}/__intm/{region_id}/{sst_file_id}/fulltext-{column_id}-{uuid}`.
100    pub(crate) fn fulltext_path(
101        &self,
102        region_id: &RegionId,
103        sst_file_id: &FileId,
104        column_id: ColumnId,
105    ) -> PathBuf {
106        let uuid = Uuid::new_v4();
107        self.base_dir
108            .join(INTERMEDIATE_DIR)
109            .join(region_id.as_u64().to_string())
110            .join(sst_file_id.to_string())
111            .join(format!("fulltext-{column_id}-{uuid}"))
112    }
113
114    /// Prunes the intermediate directory for SST files.
115    pub(crate) async fn prune_sst_dir(
116        &self,
117        region_id: &RegionId,
118        sst_file_id: &FileId,
119    ) -> Result<()> {
120        let region_id = region_id.as_u64();
121        let sst_dir = format!("{INTERMEDIATE_DIR}/{region_id}/{sst_file_id}/");
122        self.store.remove_all(&sst_dir).await
123    }
124
125    /// Prunes the intermediate directory for region files.
126    pub(crate) async fn prune_region_dir(&self, region_id: &RegionId) -> Result<()> {
127        let region_id = region_id.as_u64();
128        let region_dir = format!("{INTERMEDIATE_DIR}/{region_id}/");
129        self.store.remove_all(&region_dir).await
130    }
131}
132
133/// `IntermediateLocation` produces paths for intermediate files
134/// during external sorting.
135#[derive(Debug, Clone)]
136pub struct IntermediateLocation {
137    files_dir: String,
138}
139
140impl IntermediateLocation {
141    /// Create a new `IntermediateLocation`. Set the root directory to
142    /// `__intm/{region_id}/{sst_file_id}/{uuid}/`, incorporating
143    /// uuid to differentiate active sorting files from orphaned data due to unexpected
144    /// process termination.
145    pub fn new(region_id: &RegionId, sst_file_id: &FileId) -> Self {
146        let region_id = region_id.as_u64();
147        let uuid = Uuid::new_v4();
148        Self {
149            files_dir: format!("{INTERMEDIATE_DIR}/{region_id}/{sst_file_id}/{uuid}/"),
150        }
151    }
152
153    /// Returns the directory to clean up when the sorting is done
154    pub fn dir_to_cleanup(&self) -> &str {
155        &self.files_dir
156    }
157
158    /// Returns the directory path for intermediate files associated with `file_group`:
159    /// `__intm/{region_id}/{sst_file_id}/{uuid}/{hash(file_group)}/`.
160    ///
161    /// Index names may contain JSON paths with filesystem-unsafe characters or exceed
162    /// directory name length limits. Fixed-seed aHash maps them to 16-character hexadecimal
163    /// directory names.
164    pub fn file_group_path(&self, file_group: &str) -> String {
165        let hash = FixedRandomState::new().hash_one(file_group);
166        util::join_path(&self.files_dir, &format!("{hash:016x}/"))
167    }
168
169    /// Returns the path of the intermediate file with the given `file_group` and `im_file_id`:
170    /// `__intm/{region_id}/{sst_file_id}/{uuid}/{hash(file_group)}/{im_file_id}.im`
171    pub fn file_path(&self, file_group: &str, im_file_id: &str) -> String {
172        util::join_path(
173            &self.file_group_path(file_group),
174            &format!("{im_file_id}.im"),
175        )
176    }
177
178    /// Returns the intermediate file id from the path.
179    pub fn im_file_id_from_path(&self, path: &str) -> String {
180        path.rsplit('/')
181            .next()
182            .and_then(|s| s.strip_suffix(".im"))
183            .unwrap_or_default()
184            .to_string()
185    }
186}
187
188/// `TempFileProvider` implements `ExternalTempFileProvider`.
189/// It uses `InstrumentedStore` to create and read intermediate files.
190pub(crate) struct TempFileProvider {
191    /// Provides the location of intermediate files.
192    location: IntermediateLocation,
193    /// Provides store to access to intermediate files.
194    manager: IntermediateManager,
195}
196
197#[async_trait]
198impl ExternalTempFileProvider for TempFileProvider {
199    async fn create(
200        &self,
201        file_group: &str,
202        file_id: &str,
203    ) -> IndexResult<Box<dyn AsyncWrite + Unpin + Send>> {
204        let path = self.location.file_path(file_group, file_id);
205        let writer = self
206            .manager
207            .store()
208            .writer(
209                &path,
210                &INDEX_INTERMEDIATE_WRITE_BYTES_TOTAL,
211                &INDEX_INTERMEDIATE_WRITE_OP_TOTAL,
212                &INDEX_INTERMEDIATE_FLUSH_OP_TOTAL,
213            )
214            .await
215            .map_err(BoxedError::new)
216            .context(index_error::ExternalSnafu)?;
217        Ok(Box::new(writer))
218    }
219
220    async fn read_all(
221        &self,
222        file_group: &str,
223    ) -> IndexResult<Vec<(String, Box<dyn AsyncRead + Unpin + Send>)>> {
224        let file_group_path = self.location.file_group_path(file_group);
225        let entries = self
226            .manager
227            .store()
228            .list(&file_group_path)
229            .await
230            .map_err(BoxedError::new)
231            .context(index_error::ExternalSnafu)?;
232        let mut readers = Vec::with_capacity(entries.len());
233
234        for entry in entries {
235            if entry.metadata().is_dir() {
236                // todo(hl): we can keep this warning once we find a way to filter self in list result.
237                debug!("Unexpected entry in index creation dir: {:?}", entry.path());
238                continue;
239            }
240
241            let im_file_id = self.location.im_file_id_from_path(entry.path());
242
243            let reader = self
244                .manager
245                .store()
246                .reader(
247                    entry.path(),
248                    &INDEX_INTERMEDIATE_READ_BYTES_TOTAL,
249                    &INDEX_INTERMEDIATE_READ_OP_TOTAL,
250                    &INDEX_INTERMEDIATE_SEEK_OP_TOTAL,
251                )
252                .await
253                .map_err(BoxedError::new)
254                .context(index_error::ExternalSnafu)?;
255            readers.push((im_file_id, Box::new(reader) as _));
256        }
257
258        Ok(readers)
259    }
260}
261
262impl TempFileProvider {
263    /// Creates a new `TempFileProvider`.
264    pub fn new(location: IntermediateLocation, manager: IntermediateManager) -> Self {
265        Self { location, manager }
266    }
267
268    /// Removes all intermediate files.
269    pub async fn cleanup(&self) -> Result<()> {
270        self.manager
271            .store()
272            .remove_all(self.location.dir_to_cleanup())
273            .await
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use std::ffi::OsStr;
280
281    use common_test_util::temp_dir;
282    use futures::{AsyncReadExt, AsyncWriteExt};
283    use regex::Regex;
284    use store_api::storage::{FileId, RegionId};
285
286    use super::*;
287
288    #[tokio::test]
289    async fn test_manager() {
290        let temp_dir = temp_dir::create_temp_dir("index_intermediate");
291        let path = temp_dir.path().to_str().unwrap();
292
293        // write a garbage file
294        tokio::fs::create_dir_all(format!("{path}/{INTERMEDIATE_DIR}"))
295            .await
296            .unwrap();
297        tokio::fs::write(format!("{path}/{INTERMEDIATE_DIR}/garbage.im"), "blahblah")
298            .await
299            .unwrap();
300
301        let _manager = IntermediateManager::init_fs(path).await.unwrap();
302
303        // cleaned up by `init_fs`
304        assert!(
305            !tokio::fs::try_exists(format!("{path}/{INTERMEDIATE_DIR}"))
306                .await
307                .unwrap()
308        );
309    }
310
311    #[tokio::test]
312    async fn test_cleanup_dir() {
313        let temp_dir = temp_dir::create_temp_dir("test_cleanup_dir_");
314
315        let region_id = RegionId::new(0, 0);
316        let sst_file_id = FileId::random();
317        let region_dir = temp_dir
318            .path()
319            .join(INTERMEDIATE_DIR)
320            .join(region_id.as_u64().to_string());
321        let sst_dir = region_dir.join(sst_file_id.to_string());
322
323        let path = temp_dir.path().to_str().unwrap();
324        let manager = IntermediateManager::init_fs(path).await.unwrap();
325
326        let location = IntermediateLocation::new(&region_id, &sst_file_id);
327        let temp_file_provider = TempFileProvider::new(location, manager.clone());
328
329        let mut f1 = temp_file_provider
330            .create("sky", "000000000000")
331            .await
332            .unwrap();
333        f1.write_all(b"hello").await.unwrap();
334        f1.flush().await.unwrap();
335        f1.close().await.unwrap();
336
337        let mut f2 = temp_file_provider
338            .create("sky", "000000000001")
339            .await
340            .unwrap();
341        f2.write_all(b"world").await.unwrap();
342        f2.flush().await.unwrap();
343        f2.close().await.unwrap();
344
345        temp_file_provider.cleanup().await.unwrap();
346
347        // sst_dir and region_dir still exists
348        assert!(tokio::fs::try_exists(&sst_dir).await.unwrap());
349        assert!(tokio::fs::try_exists(&region_dir).await.unwrap());
350
351        // sst_dir should be deleted, region_dir still exists
352        manager
353            .prune_sst_dir(&region_id, &sst_file_id)
354            .await
355            .unwrap();
356        assert!(tokio::fs::try_exists(&region_dir).await.unwrap());
357        assert!(!tokio::fs::try_exists(&sst_dir).await.unwrap());
358
359        // sst_dir, region_dir should be deleted
360        manager.prune_region_dir(&region_id).await.unwrap();
361        assert!(!tokio::fs::try_exists(&sst_dir).await.unwrap());
362        assert!(!tokio::fs::try_exists(&region_dir).await.unwrap());
363    }
364
365    #[test]
366    fn test_intermediate_location() {
367        let sst_file_id = FileId::random();
368        let location = IntermediateLocation::new(&RegionId::new(0, 0), &sst_file_id);
369
370        let re = Regex::new(&format!(
371            "{INTERMEDIATE_DIR}/0/{sst_file_id}/{}/",
372            r"\w{8}-\w{4}-\w{4}-\w{4}-\w{12}"
373        ))
374        .unwrap();
375        assert!(re.is_match(&location.files_dir));
376
377        let uuid = location.files_dir.split('/').nth(3).unwrap();
378
379        let file_group = "1";
380        let group_path = location.file_group_path(file_group);
381        assert_eq!(group_path, location.file_group_path(file_group));
382        let group_dir = group_path.trim_end_matches('/').rsplit('/').next().unwrap();
383        assert_eq!(group_dir.len(), 16);
384        assert!(group_dir.bytes().all(|byte| byte.is_ascii_hexdigit()));
385        assert_eq!(
386            location.file_group_path(file_group),
387            format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{group_dir}/")
388        );
389
390        let im_file_id = "000000000010";
391        let file_path = location.file_path(file_group, im_file_id);
392        assert_eq!(
393            file_path,
394            format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{group_dir}/{im_file_id}.im")
395        );
396
397        assert_eq!(location.im_file_id_from_path(&file_path), im_file_id);
398    }
399
400    #[tokio::test]
401    async fn test_fulltext_intm_path() {
402        let temp_dir = temp_dir::create_temp_dir("test_fulltext_intm_path_");
403        let aux_path = temp_dir.path().to_string_lossy().to_string();
404
405        let manager = IntermediateManager::init_fs(&aux_path).await.unwrap();
406        let region_id = RegionId::new(0, 0);
407        let sst_file_id = FileId::random();
408        let column_id = 1;
409        let fulltext_path = manager.fulltext_path(&region_id, &sst_file_id, column_id);
410
411        let mut pi = fulltext_path.iter();
412        for a in temp_dir.path().iter() {
413            assert_eq!(a, pi.next().unwrap());
414        }
415        assert_eq!(pi.next().unwrap(), INTERMEDIATE_DIR);
416        assert_eq!(pi.next().unwrap(), "0"); // region id
417        assert_eq!(pi.next().unwrap(), OsStr::new(&sst_file_id.to_string())); // sst file id
418        assert!(
419            Regex::new(r"fulltext-1-\w{8}-\w{4}-\w{4}-\w{4}-\w{12}")
420                .unwrap()
421                .is_match(&pi.next().unwrap().to_string_lossy())
422        ); // fulltext path
423        assert!(pi.next().is_none());
424    }
425
426    #[tokio::test]
427    async fn test_temp_file_provider_basic() {
428        let temp_dir = temp_dir::create_temp_dir("intermediate");
429        let path = temp_dir.path().display().to_string();
430
431        let region_id = RegionId::new(0, 0);
432        let location = IntermediateLocation::new(&region_id, &FileId::random());
433        let store = IntermediateManager::init_fs(path).await.unwrap();
434        let provider = TempFileProvider::new(location.clone(), store);
435
436        let group0 = index::target::IndexTarget::new_json_path(1, vec!["a/\":\\".repeat(100)])
437            .unwrap()
438            .encode();
439        let group1 = index::target::IndexTarget::new_json_path(1, vec!["b/\":\\".repeat(100)])
440            .unwrap()
441            .encode();
442        assert!(group0.len() > 255 && group1.len() > 255);
443        let file_group = group0.as_str();
444        let file_id = "0000000010";
445        let mut writer = provider.create(file_group, file_id).await.unwrap();
446        writer.write_all(b"hello").await.unwrap();
447        writer.flush().await.unwrap();
448        writer.close().await.unwrap();
449
450        let file_id = "0000000100";
451        let mut writer = provider.create(file_group, file_id).await.unwrap();
452        writer.write_all(b"world").await.unwrap();
453        writer.flush().await.unwrap();
454        writer.close().await.unwrap();
455
456        let file_group = group1.as_str();
457        let file_id = "0000000010";
458        let mut writer = provider.create(file_group, file_id).await.unwrap();
459        writer.write_all(b"foo").await.unwrap();
460        writer.flush().await.unwrap();
461        writer.close().await.unwrap();
462
463        let readers = provider.read_all(&group0).await.unwrap();
464        assert_eq!(readers.len(), 2);
465        for (_, mut reader) in readers {
466            let mut buf = Vec::new();
467            reader.read_to_end(&mut buf).await.unwrap();
468            assert!(matches!(buf.as_slice(), b"hello" | b"world"));
469        }
470        let readers = provider.read_all(&group1).await.unwrap();
471        assert_eq!(readers.len(), 1);
472        let mut reader = readers.into_iter().map(|x| x.1).next().unwrap();
473        let mut buf = Vec::new();
474        reader.read_to_end(&mut buf).await.unwrap();
475        assert_eq!(buf, b"foo");
476
477        provider.cleanup().await.unwrap();
478
479        assert!(
480            provider
481                .manager
482                .store()
483                .list(location.dir_to_cleanup())
484                .await
485                .unwrap()
486                .is_empty()
487        );
488    }
489}
mito2/sst/index/intermediate.rs

mito2/sst/index/
intermediate.rs
