mysten_metrics/
guards.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use prometheus::IntGauge;
5use std::future::Future;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8
9/// Increments gauge when acquired, decrements when guard drops
10pub struct GaugeGuard<'a>(&'a IntGauge);
11
12impl<'a> GaugeGuard<'a> {
13    pub fn acquire(g: &'a IntGauge) -> Self {
14        g.inc();
15        Self(g)
16    }
17}
18
19impl Drop for GaugeGuard<'_> {
20    fn drop(&mut self) {
21        self.0.dec();
22    }
23}
24
25/// Difference vs GaugeGuard: Stores the gauge by value to avoid borrowing issues.
26pub struct InflightGuard(IntGauge);
27
28impl InflightGuard {
29    pub fn acquire(g: IntGauge) -> Self {
30        g.inc();
31        Self(g)
32    }
33}
34
35impl Drop for InflightGuard {
36    fn drop(&mut self) {
37        self.0.dec();
38    }
39}
40
41pub trait InflightGuardFutureExt: Future + Sized {
42    /// Count number of in flight futures running
43    fn count_in_flight(self, g: IntGauge) -> InflightGuardFuture<Self>;
44}
45
46impl<F: Future> InflightGuardFutureExt for F {
47    fn count_in_flight(self, g: IntGauge) -> InflightGuardFuture<Self> {
48        InflightGuardFuture {
49            f: Box::pin(self),
50            _guard: InflightGuard::acquire(g),
51        }
52    }
53}
54
55pub struct InflightGuardFuture<F: Sized> {
56    f: Pin<Box<F>>,
57    _guard: InflightGuard,
58}
59
60impl<F: Future> Future for InflightGuardFuture<F> {
61    type Output = F::Output;
62
63    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
64        self.f.as_mut().poll(cx)
65    }
66}
mysten_metrics/guards.rs

mysten_metrics/
guards.rs
