use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{AcquireError, Semaphore};
const CREDIT_TARGET: u16 = 512;
pub(crate) const DEFAULT_CREDIT_WAIT: Duration = Duration::from_secs(30);
pub(crate) struct CreditPool {
permits: Mutex<Arc<Semaphore>>,
wait_ms: AtomicU64,
}
impl CreditPool {
pub(crate) fn new() -> Self {
Self {
permits: Mutex::new(Arc::new(Semaphore::new(0))),
wait_ms: AtomicU64::new(DEFAULT_CREDIT_WAIT.as_millis() as u64),
}
}
fn current(&self) -> Arc<Semaphore> {
Arc::clone(&self.permits.lock().unwrap())
}
pub(crate) fn reset(&self) {
*self.permits.lock().unwrap() = Arc::new(Semaphore::new(0));
}
pub(crate) fn available(&self) -> u16 {
self.current().available_permits().min(u16::MAX as usize) as u16
}
pub(crate) fn grant(&self, credits: u16) {
if credits > 0 {
self.current().add_permits(credits as usize);
}
}
pub(crate) fn try_reserve(&self, charge: u16) -> bool {
match self.current().try_acquire_many(u32::from(charge)) {
Ok(permit) => {
permit.forget();
true
}
Err(_) => false,
}
}
pub(crate) async fn reserve(&self, charge: u16) -> Result<(), AcquireError> {
self.current()
.acquire_many_owned(u32::from(charge))
.await?
.forget();
Ok(())
}
pub(crate) fn refund(&self, charge: u16) {
self.grant(charge);
}
pub(crate) fn request_for(&self, charge: u16) -> u16 {
charge.saturating_add(CREDIT_TARGET.saturating_sub(self.available()))
}
pub(crate) fn wait_timeout(&self) -> Duration {
Duration::from_millis(self.wait_ms.load(Ordering::Relaxed))
}
pub(crate) fn set_wait_timeout(&self, after: Duration) {
let ms = u64::try_from(after.as_millis()).unwrap_or(u64::MAX);
self.wait_ms.store(ms, Ordering::Relaxed);
}
pub(crate) fn close(&self) {
self.current().close();
}
pub(crate) fn is_closed(&self) -> bool {
self.current().is_closed()
}
#[cfg(test)]
pub(crate) fn set_available(&self, credits: u16) {
let permits = self.current();
let have = permits.available_permits();
let want = usize::from(credits);
match want.cmp(&have) {
std::cmp::Ordering::Greater => permits.add_permits(want - have),
std::cmp::Ordering::Less => {
if let Ok(permit) = permits.try_acquire_many((have - want) as u32) {
permit.forget();
}
}
std::cmp::Ordering::Equal => {}
}
}
}
#[must_use = "dropping the reservation refunds the credits without sending"]
pub(crate) struct CreditReservation<'a> {
pool: Option<&'a CreditPool>,
charge: u16,
}
impl<'a> CreditReservation<'a> {
pub(crate) fn new(pool: &'a CreditPool, charge: u16) -> Self {
Self {
pool: Some(pool),
charge,
}
}
pub(crate) fn exempt() -> Self {
Self {
pool: None,
charge: 0,
}
}
pub(crate) fn commit(mut self) {
self.pool = None;
}
}
impl Drop for CreditReservation<'_> {
fn drop(&mut self) {
if let Some(pool) = self.pool {
pool.refund(self.charge);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_reservation_holds_credits_out_of_the_pool_until_it_is_refunded() {
let pool = CreditPool::new();
pool.set_available(10);
assert!(pool.try_reserve(4));
assert_eq!(pool.available(), 6);
let reservation = CreditReservation::new(&pool, 4);
drop(reservation);
assert_eq!(
pool.available(),
10,
"an unsent request gives its credits back"
);
}
#[test]
fn a_committed_reservation_leaves_the_credits_with_the_server() {
let pool = CreditPool::new();
pool.set_available(10);
assert!(pool.try_reserve(4));
CreditReservation::new(&pool, 4).commit();
assert_eq!(
pool.available(),
6,
"credits spent on the wire only come back as a grant"
);
}
#[test]
fn a_charge_larger_than_the_window_is_not_partially_reserved() {
let pool = CreditPool::new();
pool.set_available(3);
assert!(!pool.try_reserve(4));
assert_eq!(pool.available(), 3, "a failed reserve takes nothing");
}
#[test]
fn the_credit_request_always_covers_the_charge_and_climbs_to_the_target() {
let pool = CreditPool::new();
pool.set_available(0);
assert_eq!(pool.request_for(8), 8 + CREDIT_TARGET);
pool.set_available(CREDIT_TARGET);
assert_eq!(
pool.request_for(8),
8,
"at target, ask only for the charge back"
);
pool.set_available(u16::MAX);
assert_eq!(pool.request_for(8), 8, "never ask for less than the charge");
}
#[tokio::test]
async fn a_grant_wakes_a_waiter() {
let pool = CreditPool::new();
pool.set_available(0);
let waiting = pool.reserve(4);
tokio::pin!(waiting);
assert!(
tokio::time::timeout(Duration::from_millis(50), &mut waiting)
.await
.is_err(),
"nothing to reserve yet"
);
pool.grant(4);
waiting.await.expect("the grant satisfies the waiter");
assert_eq!(pool.available(), 0);
}
#[tokio::test]
async fn closing_the_pool_fails_waiters_instead_of_parking_them() {
let pool = CreditPool::new();
pool.set_available(0);
let waiting = pool.reserve(1);
tokio::pin!(waiting);
assert!(
tokio::time::timeout(Duration::from_millis(50), &mut waiting)
.await
.is_err()
);
pool.close();
assert!(waiting.await.is_err());
assert!(pool.is_closed());
}
#[test]
fn a_fresh_pool_funds_nothing_until_the_server_grants() {
let pool = CreditPool::new();
assert_eq!(pool.available(), 0);
assert!(
!pool.try_reserve(1),
"a credit the server never granted must not be spendable: \
NEGOTIATE is exempt from the window, and an exemption is not a \
permit anything else can take"
);
}
#[tokio::test]
async fn a_reset_pool_holds_no_ungranted_credit_and_is_no_longer_closed() {
let pool = CreditPool::new();
pool.set_available(400);
pool.close();
assert!(pool.is_closed());
pool.reset();
assert!(
!pool.is_closed(),
"a revived connection needs a live budget"
);
assert_eq!(
pool.available(),
0,
"credits granted by a dead session must not carry over -- the new \
server's window may be far smaller, and until it says so the \
client holds nothing"
);
assert!(
!pool.try_reserve(1),
"a request racing the revival must wait for a real grant rather \
than spend the handshake's exemption"
);
}
#[tokio::test]
async fn a_waiter_on_the_old_budget_is_failed_by_the_reset_rather_than_migrated() {
let pool = CreditPool::new();
pool.set_available(0);
let waiting = pool.reserve(4);
tokio::pin!(waiting);
assert!(
tokio::time::timeout(Duration::from_millis(50), &mut waiting)
.await
.is_err(),
"nothing to reserve yet"
);
pool.close();
pool.reset();
pool.grant(64);
assert!(
tokio::time::timeout(Duration::from_millis(200), waiting)
.await
.expect("the waiter must resolve, not hang")
.is_err(),
"a send queued against the old session must fail rather than \
silently continue on the new one"
);
}
}