1//! The memory the daemon has already promised to commands it let through.
2//!
3//! One free-memory reading cannot stop several agents from starting several
4//! builds at once: each is judged in the same second, each sees the same
5//! free figure, and each is told there is room. A command's memory only
6//! shows up in that figure once the command has grown into it, which for a
7//! build is tens of seconds later.
8//!
9//! So the daemon keeps a booking for every command it has judged: the
10//! megabytes Jev's load answer says it needs. The next command is judged
11//! against what is free LESS what is booked. A booking ends when its command
12//! does. Until then it counts in full while the command is young, and less as
13//! it ages, because by then the command's real use is in the free figure
14//! itself and counting both would book the same memory twice.
15//!
16//! Every session on the machine talks to this one daemon, so the ledger is
17//! the machine's, not a session's.
18
19use std::collections::HashMap;
20use std::time::{Duration, Instant};
21
22/// A booking counts in full for this long after its command was let through.
23pub const COUNTED_IN_FULL_FOR: Duration = Duration::from_secs(20);
24
25/// Past this age a booking counts for nothing: the command has taken what it
26/// is going to take, and the free figure shows it. Between the two, the
27/// share falls in a straight line.
28pub const COUNTED_UNTIL: Duration = Duration::from_secs(90);
29
30struct Booking {
31    session: String,
32    mb: f64,
33    at: Instant,
34}
35
36/// The share of a booking of this age that still counts.
37fn weight(age: Duration) -> f64 {
38    if age <= COUNTED_IN_FULL_FOR {
39        1.0
40    } else if age >= COUNTED_UNTIL {
41        0.0
42    } else {
43        let fading = (COUNTED_UNTIL - COUNTED_IN_FULL_FOR).as_secs_f64();
44        1.0 - (age - COUNTED_IN_FULL_FOR).as_secs_f64() / fading
45    }
46}
47
48#[derive(Default)]
49pub struct Ledger {
50    /// By tool call id.
51    bookings: HashMap<String, Booking>,
52}
53
54impl Ledger {
55    /// Books `mb` megabytes for the tool call `id`. Booking the same call
56    /// again replaces its entry, so a call is never counted twice.
57    pub fn book(&mut self, id: &str, session: &str, mb: f64, now: Instant) {
58        self.bookings.insert(id.to_owned(), Booking { session: session.to_owned(), mb, at: now });
59    }
60
61    /// Ends the booking of a call that has finished, failed or never ran.
62    pub fn release(&mut self, id: &str) {
63        self.bookings.remove(id);
64    }
65
66    /// Ends every booking of a session that has closed.
67    pub fn release_session(&mut self, session: &str) {
68        self.bookings.retain(|_, booking| booking.session != session);
69    }
70
71    /// The megabytes booked right now, each booking weighted by its age.
72    pub fn booked_mb(&self, now: Instant) -> f64 {
73        self.bookings.values().map(|booking| booking.mb * weight(now.saturating_duration_since(booking.at))).sum()
74    }
75
76    /// How many commands hold a booking.
77    pub fn running(&self) -> usize {
78        self.bookings.len()
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    /// The case the ledger exists for.
87    #[test]
88    fn five_builds_in_one_second_are_all_counted() {
89        // The case this exists for: 8 GB free, five agents each start a
90        // 3 GB build at the same moment. On the free figure alone all five
91        // fit; with bookings, the third already does not.
92        let now = Instant::now();
93        let mut ledger = Ledger::default();
94        let free = 8000.0;
95        let mut let_through = 0;
96        for call in 0..5 {
97            if free - ledger.booked_mb(now) >= 3000.0 {
98                ledger.book(&format!("call-{call}"), "session", 3000.0, now);
99                let_through += 1;
100            }
101        }
102        assert_eq!(let_through, 2);
103        assert_eq!(ledger.booked_mb(now), 6000.0);
104    }
105
106    /// A booking ends with its command, or with its session, so a crashed session cannot hold
107    /// memory for ever.
108    #[test]
109    fn a_finished_command_gives_its_booking_back() {
110        let now = Instant::now();
111        let mut ledger = Ledger::default();
112        ledger.book("a", "one", 3000.0, now);
113        ledger.book("b", "two", 500.0, now);
114        ledger.release("a");
115        assert_eq!(ledger.booked_mb(now), 500.0);
116        ledger.release_session("two");
117        assert_eq!(ledger.running(), 0);
118    }
119
120    /// A long-running command must not be counted twice: once here, and once in the machine's own
121    /// figure after it has taken its memory.
122    #[test]
123    fn a_booking_fades_as_the_command_grows_into_its_memory() {
124        let at = Instant::now();
125        let mut ledger = Ledger::default();
126        ledger.book("a", "one", 1000.0, at);
127        assert_eq!(ledger.booked_mb(at + COUNTED_IN_FULL_FOR), 1000.0);
128        let halfway = COUNTED_IN_FULL_FOR + (COUNTED_UNTIL - COUNTED_IN_FULL_FOR) / 2;
129        assert!((ledger.booked_mb(at + halfway) - 500.0).abs() < 1.0);
130        assert_eq!(ledger.booked_mb(at + COUNTED_UNTIL), 0.0);
131        // Still held, though it no longer counts: only the outcome ends it.
132        assert_eq!(ledger.running(), 1);
133    }
134
135    /// The mod asks about a waiting command many times; its booking is keyed by the call, so it is
136    /// one booking.
137    #[test]
138    fn booking_a_call_twice_counts_it_once() {
139        let now = Instant::now();
140        let mut ledger = Ledger::default();
141        ledger.book("a", "one", 3000.0, now);
142        ledger.book("a", "one", 3000.0, now);
143        assert_eq!(ledger.booked_mb(now), 3000.0);
144    }
145}