#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum BudgetDimension {
Token,
CostCents,
TimeSecs,
}
impl std::fmt::Display for BudgetDimension {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let name = match self {
BudgetDimension::Token => "token",
BudgetDimension::CostCents => "cost_cents",
BudgetDimension::TimeSecs => "time_secs",
};
write!(f, "{name}")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct BudgetConsumption {
pub dimension: BudgetDimension,
pub amount: u64,
}
impl BudgetConsumption {
pub fn new(dimension: BudgetDimension, amount: u64) -> Self {
Self { dimension, amount }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct BudgetAllocation {
dimension: BudgetDimension,
limit: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BudgetAllocationError {
ZeroLimit,
}
impl std::fmt::Display for BudgetAllocationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BudgetAllocationError::ZeroLimit => write!(f, "budget limit must be greater than zero"),
}
}
}
impl std::error::Error for BudgetAllocationError {}
impl BudgetAllocation {
pub fn new(dimension: BudgetDimension, limit: u64) -> Result<Self, BudgetAllocationError> {
if limit == 0 {
return Err(BudgetAllocationError::ZeroLimit);
}
Ok(Self { dimension, limit })
}
pub fn dimension(&self) -> BudgetDimension {
self.dimension
}
pub fn limit(&self) -> u64 {
self.limit
}
}
#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for BudgetAllocation {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(serde::Deserialize)]
struct Wire {
dimension: BudgetDimension,
limit: u64,
}
let wire = Wire::deserialize(deserializer)?;
BudgetAllocation::new(wire.dimension, wire.limit)
.map_err(|_| serde::de::Error::custom("budget limit must be greater than zero"))
}
}