atagen

roundabout

public
Star0
Watch0
ecc1cf8df0f5b88c6bf2cd07254c98d13827cd26
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/sched.rsRust353 lines9.2 KB
1
//! fifo with easy backfill. the first job that can't start gets a reservation,
2
//! and anyone behind it only cuts in if they'll be gone before it starts or fit
3
//! in the leftovers once it does
4
5
use std::{collections::BTreeMap, time::Duration};
6
7
use serde::{Deserialize, Serialize};
8
9
use crate::units::Bytes;
10
11
pub type JobId = u64;
12
pub type FsId = u64;
13
14
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
15
pub struct Need {
16
    pub ram: Bytes,
17
    pub disk: Bytes,
18
    pub cores: u32,
19
    pub time: Duration,
20
}
21
22
#[derive(Clone, Debug)]
23
pub struct Queued {
24
    pub id: JobId,
25
    pub need: Need,
26
    pub heavy: bool,
27
    pub fs: Vec<FsId>,
28
}
29
30
#[derive(Clone, Debug)]
31
pub struct Running {
32
    pub need: Need,
33
    pub heavy: bool,
34
    pub fs: Vec<FsId>,
35
    pub rss: Bytes,
36
    /// zero once overdue
37
    pub remaining: Duration,
38
}
39
40
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41
pub struct Disk {
42
    pub free: Bytes,
43
    pub reserve: Bytes,
44
    /// written since the fs last went from idle to busy. free space already
45
    /// reflects this, so it's netted off promises to avoid counting it twice
46
    pub consumed: Bytes,
47
}
48
49
#[derive(Clone, Debug)]
50
pub struct Capacity {
51
    pub ram_budget: Bytes,
52
    pub ram_live: Bytes,
53
    pub ram_reserve: Bytes,
54
    pub cores: u32,
55
    pub disks: BTreeMap<FsId, Disk>,
56
    pub pressure: bool,
57
}
58
59
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
60
#[serde(rename_all = "snake_case", tag = "reason", content = "fs")]
61
pub enum Reason {
62
    HeavyRunning,
63
    Pressure,
64
    Ram,
65
    Cores,
66
    Disk(FsId),
67
    Reserved,
68
}
69
70
impl std::fmt::Display for Reason {
71
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72
        f.write_str(match self {
73
            Self::HeavyRunning => "another heavy build is running",
74
            Self::Pressure => "memory pressure is high",
75
            Self::Ram => "not enough free RAM",
76
            Self::Cores => "not enough free cores",
77
            Self::Disk(_) => "not enough free disk",
78
            Self::Reserved => "waiting behind a reserved build",
79
        })
80
    }
81
}
82
83
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
84
pub enum Verdict {
85
    Admit,
86
    Wait(Reason),
87
}
88
89
#[derive(Clone, Debug)]
90
struct Free {
91
    ram: Bytes,
92
    cores: u32,
93
    disk: BTreeMap<FsId, Bytes>,
94
    heavy_slot: bool,
95
}
96
97
impl Free {
98
    fn shortfall(&self, need: &Need, fs: &[FsId], heavy: bool) -> Option<Reason> {
99
        if heavy && !self.heavy_slot {
100
            return Some(Reason::HeavyRunning);
101
        }
102
        if need.ram > self.ram {
103
            return Some(Reason::Ram);
104
        }
105
        if need.cores > self.cores {
106
            return Some(Reason::Cores);
107
        }
108
        fs.iter()
109
            .find(|id| self.disk.get(id).is_some_and(|free| need.disk > *free))
110
            .map(|id| Reason::Disk(*id))
111
    }
112
113
    fn take(&mut self, need: &Need, fs: &[FsId], heavy: bool) {
114
        self.ram = self.ram.saturating_sub(need.ram);
115
        self.cores = self.cores.saturating_sub(need.cores);
116
        for id in fs {
117
            if let Some(free) = self.disk.get_mut(id) {
118
                *free = free.saturating_sub(need.disk);
119
            }
120
        }
121
        if heavy {
122
            self.heavy_slot = false;
123
        }
124
    }
125
126
    /// no disk here. builds leave their output behind when they finish
127
    fn give_back(&mut self, ram: Bytes, cores: u32, heavy: bool) {
128
        self.ram += ram;
129
        self.cores += cores;
130
        if heavy {
131
            self.heavy_slot = true;
132
        }
133
    }
134
}
135
136
fn clamp(need: Need, cap: &Capacity) -> Need {
137
    Need {
138
        ram: need.ram.min(cap.ram_budget),
139
        cores: need.cores.clamp(1, cap.cores.max(1)),
140
        ..need
141
    }
142
}
143
144
fn free_now(running: &[Running], cap: &Capacity) -> Free {
145
    let charged: Bytes = running.iter().map(|r| r.need.ram.max(r.rss)).sum();
146
    let unrealised: Bytes = running
147
        .iter()
148
        .map(|r| r.need.ram.saturating_sub(r.rss))
149
        .sum();
150
    let by_budget = cap.ram_budget.saturating_sub(charged);
151
    let by_live = cap
152
        .ram_live
153
        .saturating_sub(cap.ram_reserve)
154
        .saturating_sub(unrealised);
155
156
    let disk = cap
157
        .disks
158
        .iter()
159
        .map(|(id, d)| {
160
            let promised: Bytes = running
161
                .iter()
162
                .filter(|r| r.fs.contains(id))
163
                .map(|r| r.need.disk)
164
                .sum();
165
            let outstanding = promised.saturating_sub(d.consumed);
166
            (
167
                *id,
168
                d.free.saturating_sub(d.reserve).saturating_sub(outstanding),
169
            )
170
        })
171
        .collect();
172
173
    Free {
174
        ram: by_budget.min(by_live),
175
        cores: cap
176
            .cores
177
            .saturating_sub(running.iter().map(|r| r.need.cores).sum()),
178
        disk,
179
        heavy_slot: !running.iter().any(|r| r.heavy),
180
    }
181
}
182
183
pub fn headroom(running: &[Running], cap: &Capacity) -> (Bytes, u32) {
184
    let free = free_now(running, cap);
185
    (free.ram, free.cores)
186
}
187
188
struct Reservation {
189
    /// `None` when nothing finishing would free enough, e.g. disk or someone
190
    /// else's ram. backfill falls back to the time cap then
191
    shadow: Option<Duration>,
192
    extra: Option<Free>,
193
}
194
195
fn reserve(
196
    job: &Queued,
197
    need: &Need,
198
    free: &Free,
199
    ends: &[(Duration, Bytes, u32, bool)],
200
) -> Reservation {
201
    let mut sim = free.clone();
202
    let mut ends = ends.to_vec();
203
    ends.sort_by_key(|e| e.0);
204
    for (at, ram, cores, heavy) in ends {
205
        sim.give_back(ram, cores, heavy);
206
        if sim.shortfall(need, &job.fs, job.heavy).is_none() {
207
            sim.take(need, &job.fs, job.heavy);
208
            return Reservation {
209
                shadow: Some(at),
210
                extra: Some(sim),
211
            };
212
        }
213
    }
214
    Reservation {
215
        shadow: None,
216
        extra: None,
217
    }
218
}
219
220
pub fn plan(
221
    queue: &[Queued],
222
    running: &[Running],
223
    cap: &Capacity,
224
    backfill_max: Duration,
225
) -> Vec<(JobId, Verdict)> {
226
    let mut free = free_now(running, cap);
227
    let mut ends: Vec<_> = running
228
        .iter()
229
        .map(|r| (r.remaining, r.need.ram.max(r.rss), r.need.cores, r.heavy))
230
        .collect();
231
    let mut reservation: Option<Reservation> = None;
232
    let mut out = Vec::with_capacity(queue.len());
233
234
    for job in queue {
235
        let need = clamp(job.need, cap);
236
237
        if job.heavy && cap.pressure {
238
            out.push((job.id, Verdict::Wait(Reason::Pressure)));
239
            continue;
240
        }
241
242
        if let Some(why) = free.shortfall(&need, &job.fs, job.heavy) {
243
            if reservation.is_none() {
244
                reservation = Some(reserve(job, &need, &free, &ends));
245
            }
246
            out.push((job.id, Verdict::Wait(why)));
247
            continue;
248
        }
249
250
        if let Some(res) = &mut reservation {
251
            let finishes_first = match res.shadow {
252
                Some(shadow) => need.time <= shadow,
253
                None => need.time <= backfill_max,
254
            };
255
            let fits_extra = res
256
                .extra
257
                .as_ref()
258
                .is_some_and(|x| x.shortfall(&need, &job.fs, job.heavy).is_none());
259
            if !finishes_first && !fits_extra {
260
                out.push((job.id, Verdict::Wait(Reason::Reserved)));
261
                continue;
262
            }
263
            if !finishes_first && let Some(extra) = &mut res.extra {
264
                extra.take(&need, &job.fs, job.heavy);
265
            }
266
        }
267
268
        free.take(&need, &job.fs, job.heavy);
269
        ends.push((need.time, need.ram, need.cores, job.heavy));
270
        out.push((job.id, Verdict::Admit));
271
    }
272
    out
273
}
274
275
#[cfg(test)]
276
mod tests {
277
    use super::*;
278
279
    const FS: FsId = 1;
280
281
    fn need(ram_gib: u64, cores: u32, disk_gib: u64, mins: u64) -> Need {
282
        Need {
283
            ram: Bytes::gib(ram_gib),
284
            disk: Bytes::gib(disk_gib),
285
            cores,
286
            time: Duration::from_secs(mins * 60),
287
        }
288
    }
289
290
    fn q(id: JobId, need: Need, heavy: bool) -> Queued {
291
        Queued {
292
            id,
293
            need,
294
            heavy,
295
            fs: vec![FS],
296
        }
297
    }
298
299
    fn cap() -> Capacity {
300
        Capacity {
301
            ram_budget: Bytes::gib(24),
302
            ram_live: Bytes::gib(28),
303
            ram_reserve: Bytes::gib(4),
304
            cores: 16,
305
            disks: BTreeMap::from([(
306
                FS,
307
                Disk {
308
                    free: Bytes::gib(100),
309
                    reserve: Bytes::gib(20),
310
                    consumed: Bytes::ZERO,
311
                },
312
            )]),
313
            pressure: false,
314
        }
315
    }
316
317
    fn verdicts(plan: &[(JobId, Verdict)]) -> Vec<Verdict> {
318
        plan.iter().map(|(_, v)| *v).collect()
319
    }
320
321
    const BF: Duration = Duration::from_secs(120);
322
323
    #[test]
324
    fn two_heavies_never_collide() {
325
        let out = plan(
326
            &[q(1, need(8, 8, 5, 10), true), q(2, need(8, 4, 5, 10), true)],
327
            &[],
328
            &cap(),
329
            BF,
330
        );
331
        assert_eq!(
332
            verdicts(&out),
333
            [Verdict::Admit, Verdict::Wait(Reason::HeavyRunning)]
334
        );
335
    }
336
337
    #[test]
338
    fn disk_never_double_booked() {
339
        let out = plan(
340
            &[
341
                q(1, need(1, 1, 50, 10), false),
342
                q(2, need(1, 1, 50, 10), false),
343
            ],
344
            &[],
345
            &cap(),
346
            BF,
347
        );
348
        assert_eq!(
349
            verdicts(&out),
350
            [Verdict::Admit, Verdict::Wait(Reason::Disk(FS))]
351
        );
352
    }
353
}