rivet_envoy_client/
async_counter.rs

1use std::sync::Arc;
2use std::sync::Weak;
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::time::Duration;
5
6use parking_lot::Mutex;
7use tokio::sync::Notify;
8
9use crate::time::Instant;
10use crate::utils::sleep;
11
12pub struct AsyncCounter {
13	value: AtomicUsize,
14	zero_notify: Notify,
15	zero_observers: Mutex<Vec<Weak<Notify>>>,
16	change_observers: Mutex<Vec<Weak<Notify>>>,
17	change_callbacks: Mutex<Vec<Arc<dyn Fn() + Send + Sync>>>,
18}
19
20impl AsyncCounter {
21	pub fn new() -> Self {
22		Self {
23			value: AtomicUsize::new(0),
24			zero_notify: Notify::new(),
25			zero_observers: Mutex::new(Vec::new()),
26			change_observers: Mutex::new(Vec::new()),
27			change_callbacks: Mutex::new(Vec::new()),
28		}
29	}
30
31	pub fn register_zero_notify(&self, notify: &Arc<Notify>) {
32		self.zero_observers.lock().push(Arc::downgrade(notify));
33	}
34
35	pub fn register_change_notify(&self, notify: &Arc<Notify>) {
36		self.change_observers.lock().push(Arc::downgrade(notify));
37	}
38
39	pub fn register_change_callback(&self, callback: Arc<dyn Fn() + Send + Sync>) {
40		self.change_callbacks.lock().push(callback);
41	}
42
43	pub fn increment(&self) {
44		self.value.fetch_add(1, Ordering::Relaxed);
45		self.notify_change();
46	}
47
48	pub fn decrement(&self) {
49		let prev = self.value.fetch_sub(1, Ordering::AcqRel);
50		debug_assert!(prev > 0, "AsyncCounter decrement below zero");
51		if prev == 1 {
52			self.zero_notify.notify_waiters();
53			let mut observers = self.zero_observers.lock();
54			observers.retain(|observer| {
55				let Some(notify) = observer.upgrade() else {
56					return false;
57				};
58				notify.notify_waiters();
59				true
60			});
61		}
62		self.notify_change();
63	}
64
65	fn notify_change(&self) {
66		let mut observers = self.change_observers.lock();
67		observers.retain(|observer| {
68			let Some(notify) = observer.upgrade() else {
69				return false;
70			};
71			notify.notify_waiters();
72			true
73		});
74		drop(observers);
75
76		let callbacks = self.change_callbacks.lock().clone();
77		for callback in callbacks {
78			callback();
79		}
80	}
81
82	pub fn load(&self) -> usize {
83		self.value.load(Ordering::Acquire)
84	}
85
86	pub async fn wait_zero(&self, deadline: Instant) -> bool {
87		loop {
88			let notified = self.zero_notify.notified();
89			tokio::pin!(notified);
90			notified.as_mut().enable();
91
92			if self.value.load(Ordering::Acquire) == 0 {
93				return true;
94			}
95
96			let timeout = deadline
97				.checked_duration_since(Instant::now())
98				.unwrap_or(Duration::ZERO);
99			tokio::select! {
100				_ = notified => {}
101				_ = sleep(timeout) => return false,
102			}
103		}
104	}
105
106	pub async fn wait_zero_unbounded(&self) {
107		loop {
108			let notified = self.zero_notify.notified();
109			tokio::pin!(notified);
110			notified.as_mut().enable();
111
112			if self.value.load(Ordering::Acquire) == 0 {
113				return;
114			}
115
116			notified.await;
117		}
118	}
119}
120
121impl Default for AsyncCounter {
122	fn default() -> Self {
123		Self::new()
124	}
125}
rivet_envoy_client/async_counter.rs

rivet_envoy_client/
async_counter.rs
