mito2/sst/
file_ref.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::collections::{HashMap, HashSet};
16use std::sync::Arc;
17
18use common_telemetry::debug;
19use dashmap::{DashMap, Entry};
20use store_api::storage::{FileRef, FileRefsManifest, RegionId};
21
22use crate::error::Result;
23use crate::metrics::GC_REF_FILE_CNT;
24use crate::region::MitoRegionRef;
25use crate::sst::file::FileMeta;
26
27/// File references grouped by their physical source region.
28#[derive(Debug, Clone, Default)]
29struct RegionFileRefs {
30    /// Counts independent metadata views of each physical file and index version.
31    files: HashMap<FileRef, usize>,
32}
33
34/// Tracks reader references to SST metadata views in one datanode for object-store GC.
35#[derive(Debug)]
36pub struct FileReferenceManager {
37    /// Datanode id. used to determine tmp ref file name.
38    node_id: Option<u64>,
39    files_per_region: DashMap<RegionId, RegionFileRefs>,
40    /// Whether global GC is enabled.
41    /// This is only meaningful when using object store (not local filesystem).
42    gc_enabled: bool,
43}
44
45pub type FileReferenceManagerRef = Arc<FileReferenceManager>;
46
47impl FileReferenceManager {
48    pub fn new(node_id: Option<u64>) -> Self {
49        Self {
50            node_id,
51            files_per_region: Default::default(),
52            gc_enabled: false,
53        }
54    }
55
56    /// Creates a new FileReferenceManager with GC configuration.
57    pub fn with_gc_enabled(node_id: Option<u64>, gc_enabled: bool) -> Self {
58        Self {
59            node_id,
60            files_per_region: Default::default(),
61            gc_enabled,
62        }
63    }
64
65    /// Returns whether global GC is enabled.
66    pub fn is_gc_enabled(&self) -> bool {
67        self.gc_enabled
68    }
69
70    fn ref_file_set(&self, region_id: RegionId) -> Option<HashSet<FileRef>> {
71        let region_refs = self.files_per_region.get(&region_id)?.clone();
72        let refs: HashSet<_> = region_refs.files.keys().cloned().collect();
73        debug!(
74            "Get file refs for region {}, node {:?}, {} files",
75            region_id,
76            self.node_id,
77            region_refs.files.len(),
78        );
79        Some(refs)
80    }
81
82    /// Gets references from live SST metadata views and related regions' manifests.
83    pub(crate) async fn get_snapshot_of_file_refs(
84        &self,
85        query_regions_for_mem: Vec<MitoRegionRef>,
86        dst_region_to_src_regions: Vec<(MitoRegionRef, HashSet<RegionId>)>,
87    ) -> Result<FileRefsManifest> {
88        // Imported SSTs are tracked under their source region even if that region is
89        // not open on this node. Capture their old readers before reading manifests.
90        let region_ids: HashSet<_> = query_regions_for_mem
91            .iter()
92            .map(|r| r.region_id())
93            .chain(
94                dst_region_to_src_regions
95                    .iter()
96                    .flat_map(|(_, sources)| sources.iter().copied()),
97            )
98            .collect();
99        let mut ref_files = HashMap::new();
100        for region_id in region_ids {
101            if let Some(files) = self.ref_file_set(region_id) {
102                ref_files.insert(region_id, files);
103            }
104        }
105
106        let mut manifest_version = HashMap::new();
107        let mut cross_region_refs = HashMap::new();
108
109        for (dst_region, src_regions) in &dst_region_to_src_regions {
110            let manifest = dst_region.manifest_ctx.manifest().await;
111            for meta in manifest.files.values() {
112                if src_regions.contains(&meta.region_id) {
113                    cross_region_refs
114                        .entry(meta.region_id)
115                        .or_insert_with(HashSet::new)
116                        .insert(dst_region.region_id());
117                    // GC and repartition hold region read locks, so staging manifests
118                    // cannot change ownership while these references are collected.
119                    ref_files
120                        .entry(meta.region_id)
121                        .or_insert_with(HashSet::new)
122                        .insert(FileRef::new(
123                            meta.region_id,
124                            meta.file_id,
125                            meta.index_version(),
126                        ));
127                }
128            }
129            manifest_version.insert(dst_region.region_id(), manifest.manifest_version);
130        }
131
132        for r in &query_regions_for_mem {
133            let manifest = r.manifest_ctx.manifest().await;
134            // The manifest only protects its current index generation; older readers
135            // still need their exact-version temporary references.
136            ref_files.entry(r.region_id()).and_modify(|refs| {
137                refs.retain(|f| {
138                    !manifest
139                        .files
140                        .get(&f.file_id)
141                        .is_some_and(|meta| meta.index_version() == f.index_version)
142                });
143            });
144            manifest_version.insert(r.region_id(), manifest.manifest_version);
145        }
146
147        Ok(FileRefsManifest {
148            file_refs: ref_files,
149            manifest_version,
150            cross_region_refs,
151        })
152    }
153
154    /// Registers a metadata view, paired with `remove_file` when its last reader leaves.
155    pub(crate) fn register_file(&self, file_meta: &FileMeta) {
156        let mut region = self
157            .files_per_region
158            .entry(file_meta.region_id)
159            .or_default();
160        let file_ref = FileRef::new(
161            file_meta.region_id,
162            file_meta.file_id,
163            file_meta.index_version(),
164        );
165        let count = region.files.entry(file_ref).or_insert_with(|| {
166            GC_REF_FILE_CNT.inc();
167            0
168        });
169        *count += 1;
170    }
171
172    /// Releases a metadata view's reference, regardless of whether the SST was retired.
173    pub(crate) fn remove_file(&self, file_meta: &FileMeta) {
174        let file_ref = FileRef::new(
175            file_meta.region_id,
176            file_meta.file_id,
177            file_meta.index_version(),
178        );
179        if let Entry::Occupied(mut region) = self.files_per_region.entry(file_meta.region_id) {
180            if let std::collections::hash_map::Entry::Occupied(mut file) =
181                region.get_mut().files.entry(file_ref)
182            {
183                *file.get_mut() -= 1;
184                if *file.get() == 0 {
185                    file.remove();
186                    GC_REF_FILE_CNT.dec();
187                }
188            }
189            if region.get().files.is_empty() {
190                region.remove();
191            }
192        }
193    }
194
195    pub fn node_id(&self) -> Option<u64> {
196        self.node_id
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use store_api::storage::FileId;
203
204    use super::*;
205    use crate::sst::file::IndexType;
206
207    #[test]
208    fn test_index_view_references_release_generations_independently() {
209        let manager = FileReferenceManager::new(None);
210        let meta = FileMeta {
211            region_id: RegionId::new(1, 1),
212            file_id: FileId::random(),
213            available_indexes: smallvec::smallvec![IndexType::InvertedIndex],
214            ..Default::default()
215        };
216        let updated = FileMeta {
217            index_version: 2,
218            ..meta.clone()
219        };
220        manager.register_file(&meta);
221        manager.register_file(&meta);
222        manager.register_file(&updated);
223        let expected: HashSet<_> = [0, 2]
224            .into_iter()
225            .map(|version| FileRef::new(meta.region_id, meta.file_id, Some(version)))
226            .collect();
227        assert_eq!(manager.ref_file_set(meta.region_id).unwrap(), expected);
228
229        // Independent readers of the same index version must both release it.
230        manager.remove_file(&meta);
231        assert_eq!(manager.ref_file_set(meta.region_id).unwrap(), expected);
232        manager.remove_file(&meta);
233        assert_eq!(
234            manager.ref_file_set(meta.region_id).unwrap(),
235            HashSet::from([FileRef::new(meta.region_id, meta.file_id, Some(2))])
236        );
237        manager.remove_file(&updated);
238        assert!(manager.ref_file_set(meta.region_id).is_none());
239    }
240}
mito2/sst/file_ref.rs

mito2/sst/
file_ref.rs
