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 #[expect(clippy::disallowed_methods, reason = "test fixture (child D4)")]
336 fn a_run_child_reads_the_requested_affinity() {
337 let available = desktop_ffi::cpu::thread_group().unwrap().mask;
338 let mask = std::num::NonZeroUsize::new(
339 usize::try_from(1_u64 << available.trailing_zeros()).unwrap(),
340 )
341 .unwrap();
342 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
343 .with_shares(Shares::CpuAndMemory {
344 limit: std::num::NonZeroU64::new(8 << 30).unwrap(),
345 })
346 .with_affinity(super::RunAffinity::Mask(mask));
347 let owner = kernel::RunId::from_bytes([9; 16]);
348 let addr = kernel::Address::parse("vault/room1").unwrap();
349 for table in [
350 backlog.clone(),
351 backlog.clone().with_affinity(super::RunAffinity::Os),
352 ] {
353 let backlog = &table;
354 let gate = backlog
355 .scratch
356 .dir(backlog.mint().unwrap())
357 .with_extension("gate");
358 let mut child = std::process::Command::new("powershell.exe");
359 child.args([
360 "-NoProfile",
361 "-NonInteractive",
362 "-Command",
363 "while (-not (Test-Path -LiteralPath $env:R05_GATE)) { Start-Sleep -Milliseconds 20 }; [Diagnostics.Process]::GetCurrentProcess().ProcessorAffinity.ToInt64(); exit 23",
364 ]);
365 child
366 .env("R05_GATE", &gate)
367 .current_dir(std::env::temp_dir());
368 let started = backlog
369 .run(owner, &addr, "affinity reader".to_owned(), child)
370 .unwrap();
371 assert!(matches!(started, crate::Started::Backgrounded { .. }));
372 let reading = backlog.processes().unwrap().remove(&owner).unwrap();
373 assert_eq!(reading.affinity, super::RunAffinity::Mask(mask));
374 assert_eq!(reading.unfollowed, 0);
375 assert_eq!(reading.share, backlog.shares());
376 std::fs::write(&gate, b"ready").unwrap();
377 let finished = harvested(backlog, owner);
378 std::fs::remove_file(gate).unwrap();
379 assert_eq!(
380 finished.exit,
381 crate::Exit::Ended { code: 23 },
382 "{}",
383 finished.stderr
384 );
385 assert_eq!(finished.stdout.trim().parse::<usize>().unwrap(), mask.get());
386 }
387 backlog.release(owner);
388 }
389
390 #[cfg(windows)]
396 #[test]
397 #[expect(clippy::disallowed_methods, reason = "test fixture (child D4)")]
398 fn a_command_refused_at_the_ceiling_reports_the_hit() {
399 let limit = std::num::NonZeroU64::new(512 << 20).unwrap();
400 let owner = kernel::RunId::from_bytes([13; 16]);
401 let addr = kernel::Address::parse("vault/room1").unwrap();
402 let quick = Backlog::with_window(crate::PollBudget::new(1_500, 20))
403 .with_shares(Shares::CpuAndMemory { limit });
404 let mut small = std::process::Command::new("powershell.exe");
405 small.args(["-NoProfile", "-NonInteractive", "-Command", "exit 0"]);
406 small.current_dir(std::env::temp_dir());
407 let settled = quick
408 .run(owner, &addr, "under the ceiling".to_owned(), small)
409 .unwrap();
410 assert!(
411 matches!(
412 settled,
413 crate::Started::Settled {
414 exit: crate::Exit::Ended { code: 0 },
415 ceiling: None,
416 ..
417 }
418 ),
419 "{settled:?}"
420 );
421 quick.release(owner);
422
423 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
424 .with_shares(Shares::CpuAndMemory { limit });
425 let gate = backlog
426 .scratch
427 .dir(backlog.mint().unwrap())
428 .with_extension("gate");
429 let mut large = std::process::Command::new("powershell.exe");
430 large.args([
431 "-NoProfile",
432 "-NonInteractive",
433 "-Command",
434 "while (-not (Test-Path -LiteralPath $env:D95_GATE)) { Start-Sleep -Milliseconds 20 }; try { $b = New-Object byte[] 1073741824; exit 0 } catch { exit 7 }",
435 ]);
436 large
437 .env("D95_GATE", &gate)
438 .current_dir(std::env::temp_dir());
439 assert!(matches!(
440 backlog
441 .run(owner, &addr, "over the ceiling".to_owned(), large)
442 .unwrap(),
443 crate::Started::Backgrounded { .. }
444 ));
445 std::fs::write(&gate, b"ready").unwrap();
446 let mut finished = None;
447 for _ in 0..1_500 {
448 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
449 finished = Some(done);
450 break;
451 }
452 std::thread::sleep(std::time::Duration::from_millis(20));
453 }
454 std::fs::remove_file(gate).unwrap();
455 let finished = finished.expect("the allocating command finishes");
456 assert_eq!(
457 (finished.exit, finished.ceiling),
458 (
459 crate::Exit::Ended { code: 7 },
460 Some(crate::Ceiling::Hit { limit })
461 ),
462 "{}",
463 finished.stderr
464 );
465 backlog.release(owner);
466 }
467
468 #[cfg(not(windows))]
470 #[test]
471 #[expect(clippy::disallowed_methods, reason = "test fixture (child D4)")]
472 fn a_run_affinity_request_elsewhere_keeps_the_command_result() {
473 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1)).with_affinity(
474 super::RunAffinity::Mask(std::num::NonZeroUsize::new(1).unwrap()),
475 );
476 let owner = kernel::RunId::from_bytes([9; 16]);
477 let addr = kernel::Address::parse("vault/room1").unwrap();
478 let mut child = std::process::Command::new("sh");
479 child.args(["-c", "sleep 1; echo affinity; exit 23"]);
480 child.current_dir(std::env::temp_dir());
481 assert!(matches!(
482 backlog
483 .run(owner, &addr, "reader".to_owned(), child)
484 .unwrap(),
485 crate::Started::Backgrounded { .. }
486 ));
487 assert_eq!(
488 backlog.processes().unwrap().get(&owner).unwrap().affinity,
489 super::RunAffinity::Os
490 );
491 for _ in 0..250 {
492 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
493 assert_eq!(
494 (done.exit, done.stdout.trim()),
495 (crate::Exit::Ended { code: 23 }, "affinity")
496 );
497 backlog.release(owner);
498 return;
499 }
500 std::thread::sleep(std::time::Duration::from_millis(20));
501 }
502 backlog.release(owner);
503 panic!("the command finishes");
504 }
505
506 #[cfg(windows)]
509 #[test]
510 #[expect(clippy::disallowed_methods, reason = "test fixture (child D4)")]
511 fn a_refused_run_affinity_keeps_the_command_result_and_share() {
512 let available = usize::try_from(desktop_ffi::cpu::thread_group().unwrap().mask).unwrap();
513 if available == usize::MAX {
514 return; }
516 let backlog = Backlog::with_window(crate::PollBudget::new(1, 1))
517 .with_shares(Shares::Cpu)
518 .with_affinity(super::RunAffinity::Mask(std::num::NonZeroUsize::MAX));
519 let owner = kernel::RunId::from_bytes([11; 16]);
520 let addr = kernel::Address::parse("vault/room1").unwrap();
521 let mut child = std::process::Command::new("powershell.exe");
522 child.args(["-NoProfile", "-NonInteractive", "-Command", "exit 23"]);
523 child.current_dir(std::env::temp_dir());
524 assert!(matches!(
525 backlog
526 .run(owner, &addr, "refused mask".to_owned(), child)
527 .unwrap(),
528 crate::Started::Backgrounded { .. }
529 ));
530 let reading = backlog.processes().unwrap().remove(&owner).unwrap();
531 assert_eq!(
532 (reading.affinity, reading.share, reading.unfollowed),
533 (super::RunAffinity::Os, Shares::Cpu, 0)
534 );
535 let done = harvested(&backlog, owner);
536 backlog.release(owner);
537 assert_eq!(
538 done.exit,
539 crate::Exit::Ended { code: 23 },
540 "{}",
541 done.stderr
542 );
543 }
544
545 #[cfg(windows)]
551 fn harvested(backlog: &Backlog, owner: kernel::RunId) -> crate::Finished {
552 loop {
553 if let Some(done) = backlog.harvest(owner).unwrap().into_iter().next() {
554 return done;
555 }
556 std::thread::sleep(std::time::Duration::from_millis(20));
557 }
558 }
559}