mito2/manifest/storage/
size_tracker.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;
16use std::fmt::Debug;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::{Arc, RwLock};
19
20use store_api::ManifestVersion;
21
22use crate::metrics::manifest_size::RegionManifestSizeMetric;
23
24/// Key to identify a manifest file.
25#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
26pub(crate) enum FileKey {
27    /// A delta file (`.json`).
28    Delta(ManifestVersion),
29    /// A checkpoint file (`.checkpoint`).
30    Checkpoint(ManifestVersion),
31}
32
33pub(crate) trait Tracker: Send + Sync + Debug {
34    fn record(&self, version: ManifestVersion, size: u64);
35}
36
37#[derive(Debug, Clone)]
38pub struct CheckpointTracker {
39    size_tracker: SizeTracker,
40}
41
42impl Tracker for CheckpointTracker {
43    fn record(&self, version: ManifestVersion, size: u64) {
44        self.size_tracker.record(FileKey::Checkpoint(version), size);
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct DeltaTracker {
50    size_tracker: SizeTracker,
51}
52
53impl Tracker for DeltaTracker {
54    fn record(&self, version: ManifestVersion, size: u64) {
55        self.size_tracker.record(FileKey::Delta(version), size);
56    }
57}
58
59#[derive(Debug, Clone)]
60pub struct NoopTracker;
61
62impl Tracker for NoopTracker {
63    fn record(&self, _version: ManifestVersion, _size: u64) {
64        // noop
65    }
66}
67
68#[derive(Debug, Clone, Default)]
69pub(crate) struct SizeTracker {
70    inventory: Arc<RwLock<Inventory>>,
71    total_size: Arc<AtomicU64>,
72    metric: Option<RegionManifestSizeMetric>,
73}
74
75#[derive(Debug, Default)]
76struct Inventory {
77    file_sizes: HashMap<FileKey, u64>,
78    checkpoint_version: Option<ManifestVersion>,
79}
80
81impl SizeTracker {
82    /// Returns a new [SizeTracker].
83    pub fn new(total_size: Arc<AtomicU64>, metric: Option<RegionManifestSizeMetric>) -> Self {
84        Self {
85            inventory: Arc::default(),
86            total_size,
87            metric,
88        }
89    }
90
91    /// Returns the manifest tracker.
92    pub(crate) fn manifest_tracker(&self) -> DeltaTracker {
93        DeltaTracker {
94            size_tracker: self.clone(),
95        }
96    }
97
98    /// Returns the checkpoint tracker.
99    pub(crate) fn checkpoint_tracker(&self) -> CheckpointTracker {
100        CheckpointTracker {
101            size_tracker: self.clone(),
102        }
103    }
104
105    /// Returns the total tracked size.
106    pub(crate) fn total(&self) -> u64 {
107        self.total_size.load(Ordering::Relaxed)
108    }
109
110    /// Whether this checkpoint's stored size is already known.
111    pub(crate) fn contains_checkpoint(&self, version: ManifestVersion) -> bool {
112        self.inventory
113            .read()
114            .unwrap()
115            .file_sizes
116            .contains_key(&FileKey::Checkpoint(version))
117    }
118
119    /// Keeps the latest observed checkpoint and deltas after it, regardless of cleanup.
120    /// Physical storage may retain older checkpoints and deltas for longer.
121    pub(crate) fn retain_after(&self, version: ManifestVersion) {
122        let mut inventory = self.inventory.write().unwrap();
123        if inventory
124            .checkpoint_version
125            .is_some_and(|current| current > version)
126        {
127            return;
128        }
129        inventory.checkpoint_version = Some(version);
130        inventory.file_sizes.retain(|key, _| match key {
131            FileKey::Delta(v) => *v > version,
132            FileKey::Checkpoint(v) => *v >= version,
133        });
134        self.total_size
135            .store(inventory.file_sizes.values().sum(), Ordering::Relaxed);
136        self.publish_locked();
137    }
138
139    /// Called with the inventory write lock held, preserving publication order.
140    fn publish_locked(&self) {
141        if let Some(metric) = &self.metric {
142            metric.publish(self.total_size.load(Ordering::Relaxed));
143        }
144    }
145
146    fn record(&self, key: FileKey, size: u64) {
147        let mut inventory = self.inventory.write().unwrap();
148        // Manifest versions are monotonic: deltas at or below the checkpoint watermark
149        // are superseded, even if a read started before publication finishes afterward.
150        // Such a late read must not resurrect old entries.
151        if inventory
152            .checkpoint_version
153            .is_some_and(|version| match key {
154                FileKey::Delta(v) => v <= version,
155                FileKey::Checkpoint(v) => v < version,
156            })
157        {
158            return;
159        }
160        if let Some(old_size) = inventory.file_sizes.insert(key, size) {
161            self.total_size.fetch_sub(old_size, Ordering::Relaxed);
162        }
163        self.total_size.fetch_add(size, Ordering::Relaxed);
164        self.publish_locked();
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use store_api::storage::RegionId;
171
172    use super::*;
173    use crate::metrics::manifest_size::remove_region_manifest_size;
174    use crate::metrics::manifest_size::tests::manifest_metric_sample;
175
176    #[test]
177    fn manifest_metric_tracks_mutations_and_activation() {
178        let id = RegionId::new(u32::MAX, 30);
179        remove_region_manifest_size(id);
180        let total = Arc::new(AtomicU64::new(0));
181        let metric = RegionManifestSizeMetric::with_threshold(id, 0);
182        let tracker = SizeTracker::new(total.clone(), Some(metric.clone()));
183        tracker.manifest_tracker().record(0, 10);
184        assert_eq!(manifest_metric_sample(id), None);
185        metric.activate(&total);
186        assert_eq!(manifest_metric_sample(id), Some(10));
187        tracker.manifest_tracker().record(0, 20);
188        tracker.checkpoint_tracker().record(1, 30);
189        assert_eq!(manifest_metric_sample(id), Some(tracker.total() as i64));
190        assert_eq!(tracker.total(), 50);
191        tracker.retain_after(1);
192        assert_eq!(manifest_metric_sample(id), Some(30));
193        // Late reads and stale retention cannot resurrect superseded bytes or gauges.
194        tracker.manifest_tracker().record(0, 999);
195        tracker.retain_after(0);
196        assert_eq!(manifest_metric_sample(id), Some(30));
197        // An unregistered opener with the same id has no right to publish.
198        let candidate = RegionManifestSizeMetric::with_threshold(id, 0);
199        let candidate_tracker = SizeTracker::new(Arc::new(AtomicU64::new(0)), Some(candidate));
200        candidate_tracker.manifest_tracker().record(2, 999);
201        assert_eq!(manifest_metric_sample(id), Some(30));
202        tracker.retain_after(2);
203        assert_eq!(manifest_metric_sample(id), Some(0));
204        metric.deactivate();
205        remove_region_manifest_size(id);
206        tracker.manifest_tracker().record(3, 99);
207        assert_eq!(manifest_metric_sample(id), None);
208    }
209
210    #[test]
211    fn manifest_metric_retention_removes_series_below_threshold() {
212        let id = RegionId::new(u32::MAX, 44);
213        let total = Arc::new(AtomicU64::new(0));
214        let metric = RegionManifestSizeMetric::with_threshold(id, 100);
215        let tracker = SizeTracker::new(total.clone(), Some(metric.clone()));
216        metric.activate(&total);
217        tracker.manifest_tracker().record(0, 80);
218        assert_eq!(tracker.total(), 80);
219        assert_eq!(manifest_metric_sample(id), None);
220        tracker.checkpoint_tracker().record(1, 30);
221        assert_eq!(manifest_metric_sample(id), Some(110));
222        tracker.retain_after(1);
223        assert_eq!(tracker.total(), 30);
224        assert_eq!(manifest_metric_sample(id), None);
225        tracker.manifest_tracker().record(0, 999);
226        assert_eq!(tracker.total(), 30);
227        assert_eq!(manifest_metric_sample(id), None);
228        metric.deactivate();
229    }
230
231    #[test]
232    fn manifest_metric_publishes_with_inventory_locked() {
233        use crate::metrics::manifest_size::TestHook;
234        let id = RegionId::new(u32::MAX, 31);
235        remove_region_manifest_size(id);
236        let total = Arc::new(AtomicU64::new(0));
237        let metric = RegionManifestSizeMetric::with_threshold(id, 0);
238        let tracker = SizeTracker::new(total.clone(), Some(metric.clone()));
239        metric.activate(&total);
240        let sizes = tracker.inventory.clone();
241        metric.set_test_hook(Some(TestHook::new(move || {
242            assert!(
243                matches!(sizes.try_write(), Err(std::sync::TryLockError::WouldBlock)),
244                "inventory must remain locked through gauge publication"
245            );
246        })));
247        tracker.manifest_tracker().record(0, 10);
248        tracker.manifest_tracker().record(0, 20);
249        tracker.checkpoint_tracker().record(1, 30);
250        tracker.retain_after(1);
251        tracker.retain_after(2);
252        assert_eq!(manifest_metric_sample(id), Some(tracker.total() as i64));
253        metric.deactivate();
254        remove_region_manifest_size(id);
255    }
256
257    #[test]
258    fn retain_checkpoint_and_later_deltas() {
259        let tracker = SizeTracker::default();
260        tracker.retain_after(0);
261        assert_eq!(tracker.total(), 0);
262        for (key, size) in [
263            (FileKey::Delta(1), 10),
264            (FileKey::Delta(2), 20),
265            (FileKey::Delta(3), 30),
266            (FileKey::Checkpoint(1), 40),
267            (FileKey::Checkpoint(2), 50),
268        ] {
269            tracker.record(key, size);
270        }
271        tracker.retain_after(2);
272        assert_eq!(tracker.total(), 80);
273        assert_eq!(
274            tracker.inventory.read().unwrap().file_sizes,
275            HashMap::from([(FileKey::Delta(3), 30), (FileKey::Checkpoint(2), 50)])
276        );
277        // Late reads and an older retention request cannot roll accounting back.
278        tracker.record(FileKey::Delta(2), 20);
279        tracker.record(FileKey::Checkpoint(1), 40);
280        tracker.retain_after(1);
281        assert_eq!(tracker.total(), 80);
282        tracker.record(FileKey::Delta(3), 35);
283        assert_eq!(tracker.total(), 85);
284    }
285}
mito2/manifest/storage/size_tracker.rs

mito2/manifest/storage/
size_tracker.rs
