use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
#[non_exhaustive]
pub enum SubscriptionRequestClass {
AuthRefresh,
Catalog,
Inference,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct SubscriptionRequestLimits {
auth_refresh: u32,
catalog: u32,
inference: u32,
max_inference_input_tokens: u64,
max_elapsed: Option<Duration>,
retries: u32,
}
impl SubscriptionRequestLimits {
pub fn new(
auth_refresh: u32,
catalog: u32,
inference: u32,
max_inference_input_tokens: u64,
max_elapsed: Option<Duration>,
retries: u32,
) -> Result<Self, SubscriptionLimitError> {
if max_elapsed.is_some_and(|elapsed| elapsed.is_zero())
|| (inference > 0 && max_inference_input_tokens == 0)
{
return Err(SubscriptionLimitError::InvalidLimits);
}
Ok(Self {
auth_refresh,
catalog,
inference,
max_inference_input_tokens,
max_elapsed,
retries,
})
}
#[must_use]
pub const fn allowance(self, class: SubscriptionRequestClass) -> u32 {
match class {
SubscriptionRequestClass::AuthRefresh => self.auth_refresh,
SubscriptionRequestClass::Catalog => self.catalog,
SubscriptionRequestClass::Inference => self.inference,
}
}
#[must_use]
pub const fn max_elapsed(self) -> Option<Duration> {
self.max_elapsed
}
#[must_use]
pub const fn max_inference_input_tokens(self) -> u64 {
self.max_inference_input_tokens
}
#[must_use]
pub const fn retries(self) -> u32 {
self.retries
}
}
pub trait SubscriptionGateClock: Send + Sync + 'static {
fn now(&self) -> Duration;
}
#[derive(Debug)]
pub struct SystemSubscriptionGateClock {
origin: Instant,
}
impl Default for SystemSubscriptionGateClock {
fn default() -> Self {
Self {
origin: Instant::now(),
}
}
}
impl SubscriptionGateClock for SystemSubscriptionGateClock {
fn now(&self) -> Duration {
self.origin.elapsed()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct SubscriptionClassSnapshot {
allowance: u32,
reserved: u32,
dispatched: u32,
abandoned: u32,
denied: u32,
input_token_upper_bound: u64,
}
impl SubscriptionClassSnapshot {
#[must_use]
pub const fn allowance(self) -> u32 {
self.allowance
}
#[must_use]
pub const fn reserved(self) -> u32 {
self.reserved
}
#[must_use]
pub const fn dispatched(self) -> u32 {
self.dispatched
}
#[must_use]
pub const fn abandoned(self) -> u32 {
self.abandoned
}
#[must_use]
pub const fn denied(self) -> u32 {
self.denied
}
#[must_use]
pub const fn input_token_upper_bound(self) -> u64 {
self.input_token_upper_bound
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct SubscriptionLimitSnapshot {
auth_refresh: SubscriptionClassSnapshot,
catalog: SubscriptionClassSnapshot,
inference: SubscriptionClassSnapshot,
elapsed: Duration,
max_elapsed: Option<Duration>,
deadline_reached: bool,
retries: u32,
}
impl SubscriptionLimitSnapshot {
#[must_use]
pub const fn class(self, class: SubscriptionRequestClass) -> SubscriptionClassSnapshot {
match class {
SubscriptionRequestClass::AuthRefresh => self.auth_refresh,
SubscriptionRequestClass::Catalog => self.catalog,
SubscriptionRequestClass::Inference => self.inference,
}
}
#[must_use]
pub const fn elapsed(self) -> Duration {
self.elapsed
}
#[must_use]
pub const fn max_elapsed(self) -> Option<Duration> {
self.max_elapsed
}
#[must_use]
pub const fn deadline_reached(self) -> bool {
self.deadline_reached
}
#[must_use]
pub const fn retries(self) -> u32 {
self.retries
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum SubscriptionLimitError {
InvalidLimits,
DeadlineExceeded,
ClassLimitExceeded {
class: SubscriptionRequestClass,
reserved: u32,
allowance: u32,
},
PermitAlreadyUsed,
InferenceBoundRequired,
InferenceInputLimitExceeded,
InferenceScopeUnavailable,
Unavailable,
}
impl fmt::Display for SubscriptionLimitError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidLimits => formatter.write_str("subscription request limits are invalid"),
Self::DeadlineExceeded => {
formatter.write_str("subscription request deadline was reached")
}
Self::ClassLimitExceeded {
class,
reserved,
allowance,
} => write!(
formatter,
"subscription {class:?} reserved {reserved} of allowance {allowance}"
),
Self::PermitAlreadyUsed => {
formatter.write_str("subscription request permit was already used")
}
Self::InferenceBoundRequired => formatter
.write_str("subscription inference dispatch requires a conservative input bound"),
Self::InferenceInputLimitExceeded => {
formatter.write_str("subscription inference input limit was reached")
}
Self::InferenceScopeUnavailable => {
formatter.write_str("subscription inference root scope is unavailable")
}
Self::Unavailable => formatter.write_str("subscription request gate is unavailable"),
}
}
}
impl std::error::Error for SubscriptionLimitError {}
pub trait SubscriptionRequestPermit: Send + 'static {
fn class(&self) -> SubscriptionRequestClass;
fn authorize_dispatch(&mut self) -> Result<(), SubscriptionLimitError>;
fn refund_not_started(&mut self) -> Result<(), SubscriptionLimitError>;
}
pub trait SubscriptionRequestGate: Send + Sync + 'static {
fn acquire(
&self,
class: SubscriptionRequestClass,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError>;
fn acquire_inference(
&self,
input_token_upper_bound: u64,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError>;
fn snapshot(&self) -> Result<SubscriptionLimitSnapshot, SubscriptionLimitError>;
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum SubscriptionInferenceScope {
Turn([u8; 16]),
SessionWorkflow([u8; 32]),
}
impl SubscriptionInferenceScope {
pub fn turn(bytes: [u8; 16]) -> Result<Self, SubscriptionLimitError> {
if bytes == [0; 16] {
Err(SubscriptionLimitError::InvalidLimits)
} else {
Ok(Self::Turn(bytes))
}
}
pub fn session_workflow(bytes: [u8; 32]) -> Result<Self, SubscriptionLimitError> {
if bytes == [0; 32] {
Err(SubscriptionLimitError::InvalidLimits)
} else {
Ok(Self::SessionWorkflow(bytes))
}
}
}
#[derive(Clone, Copy, Debug, Default)]
struct ClassState {
dispatch_attempts: u32,
reserved: u32,
dispatched: u32,
abandoned: u32,
denied: u32,
input_token_upper_bound: u64,
}
#[derive(Debug, Default)]
struct GateState {
auth_refresh: ClassState,
catalog: ClassState,
inference: ClassState,
}
impl GateState {
fn class_mut(&mut self, class: SubscriptionRequestClass) -> &mut ClassState {
match class {
SubscriptionRequestClass::AuthRefresh => &mut self.auth_refresh,
SubscriptionRequestClass::Catalog => &mut self.catalog,
SubscriptionRequestClass::Inference => &mut self.inference,
}
}
const fn class(&self, class: SubscriptionRequestClass) -> ClassState {
match class {
SubscriptionRequestClass::AuthRefresh => self.auth_refresh,
SubscriptionRequestClass::Catalog => self.catalog,
SubscriptionRequestClass::Inference => self.inference,
}
}
}
struct GateInner {
limits: SubscriptionRequestLimits,
clock: Arc<dyn SubscriptionGateClock>,
started: Duration,
state: Mutex<GateState>,
}
impl fmt::Debug for GateInner {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SubscriptionRequestGate")
.field("limits", &self.limits)
.field("started", &self.started)
.finish_non_exhaustive()
}
}
impl GateInner {
fn elapsed(&self) -> Duration {
self.clock.now().saturating_sub(self.started)
}
fn deadline_reached(&self) -> bool {
self.limits
.max_elapsed
.is_some_and(|limit| self.elapsed() >= limit)
}
}
#[derive(Clone)]
pub struct InMemorySubscriptionRequestGate {
inner: Arc<GateInner>,
}
impl fmt::Debug for InMemorySubscriptionRequestGate {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.inner.fmt(formatter)
}
}
impl InMemorySubscriptionRequestGate {
#[must_use]
pub fn new(limits: SubscriptionRequestLimits, clock: Arc<dyn SubscriptionGateClock>) -> Self {
let started = clock.now();
Self {
inner: Arc::new(GateInner {
limits,
clock,
started,
state: Mutex::new(GateState::default()),
}),
}
}
}
thread_local! {
static ACTIVE_INFERENCE_SCOPES: RefCell<Vec<(u64, InMemorySubscriptionRequestGate)>> =
const { RefCell::new(Vec::new()) };
}
static NEXT_TURN_GATE_ID: AtomicU64 = AtomicU64::new(1);
#[derive(Clone)]
pub struct SubscriptionInferenceScopeHandle {
gate_id: u64,
scope: SubscriptionInferenceScope,
inference: InMemorySubscriptionRequestGate,
global: InMemorySubscriptionRequestGate,
}
pub struct SubscriptionInferenceScopeLease {
gate: Arc<TurnScopedSubscriptionRequestGate>,
handle: SubscriptionInferenceScopeHandle,
}
impl fmt::Debug for SubscriptionInferenceScopeLease {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SubscriptionInferenceScopeLease")
.field("scope", &self.handle.scope)
.finish_non_exhaustive()
}
}
impl SubscriptionInferenceScopeLease {
#[must_use]
pub fn handle(&self) -> SubscriptionInferenceScopeHandle {
self.handle.clone()
}
#[must_use]
pub const fn scope(&self) -> SubscriptionInferenceScope {
self.handle.scope
}
}
impl Drop for SubscriptionInferenceScopeLease {
fn drop(&mut self) {
let mut scopes = match self.gate.scopes.lock() {
Ok(scopes) => scopes,
Err(poisoned) => poisoned.into_inner(),
};
scopes.live.remove(&self.handle.scope);
scopes.retired.insert(self.handle.scope);
}
}
impl fmt::Debug for SubscriptionInferenceScopeHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SubscriptionInferenceScopeHandle")
.field("scope", &self.scope)
.finish_non_exhaustive()
}
}
impl SubscriptionInferenceScopeHandle {
#[must_use]
pub const fn scope(&self) -> SubscriptionInferenceScope {
self.scope
}
#[must_use]
pub fn enter(&self) -> SubscriptionInferenceScopeGuard {
ACTIVE_INFERENCE_SCOPES.with(|active| {
active
.borrow_mut()
.push((self.gate_id, self.inference.clone()));
});
SubscriptionInferenceScopeGuard {
gate_id: self.gate_id,
}
}
pub fn snapshot(&self) -> Result<SubscriptionLimitSnapshot, SubscriptionLimitError> {
scope_snapshot(&self.global, &self.inference)
}
}
pub struct SubscriptionInferenceScopeGuard {
gate_id: u64,
}
impl Drop for SubscriptionInferenceScopeGuard {
fn drop(&mut self) {
ACTIVE_INFERENCE_SCOPES.with(|active| {
let popped = active.borrow_mut().pop();
debug_assert!(popped.is_some_and(|(gate_id, _)| gate_id == self.gate_id));
});
}
}
pub struct TurnScopedSubscriptionRequestGate {
limits: SubscriptionRequestLimits,
clock: Arc<dyn SubscriptionGateClock>,
global: InMemorySubscriptionRequestGate,
gate_id: u64,
scopes: Mutex<InferenceScopes>,
}
#[derive(Default)]
struct InferenceScopes {
live: HashMap<SubscriptionInferenceScope, InMemorySubscriptionRequestGate>,
retired: HashSet<SubscriptionInferenceScope>,
}
impl fmt::Debug for TurnScopedSubscriptionRequestGate {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("TurnScopedSubscriptionRequestGate")
.field("limits", &self.limits)
.finish_non_exhaustive()
}
}
impl TurnScopedSubscriptionRequestGate {
#[must_use]
pub fn new(limits: SubscriptionRequestLimits, clock: Arc<dyn SubscriptionGateClock>) -> Self {
Self {
limits,
clock: Arc::clone(&clock),
global: InMemorySubscriptionRequestGate::new(limits, clock),
gate_id: NEXT_TURN_GATE_ID.fetch_add(1, Ordering::Relaxed),
scopes: Mutex::new(InferenceScopes::default()),
}
}
pub fn claim_scope(
self: &Arc<Self>,
scope: SubscriptionInferenceScope,
) -> Result<SubscriptionInferenceScopeLease, SubscriptionLimitError> {
let mut scopes = self
.scopes
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
if scopes.live.contains_key(&scope) || scopes.retired.contains(&scope) {
return Err(SubscriptionLimitError::InferenceScopeUnavailable);
}
let inference = InMemorySubscriptionRequestGate::new(self.limits, Arc::clone(&self.clock));
scopes.live.insert(scope, inference.clone());
let handle = SubscriptionInferenceScopeHandle {
gate_id: self.gate_id,
scope,
inference,
global: self.global.clone(),
};
Ok(SubscriptionInferenceScopeLease {
gate: Arc::clone(self),
handle,
})
}
pub fn bind_scope(
&self,
scope: SubscriptionInferenceScope,
) -> Result<SubscriptionInferenceScopeHandle, SubscriptionLimitError> {
let scopes = self
.scopes
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
let inference = scopes
.live
.get(&scope)
.cloned()
.ok_or(SubscriptionLimitError::InferenceScopeUnavailable)?;
Ok(SubscriptionInferenceScopeHandle {
gate_id: self.gate_id,
scope,
inference,
global: self.global.clone(),
})
}
fn active_inference_gate(&self) -> Option<InMemorySubscriptionRequestGate> {
if let Some(gate) = ACTIVE_INFERENCE_SCOPES.with(|active| {
active
.borrow()
.iter()
.rev()
.find(|(gate_id, _)| *gate_id == self.gate_id)
.map(|(_, gate)| gate.clone())
}) {
return Some(gate);
}
let scopes = self.scopes.lock().ok()?;
if scopes.live.len() != 1 {
return None;
}
scopes.live.values().next().cloned()
}
}
impl SubscriptionRequestGate for TurnScopedSubscriptionRequestGate {
fn acquire(
&self,
class: SubscriptionRequestClass,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError> {
if class == SubscriptionRequestClass::Inference {
return Err(SubscriptionLimitError::InferenceBoundRequired);
}
self.active_inference_gate()
.unwrap_or_else(|| self.global.clone())
.acquire(class)
}
fn acquire_inference(
&self,
input_token_upper_bound: u64,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError> {
self.active_inference_gate()
.ok_or(SubscriptionLimitError::InferenceScopeUnavailable)?
.acquire_inference(input_token_upper_bound)
}
fn snapshot(&self) -> Result<SubscriptionLimitSnapshot, SubscriptionLimitError> {
let global = self.global.snapshot()?;
let inference = self
.active_inference_gate()
.map(|gate| gate.snapshot())
.transpose()?
.map_or_else(
|| SubscriptionClassSnapshot {
allowance: self.limits.allowance(SubscriptionRequestClass::Inference),
reserved: 0,
dispatched: 0,
abandoned: 0,
denied: 0,
input_token_upper_bound: 0,
},
|snapshot| snapshot.class(SubscriptionRequestClass::Inference),
);
Ok(SubscriptionLimitSnapshot {
auth_refresh: global.class(SubscriptionRequestClass::AuthRefresh),
catalog: global.class(SubscriptionRequestClass::Catalog),
inference,
elapsed: global.elapsed(),
max_elapsed: global.max_elapsed(),
deadline_reached: global.deadline_reached(),
retries: global.retries(),
})
}
}
fn scope_snapshot(
global_gate: &InMemorySubscriptionRequestGate,
inference_gate: &InMemorySubscriptionRequestGate,
) -> Result<SubscriptionLimitSnapshot, SubscriptionLimitError> {
let global = global_gate.snapshot()?;
let inference = inference_gate.snapshot()?;
Ok(SubscriptionLimitSnapshot {
auth_refresh: global.class(SubscriptionRequestClass::AuthRefresh),
catalog: global.class(SubscriptionRequestClass::Catalog),
inference: inference.class(SubscriptionRequestClass::Inference),
elapsed: inference.elapsed(),
max_elapsed: inference.max_elapsed(),
deadline_reached: inference.deadline_reached(),
retries: inference.retries(),
})
}
impl InMemorySubscriptionRequestGate {
pub fn dispatch_attempts(
&self,
class: SubscriptionRequestClass,
) -> Result<u32, SubscriptionLimitError> {
let state = self
.inner
.state
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
Ok(state.class(class).dispatch_attempts)
}
}
#[cfg(feature = "http-refresh")]
pub struct CatalogRequestDispatch {
gate: Option<Arc<dyn SubscriptionRequestGate>>,
permit: Option<Box<dyn SubscriptionRequestPermit>>,
}
#[cfg(feature = "http-refresh")]
impl CatalogRequestDispatch {
#[must_use]
pub fn new(gate: Option<Arc<dyn SubscriptionRequestGate>>) -> Self {
Self { gate, permit: None }
}
pub fn observe(
&mut self,
failure: Option<plakous_http::TransmissionFact>,
) -> Result<(), SubscriptionLimitError> {
if let Some(failure) = failure {
if failure == plakous_http::TransmissionFact::NotStarted
&& let Some(permit) = &mut self.permit
{
permit.refund_not_started()?;
}
} else {
self.permit = None;
if let Some(gate) = &self.gate {
let mut permit = gate.acquire(SubscriptionRequestClass::Catalog)?;
permit.authorize_dispatch()?;
self.permit = Some(permit);
}
}
Ok(())
}
}
impl SubscriptionRequestGate for InMemorySubscriptionRequestGate {
fn acquire(
&self,
class: SubscriptionRequestClass,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError> {
if class == SubscriptionRequestClass::Inference {
return Err(SubscriptionLimitError::InferenceBoundRequired);
}
self.acquire_inner(class, None)
}
fn acquire_inference(
&self,
input_token_upper_bound: u64,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError> {
if input_token_upper_bound == 0 {
return Err(SubscriptionLimitError::InferenceBoundRequired);
}
self.acquire_inner(
SubscriptionRequestClass::Inference,
Some(input_token_upper_bound),
)
}
fn snapshot(&self) -> Result<SubscriptionLimitSnapshot, SubscriptionLimitError> {
let state = self
.inner
.state
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
let snapshot = |class| {
let value = state.class(class);
SubscriptionClassSnapshot {
allowance: self.inner.limits.allowance(class),
reserved: value.reserved,
dispatched: value.dispatched,
abandoned: value.abandoned,
denied: value.denied,
input_token_upper_bound: value.input_token_upper_bound,
}
};
Ok(SubscriptionLimitSnapshot {
auth_refresh: snapshot(SubscriptionRequestClass::AuthRefresh),
catalog: snapshot(SubscriptionRequestClass::Catalog),
inference: snapshot(SubscriptionRequestClass::Inference),
elapsed: self.inner.elapsed(),
max_elapsed: self.inner.limits.max_elapsed,
deadline_reached: self.inner.deadline_reached(),
retries: self.inner.limits.retries,
})
}
}
impl InMemorySubscriptionRequestGate {
fn acquire_inner(
&self,
class: SubscriptionRequestClass,
input_token_upper_bound: Option<u64>,
) -> Result<Box<dyn SubscriptionRequestPermit>, SubscriptionLimitError> {
let mut state = self
.inner
.state
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
let class_state = state.class_mut(class);
if self.inner.deadline_reached() {
class_state.denied = class_state
.denied
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
return Err(SubscriptionLimitError::DeadlineExceeded);
}
let allowance = self.inner.limits.allowance(class);
if class_state.reserved >= allowance {
class_state.denied = class_state
.denied
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
return Err(SubscriptionLimitError::ClassLimitExceeded {
class,
reserved: class_state.reserved,
allowance,
});
}
if let Some(bound) = input_token_upper_bound {
let candidate = class_state
.input_token_upper_bound
.checked_add(bound)
.ok_or(SubscriptionLimitError::Unavailable)?;
if candidate > self.inner.limits.max_inference_input_tokens {
class_state.denied = class_state
.denied
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
return Err(SubscriptionLimitError::InferenceInputLimitExceeded);
}
class_state.input_token_upper_bound = candidate;
}
class_state.reserved = class_state
.reserved
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
Ok(Box::new(InMemoryPermit {
inner: Arc::clone(&self.inner),
class,
authorized: false,
settled: false,
refunded: false,
input_token_upper_bound: input_token_upper_bound.unwrap_or(0),
}))
}
}
struct InMemoryPermit {
inner: Arc<GateInner>,
class: SubscriptionRequestClass,
authorized: bool,
settled: bool,
refunded: bool,
input_token_upper_bound: u64,
}
impl SubscriptionRequestPermit for InMemoryPermit {
fn class(&self) -> SubscriptionRequestClass {
self.class
}
fn authorize_dispatch(&mut self) -> Result<(), SubscriptionLimitError> {
if self.authorized {
return Err(SubscriptionLimitError::PermitAlreadyUsed);
}
if self.inner.deadline_reached() {
return Err(SubscriptionLimitError::DeadlineExceeded);
}
let mut state = self
.inner
.state
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
let class = state.class_mut(self.class);
class.dispatch_attempts = class
.dispatch_attempts
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
class.dispatched = class
.dispatched
.checked_add(1)
.ok_or(SubscriptionLimitError::Unavailable)?;
self.authorized = true;
self.settled = true;
Ok(())
}
fn refund_not_started(&mut self) -> Result<(), SubscriptionLimitError> {
if !self.authorized || self.refunded {
return Err(SubscriptionLimitError::PermitAlreadyUsed);
}
let mut state = self
.inner
.state
.lock()
.map_err(|_| SubscriptionLimitError::Unavailable)?;
let class = state.class_mut(self.class);
class.reserved -= 1;
class.dispatched -= 1;
class.input_token_upper_bound -= self.input_token_upper_bound;
self.refunded = true;
Ok(())
}
}
impl Drop for InMemoryPermit {
fn drop(&mut self) {
if self.settled {
return;
}
if let Ok(mut state) = self.inner.state.lock() {
let class = state.class_mut(self.class);
class.abandoned = class.abandoned.saturating_add(1);
self.settled = true;
}
}
}
#[cfg(test)]
#[allow(
clippy::expect_used,
reason = "gate test fixtures use explicit panic messages for invariant setup"
)]
mod tests {
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread;
use super::*;
#[derive(Default)]
struct ManualClock(AtomicU64);
impl ManualClock {
fn advance(&self, seconds: u64) {
self.0.fetch_add(seconds, Ordering::SeqCst);
}
}
impl SubscriptionGateClock for ManualClock {
fn now(&self) -> Duration {
Duration::from_secs(self.0.load(Ordering::SeqCst))
}
}
fn gate(clock: Arc<ManualClock>) -> InMemorySubscriptionRequestGate {
InMemorySubscriptionRequestGate::new(
SubscriptionRequestLimits::new(0, 1, 5, 4_096, Some(Duration::from_secs(300)), 0)
.expect("fixed limits are valid"),
clock,
)
}
#[test]
fn proven_not_started_refunds_count_and_input_without_reusing_authorization() {
let gate = gate(Arc::new(ManualClock::default()));
let before = gate.snapshot().expect("snapshot").inference;
let mut permit = gate
.acquire_inference(4_096)
.expect("whole input allowance");
permit.authorize_dispatch().expect("dispatch");
permit.refund_not_started().expect("proven unsent refund");
assert_eq!(gate.snapshot().expect("snapshot").inference, before);
assert_eq!(
permit.authorize_dispatch(),
Err(SubscriptionLimitError::PermitAlreadyUsed)
);
assert_eq!(
permit.refund_not_started(),
Err(SubscriptionLimitError::PermitAlreadyUsed)
);
drop(permit);
assert_eq!(gate.snapshot().expect("snapshot").inference, before);
assert!(gate.acquire_inference(4_096).is_ok());
}
#[test]
fn zero_refresh_one_catalog_and_five_inference_are_exact() {
let gate = gate(Arc::new(ManualClock::default()));
assert_eq!(
gate.acquire(SubscriptionRequestClass::AuthRefresh).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::AuthRefresh,
reserved: 0,
allowance: 0,
})
);
let mut catalog = gate
.acquire(SubscriptionRequestClass::Catalog)
.expect("catalog permit");
catalog.authorize_dispatch().expect("catalog dispatch");
assert_eq!(
catalog.authorize_dispatch(),
Err(SubscriptionLimitError::PermitAlreadyUsed)
);
assert!(gate.acquire(SubscriptionRequestClass::Catalog).is_err());
for _ in 0..5 {
gate.acquire_inference(100)
.expect("inference permit")
.authorize_dispatch()
.expect("inference dispatch");
}
assert_eq!(
gate.acquire(SubscriptionRequestClass::Inference).err(),
Some(SubscriptionLimitError::InferenceBoundRequired)
);
assert_eq!(
gate.acquire_inference(100).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 5,
allowance: 5,
})
);
let snapshot = gate.snapshot().expect("snapshot");
assert_eq!(snapshot.retries(), 0);
assert_eq!(
snapshot.class(SubscriptionRequestClass::AuthRefresh),
SubscriptionClassSnapshot {
allowance: 0,
reserved: 0,
dispatched: 0,
abandoned: 0,
denied: 1,
input_token_upper_bound: 0,
}
);
assert_eq!(
snapshot
.class(SubscriptionRequestClass::Inference)
.dispatched(),
5
);
assert_eq!(
snapshot
.class(SubscriptionRequestClass::Inference)
.input_token_upper_bound(),
500
);
}
#[test]
fn abandoned_permit_consumes_capacity_and_cannot_mint_retry() {
let gate = gate(Arc::new(ManualClock::default()));
drop(
gate.acquire(SubscriptionRequestClass::Catalog)
.expect("catalog permit"),
);
assert_eq!(
gate.acquire(SubscriptionRequestClass::Catalog).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Catalog,
reserved: 1,
allowance: 1,
})
);
let state = gate
.snapshot()
.expect("snapshot")
.class(SubscriptionRequestClass::Catalog);
assert_eq!(state.reserved(), 1);
assert_eq!(state.dispatched(), 0);
assert_eq!(state.abandoned(), 1);
assert_eq!(state.denied(), 1);
}
#[test]
fn one_absolute_deadline_applies_to_existing_and_new_permits() {
let clock = Arc::new(ManualClock::default());
let gate = gate(clock.clone());
let mut permit = gate.acquire_inference(1).expect("pre-deadline permit");
clock.advance(300);
assert_eq!(
permit.authorize_dispatch(),
Err(SubscriptionLimitError::DeadlineExceeded)
);
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::DeadlineExceeded)
);
drop(permit);
let snapshot = gate.snapshot().expect("snapshot");
assert!(snapshot.deadline_reached());
assert_eq!(snapshot.elapsed(), Duration::from_secs(300));
assert_eq!(snapshot.max_elapsed(), Some(Duration::from_secs(300)));
}
#[test]
fn concurrent_reservations_never_overshoot() {
let gate = Arc::new(gate(Arc::new(ManualClock::default())));
let threads = (0..32)
.map(|_| {
let gate = Arc::clone(&gate);
thread::spawn(move || {
gate.acquire_inference(1)
.map(|mut permit| permit.authorize_dispatch())
})
})
.collect::<Vec<_>>();
let successes = threads
.into_iter()
.map(thread::JoinHandle::join)
.filter(|result| matches!(result, Ok(Ok(Ok(())))))
.count();
assert_eq!(successes, 5);
let state = gate
.snapshot()
.expect("snapshot")
.class(SubscriptionRequestClass::Inference);
assert_eq!(state.reserved(), 5);
assert_eq!(state.dispatched(), 5);
assert_eq!(state.denied(), 27);
}
#[test]
fn scope_handles_prevent_a_b_a_renewal_and_keep_truthful_deadlines() {
let clock = Arc::new(ManualClock::default());
let limits = SubscriptionRequestLimits::new(0, 1, 2, 10, Some(Duration::from_secs(5)), 0)
.expect("limits");
let gate = Arc::new(TurnScopedSubscriptionRequestGate::new(
limits,
clock.clone(),
));
let first = SubscriptionInferenceScope::turn([1; 16]).expect("first scope");
let second = SubscriptionInferenceScope::turn([2; 16]).expect("second scope");
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::InferenceScopeUnavailable),
"inference before a claimed root must fail closed"
);
let first_lease = gate.claim_scope(first).expect("claim first");
gate.acquire_inference(1)
.expect("the sole claimed root is unambiguous without a poll-local enter")
.authorize_dispatch()
.expect("sole-root dispatch");
let first_handle = first_lease.handle();
let first_active = first_handle.enter();
gate.acquire_inference(5)
.expect("count permit")
.authorize_dispatch()
.expect("count dispatch");
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 2,
allowance: 2,
})
);
drop(first_active);
let second_lease = gate.claim_scope(second).expect("claim second");
let second_handle = second_lease.handle();
let second_active = second_handle.enter();
gate.acquire_inference(1)
.expect("second scope is independent")
.authorize_dispatch()
.expect("second dispatch");
drop(second_active);
let first_again = gate.bind_scope(first).expect("rebind live first");
assert_eq!(first_again.scope(), first);
let first_again_active = first_again.enter();
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 2,
allowance: 2,
})
);
drop(first_again_active);
let _second_active = second_handle.enter();
assert_eq!(
gate.acquire_inference(11).err(),
Some(SubscriptionLimitError::InferenceInputLimitExceeded)
);
let mut before_deadline = gate.acquire_inference(1).expect("deadline permit");
clock.advance(5);
assert_eq!(
before_deadline.authorize_dispatch(),
Err(SubscriptionLimitError::DeadlineExceeded)
);
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::DeadlineExceeded)
);
let snapshot = second_handle.snapshot().expect("scope snapshot");
assert!(snapshot.deadline_reached());
assert_eq!(snapshot.elapsed(), Duration::from_secs(5));
let state = snapshot.class(SubscriptionRequestClass::Inference);
assert_eq!(state.reserved(), 2);
assert_eq!(state.dispatched(), 1);
assert_eq!(state.denied(), 2);
assert_eq!(state.input_token_upper_bound(), 2);
drop(first_lease);
assert_eq!(
gate.claim_scope(first).err(),
Some(SubscriptionLimitError::InferenceScopeUnavailable)
);
drop(second_lease);
}
#[test]
fn concurrent_scope_activation_never_cross_charges() {
let clock = Arc::new(ManualClock::default());
let limits = SubscriptionRequestLimits::new(0, 1, 1, 10, Some(Duration::from_secs(5)), 0)
.expect("limits");
let gate = Arc::new(TurnScopedSubscriptionRequestGate::new(limits, clock));
let first_lease = gate
.claim_scope(SubscriptionInferenceScope::turn([1; 16]).expect("first"))
.expect("claim first");
let second_lease = gate
.claim_scope(SubscriptionInferenceScope::turn([2; 16]).expect("second"))
.expect("claim second");
let first = first_lease.handle();
let second = second_lease.handle();
let threads = [first.clone(), second.clone()].map(|scope| {
let gate = Arc::clone(&gate);
thread::spawn(move || {
let _active = scope.enter();
gate.acquire_inference(3)?.authorize_dispatch()
})
});
for thread in threads {
assert_eq!(thread.join().expect("thread"), Ok(()));
}
for scope in [first, second] {
let inference = scope
.snapshot()
.expect("scope snapshot")
.class(SubscriptionRequestClass::Inference);
assert_eq!(inference.reserved(), 1);
assert_eq!(inference.dispatched(), 1);
assert_eq!(inference.input_token_upper_bound(), 3);
}
}
#[test]
fn dyn_gate_acquires_after_typed_claim_without_tls_and_after_lease_drop() {
let clock = Arc::new(ManualClock::default());
let limits =
SubscriptionRequestLimits::new(1, 1, 8, 1_000_000, Some(Duration::from_secs(30)), 0)
.expect("limits");
let typed = Arc::new(TurnScopedSubscriptionRequestGate::new(limits, clock));
let dyn_gate: Arc<dyn SubscriptionRequestGate> = typed.clone();
let scope = SubscriptionInferenceScope::turn([9; 16]).expect("scope");
let lease = typed.claim_scope(scope).expect("claim");
dyn_gate
.acquire_inference(100_000)
.expect("manager-style dyn acquire must succeed on the claimed root without TLS")
.authorize_dispatch()
.expect("dispatch");
drop(lease);
assert_eq!(
dyn_gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::InferenceScopeUnavailable)
);
}
#[test]
fn claimed_root_inference_is_independent_of_a_saturated_process_class() {
let clock = Arc::new(ManualClock::default());
let limits =
SubscriptionRequestLimits::new(1, 1, 8, 1_000_000, Some(Duration::from_secs(30)), 0)
.expect("limits");
let typed = Arc::new(TurnScopedSubscriptionRequestGate::new(limits, clock));
let dyn_gate: Arc<dyn SubscriptionRequestGate> = typed.clone();
let global_permits = (0..8)
.map(|_| {
typed
.global
.acquire_inference(1)
.expect("process class reservation")
})
.collect::<Vec<_>>();
assert_eq!(
typed.global.acquire_inference(1).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 8,
allowance: 8,
})
);
let lease = typed
.claim_scope(SubscriptionInferenceScope::turn([3; 16]).expect("scope"))
.expect("claim");
dyn_gate
.acquire_inference(1)
.expect("claimed root must not inherit a saturated process inference class")
.authorize_dispatch()
.expect("root dispatch");
drop(lease);
drop(global_permits);
}
#[test]
fn unclocked_gate_charges_refreshes_to_the_live_root_for_hours() {
let clock = Arc::new(ManualClock::default());
let limits = SubscriptionRequestLimits::new(1, 1, 8, 1_000_000, None, 0).expect("limits");
let gate = Arc::new(TurnScopedSubscriptionRequestGate::new(
limits,
clock.clone(),
));
let mut resolution = gate
.acquire(SubscriptionRequestClass::Catalog)
.expect("resolution catalog");
resolution
.authorize_dispatch()
.expect("resolution dispatch");
for (turn, seed) in [(1_u64, 1_u8), (2, 2), (3, 3)] {
clock.advance(turn * 3_600);
let lease = gate
.claim_scope(SubscriptionInferenceScope::turn([seed; 16]).expect("scope"))
.expect("claim");
for class in [
SubscriptionRequestClass::AuthRefresh,
SubscriptionRequestClass::Catalog,
] {
gate.acquire(class)
.expect("in-turn refresh is charged to the root")
.authorize_dispatch()
.expect("in-turn refresh dispatch");
}
clock.advance(3_600);
gate.acquire_inference(1)
.expect("an hour-long turn keeps inference")
.authorize_dispatch()
.expect("inference dispatch");
assert!(!gate.snapshot().expect("snapshot").deadline_reached());
drop(lease);
}
let global = gate.global.snapshot().expect("global snapshot");
assert_eq!(global.max_elapsed(), None);
assert_eq!(
global.class(SubscriptionRequestClass::Catalog).reserved(),
1
);
assert_eq!(
global
.class(SubscriptionRequestClass::AuthRefresh)
.reserved(),
0
);
}
#[test]
fn two_live_roots_do_not_cross_charge_without_tls() {
let clock = Arc::new(ManualClock::default());
let limits =
SubscriptionRequestLimits::new(1, 1, 8, 1_000_000, Some(Duration::from_secs(30)), 0)
.expect("limits");
let gate = Arc::new(TurnScopedSubscriptionRequestGate::new(limits, clock));
let first = gate
.claim_scope(SubscriptionInferenceScope::turn([1; 16]).expect("first"))
.expect("claim first");
let _second = gate
.claim_scope(SubscriptionInferenceScope::turn([2; 16]).expect("second"))
.expect("claim second");
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::InferenceScopeUnavailable)
);
let _active = first.handle().enter();
gate.acquire_inference(1)
.expect("TLS selects the entered root")
.authorize_dispatch()
.expect("dispatch");
}
#[test]
fn eight_abandoned_inference_permits_report_reserved_and_allowance() {
let limits =
SubscriptionRequestLimits::new(1, 2, 8, 1_000_000, Some(Duration::from_secs(300)), 0)
.expect("limits");
let gate = InMemorySubscriptionRequestGate::new(limits, Arc::new(ManualClock::default()));
let permits = (0..8)
.map(|_| gate.acquire_inference(1).expect("inference reservation"))
.collect::<Vec<_>>();
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 8,
allowance: 8,
})
);
let state = gate
.snapshot()
.expect("snapshot")
.class(SubscriptionRequestClass::Inference);
assert_eq!(state.reserved(), 8);
assert_eq!(state.dispatched(), 0);
assert_eq!(state.abandoned(), 0);
assert_eq!(state.denied(), 1);
drop(permits);
let abandoned = gate
.snapshot()
.expect("snapshot")
.class(SubscriptionRequestClass::Inference);
assert_eq!(abandoned.reserved(), 8);
assert_eq!(abandoned.abandoned(), 8);
assert_eq!(
gate.acquire_inference(1).err(),
Some(SubscriptionLimitError::ClassLimitExceeded {
class: SubscriptionRequestClass::Inference,
reserved: 8,
allowance: 8,
})
);
}
}
#[cfg(all(test, feature = "http-refresh"))]
#[allow(
clippy::expect_used,
reason = "test constants and accounting assertions"
)]
mod catalog_dispatch_tests {
use super::*;
#[test]
fn catalog_attempts_keep_fresh_permits_and_only_refund_not_started() {
let gate = Arc::new(InMemorySubscriptionRequestGate::new(
SubscriptionRequestLimits::new(0, 3, 0, 0, None, 0).expect("limits"),
Arc::new(SystemSubscriptionGateClock::default()),
));
let mut dispatch = CatalogRequestDispatch::new(Some(gate.clone()));
dispatch.observe(None).expect("first permit");
dispatch
.observe(Some(plakous_http::TransmissionFact::NotStarted))
.expect("refund");
dispatch.observe(None).expect("fresh permit");
dispatch
.observe(Some(plakous_http::TransmissionFact::MayHaveTransmitted))
.expect("sent");
let state = gate
.snapshot()
.expect("snapshot")
.class(SubscriptionRequestClass::Catalog);
assert_eq!(state.reserved(), 1);
assert_eq!(state.dispatched(), 1);
assert_eq!(
gate.dispatch_attempts(SubscriptionRequestClass::Catalog)
.expect("attempts"),
2
);
}
}