use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Resource {
InputBytes,
BundleEntries,
GraphElements,
MetadataBytes,
Depth,
Steps,
OutputBytes,
}
impl Resource {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::InputBytes => "input bytes",
Self::BundleEntries => "bundle entries",
Self::GraphElements => "graph elements",
Self::MetadataBytes => "metadata bytes",
Self::Depth => "depth",
Self::Steps => "steps",
Self::OutputBytes => "output bytes",
}
}
}
impl std::fmt::Display for Resource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
#[error("{resource} limit exceeded: the configured bound is {limit}")]
pub struct LimitExceeded {
pub resource: Resource,
pub limit: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResourceLimits {
pub input_bytes: u64,
pub bundle_entries: u64,
pub graph_elements: u64,
pub metadata_bytes: u64,
pub depth: u64,
pub steps: u64,
pub output_bytes: u64,
}
impl ResourceLimits {
#[must_use]
pub const fn defaults() -> Self {
Self {
input_bytes: 64 * 1024 * 1024,
bundle_entries: 4_096,
graph_elements: 1_000_000,
metadata_bytes: 16 * 1024 * 1024,
depth: 128,
steps: 10_000_000,
output_bytes: 64 * 1024 * 1024,
}
}
#[must_use]
pub const fn unbounded() -> Self {
Self {
input_bytes: 0,
bundle_entries: 0,
graph_elements: 0,
metadata_bytes: 0,
depth: 0,
steps: 0,
output_bytes: 0,
}
}
#[must_use]
pub const fn get(&self, resource: Resource) -> u64 {
match resource {
Resource::InputBytes => self.input_bytes,
Resource::BundleEntries => self.bundle_entries,
Resource::GraphElements => self.graph_elements,
Resource::MetadataBytes => self.metadata_bytes,
Resource::Depth => self.depth,
Resource::Steps => self.steps,
Resource::OutputBytes => self.output_bytes,
}
}
#[must_use]
pub fn budget(self) -> Budget {
Budget::new(self)
}
}
impl Default for ResourceLimits {
fn default() -> Self {
Self::defaults()
}
}
#[derive(Clone, Debug)]
pub struct Budget {
limits: ResourceLimits,
consumed: Arc<Consumed>,
}
#[derive(Debug, Default)]
struct Consumed {
input_bytes: AtomicU64,
bundle_entries: AtomicU64,
graph_elements: AtomicU64,
metadata_bytes: AtomicU64,
steps: AtomicU64,
output_bytes: AtomicU64,
}
impl Budget {
#[must_use]
pub fn new(limits: ResourceLimits) -> Self {
Self {
limits,
consumed: Arc::new(Consumed::default()),
}
}
#[must_use]
pub fn with_defaults() -> Self {
Self::new(ResourceLimits::defaults())
}
#[must_use]
pub const fn limits(&self) -> &ResourceLimits {
&self.limits
}
pub fn charge(&self, resource: Resource, amount: u64) -> Result<(), LimitExceeded> {
let limit = self.limits.get(resource);
let Some(counter) = self.counter(resource) else {
return self.check_depth(amount);
};
let total = counter.fetch_add(amount, Ordering::Relaxed) + amount;
if limit != 0 && total > limit {
return Err(LimitExceeded { resource, limit });
}
Ok(())
}
pub const fn enter(&self, level: u64) -> Result<(), LimitExceeded> {
self.check_depth(level)
}
const fn check_depth(&self, level: u64) -> Result<(), LimitExceeded> {
let limit = self.limits.depth;
if limit != 0 && level > limit {
return Err(LimitExceeded {
resource: Resource::Depth,
limit,
});
}
Ok(())
}
#[must_use]
pub fn consumed(&self, resource: Resource) -> u64 {
self.counter(resource)
.map_or(0, |c| c.load(Ordering::Relaxed))
}
fn counter(&self, resource: Resource) -> Option<&AtomicU64> {
match resource {
Resource::InputBytes => Some(&self.consumed.input_bytes),
Resource::BundleEntries => Some(&self.consumed.bundle_entries),
Resource::GraphElements => Some(&self.consumed.graph_elements),
Resource::MetadataBytes => Some(&self.consumed.metadata_bytes),
Resource::Steps => Some(&self.consumed.steps),
Resource::OutputBytes => Some(&self.consumed.output_bytes),
Resource::Depth => None,
}
}
}
impl Default for Budget {
fn default() -> Self {
Self::with_defaults()
}
}