use crate::resource::BudgetGroupError;
use crate::resource::InsufficientBudgetError;
use crate::resource::MeasuredBudgetError;
use crate::resource::ResourceLimit;
use crate::resource::ResourceQuantity;
mod string_output;
#[derive(Debug, PartialEq, Eq)]
pub struct ResourceBudget<R, Q = u64>
where
Q: ResourceQuantity,
{
limit: ResourceLimit<R, Q>,
remaining: Q,
}
impl<R, Q> ResourceBudget<R, Q>
where
Q: ResourceQuantity,
{
#[inline]
#[must_use]
pub const fn new(resource: R, limit: Q) -> Self {
Self {
limit: ResourceLimit::new(resource, limit),
remaining: limit,
}
}
#[inline]
#[must_use]
pub fn from_limit(limit: ResourceLimit<R, Q>) -> Self {
Self {
remaining: limit.maximum(),
limit,
}
}
#[must_use = "the availability check result must be handled"]
pub fn check_available(&self, amount: Q) -> Result<(), InsufficientBudgetError<R, Q>>
where
R: Clone,
{
if amount <= self.remaining {
Ok(())
} else {
Err(InsufficientBudgetError {
resource: self.limit.resource().clone(),
limit: self.limit.maximum(),
remaining: self.remaining,
requested: amount,
})
}
}
#[inline]
pub fn try_consume(&mut self, amount: Q) -> Result<(), InsufficientBudgetError<R, Q>>
where
R: Clone,
{
self.check_available(amount)?;
self.remaining = self.remaining - amount;
Ok(())
}
#[inline]
#[must_use = "the availability check result must be handled"]
pub fn check_available_usize(&self, amount: usize) -> Result<(), MeasuredBudgetError<R, Q>>
where
R: Clone,
{
let amount = Q::try_from_usize(amount)
.map_err(|source| MeasuredBudgetError::quantity(self.resource().clone(), source))?;
self.check_available(amount).map_err(MeasuredBudgetError::from)
}
#[inline]
#[must_use = "the availability check result must be handled"]
pub fn check_available_u64(&self, amount: u64) -> Result<(), MeasuredBudgetError<R, Q>>
where
R: Clone,
{
let amount =
Q::try_from_u64(amount).map_err(|source| MeasuredBudgetError::quantity(self.resource().clone(), source))?;
self.check_available(amount).map_err(MeasuredBudgetError::from)
}
#[inline]
pub fn try_consume_usize(&mut self, amount: usize) -> Result<(), MeasuredBudgetError<R, Q>>
where
R: Clone,
{
let amount = Q::try_from_usize(amount)
.map_err(|source| MeasuredBudgetError::quantity(self.resource().clone(), source))?;
self.try_consume(amount).map_err(MeasuredBudgetError::from)
}
#[inline]
pub fn try_consume_u64(&mut self, amount: u64) -> Result<(), MeasuredBudgetError<R, Q>>
where
R: Clone,
{
let amount =
Q::try_from_u64(amount).map_err(|source| MeasuredBudgetError::quantity(self.resource().clone(), source))?;
self.try_consume(amount).map_err(MeasuredBudgetError::from)
}
pub fn try_consume_group(budgets: &mut [&mut Self], amount: Q) -> Result<(), BudgetGroupError<R, Q>>
where
R: Clone,
{
for (index, budget) in budgets.iter().enumerate() {
budget
.check_available(amount)
.map_err(|source| BudgetGroupError::new(index, source))?;
}
for budget in budgets {
let consumed = budget.consume_available(amount);
debug_assert_eq!(consumed, amount);
}
Ok(())
}
#[inline]
#[must_use]
pub fn consume_available(&mut self, requested: Q) -> Q {
let consumed = requested.min(self.remaining);
self.remaining = self.remaining - consumed;
consumed
}
#[must_use]
#[inline(always)]
pub const fn resource(&self) -> &R {
self.limit.resource()
}
#[must_use]
#[inline(always)]
pub const fn resource_limit(&self) -> &ResourceLimit<R, Q> {
&self.limit
}
#[must_use]
#[inline(always)]
pub const fn limit(&self) -> Q {
self.limit.maximum()
}
#[must_use]
#[inline(always)]
pub const fn remaining(&self) -> Q {
self.remaining
}
#[must_use]
#[inline(always)]
pub fn used(&self) -> Q {
self.limit.maximum() - self.remaining
}
}