gwr_components/
capacity_allocator.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! Capacity accounting and scoped reservations.
4
5use std::cell::RefCell;
6use std::rc::Rc;
7
8use gwr_engine::events::repeated::Repeated;
9use gwr_engine::sim_error;
10use gwr_engine::traits::Event;
11use gwr_engine::types::{SimError, SimResult};
12use gwr_model_builder::{EntityDisplay, EntityGet};
13use gwr_track::entity::Entity;
14
15/// A reusable capacity allocator for object counts, bytes, or other units.
16#[derive(Clone, EntityGet, EntityDisplay)]
17pub struct CapacityAllocator {
18    entity: Rc<Entity>,
19    capacity: usize,
20    capacity_unit: Rc<String>,
21    used: Rc<RefCell<usize>>,
22    level_change: Repeated<usize>,
23}
24
25#[must_use = "the reservation is released when dropped"]
26pub struct CapacityReservation {
27    allocator: CapacityAllocator,
28    units: usize,
29}
30
31impl Drop for CapacityReservation {
32    fn drop(&mut self) {
33        self.allocator.release(self.units);
34    }
35}
36
37impl CapacityAllocator {
38    /// Create a standalone allocator with its own child entity.
39    pub fn new(
40        parent: &Rc<Entity>,
41        name: &str,
42        capacity: usize,
43        capacity_unit: impl Into<String>,
44    ) -> Result<Self, SimError> {
45        let entity = Rc::new(Entity::new(parent, name));
46        Self::for_entity(&entity, capacity, capacity_unit)
47    }
48
49    /// Create allocator bookkeeping on an existing entity.
50    pub fn for_entity(
51        entity: &Rc<Entity>,
52        capacity: usize,
53        capacity_unit: impl Into<String>,
54    ) -> Result<Self, SimError> {
55        if capacity == 0 {
56            return sim_error!("Unsupported CapacityAllocator with capacity of 0");
57        }
58        let capacity_unit = capacity_unit.into();
59        entity.track_capacity(capacity, &capacity_unit);
60
61        Ok(Self {
62            entity: entity.clone(),
63            capacity,
64            capacity_unit: Rc::new(capacity_unit),
65            used: Rc::new(RefCell::new(0)),
66            level_change: Repeated::new(usize::default()),
67        })
68    }
69
70    #[must_use]
71    pub fn used(&self) -> usize {
72        *self.used.borrow()
73    }
74
75    #[must_use]
76    fn has_capacity_for(&self, units: usize) -> bool {
77        units <= self.capacity - self.used()
78    }
79
80    fn validate_request_size(&self, units: usize) -> SimResult {
81        if units > self.capacity {
82            return sim_error!(
83                "Cannot allocate {units} {} in {:?} with capacity {}",
84                self.capacity_unit,
85                self.entity.full_name(),
86                self.capacity
87            );
88        }
89        Ok(())
90    }
91
92    #[must_use]
93    pub fn capacity(&self) -> usize {
94        self.capacity
95    }
96
97    #[must_use]
98    pub fn capacity_unit(&self) -> &str {
99        self.capacity_unit.as_str()
100    }
101
102    #[must_use]
103    pub fn level_change_event(&self) -> Repeated<usize> {
104        self.level_change.clone()
105    }
106
107    pub async fn wait_for_capacity(&self, units: usize) -> SimResult {
108        self.validate_request_size(units)?;
109        let level_change = self.level_change_event();
110        while !self.has_capacity_for(units) {
111            level_change.listen().await;
112        }
113        Ok(())
114    }
115
116    pub fn allocate(&self, units: usize) -> SimResult {
117        self.validate_request_size(units)?;
118        if !self.has_capacity_for(units) {
119            return sim_error!("Overflow in {:?}", self.entity.full_name());
120        }
121
122        let used = {
123            let mut used = self.used.borrow_mut();
124            *used += units;
125            *used
126        };
127        self.level_change.notify_result(used);
128        Ok(())
129    }
130
131    pub fn release(&self, units: usize) {
132        let used = {
133            let mut used = self.used.borrow_mut();
134            *used = used
135                .checked_sub(units)
136                .expect("capacity allocator underflow");
137            *used
138        };
139        self.level_change.notify_result(used);
140    }
141
142    pub async fn reserve(&self, units: usize) -> Result<CapacityReservation, SimError> {
143        self.wait_for_capacity(units).await?;
144        self.allocate(units)?;
145        Ok(CapacityReservation {
146            allocator: self.clone(),
147            units,
148        })
149    }
150}
gwr_components/capacity_allocator.rs

gwr_components/
capacity_allocator.rs
