boxmaker/crates/loopd/src/llama/gate.rs

//! One request in flight per server slot at a time, in `loopd`.
//!
//! The inference server is shared, so two requests on the same slot would pile up on it and queue
//! invisibly. The gate bounds that queue here, where the caller is told it is waiting. The gate is
//! held for one request: from before it is sent until `chat` returns.
//!
//! A slot is free only when nobody holds it and nobody is waiting for it; a free slot is taken at
//! once. Otherwise the acquirer joins a first-in-first-out queue, is told how many are ahead, and
//! waits. The queue is bounded: past `max_queue` waiters the slot is reported full at once.

use std::collections::{HashMap, VecDeque};
use std::sync::{Condvar, Mutex};

/// Per-slot state, behind the mutex.
#[derive(Default)]
struct State {
    slots: HashMap<u32, Slot>,
}

#[derive(Default)]
struct Slot {
    /// Held by exactly one permit at a time.
    holder: bool,
    /// Tickets of the waiters, in the order they arrived.
    queue: VecDeque<u32>,
    /// A fresh ticket for each waiter that joins this slot's queue.
    tickets: u32,
}

impl State {
    fn slot_mut(&mut self, slot: u32) -> &mut Slot {
        self.slots.entry(slot).or_default()
    }
}

#[derive(Default)]
pub struct SlotGate {
    state: Mutex<State>,
    cond: Condvar,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GateFull;

/// The right to have a request in flight on the slot. Dropping it passes the slot on.
pub struct Permit<'a> {
    gate: &'a SlotGate,
    slot: u32,
}

impl SlotGate {
    pub fn new() -> Self {
        Self {
            state: Mutex::new(State::default()),
            cond: Condvar::new(),
        }
    }

    pub fn acquire(
        &self,
        slot: u32,
        max_queue: usize,
        on_queued: &mut dyn FnMut(usize),
    ) -> Result<Permit<'_>, GateFull> {
        // Rule 1: a free slot with nobody waiting is taken at once.
        {
            let mut state = self.lock();
            let slot_ref = state.slot_mut(slot);
            if !slot_ref.holder && slot_ref.queue.is_empty() {
                slot_ref.holder = true;
                return Ok(Permit { gate: self, slot });
            }
        }

        // Rules 2 and 3: the slot is taken or someone is waiting. Join the queue, bounded.
        let (ticket, ahead): (u32, usize);
        {
            let mut state = self.lock();
            let slot_ref = state.slot_mut(slot);
            if slot_ref.queue.len() >= max_queue {
                return Err(GateFull);
            }
            slot_ref.tickets = slot_ref.tickets.wrapping_add(1);
            ticket = slot_ref.tickets;
            ahead = usize::from(slot_ref.holder) + slot_ref.queue.len();
            slot_ref.queue.push_back(ticket);
        }
        on_queued(ahead);

        // Rule 4: wait, in arrival order, for the slot to reach the front of the queue.
        loop {
            let mut state = self.lock();
            if self.take_if_front(&mut state, slot, ticket) {
                return Ok(Permit { gate: self, slot });
            }
            state = self
                .cond
                .wait(state)
                .unwrap_or_else(|poisoned| poisoned.into_inner());
        }
    }
}

impl Drop for Permit<'_> {
    fn drop(&mut self) {
        let mut state = self.gate.lock();
        state.slot_mut(self.slot).holder = false;
        // Wake every waiter; only the front one with a free slot takes it (rule 4).
        self.gate.cond.notify_all();
    }
}

impl SlotGate {
    fn lock(&self) -> std::sync::MutexGuard<'_, State> {
        self.state
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    fn take_if_front(
        &self,
        state: &mut std::sync::MutexGuard<'_, State>,
        slot: u32,
        ticket: u32,
    ) -> bool {
        let slot_ref = state.slot_mut(slot);
        if slot_ref.holder {
            return false;
        }
        match slot_ref.queue.front() {
            Some(&front) if front == ticket => {
                slot_ref.queue.pop_front();
                slot_ref.holder = true;
                true
            }
            _ => false,
        }
    }
}