use crate::identity::BillingSubject;
use crate::observation::MeterKind;
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct QuotaRequest {
pub subject: BillingSubject,
pub requested: UsageAmount,
pub _placeholder: (),
}
#[derive(Debug, Clone)]
pub enum QuotaDecision {
Allowed {
remaining: Option<u64>,
},
Denied {
reason: String,
},
}
#[derive(Debug, Clone)]
pub struct ReservationRequest {
pub subject: BillingSubject,
pub amount: UsageAmount,
pub _placeholder: (),
}
#[derive(Debug, Clone)]
pub struct Reservation {
pub id: String,
pub amount: UsageAmount,
pub _placeholder: (),
}
#[derive(Debug, Clone, Default)]
pub struct UsageAmount {
pub meters: HashMap<MeterKind, u64>,
}
impl UsageAmount {
pub fn new() -> Self {
Self {
meters: HashMap::new(),
}
}
pub fn insert(&mut self, kind: MeterKind, quantity: u64) {
self.meters
.entry(kind)
.and_modify(|v| *v += quantity)
.or_insert(quantity);
}
pub fn get(&self, kind: &MeterKind) -> u64 {
self.meters.get(kind).copied().unwrap_or(0)
}
}
pub trait QuotaAuthorizer: Send + Sync {
fn authorize(&self, request: &QuotaRequest) -> Result<QuotaDecision, QuotaError>;
}
pub trait QuotaReservator: Send + Sync {
fn reserve(
&self,
reservation: &ReservationRequest,
) -> Result<Reservation, QuotaError>;
fn commit(&self, reservation_id: &str, amount: &UsageAmount)
-> Result<(), QuotaError>;
fn refund(&self, reservation_id: &str, unused: &UsageAmount)
-> Result<(), QuotaError>;
}
#[derive(Debug, Clone)]
pub enum QuotaError {
Exceeded {
subject: String,
limit: u64,
requested: u64,
},
ConnectionError(String),
ReservationNotFound(String),
InvalidState(String),
Other(String),
}
impl std::fmt::Display for QuotaError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
QuotaError::Exceeded {
subject,
limit,
requested,
} => write!(
f,
"Quota exceeded for {subject}: limit={limit}, requested={requested}"
),
QuotaError::ConnectionError(e) => {
write!(f, "Quota connection error: {e}")
}
QuotaError::ReservationNotFound(id) => {
write!(f, "Reservation not found: {id}")
}
QuotaError::InvalidState(e) => write!(f, "Invalid reservation state: {e}"),
QuotaError::Other(e) => write!(f, "Quota error: {e}"),
}
}
}
impl std::error::Error for QuotaError {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn usage_amount_insert_and_get() {
let mut ua = UsageAmount::new();
ua.insert(MeterKind::InputTokens, 100);
ua.insert(MeterKind::InputTokens, 50);
assert_eq!(ua.get(&MeterKind::InputTokens), 150);
}
#[test]
fn usage_amount_default_empty() {
let ua = UsageAmount::default();
assert_eq!(ua.get(&MeterKind::InputTokens), 0);
}
#[test]
fn quota_error_display() {
let err = QuotaError::Exceeded {
subject: "user-1".to_string(),
limit: 1000,
requested: 1500,
};
assert!(err.to_string().contains("exceeded"));
}
}