iota_metrics/
guards.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    future::Future,
7    pin::Pin,
8    task::{Context, Poll},
9};
10
11use prometheus::IntGauge;
12
13/// Increments gauge when acquired, decrements when guard drops
14pub struct GaugeGuard<'a>(&'a IntGauge);
15
16impl<'a> GaugeGuard<'a> {
17    /// Acquires an `IntGauge` by incrementing its value and creating a new
18    /// `IntGaugeGuard` instance that holds a reference to the gauge.
19    pub fn acquire(g: &'a IntGauge) -> Self {
20        g.inc();
21        Self(g)
22    }
23}
24
25impl Drop for GaugeGuard<'_> {
26    /// Decrements the value of the `IntGauge` when the `IntGaugeGuard` is
27    /// dropped.
28    fn drop(&mut self) {
29        self.0.dec();
30    }
31}
32
33pub trait GaugeGuardFutureExt: Future + Sized {
34    /// Count number of in flight futures running
35    fn count_in_flight(self, g: &IntGauge) -> GaugeGuardFuture<Self>;
36}
37
38impl<F: Future> GaugeGuardFutureExt for F {
39    /// Count number of in flight futures running.
40    fn count_in_flight(self, g: &IntGauge) -> GaugeGuardFuture<Self> {
41        GaugeGuardFuture {
42            f: Box::pin(self),
43            _guard: GaugeGuard::acquire(g),
44        }
45    }
46}
47
48/// A struct that wraps a future (`f`) with a `GaugeGuard`. The
49/// `GaugeGuardFuture` is used to manage the lifecycle of a future while
50/// ensuring the associated `GaugeGuard` properly tracks the resource usage
51/// during the future's execution. The guard increments the gauge
52/// when the future starts and decrements it when the `GaugeGuardFuture` is
53/// dropped.
54pub struct GaugeGuardFuture<'a, F: Sized> {
55    f: Pin<Box<F>>,
56    _guard: GaugeGuard<'a>,
57}
58
59impl<F: Future> Future for GaugeGuardFuture<'_, F> {
60    type Output = F::Output;
61
62    /// Polls the wrapped future (`f`) to determine its readiness. This function
63    /// forwards the poll operation to the inner future, allowing the
64    /// `GaugeGuardFuture` to manage the polling lifecycle.
65    /// Returns `Poll::Pending` if the future is not ready or `Poll::Ready` with
66    /// the future's result if complete.
67    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
68        self.f.as_mut().poll(cx)
69    }
70}
iota_metrics/guards.rs

iota_metrics/
guards.rs
