use std::{collections::VecDeque, time::Instant};
#[derive(Debug)]
pub struct SlotPool {
slot_bytes: u64,
free: Vec<usize>,
quarantine: VecDeque<(Instant, usize)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Oversized {
pub bytes: u64,
pub slot_bytes: u64,
}
impl SlotPool {
pub fn new(slots: usize, slot_bytes: u64) -> Self {
Self {
slot_bytes,
free: (0..slots).rev().collect(),
quarantine: VecDeque::new(),
}
}
pub fn slot_bytes(&self) -> u64 {
self.slot_bytes
}
pub fn lease(&mut self, bytes: u64, now: Instant) -> Result<Option<usize>, Oversized> {
if bytes > self.slot_bytes {
return Err(Oversized {
bytes,
slot_bytes: self.slot_bytes,
});
}
while self
.quarantine
.front()
.is_some_and(|&(until, _)| until <= now)
{
let (_, slot) = self.quarantine.pop_front().expect("checked");
self.free.push(slot);
}
Ok(self.free.pop())
}
pub fn release(&mut self, slot: usize) {
self.free.push(slot);
}
pub fn quarantine(&mut self, slot: usize, until: Instant) {
self.quarantine.push_back((until, slot));
}
pub fn free(&self) -> usize {
self.free.len()
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
#[test]
fn slots_are_leased_released_and_quarantined() {
let now = Instant::now();
let mut pool = SlotPool::new(2, 100);
assert_eq!(
pool.lease(101, now),
Err(Oversized {
bytes: 101,
slot_bytes: 100
})
);
assert_eq!(pool.lease(100, now), Ok(Some(0)));
assert_eq!(pool.lease(1, now), Ok(Some(1)));
assert_eq!(pool.lease(1, now), Ok(None));
pool.release(1);
assert_eq!(pool.lease(1, now), Ok(Some(1)));
pool.quarantine(0, now + Duration::from_secs(10));
assert_eq!(pool.lease(1, now + Duration::from_secs(9)), Ok(None));
assert_eq!(pool.lease(1, now + Duration::from_secs(10)), Ok(Some(0)));
}
}