use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
pub struct InFlightGauge {
count: AtomicU64,
idle: tokio::sync::Notify,
}
impl InFlightGauge {
pub fn new() -> Self {
Self {
count: AtomicU64::new(0),
idle: tokio::sync::Notify::new(),
}
}
pub(crate) fn inc(&self) {
self.count.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn dec(&self) {
let prev = self.count.fetch_sub(1, Ordering::Release);
if prev == 1 {
self.idle.notify_waiters();
}
}
pub fn total(&self) -> u64 {
self.count.load(Ordering::Acquire)
}
pub fn idle(&self) -> &tokio::sync::Notify {
&self.idle
}
}
impl Default for InFlightGauge {
fn default() -> Self {
Self::new()
}
}
pub struct InFlightClaim(Arc<InFlightGauge>);
impl InFlightClaim {
pub fn attach(gauge: &Arc<InFlightGauge>) -> Self {
gauge.inc();
Self(Arc::clone(gauge))
}
pub fn split(&self) -> Self {
self.0.inc();
Self(Arc::clone(&self.0))
}
}
impl Drop for InFlightClaim {
fn drop(&mut self) {
self.0.dec();
}
}
impl fmt::Debug for InFlightClaim {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("InFlightClaim")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn claim_attach_increments_and_drop_decrements() {
let gauge = Arc::new(InFlightGauge::new());
let claim = InFlightClaim::attach(&gauge);
assert_eq!(gauge.total(), 1);
drop(claim);
assert_eq!(gauge.total(), 0);
}
#[test]
fn claim_split_adds_one_sibling() {
let gauge = Arc::new(InFlightGauge::new());
let original = InFlightClaim::attach(&gauge);
assert_eq!(gauge.total(), 1);
let sibling = original.split();
assert_eq!(gauge.total(), 2);
drop(sibling);
assert_eq!(gauge.total(), 1);
drop(original);
assert_eq!(gauge.total(), 0);
}
#[test]
fn claim_debug_does_not_leak_pointer() {
let gauge = Arc::new(InFlightGauge::new());
let claim = InFlightClaim::attach(&gauge);
assert_eq!(format!("{claim:?}"), "InFlightClaim");
}
#[test]
fn gauge_counts_claim_lifecycle() {
let gauge = Arc::new(InFlightGauge::new());
let first = InFlightClaim::attach(&gauge);
let second = InFlightClaim::attach(&gauge);
assert_eq!(gauge.total(), 2);
drop(first);
assert_eq!(gauge.total(), 1);
drop(second);
assert_eq!(gauge.total(), 0);
}
#[tokio::test]
async fn gauge_notifies_on_last_release_only() {
let gauge = Arc::new(InFlightGauge::new());
let mut idle = std::pin::pin!(gauge.idle().notified());
idle.as_mut().enable();
let first = InFlightClaim::attach(&gauge);
let second = InFlightClaim::attach(&gauge);
drop(first);
let notified_early =
tokio::time::timeout(std::time::Duration::from_millis(100), &mut idle).await;
assert!(
notified_early.is_err(),
"a non-final release must not notify idle waiters"
);
drop(second);
tokio::time::timeout(std::time::Duration::from_secs(1), &mut idle)
.await
.expect("the last release must resolve the enabled waiter");
}
#[tokio::test]
async fn gauge_notify_wakes_all_enabled_waiters() {
let gauge = Arc::new(InFlightGauge::new());
let mut first = std::pin::pin!(gauge.idle().notified());
first.as_mut().enable();
let mut second = std::pin::pin!(gauge.idle().notified());
second.as_mut().enable();
let claim = InFlightClaim::attach(&gauge);
drop(claim);
tokio::time::timeout(std::time::Duration::from_secs(1), &mut first)
.await
.expect("first enabled waiter must resolve on last release");
tokio::time::timeout(std::time::Duration::from_secs(1), &mut second)
.await
.expect("second enabled waiter must resolve on last release");
}
}