mito2/metrics/
manifest_size.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::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, LazyLock, Mutex, MutexGuard};
17
18use common_base::readable_size::ReadableSize;
19use lazy_static::lazy_static;
20use prometheus::{IntGaugeVec, register_int_gauge_vec};
21use store_api::metrics::REGION_LABEL;
22use store_api::storage::RegionId;
23
24const MANIFEST_SIZE_METRIC_THRESHOLD_ENV: &str =
25    "GREPTIME_REGION_MANIFEST_SIZE_METRIC_THRESHOLD_BYTES";
26const DEFAULT_MANIFEST_SIZE_METRIC_THRESHOLD: ReadableSize = ReadableSize::mb(512);
27
28/// Manifest metrics are enabled independently of the logging query-metrics switch.
29/// The process-wide env threshold is read once on first use; changes require a restart.
30/// Zero exposes all live regions.
31static MANIFEST_SIZE_METRIC_THRESHOLD: LazyLock<u64> = LazyLock::new(|| {
32    let value = std::env::var(MANIFEST_SIZE_METRIC_THRESHOLD_ENV).ok();
33    parse_manifest_size_metric_threshold(value.as_deref())
34});
35
36fn parse_manifest_size_metric_threshold(value: Option<&str>) -> u64 {
37    match value {
38        None => DEFAULT_MANIFEST_SIZE_METRIC_THRESHOLD.as_bytes(),
39        Some(value) => value.parse().unwrap_or_else(|error| {
40            common_telemetry::warn!(
41                "Invalid {MANIFEST_SIZE_METRIC_THRESHOLD_ENV}={value:?}: {error}; expected non-negative bytes; using {DEFAULT_MANIFEST_SIZE_METRIC_THRESHOLD}"
42            );
43            DEFAULT_MANIFEST_SIZE_METRIC_THRESHOLD.as_bytes()
44        }),
45    }
46}
47
48lazy_static! {
49    /// Manifest size of live regions above the threshold, from the region statistics counter.
50    /// Falling to or below the threshold removes the series; it may reappear later.
51    /// An absent series does not imply a zero size.
52    pub(crate) static ref REGION_MANIFEST_SIZE: IntGaugeVec = register_int_gauge_vec!(
53        "greptime_mito_region_manifest_size_bytes",
54        "mito region manifest size in bytes",
55        &[REGION_LABEL]
56    ).unwrap();
57}
58
59/// Publishes the manifest size of `region_id`.
60fn set_region_manifest_size(region_id: RegionId, size: u64) {
61    let region_id = region_id.as_u64().to_string();
62    REGION_MANIFEST_SIZE
63        .with_label_values(&[&region_id])
64        .set(size as i64);
65}
66
67/// Removes the per-region series. Idempotent.
68pub(crate) fn remove_region_manifest_size(region_id: RegionId) {
69    let region_id = region_id.as_u64().to_string();
70    let _ = REGION_MANIFEST_SIZE.remove_label_values(&[&region_id]);
71}
72
73/// Shared publishing state of one region's manifest size gauge.
74///
75/// Publishing can run on the worker or the background checkpoint task
76/// (`common_runtime::spawn_global`). Activation, deactivation and publishing share one
77/// mutex, so a late publish can neither resurrect a series removed by
78/// close/drop nor overwrite a successor instance's value.
79#[derive(Debug)]
80struct MetricState {
81    active: bool,
82    published: bool,
83    threshold: u64,
84    /// Test-only pause point between the active check and the gauge write.
85    #[cfg(test)]
86    hook: Option<TestHook>,
87}
88
89#[cfg(test)]
90#[derive(Clone)]
91pub(crate) struct TestHook(Arc<dyn Fn() + Send + Sync>);
92
93#[cfg(test)]
94impl std::fmt::Debug for TestHook {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.write_str("TestHook")
97    }
98}
99
100#[cfg(test)]
101impl TestHook {
102    pub(crate) fn new(hook: impl Fn() + Send + Sync + 'static) -> Self {
103        Self(Arc::new(hook))
104    }
105
106    fn call(&self) {
107        (self.0)();
108    }
109}
110
111/// Publishing state of one region's manifest size gauge.
112///
113/// Clones share publication rights. Only the instance in the live map may publish;
114/// a replacement revokes the old instance's rights before publishing its size.
115#[derive(Debug, Clone)]
116pub(crate) struct RegionManifestSizeMetric {
117    region_id: RegionId,
118    state: Arc<Mutex<MetricState>>,
119}
120
121impl RegionManifestSizeMetric {
122    pub(crate) fn new(region_id: RegionId) -> Self {
123        Self::with_threshold(region_id, *MANIFEST_SIZE_METRIC_THRESHOLD)
124    }
125
126    pub(crate) fn with_threshold(region_id: RegionId, threshold: u64) -> Self {
127        Self {
128            region_id,
129            state: Arc::new(Mutex::new(MetricState {
130                active: false,
131                published: false,
132                threshold,
133                #[cfg(test)]
134                hook: None,
135            })),
136        }
137    }
138
139    /// Overrides the threshold without changing the process environment.
140    #[cfg(test)]
141    pub(crate) fn set_test_threshold(&self, threshold: u64) {
142        let mut state = self.lock();
143        assert!(state.active, "test region must already be live");
144        state.threshold = threshold;
145    }
146
147    /// Installs a pause hook used by the deterministic concurrency tests.
148    #[cfg(test)]
149    pub(crate) fn set_test_hook(&self, hook: Option<TestHook>) {
150        self.lock().hook = hook;
151    }
152
153    fn lock(&self) -> MutexGuard<'_, MetricState> {
154        // Metrics must keep working after an unrelated panic poisons the lock.
155        self.state.lock().unwrap_or_else(|e| e.into_inner())
156    }
157
158    /// Activates publication and applies the threshold to the current total.
159    ///
160    /// The total is read while holding the lock: a concurrent writer that
161    /// skipped publishing because the gauge was not active yet cannot be lost.
162    /// Does not call [`Self::publish`] to avoid taking the lock twice.
163    pub(crate) fn activate(&self, total: &AtomicU64) {
164        let mut state = self.lock();
165        state.active = true;
166        let total = total.load(Ordering::Relaxed);
167        self.write(&mut state, total);
168    }
169
170    /// Stops publishing. Called when the instance leaves the live map or is replaced.
171    ///
172    /// Removes the series before a successor instance can publish.
173    pub(crate) fn deactivate(&self) {
174        let mut state = self.lock();
175        state.active = false;
176        if state.published {
177            remove_region_manifest_size(self.region_id);
178            state.published = false;
179        }
180    }
181
182    /// Writes `total` when active.
183    ///
184    /// Callers must hold the size tracker's write lock so a concurrent publisher
185    /// cannot publish an older total afterwards.
186    pub(crate) fn publish(&self, total: u64) {
187        let mut state = self.lock();
188        if state.active {
189            self.write(&mut state, total);
190        }
191    }
192
193    /// Writes the gauge. The caller keeps the state lock alive across this call,
194    /// so the pause hook (and the write) cannot interleave with activate/deactivate.
195    fn write(&self, state: &mut MetricState, total: u64) {
196        #[cfg(test)]
197        if let Some(hook) = &state.hook {
198            hook.call();
199        }
200        if state.threshold == 0 || total > state.threshold {
201            set_region_manifest_size(self.region_id, total);
202            state.published = true;
203        } else if state.published {
204            remove_region_manifest_size(self.region_id);
205            state.published = false;
206        }
207    }
208}
209
210#[cfg(test)]
211pub(crate) mod tests {
212    use std::sync::mpsc;
213    use std::time::Duration;
214
215    use super::*;
216
217    /// Looks up a sample without creating a series. Region ids must be unique per test.
218    ///
219    /// Shared by the engine lifecycle, size tracker, and metrics tests.
220    pub(crate) fn manifest_metric_sample(region_id: RegionId) -> Option<i64> {
221        use prometheus::core::Collector;
222        let label = region_id.as_u64().to_string();
223        let samples: Vec<_> = REGION_MANIFEST_SIZE
224            .collect()
225            .into_iter()
226            .flat_map(|family| family.get_metric().to_vec())
227            .filter(|metric| {
228                metric
229                    .get_label()
230                    .iter()
231                    .any(|pair| pair.name() == REGION_LABEL && pair.value() == label)
232            })
233            .collect();
234        assert!(samples.len() <= 1);
235        samples
236            .first()
237            .map(|metric| metric.get_gauge().as_ref().unwrap().value() as i64)
238    }
239
240    #[test]
241    fn manifest_metric_threshold_parsing() {
242        for (input, expected) in [
243            (None, 536_870_912),
244            (Some("536870912"), 536_870_912),
245            (Some("123"), 123),
246            (Some("0"), 0),
247            (Some(""), 536_870_912),
248            (Some("-1"), 536_870_912),
249            (Some("512MiB"), 536_870_912),
250            (Some("18446744073709551616"), 536_870_912),
251        ] {
252            assert_eq!(parse_manifest_size_metric_threshold(input), expected);
253        }
254    }
255
256    #[rstest::rstest]
257    #[case::filtered(100, 40)]
258    #[case::unfiltered(0, 41)]
259    fn manifest_metric_threshold_transitions(#[case] threshold: u64, #[case] sequence: u32) {
260        let id = RegionId::new(u32::MAX, sequence);
261        let metric = RegionManifestSizeMetric::with_threshold(id, threshold);
262        metric.publish(101);
263        assert_eq!(manifest_metric_sample(id), None);
264        metric.activate(&AtomicU64::new(99));
265        assert_eq!(manifest_metric_sample(id), (threshold == 0).then_some(99));
266        for total in [100, 101, 150, 100, 99, 0, 101] {
267            metric.publish(total);
268            let expected = (threshold == 0 || total > threshold).then_some(total as i64);
269            assert_eq!(manifest_metric_sample(id), expected);
270            assert_eq!(metric.lock().published, expected.is_some());
271        }
272        metric.deactivate();
273        assert_eq!(manifest_metric_sample(id), None);
274        metric.publish(200);
275        assert_eq!(manifest_metric_sample(id), None);
276    }
277
278    #[test]
279    fn manifest_metric_replacement_below_threshold() {
280        let id = RegionId::new(u32::MAX, 42);
281        let old = RegionManifestSizeMetric::with_threshold(id, 100);
282        old.activate(&AtomicU64::new(200));
283        assert_eq!(manifest_metric_sample(id), Some(200));
284        let new = RegionManifestSizeMetric::with_threshold(id, 100);
285        // An inactive opener must not remove the live instance's series.
286        new.publish(50);
287        new.deactivate();
288        assert_eq!(manifest_metric_sample(id), Some(200));
289        old.deactivate();
290        new.activate(&AtomicU64::new(50));
291        assert_eq!(manifest_metric_sample(id), None);
292        new.publish(150);
293        old.publish(300);
294        old.deactivate();
295        assert_eq!(manifest_metric_sample(id), Some(150));
296        new.deactivate();
297    }
298
299    #[rstest::rstest]
300    #[case::remove(false, false)]
301    #[case::replace(true, false)]
302    #[case::threshold_removal(true, true)]
303    fn publishing_holds_lock_through_gauge_write(
304        #[case] replace: bool,
305        #[case] below_threshold: bool,
306    ) {
307        let id = RegionId::new(
308            u32::MAX,
309            if below_threshold {
310                43
311            } else if replace {
312                21
313            } else {
314                20
315            },
316        );
317        remove_region_manifest_size(id);
318        let old = RegionManifestSizeMetric::with_threshold(id, 5);
319        old.activate(&AtomicU64::new(10));
320        let new = RegionManifestSizeMetric::with_threshold(id, 0);
321        let (paused_tx, paused_rx) = mpsc::channel();
322        let (go_tx, go_rx) = mpsc::channel();
323        let go_rx = Mutex::new(go_rx);
324        old.set_test_hook(Some(TestHook::new(move || {
325            paused_tx.send(()).unwrap();
326            go_rx
327                .lock()
328                .unwrap()
329                .recv_timeout(Duration::from_secs(5))
330                .unwrap();
331        })));
332        let publisher = std::thread::spawn({
333            let old = old.clone();
334            move || old.publish(if below_threshold { 5 } else { 11 })
335        });
336        paused_rx.recv_timeout(Duration::from_secs(5)).unwrap();
337        // Directly prove the lock is held at the gauge write boundary. This fails
338        // deterministically if publication releases its guard before writing.
339        let held = matches!(
340            old.state.try_lock(),
341            Err(std::sync::TryLockError::WouldBlock)
342        );
343        let (started_tx, started_rx) = mpsc::channel();
344        let successor = std::thread::spawn({
345            let old = old.clone();
346            let new = new.clone();
347            move || {
348                started_tx.send(()).unwrap();
349                old.deactivate();
350                if replace {
351                    new.activate(&AtomicU64::new(20));
352                }
353            }
354        });
355        let started = started_rx.recv_timeout(Duration::from_secs(5));
356        // Release before assertions so a failing assertion cannot strand threads.
357        go_tx.send(()).unwrap();
358        publisher.join().unwrap();
359        successor.join().unwrap();
360        assert!(
361            held,
362            "publication must hold the descriptor lock through the write"
363        );
364        started.unwrap();
365        old.set_test_hook(None);
366        old.publish(1);
367        assert_eq!(manifest_metric_sample(id), replace.then_some(20));
368        new.deactivate();
369        remove_region_manifest_size(id);
370    }
371}
mito2/metrics/manifest_size.rs

mito2/metrics/
manifest_size.rs
