use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio::time::Instant;
pub const MAX_SEAM_LIMIT: u32 = 65_536;
pub const QUEUE_FULL: &str = "remote_seam_queue_full";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SeamCounters {
pub depth: u64,
pub refused: u64,
pub late_responses_discarded: u64,
}
#[derive(Default)]
struct Counters {
refused: AtomicU64,
late: AtomicU64,
}
#[derive(Clone)]
pub struct SeamQueue {
active: Arc<Semaphore>,
slots: Arc<Semaphore>,
total: usize,
counters: Arc<Counters>,
}
impl std::fmt::Debug for SeamQueue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SeamQueue")
.field("counters", &self.counters())
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QueueFull;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QueueDeadline;
pub struct Slot {
slot: OwnedSemaphorePermit,
active: Arc<Semaphore>,
}
pub struct Running {
_slot: OwnedSemaphorePermit,
_active: OwnedSemaphorePermit,
}
impl SeamQueue {
pub fn new(in_flight: u32, capacity: u32) -> Result<Self, String> {
for (name, value) in [("in-flight limit", in_flight), ("queue capacity", capacity)] {
if !(1..=MAX_SEAM_LIMIT).contains(&value) {
return Err(format!("{name} must be 1 to {MAX_SEAM_LIMIT}"));
}
}
let total = (in_flight + capacity) as usize;
Ok(SeamQueue {
active: Arc::new(Semaphore::new(in_flight as usize)),
slots: Arc::new(Semaphore::new(total)),
total,
counters: Arc::default(),
})
}
pub fn try_enter(&self) -> Result<Slot, QueueFull> {
match self.slots.clone().try_acquire_owned() {
Ok(slot) => Ok(Slot {
slot,
active: self.active.clone(),
}),
Err(_) => {
self.counters.refused.fetch_add(1, Ordering::Relaxed);
Err(QueueFull)
}
}
}
pub fn count_late(&self) {
self.counters.late.fetch_add(1, Ordering::Relaxed);
}
pub fn counters(&self) -> SeamCounters {
SeamCounters {
depth: (self.total - self.slots.available_permits()) as u64,
refused: self.counters.refused.load(Ordering::Relaxed),
late_responses_discarded: self.counters.late.load(Ordering::Relaxed),
}
}
}
impl Slot {
pub async fn start(self, deadline: Instant) -> Result<Running, QueueDeadline> {
match tokio::time::timeout_at(deadline, self.active.clone().acquire_owned()).await {
Ok(Ok(active)) => Ok(Running {
_slot: self.slot,
_active: active,
}),
_ => Err(QueueDeadline),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[test]
fn limits_are_declared_and_bounded() {
assert!(SeamQueue::new(0, 1).is_err());
assert!(SeamQueue::new(1, 0).is_err());
assert!(SeamQueue::new(MAX_SEAM_LIMIT + 1, 1).is_err());
assert!(SeamQueue::new(1, MAX_SEAM_LIMIT + 1).is_err());
assert!(SeamQueue::new(MAX_SEAM_LIMIT, MAX_SEAM_LIMIT).is_ok());
}
#[tokio::test(start_paused = true)]
async fn a_full_queue_refuses_and_a_queued_call_honours_its_deadline() {
let q = SeamQueue::new(1, 1).unwrap();
let running = q.try_enter().unwrap();
let running = running.start(Instant::now()).await.unwrap();
let queued = q.try_enter().unwrap();
assert_eq!(q.try_enter().err(), Some(QueueFull));
assert_eq!(
q.counters(),
SeamCounters {
depth: 2,
refused: 1,
late_responses_discarded: 0
}
);
let deadline = Instant::now() + Duration::from_millis(5);
assert_eq!(queued.start(deadline).await.err(), Some(QueueDeadline));
assert_eq!(q.counters().depth, 1);
drop(running);
assert_eq!(q.counters().depth, 0);
q.count_late();
assert_eq!(q.counters().late_responses_discarded, 1);
}
}