1use std::collections::{BTreeMap, BTreeSet};
24use std::num::{NonZeroU64, NonZeroUsize};
25
26use kernel::{AxError, RunId};
27
28use super::ceiling::{Ceiling, Mark, Unapplied};
29#[cfg(windows)]
30use super::run_job::{RUN_CPU_WEIGHT, RunJob};
31use super::{Backlog, BacklogId, Body, Claim, Member};
32
33#[derive(Debug, Clone, Default, PartialEq, Eq)]
35pub struct RunProcesses {
36 pub pids: BTreeSet<u32>,
40 pub unfollowed: u32,
46 pub share: Shares,
49 pub affinity: RunAffinity,
52}
53
54#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
57pub enum RunAffinity {
58 #[default]
60 Os,
61 Mask(NonZeroUsize),
63}
64
65#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
71pub enum Shares {
72 #[default]
75 Unset,
76 Cpu,
78 CpuAndMemory {
81 limit: NonZeroU64,
83 },
84}
85
86impl Shares {
87 pub(super) fn memory(self) -> Option<NonZeroU64> {
89 match self {
90 Shares::Unset | Shares::Cpu => None,
91 Shares::CpuAndMemory { limit } => Some(limit),
92 }
93 }
94}
95
96#[derive(Default)]
102pub(super) struct Jobs {
103 #[cfg(windows)]
104 runs: BTreeMap<RunId, RunJob>,
105 #[cfg(not(windows))]
106 cgroups: super::cgroup::Cgroups,
107}
108
109impl Backlog {
110 pub fn processes(&self) -> Result<BTreeMap<RunId, RunProcesses>, AxError> {
116 let table = self.hold()?;
117 let mut runs: BTreeMap<RunId, RunProcesses> = BTreeMap::new();
118 for member in table.members.values() {
119 if let Body::Command {
120 child,
121 claim: Claim::Window(owner) | Claim::Run(owner),
122 ..
123 } = &member.body
124 && let Some(child) = child.client()
125 {
126 runs.entry(*owner).or_default().pids.insert(child.id());
127 }
128 }
129 table.jobs.follow(&mut runs);
130 Ok(runs)
131 }
132
133 pub(super) fn enrol(&self, id: BacklogId, mut member: Member) -> Result<(), AxError> {
136 let mut table = self.hold()?;
137 let mark = table.jobs.enter(&member, self.shares, self.affinity);
138 if let Body::Command { ceiling, .. } = &mut member.body {
139 *ceiling = mark;
140 }
141 table.members.insert(id, member);
142 Ok(())
143 }
144}
145
146impl Jobs {
147 pub(super) fn ceiling(&mut self, owner: RunId, mark: Mark) -> Option<Ceiling> {
151 let count = if mark.watching() {
152 self.refusals(owner)
153 } else {
154 None
155 };
156 super::ceiling::verdict(mark, count)
157 }
158
159 #[cfg(windows)]
160 fn refusals(&mut self, owner: RunId) -> Option<u64> {
161 self.runs.get_mut(&owner)?.refusals()
162 }
163
164 #[cfg(not(windows))]
165 fn refusals(&mut self, owner: RunId) -> Option<u64> {
166 self.cgroups.refusals(owner)
167 }
168
169 #[cfg(windows)]
174 pub(super) fn native_job(
175 &mut self,
176 owner: RunId,
177 memory: Option<NonZeroUsize>,
178 affinity: RunAffinity,
179 ) -> Result<(NonZeroUsize, Mark), AxError> {
180 use crate::tools::native_windows::denied;
181 let run = self.runs.entry(owner).or_default();
182 for retiring in &mut run.retiring {
183 retiring.cleanup()?;
184 }
185 run.retiring.clear();
186 if run.job.is_none() {
187 run.job = Some(win32job::Job::create().map_err(|err| denied("create run Job", err))?);
188 }
189 let job = run
190 .job
191 .as_ref()
192 .ok_or_else(|| denied("create run Job", "Job owner disappeared"))?;
193 desktop_ffi::cpu::job_share(
194 job.handle(),
195 RUN_CPU_WEIGHT,
196 memory.map_or(0, NonZeroUsize::get),
197 )
198 .map_err(|err| denied("set run memory ceiling", format!("{err:?}")))?;
199 if memory.is_some() && run.watch.is_none() {
200 run.watch = Some(
201 desktop_ffi::cpu::JobWatch::attach(job.handle())
202 .map_err(|err| denied("watch run memory ceiling", format!("{err:?}")))?,
203 );
204 }
205 if let RunAffinity::Mask(mask) = affinity {
206 let mut info = job
207 .query_extended_limit_info()
208 .map_err(|err| denied("read run Job affinity", err))?;
209 info.limit_affinity(mask.get());
210 job.set_extended_limit_info(&info)
211 .map_err(|err| denied("set run Job affinity", err))?;
212 }
213 run.share = match memory {
214 None => Shares::Cpu,
215 Some(memory) => Shares::CpuAndMemory {
216 limit: NonZeroU64::new(
217 u64::try_from(memory.get())
218 .map_err(|err| denied("read run Job memory", err))?,
219 )
220 .ok_or_else(|| denied("read run Job memory", "zero memory"))?,
221 },
222 };
223 run.affinity = affinity;
224 let lent = NonZeroUsize::new(
225 usize::try_from(job.handle()).map_err(|err| denied("lend run Job", err))?,
226 )
227 .ok_or_else(|| denied("lend run Job", "invalid Job handle"))?;
228 let asked = run.share.memory();
229 Ok((lent, run.mark(asked)))
230 }
231
232 #[cfg(windows)]
233 pub(super) fn keep_failed_native(
234 &mut self,
235 owner: RunId,
236 resources: super::native_windows::RetiringNative,
237 ) {
238 self.runs.entry(owner).or_default().retiring.push(resources);
239 }
240
241 #[cfg(windows)]
243 pub(super) fn forget(&mut self, owner: RunId) {
244 self.runs.remove(&owner);
245 }
246
247 #[cfg(not(windows))]
248 pub(super) fn forget(&mut self, owner: RunId) {
249 self.cgroups.forget(owner);
250 }
251
252 #[cfg(windows)]
253 fn enter(&mut self, member: &Member, shares: Shares, affinity: RunAffinity) -> Mark {
254 let Body::Command {
255 child,
256 claim: Claim::Window(owner) | Claim::Run(owner),
257 ..
258 } = &member.body
259 else {
260 return Mark::NotAsked;
261 };
262 let Some(child) = child.client() else {
263 return Mark::outside(shares.memory(), Unapplied::Unjoined);
264 };
265 let run = self.runs.entry(*owner).or_default();
266 match run.join(child, shares, affinity) {
267 Ok(()) => run.mark(shares.memory()),
268 Err(why) => {
269 run.unjoined = run.unjoined.saturating_add(1);
270 Mark::outside(shares.memory(), why)
271 }
272 }
273 }
274
275 #[cfg(not(windows))]
276 fn enter(&mut self, member: &Member, shares: Shares, _affinity: RunAffinity) -> Mark {
277 let Body::Command {
278 child,
279 claim: Claim::Window(owner) | Claim::Run(owner),
280 ..
281 } = &member.body
282 else {
283 return Mark::NotAsked;
284 };
285 match child.client() {
286 Some(child) => self.cgroups.enter(*owner, child.id(), shares),
287 None => Mark::outside(shares.memory(), Unapplied::Unjoined),
288 }
289 }
290
291 #[cfg(windows)]
292 fn follow(&self, readings: &mut BTreeMap<RunId, RunProcesses>) {
293 for (owner, run) in &self.runs {
294 let reading = readings.entry(*owner).or_default();
295 let listed = run.job.as_ref().map(win32job::Job::query_process_id_list);
296 if let Some(Ok(pids)) = listed {
297 reading
299 .pids
300 .extend(pids.into_iter().filter_map(|pid| u32::try_from(pid).ok()));
301 reading.unfollowed = reading.unfollowed.saturating_add(run.unjoined);
302 reading.share = run.share;
303 reading.affinity = run.affinity;
304 } else {
305 reading.unfollowed = u32::try_from(reading.pids.len()).unwrap_or(u32::MAX);
306 }
307 }
308 }
309
310 #[cfg(not(windows))]
311 fn follow(&self, readings: &mut BTreeMap<RunId, RunProcesses>) {
312 for (owner, reading) in readings.iter_mut() {
313 reading.unfollowed = u32::try_from(reading.pids.len()).unwrap_or(u32::MAX);
314 reading.share = self.cgroups.held(*owner);
315 }
316 }
317}
318
319#[cfg(test)]
320#[allow(
321 clippy::unwrap_used,
322 clippy::expect_used,
323 clippy::panic,
324 reason = "test code"
325)]
326mod tests {
327 use super::super::Backlog;
328 #[cfg(windows)]
329 use super::Shares;
330
331 #[cfg(windows)]
334 #[test]
335 fn a_run_child_reads_the_requested_affinity() {
336 let available = desktop_ffi::cpu::thread_group().unwrap().mask;
337 let mask = std::num::NonZeroUsize::new(
338 usize::try_from(1_u64 << available.trailing_zeros()).unwrap(),
339 )
340 .unwrap();
341 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
342 .with_shares(Shares::CpuAndMemory {
343 limit: std::num::NonZeroU64::new(8 << 30).unwrap(),
344 })
345 .with_affinity(super::RunAffinity::Mask(mask));
346 let owner = kernel::RunId::from_bytes([9; 16]);
347 let addr = kernel::Address::parse("vault/room1").unwrap();
348 for table in [
349 backlog.clone(),
350 backlog.clone().with_affinity(super::RunAffinity::Os),
351 ] {
352 let backlog = &table;
353 let gate = backlog
354 .scratch
355 .dir(backlog.mint().unwrap())
356 .with_extension("gate");
357 let mut child = std::process::Command::new("powershell.exe");
358 child.args([
359 "-NoProfile",
360 "-NonInteractive",
361 "-Command",
362 "while (-not (Test-Path -LiteralPath $env:R05_GATE)) { Start-Sleep -Milliseconds 20 }; [Diagnostics.Process]::GetCurrentProcess().ProcessorAffinity.ToInt64(); exit 23",
363 ]);
364 child
365 .env("R05_GATE", &gate)
366 .current_dir(std::env::temp_dir());
367 let started = backlog
368 .run(owner, &addr, "affinity reader".to_owned(), child)
369 .unwrap();
370 assert!(matches!(started, crate::Started::Backgrounded { .. }));
371 let reading = backlog.processes().unwrap().remove(&owner).unwrap();
372 assert_eq!(reading.affinity, super::RunAffinity::Mask(mask));
373 assert_eq!(reading.unfollowed, 0);
374 assert_eq!(reading.share, backlog.shares());
375 std::fs::write(&gate, b"ready").unwrap();
376 let finished = harvested(backlog, owner);
377 std::fs::remove_file(gate).unwrap();
378 assert_eq!(
379 finished.exit,
380 crate::Exit::Ended { code: 23 },
381 "{}",
382 finished.stderr
383 );
384 assert_eq!(finished.stdout.trim().parse::<usize>().unwrap(), mask.get());
385 }
386 backlog.release(owner);
387 }
388
389 #[cfg(windows)]
395 #[test]
396 fn a_command_refused_at_the_ceiling_reports_the_hit() {
397 let limit = std::num::NonZeroU64::new(512 << 20).unwrap();
398 let owner = kernel::RunId::from_bytes([13; 16]);
399 let addr = kernel::Address::parse("vault/room1").unwrap();
400 let quick = Backlog::with_window(crate::PollBudget::new(1_500, 20))
401 .with_shares(Shares::CpuAndMemory { limit });
402 let mut small = std::process::Command::new("powershell.exe");
403 small.args(["-NoProfile", "-NonInteractive", "-Command", "exit 0"]);
404 small.current_dir(std::env::temp_dir());
405 let settled = quick
406 .run(owner, &addr, "under the ceiling".to_owned(), small)
407 .unwrap();
408 assert!(
409 matches!(
410 settled,
411 crate::Started::Settled {
412 exit: crate::Exit::Ended { code: 0 },
413 ceiling: None,
414 ..
415 }
416 ),
417 "{settled:?}"
418 );
419 quick.release(owner);
420
421 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
422 .with_shares(Shares::CpuAndMemory { limit });
423 let gate = backlog
424 .scratch
425 .dir(backlog.mint().unwrap())
426 .with_extension("gate");
427 let mut large = std::process::Command::new("powershell.exe");
428 large.args([
429 "-NoProfile",
430 "-NonInteractive",
431 "-Command",
432 "while (-not (Test-Path -LiteralPath $env:D95_GATE)) { Start-Sleep -Milliseconds 20 }; try { $b = New-Object byte[] 1073741824; exit 0 } catch { exit 7 }",
433 ]);
434 large
435 .env("D95_GATE", &gate)
436 .current_dir(std::env::temp_dir());
437 assert!(matches!(
438 backlog
439 .run(owner, &addr, "over the ceiling".to_owned(), large)
440 .unwrap(),
441 crate::Started::Backgrounded { .. }
442 ));
443 std::fs::write(&gate, b"ready").unwrap();
444 let mut finished = None;
445 for _ in 0..1_500 {
446 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
447 finished = Some(done);
448 break;
449 }
450 std::thread::sleep(std::time::Duration::from_millis(20));
451 }
452 std::fs::remove_file(gate).unwrap();
453 let finished = finished.expect("the allocating command finishes");
454 assert_eq!(
455 (finished.exit, finished.ceiling),
456 (
457 crate::Exit::Ended { code: 7 },
458 Some(crate::Ceiling::Hit { limit })
459 ),
460 "{}",
461 finished.stderr
462 );
463 backlog.release(owner);
464 }
465
466 #[cfg(not(windows))]
468 #[test]
469 fn a_run_affinity_request_elsewhere_keeps_the_command_result() {
470 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1)).with_affinity(
471 super::RunAffinity::Mask(std::num::NonZeroUsize::new(1).unwrap()),
472 );
473 let owner = kernel::RunId::from_bytes([9; 16]);
474 let addr = kernel::Address::parse("vault/room1").unwrap();
475 let mut child = std::process::Command::new("sh");
476 child.args(["-c", "sleep 1; echo affinity; exit 23"]);
477 child.current_dir(std::env::temp_dir());
478 assert!(matches!(
479 backlog
480 .run(owner, &addr, "reader".to_owned(), child)
481 .unwrap(),
482 crate::Started::Backgrounded { .. }
483 ));
484 assert_eq!(
485 backlog.processes().unwrap().get(&owner).unwrap().affinity,
486 super::RunAffinity::Os
487 );
488 for _ in 0..250 {
489 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
490 assert_eq!(
491 (done.exit, done.stdout.trim()),
492 (crate::Exit::Ended { code: 23 }, "affinity")
493 );
494 backlog.release(owner);
495 return;
496 }
497 std::thread::sleep(std::time::Duration::from_millis(20));
498 }
499 backlog.release(owner);
500 panic!("the command finishes");
501 }
502
503 #[cfg(windows)]
506 #[test]
507 fn a_refused_run_affinity_keeps_the_command_result_and_share() {
508 let available = usize::try_from(desktop_ffi::cpu::thread_group().unwrap().mask).unwrap();
509 if available == usize::MAX {
510 return; }
512 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
513 .with_shares(Shares::Cpu)
514 .with_affinity(super::RunAffinity::Mask(std::num::NonZeroUsize::MAX));
515 let owner = kernel::RunId::from_bytes([11; 16]);
516 let addr = kernel::Address::parse("vault/room1").unwrap();
517 let mut child = std::process::Command::new("powershell.exe");
518 child.args(["-NoProfile", "-NonInteractive", "-Command", "exit 23"]);
519 child.current_dir(std::env::temp_dir());
520 assert!(matches!(
521 backlog
522 .run(owner, &addr, "refused mask".to_owned(), child)
523 .unwrap(),
524 crate::Started::Backgrounded { .. }
525 ));
526 let reading = backlog.processes().unwrap().remove(&owner).unwrap();
527 assert_eq!(
528 (reading.affinity, reading.share, reading.unfollowed),
529 (super::RunAffinity::Os, Shares::Cpu, 0)
530 );
531 let done = harvested(&backlog, owner);
532 backlog.release(owner);
533 assert_eq!(
534 done.exit,
535 crate::Exit::Ended { code: 23 },
536 "{}",
537 done.stderr
538 );
539 }
540
541 #[cfg(windows)]
547 fn harvested(backlog: &Backlog, owner: kernel::RunId) -> crate::Finished {
548 loop {
549 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
550 return done;
551 }
552 std::thread::sleep(std::time::Duration::from_millis(20));
553 }
554 }
555}