cherenkov/server/
registry.rs

1//! Bounded request registration and cancellation, independent of the GPU worker.
2
3use super::failure::Failure;
4use anyhow::{Result, ensure};
5use rand::Rng;
6use serde_json::{Value, json};
7use std::{
8    collections::HashMap,
9    sync::{
10        Arc, Mutex,
11        atomic::{AtomicU8, Ordering},
12    },
13};
14
15const LIVE: u8 = 0;
16const CANCELLED: u8 = 1;
17const COMPLETED: u8 = 2;
18
19pub(super) fn id(prefix: &str) -> String {
20    format!("{prefix}-{:032x}", rand::rng().random::<u128>())
21}
22
23/// OpenAI-style identifier minted for one decoded model tool call.
24pub(super) fn tool_call_id() -> String {
25    format!("call_{:016x}", rand::rng().random::<u64>())
26}
27
28#[derive(Default)]
29pub(super) struct Registry {
30    entries: Mutex<HashMap<String, Arc<AtomicU8>>>,
31}
32
33impl Registry {
34    pub(super) fn register(
35        self: &Arc<Self>,
36        requested: Option<&str>,
37        limit: usize,
38    ) -> Result<Arc<Ticket>> {
39        let id = requested.map(str::to_owned).unwrap_or_else(|| id("req"));
40
41        ensure!(
42            !id.is_empty()
43                && id.len() <= 80
44                && id
45                    .bytes()
46                    .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_'),
47            "request_id must be 1..80 ASCII letters, digits, '-' or '_'"
48        );
49
50        let mut entries = self.entries.lock().unwrap();
51
52        ensure!(
53            !entries.contains_key(&id),
54            Failure(409, "request_id is already active")
55        );
56        ensure!(
57            entries.len() < limit,
58            Failure(503, "request capacity exhausted")
59        );
60
61        let flag = Arc::new(AtomicU8::new(LIVE));
62
63        entries.insert(id.clone(), flag.clone());
64
65        Ok(Arc::new(Ticket {
66            id,
67            flag,
68            registry: self.clone(),
69        }))
70    }
71
72    pub(super) fn cancel(&self, id: &str) -> bool {
73        let entries = self.entries.lock().unwrap();
74        let Some(flag) = entries.get(id) else {
75            return false;
76        };
77
78        matches!(
79            flag.compare_exchange(LIVE, CANCELLED, Ordering::AcqRel, Ordering::Acquire),
80            Ok(_) | Err(CANCELLED)
81        )
82    }
83
84    pub(super) fn list(&self) -> Value {
85        let entries = self.entries.lock().unwrap();
86        let mut ids: Vec<_> = entries.keys().collect();
87
88        ids.sort();
89
90        let requests: Vec<_> = ids
91            .into_iter()
92            .map(|id| {
93                json!({
94                    "id": id,
95                    "cancel_requested": entries[id].load(Ordering::Acquire) == CANCELLED,
96                })
97            })
98            .collect();
99
100        json!({"object": "list", "data": requests})
101    }
102}
103
104pub(super) struct Ticket {
105    pub id: String,
106    flag: Arc<AtomicU8>,
107    registry: Arc<Registry>,
108}
109
110impl Ticket {
111    pub(super) fn cancelled(&self) -> bool {
112        self.flag.load(Ordering::Acquire) == CANCELLED
113    }
114
115    pub(super) fn cancel(&self) {
116        let _ = self.registry.cancel(&self.id);
117    }
118
119    /// Cancellation and final publication have one atomic ordering boundary.
120    pub(super) fn complete(&self) -> bool {
121        self.flag
122            .compare_exchange(LIVE, COMPLETED, Ordering::AcqRel, Ordering::Acquire)
123            .is_ok()
124    }
125}
126
127impl Drop for Ticket {
128    fn drop(&mut self) {
129        self.registry.entries.lock().unwrap().remove(&self.id);
130    }
131}
132
133#[cfg(test)]
134#[path = "../../tests/unit/server/registry.rs"]
135mod tests;
cherenkov/server/registry.rs

cherenkov/server/
registry.rs
