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_filtered::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
33/// Difference vs `GaugeGuard`: stores the gauge by value to avoid borrowing
34/// issues. Increments the gauge when acquired, decrements when the guard drops.
35pub struct InflightGuard(IntGauge);
36
37impl InflightGuard {
38    /// Acquires an `IntGauge` by incrementing its value and taking ownership of
39    /// the gauge so it can be decremented on drop.
40    pub fn acquire(g: IntGauge) -> Self {
41        g.inc();
42        Self(g)
43    }
44}
45
46impl Drop for InflightGuard {
47    /// Decrements the value of the `IntGauge` when the guard is dropped.
48    fn drop(&mut self) {
49        self.0.dec();
50    }
51}
52
53pub trait InflightGuardFutureExt: Future + Sized {
54    /// Count number of in flight futures running
55    fn count_in_flight(self, g: IntGauge) -> InflightGuardFuture<Self>;
56}
57
58impl<F: Future> InflightGuardFutureExt for F {
59    /// Count number of in flight futures running.
60    fn count_in_flight(self, g: IntGauge) -> InflightGuardFuture<Self> {
61        InflightGuardFuture {
62            f: Box::pin(self),
63            _guard: InflightGuard::acquire(g),
64        }
65    }
66}
67
68/// A struct that wraps a future (`f`) with an `InflightGuard`. The
69/// `InflightGuardFuture` is used to manage the lifecycle of a future while
70/// ensuring the associated `InflightGuard` properly tracks the resource usage
71/// during the future's execution. The guard increments the gauge
72/// when the future starts and decrements it when the `InflightGuardFuture` is
73/// dropped.
74pub struct InflightGuardFuture<F: Sized> {
75    f: Pin<Box<F>>,
76    _guard: InflightGuard,
77}
78
79impl<F: Future> Future for InflightGuardFuture<F> {
80    type Output = F::Output;
81
82    /// Polls the wrapped future (`f`) to determine its readiness. This function
83    /// forwards the poll operation to the inner future, allowing the
84    /// `InflightGuardFuture` to manage the polling lifecycle.
85    /// Returns `Poll::Pending` if the future is not ready or `Poll::Ready` with
86    /// the future's result if complete.
87    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
88        self.f.as_mut().poll(cx)
89    }
90}
iota_metrics/guards.rs

iota_metrics/
guards.rs
