use std::mem::ManuallyDrop;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use jiff::Timestamp;
use crate::deny::DenyReason;
use crate::ids::{AccountId, FencingToken, LeaseId};
use crate::sharding::{LocalSharding, Locality};
use crate::units::CostUnits;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct LeaseGrant {
pub lease_id: LeaseId,
pub account_id: AccountId,
pub fencing_token: FencingToken,
pub units: CostUnits,
pub expires_at: Timestamp,
}
#[derive(Debug)]
pub struct AccountOverage {
account_id: AccountId,
spent: AtomicU64,
committed: AtomicU64,
commit_publications: AtomicUsize,
contended: AtomicU64,
}
struct OverageCommitPublication<'a> {
overage: &'a AccountOverage,
}
impl OverageCommitPublication<'_> {
fn publish(self, units: CostUnits) {
let prior = self
.overage
.committed
.fetch_add(units.get(), Ordering::AcqRel);
debug_assert!(
prior
.checked_add(units.get())
.is_some_and(|committed| committed <= self.overage.spent.load(Ordering::Acquire)),
"committed overage exceeds total recorded spend"
);
}
}
impl Drop for OverageCommitPublication<'_> {
fn drop(&mut self) {
let prior = self
.overage
.commit_publications
.fetch_sub(1, Ordering::AcqRel);
debug_assert!(prior > 0, "overage commit publication count underflowed");
}
}
impl AccountOverage {
#[must_use]
pub fn new(account_id: AccountId) -> Self {
AccountOverage {
account_id,
spent: AtomicU64::new(0),
committed: AtomicU64::new(0),
commit_publications: AtomicUsize::new(0),
contended: AtomicU64::new(0),
}
}
#[must_use]
pub fn contended_debits(&self) -> u64 {
self.contended.load(Ordering::Relaxed)
}
#[must_use]
pub fn account_id(&self) -> AccountId {
self.account_id
}
#[must_use]
pub fn spent(&self) -> CostUnits {
CostUnits(self.spent.load(Ordering::Acquire))
}
#[must_use]
pub fn headroom(&self, cap: CostUnits) -> CostUnits {
CostUnits(cap.get().saturating_sub(self.spent.load(Ordering::Acquire)))
}
#[inline]
pub(crate) fn try_debit(&self, units: CostUnits, cap: CostUnits) -> Result<(), DenyReason> {
let want = units.get();
let mut current = self.spent.load(Ordering::Acquire);
let mut lost = 0u64;
let note = |lost: u64| {
if lost != 0 {
self.contended.fetch_add(lost, Ordering::Relaxed);
}
};
loop {
let refused = || {
if self.commit_publications.fetch_add(0, Ordering::AcqRel) != 0 {
return DenyReason::OverageCommitInProgress {
spent: CostUnits(current),
overage_cap: cap,
};
}
let committed = self.committed.load(Ordering::Acquire);
if committed
.checked_add(want)
.is_some_and(|next| next <= cap.get())
{
DenyReason::OverageCapTemporarilyExhausted {
spent: CostUnits(current),
overage_cap: cap,
}
} else {
DenyReason::OverageCapExhausted {
spent: CostUnits(current),
overage_cap: cap,
}
}
};
let Some(next) = current.checked_add(want) else {
note(lost);
return Err(refused());
};
if next > cap.get() {
note(lost);
return Err(refused());
}
match self.spent.compare_exchange_weak(
current,
next,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
note(lost);
return Ok(());
}
Err(observed) => {
lost += 1;
current = observed;
}
}
}
}
#[inline]
pub(crate) fn publish_claim<T, E>(
&self,
units: CostUnits,
claim: impl FnOnce() -> Result<T, E>,
) -> Result<T, E> {
let prior = self.commit_publications.fetch_add(1, Ordering::AcqRel);
debug_assert_ne!(
prior,
usize::MAX,
"live overage commit publications exceed the address space"
);
let publication = OverageCommitPublication { overage: self };
match claim() {
Ok(value) => {
publication.publish(units);
Ok(value)
}
Err(error) => Err(error),
}
}
#[inline]
pub(crate) fn credit(&self, units: CostUnits) {
let prior = self.spent.fetch_sub(units.get(), Ordering::AcqRel);
debug_assert!(
prior >= units.get(),
"overage credit of {} exceeds recorded spend {prior}",
units.get()
);
}
#[inline]
pub(crate) fn debit_tentatively(
&self,
units: CostUnits,
cap: CostUnits,
) -> Result<TentativeOverage<'_>, DenyReason> {
self.try_debit(units, cap)?;
Ok(TentativeOverage {
overage: self,
units,
})
}
}
#[derive(Debug)]
pub(crate) struct TentativeOverage<'a> {
overage: &'a AccountOverage,
units: CostUnits,
}
impl TentativeOverage<'_> {
#[inline]
pub(crate) fn publish_commit<T, E>(self, claim: impl FnOnce() -> Result<T, E>) -> Result<T, E> {
let this = ManuallyDrop::new(self);
match this.overage.publish_claim(this.units, claim) {
Ok(value) => Ok(value),
Err(error) => {
this.overage.credit(this.units);
Err(error)
}
}
}
}
impl Drop for TentativeOverage<'_> {
fn drop(&mut self) {
self.overage.credit(self.units);
}
}
pub trait RefillSignal: Send + Sync + core::fmt::Debug {
fn request_refill(&self);
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RefillVerdict {
Idle,
Draining,
Refused,
}
#[repr(align(128))]
#[derive(Debug)]
struct LeaseShard {
remaining: AtomicU64,
low_water: u64,
signalled: AtomicBool,
contended: AtomicU64,
}
impl LeaseShard {
fn new(remaining: u64, low_water: u64) -> Self {
Self {
remaining: AtomicU64::new(remaining),
low_water,
signalled: AtomicBool::new(false),
contended: AtomicU64::new(0),
}
}
#[inline]
fn note_contention(&self, lost: u64) {
if lost != 0 {
self.contended.fetch_add(lost, Ordering::Relaxed);
}
}
}
#[derive(Debug)]
enum LeaseBalance {
Single(LeaseShard),
Sharded(Box<[LeaseShard]>),
}
impl LeaseBalance {
fn as_slice(&self) -> &[LeaseShard] {
match self {
Self::Single(shard) => std::slice::from_ref(shard),
Self::Sharded(shards) => shards,
}
}
}
#[derive(Debug)]
pub(crate) struct LeaseDebit {
shard: usize,
units: u64,
}
#[derive(Debug)]
struct LeaseInner {
grant: LeaseGrant,
balance: LeaseBalance,
refill: OnceLock<Arc<dyn RefillSignal>>,
low_water: u64,
refused: AtomicBool,
refused_quote: AtomicU64,
contention_reported: AtomicU64,
usable_until: Timestamp,
}
#[derive(Debug, Clone)]
#[repr(align(128))]
pub struct LocalLease {
inner: Arc<LeaseInner>,
}
fn usable_until(expires_at: Timestamp, margin: jiff::SignedDuration) -> Timestamp {
if margin < jiff::SignedDuration::ZERO {
Timestamp::MIN
} else {
expires_at
.checked_sub(margin)
.unwrap_or(Timestamp::MIN)
}
}
fn partition(total: u64, count: usize, index: usize) -> u64 {
let count = u64::try_from(count).expect("local shard count fits u64");
let index = u64::try_from(index).expect("local shard index fits u64");
total / count + u64::from(index < total % count)
}
impl LocalLease {
#[must_use]
pub fn new(grant: LeaseGrant, low_water: CostUnits) -> Self {
Self::with_safety_margin(grant, low_water, jiff::SignedDuration::ZERO)
}
#[must_use]
pub fn with_safety_margin(
grant: LeaseGrant,
low_water: CostUnits,
margin: jiff::SignedDuration,
) -> Self {
let usable_until = usable_until(grant.expires_at, margin);
LocalLease {
inner: Arc::new(LeaseInner {
balance: LeaseBalance::Single(LeaseShard::new(grant.units.get(), low_water.get())),
low_water: low_water.get(),
refused: AtomicBool::new(false),
refused_quote: AtomicU64::new(0),
contention_reported: AtomicU64::new(0),
usable_until,
refill: OnceLock::new(),
grant,
}),
}
}
#[must_use]
pub fn with_sharding(
grant: LeaseGrant,
low_water: CostUnits,
margin: jiff::SignedDuration,
sharding: LocalSharding,
) -> Self {
if sharding == LocalSharding::SINGLE {
return Self::with_safety_margin(grant, low_water, margin);
}
let usable_until = usable_until(grant.expires_at, margin);
let count = sharding.get();
let shards = (0..count)
.map(|index| {
LeaseShard::new(
partition(grant.units.get(), count, index),
partition(low_water.get(), count, index),
)
})
.collect::<Vec<_>>()
.into_boxed_slice();
Self {
inner: Arc::new(LeaseInner {
grant,
balance: LeaseBalance::Sharded(shards),
refill: OnceLock::new(),
low_water: low_water.get(),
refused: AtomicBool::new(false),
refused_quote: AtomicU64::new(0),
contention_reported: AtomicU64::new(0),
usable_until,
}),
}
}
#[must_use]
pub fn with_refill(self, signal: Arc<dyn RefillSignal>) -> Self {
let _already_attached = self.inner.refill.set(signal);
self
}
#[must_use]
pub fn usable_until(&self) -> Timestamp {
self.inner.usable_until
}
#[must_use]
pub fn grant(&self) -> &LeaseGrant {
&self.inner.grant
}
#[must_use]
pub fn is_only_local_view(&self) -> bool {
Arc::strong_count(&self.inner) == 1
}
#[must_use]
pub fn largest_refused_quote(&self) -> CostUnits {
CostUnits(self.inner.refused_quote.load(Ordering::Acquire))
}
#[must_use]
pub fn contended_debits(&self) -> u64 {
self.inner
.balance
.as_slice()
.iter()
.fold(0u64, |total, shard| {
total.saturating_add(shard.contended.load(Ordering::Relaxed))
})
}
#[must_use]
pub fn take_unreported_contention(&self) -> u64 {
let total = self.contended_debits();
let reported = self
.inner
.contention_reported
.fetch_max(total, Ordering::AcqRel);
total.saturating_sub(reported)
}
#[must_use]
pub fn unreported_contention(&self) -> u64 {
self.contended_debits()
.saturating_sub(self.inner.contention_reported.load(Ordering::Acquire))
}
#[must_use]
pub fn remaining(&self) -> CostUnits {
let remaining = self
.inner
.balance
.as_slice()
.iter()
.fold(0u64, |total, shard| {
total.saturating_add(shard.remaining.load(Ordering::Acquire))
})
.min(self.inner.grant.units.get());
CostUnits(remaining)
}
#[must_use]
pub fn needs_refill(&self) -> bool {
self.remaining().get() <= self.inner.low_water
}
#[must_use]
pub fn refill_due_or_rearm(&self) -> RefillVerdict {
if self.inner.refused.swap(false, Ordering::AcqRel) {
return RefillVerdict::Refused;
}
if self.needs_refill() {
return RefillVerdict::Draining;
}
for shard in self.inner.balance.as_slice() {
shard.signalled.swap(false, Ordering::AcqRel);
}
if self.needs_refill() {
RefillVerdict::Draining
} else {
RefillVerdict::Idle
}
}
#[inline]
pub fn try_debit(&self, units: CostUnits, now: Timestamp) -> Result<(), DenyReason> {
self.try_reserve_at(units, now, Locality::current())
.map(|_| ())
}
#[inline]
pub(crate) fn try_reserve_at(
&self,
units: CostUnits,
now: Timestamp,
locality: Locality,
) -> Result<LeaseDebit, DenyReason> {
if now >= self.inner.usable_until {
self.signal_refusal();
return Err(DenyReason::LeaseExpired);
}
let want = units.get();
let shards = self.inner.balance.as_slice();
let first = locality.index(LocalSharding::new(
std::num::NonZeroUsize::new(shards.len()).expect("lease has at least one shard"),
));
for offset in 0..shards.len() {
let index = (first + offset) % shards.len();
if let Some(part) = self.try_whole(index, want) {
return Ok(part);
}
}
let mut needed = want;
for offset in 0..shards.len() {
let index = (first + offset) % shards.len();
needed -= self.take_up_to(index, needed);
if needed == 0 {
let next = shards[first].remaining.load(Ordering::Acquire);
self.maybe_signal_refill(first, next);
return Ok(LeaseDebit {
shard: first,
units: want,
});
}
}
self.credit_to(first, want - needed);
self.inner.refused_quote.fetch_max(want, Ordering::Relaxed);
self.signal_refusal();
Err(DenyReason::LeaseExhausted {
remaining: self.remaining(),
})
}
#[inline]
fn try_whole(&self, shard_index: usize, want: u64) -> Option<LeaseDebit> {
let shard = &self.inner.balance.as_slice()[shard_index];
let mut current = shard.remaining.load(Ordering::Acquire);
let mut lost = 0u64;
loop {
let Some(next) = current.checked_sub(want) else {
shard.note_contention(lost);
return None;
};
match shard.remaining.compare_exchange_weak(
current,
next,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
shard.note_contention(lost);
self.maybe_signal_refill(shard_index, next);
return Some(LeaseDebit {
shard: shard_index,
units: want,
});
}
Err(observed) => {
lost += 1;
current = observed;
}
}
}
}
#[cold]
fn take_up_to(&self, shard_index: usize, want: u64) -> u64 {
let shard = &self.inner.balance.as_slice()[shard_index];
let mut current = shard.remaining.load(Ordering::Acquire);
let mut lost = 0u64;
loop {
if current == 0 {
shard.note_contention(lost);
return 0;
}
let taken = current.min(want);
let next = current - taken;
match shard.remaining.compare_exchange_weak(
current,
next,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
shard.note_contention(lost);
return taken;
}
Err(observed) => {
lost += 1;
current = observed;
}
}
}
}
#[inline]
fn maybe_signal_refill(&self, shard_index: usize, next: u64) {
if next <= self.inner.balance.as_slice()[shard_index].low_water {
self.signal_refill(shard_index);
}
}
#[cold]
#[inline(never)]
fn signal_refusal(&self) {
if self.inner.refused.swap(true, Ordering::AcqRel) {
return;
}
if let Some(signal) = self.inner.refill.get() {
signal.request_refill();
}
}
#[cold]
#[inline(never)]
fn signal_refill(&self, shard_index: usize) {
let Some(signal) = self.inner.refill.get() else {
return;
};
if !self.inner.balance.as_slice()[shard_index]
.signalled
.swap(true, Ordering::Release)
{
signal.request_refill();
}
}
#[inline]
pub(crate) fn credit(&self, debit: &LeaseDebit) {
self.credit_to(debit.shard, debit.units);
}
fn credit_to(&self, shard: usize, units: u64) {
self.inner.balance.as_slice()[shard]
.remaining
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
current.checked_add(units)
})
.expect("a debit receipt cannot credit beyond its lease grant");
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::num::NonZeroUsize;
fn t(secs: i64) -> Timestamp {
Timestamp::from_second(secs).unwrap()
}
fn overage() -> AccountOverage {
AccountOverage::new(AccountId(1))
}
#[test]
fn dropping_an_unresolved_tentative_debit_returns_the_credit() {
let o = overage();
{
let tentative = o.debit_tentatively(CostUnits(40), CostUnits(100)).unwrap();
assert_eq!(o.spent(), CostUnits(40));
drop(tentative);
}
assert_eq!(o.spent(), CostUnits::ZERO);
assert_eq!(o.headroom(CostUnits(100)), CostUnits(100));
}
#[test]
fn a_tentative_debit_whose_claim_loses_returns_the_credit() {
let o = overage();
let tentative = o.debit_tentatively(CostUnits(40), CostUnits(100)).unwrap();
assert_eq!(o.spent(), CostUnits(40));
let lost: Result<(), ()> = tentative.publish_commit(|| Err(()));
assert!(lost.is_err());
assert_eq!(o.spent(), CostUnits::ZERO);
assert_eq!(o.headroom(CostUnits(100)), CostUnits(100));
}
#[test]
fn a_tentative_debit_whose_claim_wins_becomes_irrevocable() {
let o = overage();
let tentative = o.debit_tentatively(CostUnits(40), CostUnits(100)).unwrap();
let won: Result<(), ()> = tentative.publish_commit(|| Ok(()));
assert!(won.is_ok());
assert_eq!(o.spent(), CostUnits(40));
assert_eq!(
o.try_debit(CostUnits(61), CostUnits(100)),
Err(DenyReason::OverageCapExhausted {
spent: CostUnits(40),
overage_cap: CostUnits(100),
})
);
}
#[test]
fn committed_overage_accumulates_up_to_the_cap_and_then_refuses() {
let o = overage();
o.try_debit(CostUnits(40), CostUnits(100)).unwrap();
o.publish_claim(CostUnits(40), || Ok::<_, ()>(())).unwrap();
o.try_debit(CostUnits(60), CostUnits(100)).unwrap();
o.publish_claim(CostUnits(60), || Ok::<_, ()>(())).unwrap();
assert_eq!(o.spent(), CostUnits(100));
assert_eq!(o.headroom(CostUnits(100)), CostUnits::ZERO);
assert_eq!(
o.try_debit(CostUnits(1), CostUnits(100)),
Err(DenyReason::OverageCapExhausted {
spent: CostUnits(100),
overage_cap: CostUnits(100),
})
);
assert_eq!(o.spent(), CostUnits(100), "a refusal claims nothing");
}
#[test]
fn pending_overage_is_transient_only_when_its_refund_would_make_room() {
let o = overage();
o.try_debit(CostUnits(60), CostUnits(100)).unwrap();
let pending_saturation = o.try_debit(CostUnits(50), CostUnits(100)).unwrap_err();
assert_eq!(
pending_saturation,
DenyReason::OverageCapTemporarilyExhausted {
spent: CostUnits(60),
overage_cap: CostUnits(100),
}
);
assert_eq!(pending_saturation.retry(), crate::deny::Retry::Transient);
let request_exceeds_cap = o.try_debit(CostUnits(101), CostUnits(100)).unwrap_err();
assert_eq!(
request_exceeds_cap,
DenyReason::OverageCapExhausted {
spent: CostUnits(60),
overage_cap: CostUnits(100),
}
);
assert_eq!(request_exceeds_cap.retry(), crate::deny::Retry::Transient);
}
#[test]
fn lowering_the_cap_below_current_spend_refuses_at_once() {
let o = overage();
o.try_debit(CostUnits(80), CostUnits(100)).unwrap();
assert_eq!(o.headroom(CostUnits(50)), CostUnits::ZERO);
assert!(o.try_debit(CostUnits(1), CostUnits(50)).is_err());
o.try_debit(CostUnits(1), CostUnits(100)).unwrap();
}
#[test]
fn an_unrepresentable_total_refuses_rather_than_wrapping() {
let o = overage();
o.try_debit(CostUnits(u64::MAX - 1), CostUnits(u64::MAX))
.unwrap();
assert!(o.try_debit(CostUnits(2), CostUnits(u64::MAX)).is_err());
assert_eq!(o.spent(), CostUnits(u64::MAX - 1));
}
#[test]
fn credit_returns_headroom_to_the_cap() {
let o = overage();
o.try_debit(CostUnits(100), CostUnits(100)).unwrap();
o.credit(CostUnits(30));
assert_eq!(o.spent(), CostUnits(70));
assert_eq!(o.headroom(CostUnits(100)), CostUnits(30));
o.try_debit(CostUnits(30), CostUnits(100)).unwrap();
assert!(o.try_debit(CostUnits(1), CostUnits(100)).is_err());
}
#[test]
fn concurrent_debits_never_exceed_the_cap() {
const THREADS: usize = 8;
const EACH: usize = 500;
const CAP: u64 = 1_000;
let o = Arc::new(overage());
let admitted = Arc::new(AtomicU64::new(0));
std::thread::scope(|scope| {
for _ in 0..THREADS {
let o = Arc::clone(&o);
let admitted = Arc::clone(&admitted);
scope.spawn(move || {
for _ in 0..EACH {
if o.try_debit(CostUnits(1), CostUnits(CAP)).is_ok() {
admitted.fetch_add(1, Ordering::Relaxed);
}
}
});
}
});
assert_eq!(o.spent(), CostUnits(CAP));
assert_eq!(
admitted.load(Ordering::Relaxed),
CAP,
"every admitted unit is one the counter recorded, and vice versa"
);
}
fn lease(units: u64, expires: i64, low_water: u64) -> LocalLease {
LocalLease::new(
LeaseGrant {
lease_id: LeaseId(7),
account_id: AccountId(1),
fencing_token: FencingToken(3),
units: CostUnits(units),
expires_at: t(expires),
},
CostUnits(low_water),
)
}
fn sharded_lease(units: u64, low_water: u64, shards: usize) -> LocalLease {
LocalLease::with_sharding(
LeaseGrant {
lease_id: LeaseId(7),
account_id: AccountId(1),
fencing_token: FencingToken(3),
units: CostUnits(units),
expires_at: t(1_000),
},
CostUnits(low_water),
jiff::SignedDuration::ZERO,
LocalSharding::new(NonZeroUsize::new(shards).unwrap()),
)
}
#[test]
fn sharded_grant_and_low_water_partitions_are_exact() {
assert_eq!(align_of::<LeaseShard>(), 128);
assert_eq!(size_of::<LeaseShard>(), 128);
assert_eq!(align_of::<LocalLease>(), 128);
let l = sharded_lease(10, 3, 4);
let shards = l.inner.balance.as_slice();
assert_eq!(shards.len(), 4);
assert_eq!(
shards
.iter()
.map(|shard| shard.remaining.load(Ordering::Relaxed))
.collect::<Vec<_>>(),
[3, 3, 2, 2]
);
assert_eq!(
shards
.iter()
.map(|shard| shard.low_water)
.collect::<Vec<_>>(),
[1, 1, 1, 0]
);
assert_eq!(l.remaining(), CostUnits(10));
}
#[test]
fn cloned_local_views_are_visible_to_quiescence_detection() {
let lease = sharded_lease(10, 0, 2);
assert!(lease.is_only_local_view());
let sibling = lease.clone();
assert!(!lease.is_only_local_view());
drop(sibling);
assert!(lease.is_only_local_view());
}
#[test]
fn fragmented_reservation_refunds_without_stranding_capacity() {
let l = Arc::new(sharded_lease(10, 0, 4));
assert_eq!(size_of::<LeaseDebit>(), 16, "the receipt is fixed-size");
let reservation = crate::Reservation::reserve(&l, CostUnits(8), t(0)).unwrap();
assert_eq!(l.remaining(), CostUnits(2));
assert_eq!(reservation.cancel(), crate::CancelOutcome::ZeroCharged);
assert_eq!(l.remaining(), CostUnits(10));
}
#[test]
fn failed_fragmented_debit_reports_true_remaining_and_rolls_back() {
let l = sharded_lease(10, 0, 4);
assert_eq!(
l.try_debit(CostUnits(11), t(0)),
Err(DenyReason::LeaseExhausted {
remaining: CostUnits(10)
})
);
assert_eq!(l.remaining(), CostUnits(10));
}
#[test]
fn rebalanced_refunds_are_exact_at_the_u64_boundary() {
let l = Arc::new(sharded_lease(u64::MAX, 0, 2));
let reservation = crate::Reservation::reserve(&l, CostUnits(u64::MAX), t(0)).unwrap();
assert_eq!(l.remaining(), CostUnits::ZERO);
assert_eq!(reservation.cancel(), crate::CancelOutcome::ZeroCharged);
assert_eq!(l.remaining(), CostUnits(u64::MAX));
let reservation = crate::Reservation::reserve(&l, CostUnits(u64::MAX), t(0)).unwrap();
assert_eq!(reservation.cancel(), crate::CancelOutcome::ZeroCharged);
assert_eq!(l.remaining(), CostUnits(u64::MAX));
}
#[cfg(any(
target_arch = "x86_64",
all(target_arch = "aarch64", target_feature = "lse")
))]
#[test]
fn an_uncontended_debit_records_no_contention() {
for l in [lease(1_000, 1_000, 0), sharded_lease(1_000, 0, 4)] {
for _ in 0..100 {
let debit = l
.try_reserve_at(CostUnits(3), t(0), Locality::current())
.unwrap();
l.credit(&debit);
}
let fragmented = l
.try_reserve_at(CostUnits(1_000), t(0), Locality::current())
.unwrap();
l.credit(&fragmented);
assert_eq!(l.contended_debits(), 0);
}
}
#[test]
fn contended_debits_are_recorded_without_disturbing_the_balance() {
const THREADS: u64 = 8;
const DEBITS: u64 = 20_000;
let l = lease(THREADS * DEBITS, 1_000, 0);
for _ in 0..50 {
std::thread::scope(|scope| {
for _ in 0..THREADS {
scope.spawn(|| {
for _ in 0..DEBITS {
let debit = l
.try_reserve_at(CostUnits(1), t(0), Locality::current())
.unwrap();
l.credit(&debit);
}
});
}
});
if l.contended_debits() > 0 {
break;
}
}
assert!(
l.contended_debits() > 0,
"eight writers on one line never lost a race"
);
assert_eq!(
l.remaining(),
CostUnits(THREADS * DEBITS),
"credits restored every debit"
);
}
#[test]
fn unreported_contention_is_handed_out_exactly_once() {
let l = sharded_lease(100, 0, 2);
let shards = l.inner.balance.as_slice();
shards[0].note_contention(3);
shards[1].note_contention(4);
assert_eq!(l.contended_debits(), 7);
assert_eq!(l.unreported_contention(), 7);
assert_eq!(l.take_unreported_contention(), 7);
assert_eq!(
l.take_unreported_contention(),
0,
"reinstalled and taken again"
);
assert_eq!(l.unreported_contention(), 0);
shards[1].note_contention(2);
assert_eq!(l.unreported_contention(), 2);
assert_eq!(l.take_unreported_contention(), 2);
assert_eq!(
l.contended_debits(),
9,
"the lease's own total never resets"
);
assert_eq!(l.clone().take_unreported_contention(), 0);
shards[0].note_contention(0);
assert_eq!(l.contended_debits(), 9, "recording zero writes nothing");
}
#[cfg(any(
target_arch = "x86_64",
all(target_arch = "aarch64", target_feature = "lse")
))]
#[test]
fn an_uncontended_overage_debit_records_no_contention() {
let overage = AccountOverage::new(AccountId(1));
for _ in 0..100 {
overage.try_debit(CostUnits(1), CostUnits(1_000)).unwrap();
}
assert!(
overage
.try_debit(CostUnits(1_000), CostUnits(1_000))
.is_err()
);
assert_eq!(overage.contended_debits(), 0);
}
#[test]
fn contended_overage_debits_are_recorded_without_disturbing_spend() {
const THREADS: u64 = 8;
const DEBITS: u64 = 20_000;
let overage = AccountOverage::new(AccountId(1));
let mut rounds = 0;
while overage.contended_debits() == 0 && rounds < 50 {
rounds += 1;
std::thread::scope(|scope| {
for _ in 0..THREADS {
scope.spawn(|| {
for _ in 0..DEBITS {
overage
.try_debit(CostUnits(1), CostUnits(u64::MAX))
.unwrap();
}
});
}
});
}
assert!(
overage.contended_debits() > 0,
"eight writers on one line never lost a race"
);
assert_eq!(
overage.spent.load(Ordering::Relaxed),
rounds * THREADS * DEBITS,
"every debit counted exactly once"
);
}
#[test]
fn sharded_lease_spends_to_exact_exhaustion_without_stranding() {
let l = sharded_lease(17, 0, 8);
for _ in 0..17 {
l.try_debit(CostUnits(1), t(0)).unwrap();
}
assert_eq!(l.remaining(), CostUnits::ZERO);
assert_eq!(
l.try_debit(CostUnits(1), t(0)),
Err(DenyReason::LeaseExhausted {
remaining: CostUnits::ZERO
})
);
}
#[test]
fn debit_search_starts_at_the_supplied_locality_and_wraps_in_order() {
let l = sharded_lease(16, 0, 4);
let debit = l
.try_reserve_at(CostUnits(1), t(0), Locality::for_test(3))
.unwrap();
let remaining: Vec<_> = l
.inner
.balance
.as_slice()
.iter()
.map(|shard| shard.remaining.load(Ordering::Relaxed))
.collect();
assert_eq!(remaining, [4, 4, 4, 3]);
l.credit(&debit);
let fragmented = l
.try_reserve_at(CostUnits(5), t(0), Locality::for_test(3))
.unwrap();
let remaining: Vec<_> = l
.inner
.balance
.as_slice()
.iter()
.map(|shard| shard.remaining.load(Ordering::Relaxed))
.collect();
assert_eq!(remaining, [3, 4, 4, 0]);
l.credit(&fragmented);
assert_eq!(l.remaining(), CostUnits(16));
}
#[test]
fn a_whole_sibling_is_used_before_fragmenting() {
let l = sharded_lease(16, 0, 4);
let first = l
.try_reserve_at(CostUnits(1), t(0), Locality::for_test(3))
.unwrap();
let sibling = l
.try_reserve_at(CostUnits(4), t(0), Locality::for_test(3))
.unwrap();
let remaining: Vec<_> = l
.inner
.balance
.as_slice()
.iter()
.map(|shard| shard.remaining.load(Ordering::Relaxed))
.collect();
assert_eq!(
remaining,
[0, 4, 4, 3],
"the whole sibling pays; the undersized local shard is untouched"
);
l.credit(&sibling);
l.credit(&first);
assert_eq!(l.remaining(), CostUnits(16));
}
#[test]
fn a_torn_aggregate_above_the_grant_reads_as_the_grant() {
let l = sharded_lease(4, 0, 2);
let LeaseBalance::Sharded(shards) = &l.inner.balance else {
panic!("a two-shard lease is sharded");
};
shards[1].remaining.store(4, Ordering::Release);
assert_eq!(
l.remaining(),
CostUnits(4),
"an over-counted walk is clamped to the grant, never asserted"
);
}
#[test]
fn a_fragmenting_rollback_never_panics_a_concurrent_aggregate_read() {
let l = Arc::new(sharded_lease(6_400, 0, 64));
let stop = Arc::new(AtomicBool::new(false));
let spenders: Vec<_> = (48..52)
.map(|locality| {
let l = Arc::clone(&l);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
while !stop.load(Ordering::Relaxed) {
drop(l.try_reserve_at(
CostUnits(6_401),
t(0),
Locality::for_test(locality),
));
}
})
})
.collect();
for _ in 0..200_000 {
assert!(l.remaining() <= CostUnits(6_400));
}
stop.store(true, Ordering::Relaxed);
for spender in spenders {
spender.join().unwrap();
}
assert_eq!(
l.remaining(),
CostUnits(6_400),
"no units were lost or created"
);
}
#[test]
fn an_early_shard_signal_is_rearmed_until_the_aggregate_crosses() {
let signal = Arc::new(CountingSignal::default());
let l = sharded_lease(100, 20, 2).with_refill(signal.clone());
l.try_debit(CostUnits(40), t(0)).unwrap();
assert_eq!(signal.count(), 1, "the first shard crossed its share");
assert_eq!(
l.refill_due_or_rearm(),
RefillVerdict::Idle,
"sixty aggregate units remain"
);
l.try_debit(CostUnits(1), t(0)).unwrap();
assert_eq!(signal.count(), 2, "the control-plane check rearmed it");
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Idle);
l.try_debit(CostUnits(40), t(0)).unwrap();
assert_eq!(l.remaining(), CostUnits(19));
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Draining);
}
#[test]
fn debit_decrements_and_credit_restores() {
let l = lease(100, 1_000, 25);
let debit = l
.try_reserve_at(CostUnits(60), t(0), Locality::current())
.unwrap();
assert_eq!(l.remaining(), CostUnits(40));
l.credit(&debit);
assert_eq!(l.remaining(), CostUnits(100));
}
#[test]
fn exhaustion_denies_with_remaining() {
let l = lease(10, 1_000, 0);
assert_eq!(
l.try_debit(CostUnits(11), t(0)),
Err(DenyReason::LeaseExhausted {
remaining: CostUnits(10)
})
);
l.try_debit(CostUnits(10), t(0)).unwrap();
assert_eq!(l.remaining(), CostUnits::ZERO);
}
#[test]
fn expiry_boundary_is_exclusive_of_expires_at() {
let l = lease(10, 500, 0);
assert_eq!(
l.try_debit(CostUnits(1), t(500)),
Err(DenyReason::LeaseExpired)
);
l.try_debit(CostUnits(1), t(499)).unwrap();
}
#[test]
fn low_water_triggers_refill_signal() {
let l = lease(100, 1_000, 25);
assert!(!l.needs_refill());
l.try_debit(CostUnits(75), t(0)).unwrap();
assert!(l.needs_refill());
}
#[derive(Debug, Default)]
struct CountingSignal(AtomicU64);
impl RefillSignal for CountingSignal {
fn request_refill(&self) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
impl CountingSignal {
fn count(&self) -> u64 {
self.0.load(Ordering::Relaxed)
}
}
#[test]
fn the_crossing_debit_raises_the_signal() {
let signal = Arc::new(CountingSignal::default());
let l = lease(100, 1_000, 25).with_refill(signal.clone());
l.try_debit(CostUnits(74), t(0)).unwrap();
assert_eq!(signal.count(), 0, "26 remaining is above low water");
l.try_debit(CostUnits(1), t(0)).unwrap();
assert_eq!(signal.count(), 1, "landing exactly on low water crosses it");
}
#[test]
fn a_refused_debit_tells_the_refill_plane_rather_than_waiting_to_be_polled() {
let signal = Arc::new(CountingSignal::default());
let l = lease(49, 1_000, 25).with_refill(signal.clone());
assert!(matches!(
l.try_debit(CostUnits(51), t(0)),
Err(DenyReason::LeaseExhausted { .. })
));
assert_eq!(signal.count(), 1, "the refusal rang the doorbell");
assert!(
!l.needs_refill(),
"and it rang from above the mark, which is the whole point"
);
assert_eq!(
l.refill_due_or_rearm(),
RefillVerdict::Refused,
"the plane is told the grant is mis-sized, not that it is draining"
);
assert_eq!(
l.refill_due_or_rearm(),
RefillVerdict::Idle,
"reading the verdict consumes it; one refusal is one rotation"
);
assert_eq!(
l.remaining(),
CostUnits(49),
"a refusal is still zero-charge"
);
}
#[test]
fn a_refused_debit_records_its_largest_quote() {
for l in [lease(49, 1_000, 25), sharded_lease(49, 25, 4)] {
assert_eq!(l.largest_refused_quote(), CostUnits::ZERO);
for quote in [51, 90, 60] {
assert!(l.try_debit(CostUnits(quote), t(0)).is_err());
}
assert_eq!(l.largest_refused_quote(), CostUnits(90));
l.try_debit(CostUnits(10), t(0)).unwrap();
assert_eq!(
l.largest_refused_quote(),
CostUnits(90),
"a funded debit is not demand the grant failed"
);
}
let expired = lease(49, 1_000, 25);
assert!(matches!(
expired.try_debit(CostUnits(51), t(1_000)),
Err(DenyReason::LeaseExpired)
));
assert_eq!(expired.largest_refused_quote(), CostUnits::ZERO);
}
#[test]
fn a_refusal_storm_rings_once_and_reports_once() {
let signal = Arc::new(CountingSignal::default());
let l = lease(49, 1_000, 25).with_refill(signal.clone());
for _ in 0..1_000 {
assert!(l.try_debit(CostUnits(51), t(0)).is_err());
}
assert_eq!(signal.count(), 1);
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
assert!(l.try_debit(CostUnits(51), t(0)).is_err());
assert_eq!(signal.count(), 2);
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
}
#[test]
fn a_refusal_outranks_a_low_water_crossing() {
let l = lease(100, 1_000, 90);
l.try_debit(CostUnits(20), t(0)).unwrap();
assert!(l.needs_refill(), "80 remaining is below the 90 mark");
assert!(l.try_debit(CostUnits(81), t(0)).is_err());
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
assert_eq!(
l.refill_due_or_rearm(),
RefillVerdict::Draining,
"the crossing is still there once the refusal has been answered"
);
}
#[test]
fn an_expired_lease_reports_its_refusal_rather_than_waiting_for_the_tick() {
let signal = Arc::new(CountingSignal::default());
let l = lease(100, 1_000, 25).with_refill(signal.clone());
assert_eq!(
l.try_debit(CostUnits(1), t(2_000)),
Err(DenyReason::LeaseExpired)
);
assert_eq!(signal.count(), 1);
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
}
#[test]
fn a_lease_with_no_doorbell_still_records_its_refusal() {
let l = lease(49, 1_000, 25);
assert!(l.try_debit(CostUnits(51), t(0)).is_err());
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
}
#[test]
fn a_sharded_refusal_reports_once_for_the_whole_grant() {
let signal = Arc::new(CountingSignal::default());
let l = sharded_lease(40, 4, 4).with_refill(signal.clone());
assert!(matches!(
l.try_debit(CostUnits(41), t(0)),
Err(DenyReason::LeaseExhausted { .. })
));
assert_eq!(signal.count(), 1);
assert_eq!(l.remaining(), CostUnits(40), "the rollback restored it all");
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
}
#[test]
fn a_lease_signals_at_most_once_however_long_it_drains() {
let signal = Arc::new(CountingSignal::default());
let l = lease(100, 1_000, 25).with_refill(signal.clone());
for _ in 0..80 {
l.try_debit(CostUnits(1), t(0)).unwrap();
}
assert_eq!(l.remaining(), CostUnits(20));
assert_eq!(
signal.count(),
1,
"one crossing, however many debits followed it"
);
let next = lease(100, 1_000, 25).with_refill(signal.clone());
next.try_debit(CostUnits(80), t(0)).unwrap();
assert_eq!(signal.count(), 2);
}
#[test]
fn a_refused_debit_reports_a_refusal_and_never_a_crossing() {
let signal = Arc::new(CountingSignal::default());
let l = lease(100, 1_000, 25).with_refill(signal.clone());
assert!(l.try_debit(CostUnits(500), t(0)).is_err(), "exhausted");
assert_eq!(signal.count(), 1, "the plane is told once");
assert!(
l.try_debit(CostUnits(10), t(10_000)).is_err(),
"past the usability window"
);
assert_eq!(
signal.count(),
1,
"and not again while that report still stands"
);
assert_eq!(l.remaining(), CostUnits(100), "still zero-charge");
assert_eq!(l.refill_due_or_rearm(), RefillVerdict::Refused);
assert_eq!(
l.refill_due_or_rearm(),
RefillVerdict::Idle,
"seventy-five units above the mark: no crossing was ever claimed"
);
}
#[test]
fn a_lease_without_a_signal_still_reports_the_crossing() {
let l = lease(100, 1_000, 25);
l.try_debit(CostUnits(80), t(0)).unwrap();
assert!(l.needs_refill());
}
#[test]
fn negative_safety_margin_fails_closed() {
let grant = LeaseGrant {
lease_id: LeaseId(8),
account_id: AccountId(1),
fencing_token: FencingToken(4),
units: CostUnits(10),
expires_at: t(100),
};
let l = LocalLease::with_safety_margin(
grant,
CostUnits::ZERO,
jiff::SignedDuration::from_secs(-10),
);
assert_eq!(
l.try_debit(CostUnits(1), t(99)),
Err(DenyReason::LeaseExpired)
);
}
#[test]
fn concurrent_debits_never_overspend() {
use std::sync::Arc;
let l = Arc::new(lease(1_000, 1_000, 0));
let mut handles = Vec::new();
for _ in 0..8 {
let l = Arc::clone(&l);
handles.push(std::thread::spawn(move || {
let mut granted = 0u64;
for _ in 0..1_000 {
if l.try_debit(CostUnits(1), t(0)).is_ok() {
granted += 1;
}
}
granted
}));
}
let total: u64 = handles.into_iter().map(|h| h.join().unwrap()).sum();
assert_eq!(total, 1_000);
assert_eq!(l.remaining(), CostUnits::ZERO);
}
}