Skip to main content

macula_rust/station_link/
admission.rs

1//! Admission of the CALLs a provider node receives, as
2//! macula_request_admission judges them: a request runs once, whichever of
3//! the node's links it arrives on. Its deadline must lie between the
4//! provider's clock minus 5 minutes and plus 10 minutes, and (caller,
5//! request_id) must be new; the entry is kept until the deadline plus 5
6//! minutes. A copy with the same request hash gets the stored reply, or
7//! request_copy while the first still runs; one with another hash is refused.
8//! The entries are bounded, and a full bound refuses rather than evicts, in
9//! this order: each caller holds at most `caller_quota` entries, each share
10//! (one link's place: the station it dialed) at most `share`, and the
11//! admission at most `cap`; stored replies take at most `reply_bytes` per
12//! caller and `reply_bytes_total` in all. A reply past either is not kept, and
13//! a copy of its request is refused reply_not_kept.
14
15use std::collections::HashMap;
16use std::hash::Hash;
17use std::sync::{Arc, Mutex, MutexGuard};
18
19use crate::frame::VerifiedRequest;
20
21use super::LinkError;
22
23const DEADLINE_PAST_TOLERANCE_MS: i64 = 5 * 60_000;
24const DEADLINE_AHEAD_MAX_MS: i64 = 10 * 60_000;
25const KEPT_PAST_DEADLINE_MS: i64 = 5 * 60_000;
26
27/// An admission's bounds. The last four bound the streaming sessions a node
28/// serves, as macula_stream_sessions does: sessions at once per caller and in
29/// all, and the bytes their inboxes hold per caller and in all.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct AdmissionLimits {
32    pub caller_quota: usize,
33    pub share: usize,
34    pub cap: usize,
35    pub reply_bytes: usize,
36    pub reply_bytes_total: usize,
37    pub sessions_per_caller: usize,
38    pub sessions: usize,
39    pub inbox_bytes_per_caller: usize,
40    pub inbox_bytes: usize,
41}
42
43impl Default for AdmissionLimits {
44    /// macula's defaults, with `cap` one share's worth: the bound of an
45    /// admission a single link holds. A pool sets `cap` to `share` times the
46    /// most links it holds.
47    fn default() -> Self {
48        AdmissionLimits {
49            caller_quota: 256,
50            share: 1024,
51            cap: 1024,
52            reply_bytes: 256 * 1024,
53            reply_bytes_total: 16 * 1024 * 1024,
54            sessions_per_caller: 16,
55            sessions: 1000,
56            inbox_bytes_per_caller: 16 * 1024 * 1024,
57            inbox_bytes: 256 * 1024 * 1024,
58        }
59    }
60}
61
62impl AdmissionLimits {
63    /// Whether macula would start with these limits: every bound positive,
64    /// and each per-caller bound within its total.
65    pub fn validate(&self) -> Result<(), LinkError> {
66        let l = self;
67        let valid = l.caller_quota > 0
68            && l.share > 0
69            && l.cap > 0
70            && l.reply_bytes > 0
71            && l.reply_bytes_total > 0
72            && l.caller_quota <= l.share
73            && l.reply_bytes <= l.reply_bytes_total
74            && l.sessions_per_caller > 0
75            && l.sessions >= l.sessions_per_caller
76            && l.inbox_bytes_per_caller > 0
77            && l.inbox_bytes >= l.inbox_bytes_per_caller;
78        if valid {
79            Ok(())
80        } else {
81            Err(LinkError::InvalidConfig(
82                "admission limits must be positive, each per-caller bound within its total".into(),
83            ))
84        }
85    }
86}
87
88/// The admission's judgement of a request: refused with a code, a copy with
89/// its stored reply (`None` while the first still runs), or new.
90#[derive(Debug, Clone, PartialEq, Eq)]
91pub(super) enum Verdict {
92    Refused(&'static str),
93    Copy(Option<Vec<u8>>),
94    New,
95}
96
97struct Entry {
98    hash: [u8; 48],
99    expires_at: i64,
100    share: String,
101    answered: bool,
102    /// `None` when answered but not kept.
103    reply: Option<Vec<u8>>,
104}
105
106impl Entry {
107    /// Judges a copy of the request this entry admitted, by its request hash:
108    /// another hash is refused, an answer not kept is refused, and otherwise
109    /// the copy gets the stored reply, `None` while the first still runs.
110    fn judge_copy(&self, request_hash: &[u8; 48]) -> Verdict {
111        if self.hash != *request_hash {
112            return Verdict::Refused("request_id_reused");
113        }
114        if self.answered && self.reply.is_none() {
115            return Verdict::Refused("reply_not_kept");
116        }
117        Verdict::Copy(self.reply.clone())
118    }
119}
120
121#[derive(Default)]
122struct Held {
123    entries: HashMap<([u8; 32], [u8; 16]), Entry>,
124    callers: HashMap<[u8; 32], usize>,
125    shares: HashMap<String, usize>,
126    reply_bytes: HashMap<[u8; 32], usize>,
127    reply_total: usize,
128    sessions: HashMap<[u8; 32], usize>,
129    sessions_total: usize,
130    inbox: HashMap<[u8; 32], usize>,
131    inbox_total: usize,
132}
133
134/// One provider node's request admission, shared by all its links.
135pub struct Admission {
136    limits: AdmissionLimits,
137    held: Mutex<Held>,
138}
139
140impl Admission {
141    /// An empty admission with `limits`, which must validate: a link given
142    /// limits that do not is refused at dial.
143    pub fn new(limits: AdmissionLimits) -> Admission {
144        Admission {
145            limits,
146            held: Mutex::new(Held::default()),
147        }
148    }
149
150    /// The admission's bounds.
151    pub fn limits(&self) -> AdmissionLimits {
152        self.limits
153    }
154
155    fn lock(&self) -> MutexGuard<'_, Held> {
156        self.held.lock().unwrap_or_else(|p| p.into_inner())
157    }
158
159    /// Judges `request` arriving on `share` at `now_ms`, sweeping the entries
160    /// that expired first.
161    pub(super) fn admit(&self, request: &VerifiedRequest, share: &str, now_ms: i64) -> Verdict {
162        let deadline = request.deadline as i64;
163        if deadline < now_ms - DEADLINE_PAST_TOLERANCE_MS {
164            return Verdict::Refused("expired");
165        }
166        if deadline > now_ms + DEADLINE_AHEAD_MAX_MS {
167            return Verdict::Refused("not_yet_valid");
168        }
169        let mut held = self.lock();
170        held.sweep(now_ms);
171        let key = (request.caller, request.request_id);
172        if let Some(entry) = held.entries.get(&key) {
173            return entry.judge_copy(&request.request_hash);
174        }
175        if held.callers.get(&request.caller).copied().unwrap_or(0) >= self.limits.caller_quota {
176            return Verdict::Refused("caller_quota");
177        }
178        if held.shares.get(share).copied().unwrap_or(0) >= self.limits.share {
179            return Verdict::Refused("share_full");
180        }
181        if held.entries.len() >= self.limits.cap {
182            return Verdict::Refused("admission_full");
183        }
184        held.entries.insert(
185            key,
186            Entry {
187                hash: request.request_hash,
188                expires_at: deadline + KEPT_PAST_DEADLINE_MS,
189                share: share.to_string(),
190                answered: false,
191                reply: None,
192            },
193        );
194        *held.callers.entry(request.caller).or_default() += 1;
195        *held.shares.entry(share.to_string()).or_default() += 1;
196        Verdict::New
197    }
198
199    /// Keeps the encoded reply of an admitted request for its copies, when
200    /// the byte bounds leave room for it.
201    pub(super) fn store(&self, request: &VerifiedRequest, reply: Vec<u8>) {
202        let mut held = self.lock();
203        let held = &mut *held;
204        let Some(entry) = held.entries.get_mut(&(request.caller, request.request_id)) else {
205            return;
206        };
207        if entry.answered || entry.hash != request.request_hash {
208            return;
209        }
210        entry.answered = true;
211        let caller_bytes = held.reply_bytes.get(&request.caller).copied().unwrap_or(0);
212        if caller_bytes + reply.len() > self.limits.reply_bytes
213            || held.reply_total + reply.len() > self.limits.reply_bytes_total
214        {
215            return;
216        }
217        *held.reply_bytes.entry(request.caller).or_default() += reply.len();
218        held.reply_total += reply.len();
219        entry.reply = Some(reply);
220    }
221
222    /// Takes a streaming session's place for `caller`, or `None` when the
223    /// per-caller or node bound is full. Dropping the place gives it back.
224    pub(super) fn open_session(self: &Arc<Self>, caller: [u8; 32]) -> Option<SessionPlace> {
225        let mut held = self.lock();
226        if held.sessions.get(&caller).copied().unwrap_or(0) >= self.limits.sessions_per_caller
227            || held.sessions_total >= self.limits.sessions
228        {
229            return None;
230        }
231        *held.sessions.entry(caller).or_default() += 1;
232        held.sessions_total += 1;
233        Some(SessionPlace {
234            admission: self.clone(),
235            caller,
236        })
237    }
238
239    /// Counts `n` more bytes that `caller`'s served streams hold unread,
240    /// refusing a charge past either bound.
241    pub(super) fn charge_inbox(&self, caller: [u8; 32], n: usize) -> bool {
242        let mut held = self.lock();
243        if held.inbox.get(&caller).copied().unwrap_or(0) + n > self.limits.inbox_bytes_per_caller
244            || held.inbox_total + n > self.limits.inbox_bytes
245        {
246            return false;
247        }
248        *held.inbox.entry(caller).or_default() += n;
249        held.inbox_total += n;
250        true
251    }
252
253    /// Gives back `n` bytes `caller`'s served streams no longer hold.
254    pub(super) fn release_inbox(&self, caller: [u8; 32], n: usize) {
255        let mut held = self.lock();
256        decrement(&mut held.inbox, caller, n);
257        held.inbox_total = held.inbox_total.saturating_sub(n);
258    }
259}
260
261/// A streaming session's place in the admission, given back when dropped.
262pub(super) struct SessionPlace {
263    admission: Arc<Admission>,
264    caller: [u8; 32],
265}
266
267impl Drop for SessionPlace {
268    fn drop(&mut self) {
269        let mut held = self.admission.lock();
270        decrement(&mut held.sessions, self.caller, 1);
271        held.sessions_total = held.sessions_total.saturating_sub(1);
272    }
273}
274
275impl Held {
276    /// Removes the entries whose deadline plus 5 minutes passed before
277    /// `now_ms`.
278    fn sweep(&mut self, now_ms: i64) {
279        let expired: Vec<_> = self
280            .entries
281            .iter()
282            .filter(|(_, e)| e.expires_at < now_ms)
283            .map(|(k, _)| *k)
284            .collect();
285        for key in expired {
286            self.remove_entry(key);
287        }
288    }
289
290    /// Removes one entry, giving back its caller's and share's places and
291    /// the bytes of its stored reply.
292    fn remove_entry(&mut self, key: ([u8; 32], [u8; 16])) {
293        let Some(entry) = self.entries.remove(&key) else {
294            return;
295        };
296        decrement(&mut self.callers, key.0, 1);
297        decrement(&mut self.shares, entry.share, 1);
298        if let Some(reply) = entry.reply {
299            decrement(&mut self.reply_bytes, key.0, reply.len());
300            self.reply_total = self.reply_total.saturating_sub(reply.len());
301        }
302    }
303}
304
305/// Takes `n` from `m[k]`, dropping the key at zero.
306fn decrement<K: Eq + Hash>(m: &mut HashMap<K, usize>, k: K, n: usize) {
307    if let Some(v) = m.get_mut(&k) {
308        *v = v.saturating_sub(n);
309        if *v == 0 {
310            m.remove(&k);
311        }
312    }
313}