refactor(rate_limit): one balance() behind tokens() and is_idle()

Both took the lock, refilled and read the field, differing only in the
comparison.
This commit is contained in:
2026-09-21 17:34:25 +08:00
parent 2d50440bf1
commit 676b4bc2ba
+10 -7
View File
@@ -91,21 +91,24 @@ impl TokenBucket {
tokio::time::sleep(Duration::from_secs_f64(wait)).await; tokio::time::sleep(Duration::from_secs_f64(wait)).await;
} }
/// Current balance, for the tests that assert a call site charged the /// Current balance, refilled to now.
/// bucket (a charge is otherwise only observable as a delay). fn balance(&self) -> f64 {
#[cfg(test)]
pub(crate) fn tokens(&self) -> f64 {
let mut state = self.state.lock(); let mut state = self.state.lock();
self.refill(&mut state); self.refill(&mut state);
state.tokens state.tokens
} }
/// Current balance, for the tests that assert a call site charged the
/// bucket (a charge is otherwise only observable as a delay).
#[cfg(test)]
pub(crate) fn tokens(&self) -> f64 {
self.balance()
}
/// True when the bucket has refilled to capacity: no debt outstanding, so /// True when the bucket has refilled to capacity: no debt outstanding, so
/// the chat has not sent anything recently. /// the chat has not sent anything recently.
fn is_idle(&self) -> bool { fn is_idle(&self) -> bool {
let mut state = self.state.lock(); self.balance() >= self.capacity
self.refill(&mut state);
state.tokens >= self.capacity
} }
} }