lonca_daemon/
limits.rs

1//! What keeps one client from tying the daemon up. All of it sits in front
2//! of the single writer and none of it reaches `decide()`: a request refused
3//! here was never logged, so a replay does not see it.
4
5use std::collections::BTreeMap;
6use std::os::unix::net::UnixStream;
7use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
8use std::sync::{Arc, Mutex, MutexGuard};
9use std::time::{Duration, Instant};
10
11use serde::de::{Deserializer, IgnoredAny, MapAccess, Visitor};
12use serde_json::{Value, json};
13
14use crate::protocol::MAX_LINE;
15
16/// A request line of any method up to here; past it, see [`Head`].
17pub const SHORT_LINE: usize = 64 * 1024;
18/// File descriptors the daemon asks the system for.
19const WANTED_FILES: u64 = 4096;
20
21#[derive(Clone, Debug)]
22pub struct Limits {
23    /// Connections at once, every uid together.
24    pub connections: usize,
25    /// Connections at once of one uid.
26    pub per_uid: usize,
27    /// How long a connection may send nothing.
28    pub idle: Duration,
29    /// How long a line may take, from its first byte to its newline.
30    pub line: Duration,
31    /// How long a response may take to be read.
32    pub write: Duration,
33    /// A line past this must be of a method that may send more.
34    pub short_line: usize,
35    pub long_line: usize,
36    /// Bytes of long lines held in memory at once, every connection together.
37    pub buffered: usize,
38}
39
40impl Default for Limits {
41    fn default() -> Self {
42        Self {
43            connections: 1024,
44            per_uid: 512,
45            idle: Duration::from_secs(15 * 60),
46            line: Duration::from_secs(60),
47            write: Duration::from_secs(60),
48            short_line: SHORT_LINE,
49            long_line: MAX_LINE,
50            buffered: 2 * MAX_LINE,
51        }
52    }
53}
54
55impl Limits {
56    /// What a daemon with `files` descriptors ([`open_files`]) runs with: half
57    /// of them for connections, the rest for repositories and children.
58    pub fn for_files(files: u64) -> Self {
59        let d = Self::default();
60        let connections = d.connections.min(files as usize / 2);
61        Self {
62            connections,
63            per_uid: d.per_uid.min(connections),
64            ..d
65        }
66    }
67}
68
69/// Asks for [`WANTED_FILES`] descriptors, as far as the hard limit goes
70/// (a launchd daemon starts with 256), and says what it has.
71pub fn open_files() -> u64 {
72    crate::sys::raise_open_files(WANTED_FILES)
73}
74
75/// How often each limit said no. The daemon's own count: nothing refused
76/// here is in a log. `ping` shows it.
77#[derive(Debug, Default)]
78pub struct Counters {
79    /// Connections refused, or long lines with no room left.
80    pub busy: AtomicU64,
81    /// Quiet connections closed to let a new one in.
82    pub evicted: AtomicU64,
83    /// Connections closed after sending nothing for too long.
84    pub idle: AtomicU64,
85    /// Connections closed in the middle of a line or a response.
86    pub slow: AtomicU64,
87    pub too_large: AtomicU64,
88    pub rate_limited: AtomicU64,
89}
90
91impl Counters {
92    pub fn add(counter: &AtomicU64) {
93        counter.fetch_add(1, Ordering::Relaxed);
94    }
95
96    pub fn to_json(&self) -> Value {
97        let n = |c: &AtomicU64| c.load(Ordering::Relaxed);
98        json!({
99            "busy": n(&self.busy),
100            "evicted": n(&self.evicted),
101            "idle": n(&self.idle),
102            "slow": n(&self.slow),
103            "too_large": n(&self.too_large),
104            "rate_limited": n(&self.rate_limited),
105        })
106    }
107}
108
109/// A connection the server holds.
110struct Conn {
111    uid: u32,
112    stream: Arc<UnixStream>,
113    /// A request is being answered: the only state that can not be cut.
114    working: bool,
115    quiet_since: Instant,
116}
117
118/// The connections alive, and the bytes of long lines they hold.
119pub struct Registry {
120    pub limits: Limits,
121    conns: Mutex<BTreeMap<u64, Conn>>,
122    next: AtomicU64,
123    buffered: AtomicUsize,
124}
125
126impl Registry {
127    pub fn new(limits: Limits) -> Arc<Self> {
128        Arc::new(Self {
129            limits,
130            conns: Mutex::default(),
131            next: AtomicU64::new(0),
132            buffered: AtomicUsize::new(0),
133        })
134    }
135
136    fn lock(&self) -> MutexGuard<'_, BTreeMap<u64, Conn>> {
137        self.conns.lock().unwrap_or_else(|e| e.into_inner())
138    }
139
140    /// A place for a new connection. When there is none the connection that
141    /// has been quiet longest gives up its own, so connections that do
142    /// nothing can not keep a client out; `None` when all of them are at
143    /// work. The second value says whether one was evicted.
144    pub fn admit(self: &Arc<Self>, uid: u32, stream: Arc<UnixStream>) -> Option<(Slot, bool)> {
145        let mut conns = self.lock();
146        let of_uid = conns.values().filter(|c| c.uid == uid).count();
147        let crowded = of_uid >= self.limits.per_uid;
148        let full = crowded || conns.len() >= self.limits.connections;
149        if full {
150            let quiet = conns.iter().filter(|(_, c)| !c.working);
151            let fair = quiet.filter(|(_, c)| !crowded || c.uid == uid);
152            let (&id, _) = fair.min_by_key(|(_, c)| c.quiet_since)?;
153            let gone = conns.remove(&id)?;
154            // Its thread wakes to end of stream and says why it is closing.
155            let _ = gone.stream.shutdown(std::net::Shutdown::Read);
156        }
157        let id = self.next.fetch_add(1, Ordering::Relaxed);
158        let conn = Conn {
159            uid,
160            stream,
161            working: false,
162            quiet_since: Instant::now(),
163        };
164        conns.insert(id, conn);
165        let slot = Slot {
166            registry: self.clone(),
167            id,
168        };
169        Some((slot, full))
170    }
171}
172
173/// A connection's place in the [`Registry`], given up when it is dropped.
174pub struct Slot {
175    registry: Arc<Registry>,
176    id: u64,
177}
178
179impl Slot {
180    pub fn limits(&self) -> &Limits {
181        &self.registry.limits
182    }
183
184    /// The uid of the connection's peer (`u32::MAX` when the socket did not
185    /// tell), for the security log.
186    pub fn uid(&self) -> Option<u32> {
187        self.registry.lock().get(&self.id).map(|c| c.uid)
188    }
189
190    /// Says a request is being answered. False if the connection was evicted
191    /// first: the request must not run, whoever sent it is told to send it
192    /// again.
193    pub fn begin(&self) -> bool {
194        let mut conns = self.registry.lock();
195        let conn = conns.get_mut(&self.id);
196        conn.map(|c| c.working = true).is_some()
197    }
198
199    /// The request was answered.
200    pub fn end(&self) {
201        if let Some(c) = self.registry.lock().get_mut(&self.id) {
202            c.working = false;
203            c.quiet_since = Instant::now();
204        }
205    }
206
207    pub fn evicted(&self) -> bool {
208        !self.registry.lock().contains_key(&self.id)
209    }
210
211    pub fn hold(&self) -> Held<'_> {
212        Held {
213            registry: &self.registry,
214            bytes: 0,
215        }
216    }
217}
218
219impl Drop for Slot {
220    fn drop(&mut self) {
221        self.registry.lock().remove(&self.id);
222    }
223}
224
225/// The bytes of one long line, counted against [`Limits::buffered`] until
226/// the request is done with.
227pub struct Held<'a> {
228    registry: &'a Registry,
229    bytes: usize,
230}
231
232impl Held<'_> {
233    /// Grows what is held to `total` bytes, if there is room.
234    pub fn grow(&mut self, total: usize) -> bool {
235        let more = total.saturating_sub(self.bytes);
236        let was = self.registry.buffered.fetch_add(more, Ordering::Relaxed);
237        if was + more > self.registry.limits.buffered {
238            self.registry.buffered.fetch_sub(more, Ordering::Relaxed);
239            return false;
240        }
241        self.bytes += more;
242        true
243    }
244}
245
246impl Drop for Held<'_> {
247    fn drop(&mut self) {
248        let bytes = self.bytes;
249        self.registry.buffered.fetch_sub(bytes, Ordering::Relaxed);
250    }
251}
252
253/// What the first [`SHORT_LINE`] bytes of a line say about it. A line that
254/// goes on past them must have named its method and token by then, so a
255/// long line is only held for a method that may send one and a caller the
256/// daemon knows.
257#[derive(Debug, Default, PartialEq, Eq)]
258pub struct Head {
259    pub id: u64,
260    pub method: Option<String>,
261    pub repo: Option<String>,
262    pub token: Option<String>,
263}
264
265impl<'de> Visitor<'de> for &mut Head {
266    type Value = ();
267
268    fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269        f.write_str("a request")
270    }
271
272    fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<(), A::Error> {
273        while let Some(key) = map.next_key::<String>()? {
274            match key.as_str() {
275                "id" => self.id = map.next_value()?,
276                "method" => self.method = map.next_value()?,
277                "repo" => self.repo = map.next_value()?,
278                "token" => self.token = map.next_value()?,
279                _ => {
280                    map.next_value::<IgnoredAny>()?;
281                }
282            }
283        }
284        Ok(())
285    }
286}
287
288/// Reads as far as the bytes go: the line is cut off, so the parse ends in
289/// an error and what came before it stands.
290pub fn head_of(bytes: &[u8]) -> Head {
291    let mut head = Head::default();
292    let mut de = serde_json::Deserializer::from_slice(bytes);
293    let _ = (&mut de).deserialize_map(&mut head);
294    head
295}
296
297/// One operation, in the units of a [`Bucket`]: the level grows by the rate
298/// in operations a minute every millisecond.
299const OP: u64 = 60_000;
300
301/// A token bucket holding a minute's worth of operations.
302#[derive(Clone, Copy, Debug)]
303pub struct Bucket {
304    level: u64,
305    at_ms: i64,
306}
307
308impl Bucket {
309    pub fn full(per_minute: u64, now_ms: i64) -> Self {
310        Self {
311            level: per_minute * OP,
312            at_ms: now_ms,
313        }
314    }
315
316    /// Takes one operation; `Err` is how many milliseconds until there is one.
317    pub fn take(&mut self, per_minute: u64, now_ms: i64) -> Result<(), u64> {
318        let passed = now_ms.saturating_sub(self.at_ms).max(0) as u64;
319        let grown = self.level.saturating_add(passed.saturating_mul(per_minute));
320        self.level = grown.min(per_minute * OP);
321        self.at_ms = now_ms;
322        match self.level.checked_sub(OP) {
323            Some(left) => {
324                self.level = left;
325                Ok(())
326            }
327            None => Err((OP - self.level).div_ceil(per_minute)),
328        }
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335
336    fn stream() -> Arc<UnixStream> {
337        Arc::new(UnixStream::pair().unwrap().0)
338    }
339
340    fn limits(connections: usize, per_uid: usize) -> Limits {
341        Limits {
342            connections,
343            per_uid,
344            ..Limits::default()
345        }
346    }
347
348    #[test]
349    fn file_descriptors_decide_the_connections() {
350        let few = Limits::for_files(256);
351        assert_eq!((few.connections, few.per_uid), (128, 128));
352        let some = Limits::for_files(1600);
353        assert_eq!((some.connections, some.per_uid), (800, 512));
354        let many = Limits::for_files(1 << 20);
355        assert_eq!((many.connections, many.per_uid), (1024, 512));
356    }
357
358    #[test]
359    fn the_daemon_asks_for_descriptors() {
360        let files = open_files();
361        assert!(files >= 256, "{files}");
362    }
363
364    #[test]
365    fn the_quietest_connection_makes_room() {
366        let r = Registry::new(limits(2, 2));
367        let a = stream();
368        let (first, cut) = r.admit(1, a.clone()).unwrap();
369        assert!(!cut);
370        let (second, cut) = r.admit(1, stream()).unwrap();
371        assert!(!cut);
372        std::thread::sleep(Duration::from_millis(2));
373        first.begin();
374        first.end(); // the first has been quiet for less time now
375        let (_third, cut) = r.admit(1, stream()).unwrap();
376        assert!(cut);
377        assert!(!first.evicted());
378        assert!(second.evicted());
379        assert!(!second.begin(), "evicted: its request must not run");
380    }
381
382    #[test]
383    fn a_connection_at_work_is_never_cut() {
384        let r = Registry::new(limits(1, 1));
385        let (slot, _) = r.admit(1, stream()).unwrap();
386        assert!(slot.begin());
387        assert!(r.admit(1, stream()).is_none());
388        slot.end();
389        assert!(r.admit(1, stream()).is_some());
390    }
391
392    #[test]
393    fn a_uid_at_its_cap_gives_up_its_own() {
394        let r = Registry::new(limits(8, 1));
395        let (other, _) = r.admit(2, stream()).unwrap();
396        std::thread::sleep(Duration::from_millis(2));
397        let (mine, cut) = r.admit(1, stream()).unwrap();
398        assert!(!cut);
399        std::thread::sleep(Duration::from_millis(2));
400        let (_next, cut) = r.admit(1, stream()).unwrap();
401        assert!(cut);
402        assert!(mine.evicted());
403        assert!(!other.evicted(), "older, but another uid's");
404    }
405
406    #[test]
407    fn a_full_daemon_takes_the_quietest_of_any_uid() {
408        let r = Registry::new(limits(1, 8));
409        let (other, _) = r.admit(2, stream()).unwrap();
410        let (_mine, cut) = r.admit(1, stream()).unwrap();
411        assert!(cut);
412        assert!(other.evicted());
413    }
414
415    #[test]
416    fn eviction_ends_the_reads_of_the_connection() {
417        use std::io::Read;
418        let (a, _keep) = UnixStream::pair().unwrap();
419        let a = Arc::new(a);
420        let r = Registry::new(limits(1, 1));
421        let (_slot, _) = r.admit(1, a.clone()).unwrap();
422        let _next = r.admit(1, stream()).unwrap();
423        assert_eq!((&*a).read(&mut [0; 8]).unwrap(), 0);
424    }
425
426    #[test]
427    fn a_slot_gives_its_place_back() {
428        let r = Registry::new(limits(1, 1));
429        let (slot, _) = r.admit(1, stream()).unwrap();
430        slot.begin();
431        drop(slot);
432        let (_slot, cut) = r.admit(1, stream()).unwrap();
433        assert!(!cut);
434    }
435
436    #[test]
437    fn long_lines_share_a_budget() {
438        let r = Registry::new(Limits {
439            buffered: 100,
440            ..Limits::default()
441        });
442        let (slot, _) = r.admit(1, stream()).unwrap();
443        let mut a = slot.hold();
444        assert!(a.grow(60));
445        assert!(a.grow(100), "its own bytes are not counted twice");
446        let mut b = slot.hold();
447        assert!(!b.grow(1));
448        drop(a);
449        assert!(b.grow(100), "a line done with gives its room back");
450        assert!(!b.grow(101));
451        assert!(b.grow(100));
452    }
453
454    #[test]
455    fn a_head_reads_what_comes_before_the_cut() {
456        let line = br#"{"id":7,"method":"tree","repo":"r","token":"t","params":{"paths":["a","#;
457        let want = Head {
458            id: 7,
459            method: Some("tree".into()),
460            repo: Some("r".into()),
461            token: Some("t".into()),
462        };
463        assert_eq!(head_of(line), want);
464        // Sorted keys, as a generic JSON map writes them: no token yet.
465        let sorted = br#"{"id":1,"method":"blobs","params":{"paths":["a","#;
466        let head = head_of(sorted);
467        assert_eq!(head.method.as_deref(), Some("blobs"));
468        assert_eq!(head.token, None);
469        assert_eq!(head_of(b"nonsense"), Head::default());
470        assert_eq!(head_of(br#"{"params":{"x":"#), Head::default());
471    }
472
473    #[test]
474    fn what_a_head_expects_is_a_request() {
475        let mut head = Head::default();
476        let mut de = serde_json::Deserializer::from_slice(b"1");
477        let e = (&mut de).deserialize_map(&mut head).unwrap_err();
478        assert!(e.to_string().contains("expected a request"), "{e}");
479    }
480
481    #[test]
482    fn a_bucket_lets_a_minute_through_and_then_the_rate() {
483        let mut b = Bucket::full(60, 0);
484        for _ in 0..60 {
485            assert_eq!(b.take(60, 0), Ok(()));
486        }
487        assert_eq!(b.take(60, 0), Err(1000));
488        assert_eq!(b.take(60, 400), Err(600));
489        assert_eq!(b.take(60, 1000), Ok(()));
490        assert_eq!(b.take(60, 1000), Err(1000));
491    }
492
493    #[test]
494    fn a_bucket_holds_no_more_than_a_minute() {
495        let mut b = Bucket::full(2, 0);
496        assert_eq!(b.take(2, 10_000_000), Ok(()));
497        assert_eq!(b.take(2, 10_000_000), Ok(()));
498        assert!(b.take(2, 10_000_000).is_err());
499    }
500
501    #[test]
502    fn a_clock_that_goes_back_adds_nothing() {
503        let mut b = Bucket::full(1, 1_000_000);
504        assert_eq!(b.take(1, 1_000_000), Ok(()));
505        assert_eq!(b.take(1, 0), Err(60_000));
506    }
507
508    #[test]
509    fn a_lower_rate_cuts_what_was_saved() {
510        let mut b = Bucket::full(600, 0);
511        assert_eq!(b.take(1, 0), Ok(()));
512        assert_eq!(b.take(1, 0), Err(60_000));
513    }
514
515    #[test]
516    fn counters_show_each_refusal() {
517        let c = Counters::default();
518        Counters::add(&c.busy);
519        Counters::add(&c.too_large);
520        Counters::add(&c.too_large);
521        let shown = c.to_json();
522        assert_eq!(shown["busy"], 1);
523        assert_eq!(shown["too_large"], 2);
524        assert_eq!(shown["rate_limited"], 0);
525    }
526}
lonca_daemon/limits.rs

lonca_daemon/
limits.rs
