1use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
13use std::fs::DirEntry;
14use std::os::unix::fs::MetadataExt;
15use std::os::unix::io::{AsFd, BorrowedFd};
16use std::path::{Path, PathBuf};
17use std::sync::atomic::{AtomicBool, Ordering};
18
19use anyhow::{bail, Context, Result};
20use inotify::{Inotify, WatchMask};
21use scx_utils::Cpumask;
22use tracing::{debug, info};
23
24fn cgroup_root_relative(path: &Path) -> String {
27 path.strip_prefix("/sys/fs/cgroup")
28 .map(|rel| format!("/{}", rel.to_string_lossy()))
29 .unwrap_or_else(|_| path.to_string_lossy().into_owned())
30}
31
32#[derive(Debug)]
34pub struct CellInfo {
35 pub cell_id: u32,
36 pub cgroup_path: Option<PathBuf>,
37 pub cgid: Option<u64>,
38 pub cpuset: Option<Cpumask>,
40}
41
42#[derive(Debug, Clone, Copy, Default)]
49struct CpuMinimum {
50 protected: usize,
51 requested: usize,
52}
53
54#[derive(Debug, Clone)]
61struct CpuRecipient {
62 id: u32,
63 weight: f64,
64 allowed: Cpumask,
65 claimed: Option<Cpumask>,
66 minimum: CpuMinimum,
67}
68
69#[derive(Debug)]
71pub struct CpuAssignment {
72 pub id: u32,
73 pub primary: Cpumask,
74 pub borrowable: Option<Cpumask>,
75}
76
77struct CpuManager<'a> {
79 domain: &'a Cpumask,
80 cpu_to_partition: Option<&'a HashMap<usize, usize>>,
83}
84
85impl<'a> CpuManager<'a> {
86 #[cfg(test)]
87 fn new(domain: &'a Cpumask) -> Self {
88 Self {
89 domain,
90 cpu_to_partition: None,
91 }
92 }
93
94 fn with_partitions(domain: &'a Cpumask, cpu_to_partition: &'a HashMap<usize, usize>) -> Self {
95 Self {
96 domain,
97 cpu_to_partition: Some(cpu_to_partition),
98 }
99 }
100
101 fn compute_targets(
106 total_cpus: usize,
107 recipients: &[(u32, f64)],
108 ) -> Result<HashMap<u32, usize>> {
109 if recipients.is_empty() {
110 bail!("compute_targets called with no cells");
111 }
112 if total_cpus < recipients.len() {
113 bail!(
114 "Not enough CPUs ({}) for {} cells (need at least 1 each)",
115 total_cpus,
116 recipients.len()
117 );
118 }
119
120 let total_weight: f64 = recipients.iter().map(|(_, w)| w).sum();
121 let num_recipients = recipients.len();
122
123 if total_weight <= 0.0 {
124 let per = total_cpus / num_recipients;
126 let remainder = total_cpus % num_recipients;
127 return Ok(recipients
128 .iter()
129 .enumerate()
130 .map(|(i, (id, _))| {
131 let extra = if i < remainder { 1 } else { 0 };
132 (*id, per + extra)
133 })
134 .collect());
135 }
136
137 let distributable = total_cpus - num_recipients;
138 let mut assigned = num_recipients;
139 let mut raw: Vec<(u32, f64, usize)> = recipients
140 .iter()
141 .map(|(id, weight)| {
142 let frac = weight / total_weight * distributable as f64;
143 let floored = frac.floor() as usize;
144 assigned += floored;
145 (*id, frac - frac.floor(), 1 + floored)
146 })
147 .collect();
148
149 let mut remainder = total_cpus - assigned;
150 raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
151 for entry in raw.iter_mut() {
152 if remainder == 0 {
153 break;
154 }
155 entry.2 += 1;
156 remainder -= 1;
157 }
158
159 Ok(raw.iter().map(|(id, _, count)| (*id, *count)).collect())
160 }
161
162 fn distribute_cpus_proportional(
169 cpus: &[usize],
170 recipients: &[(u32, f64)],
171 ) -> Result<HashMap<u32, Vec<usize>>> {
172 if cpus.is_empty() {
173 bail!("distribute_cpus_proportional called with no CPUs");
174 }
175 if recipients.is_empty() {
176 bail!("distribute_cpus_proportional called with no recipients");
177 }
178
179 let total_weight: f64 = recipients.iter().map(|(_, w)| w).sum();
180 let n = cpus.len();
181
182 let mut allocs: Vec<(u32, usize)> = if total_weight <= 0.0 {
183 let per = n / recipients.len();
185 let remainder = n % recipients.len();
186 recipients
187 .iter()
188 .enumerate()
189 .map(|(i, (id, _))| {
190 let extra = if i < remainder { 1 } else { 0 };
191 (*id, per + extra)
192 })
193 .collect()
194 } else {
195 let mut assigned = 0usize;
197 let mut raw: Vec<(u32, f64, usize, bool)> = recipients
198 .iter()
199 .map(|(id, weight)| {
200 let frac = weight / total_weight * n as f64;
201 let floored = frac.floor() as usize;
202 assigned += floored;
203 (*id, frac - frac.floor(), floored, *weight > 0.0)
204 })
205 .collect();
206
207 let mut remainder = n - assigned;
208 raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
209 for entry in raw.iter_mut() {
210 if remainder == 0 {
211 break;
212 }
213 entry.2 += 1;
214 remainder -= 1;
215 }
216
217 loop {
222 let Some(starved_idx) = raw
223 .iter()
224 .position(|(_, _, count, pos_weight)| *pos_weight && *count == 0)
225 else {
226 break;
227 };
228 let Some(donor_idx) = raw
229 .iter()
230 .enumerate()
231 .filter(|(_, (_, _, count, _))| *count > 1)
232 .max_by_key(|(_, (_, _, count, _))| *count)
233 .map(|(i, _)| i)
234 else {
235 break; };
237 raw[donor_idx].2 -= 1;
238 raw[starved_idx].2 += 1;
239 }
240
241 raw.iter().map(|(id, _, count, _)| (*id, *count)).collect()
242 };
243
244 let mut cpu_iter = cpus.iter().copied();
246 allocs.sort_by_key(|(id, _)| *id);
247 let mut result: HashMap<u32, Vec<usize>> = HashMap::new();
248 for (id, count) in allocs {
249 let cpus_for_recipient: Vec<usize> = cpu_iter.by_ref().take(count).collect();
250 if cpus_for_recipient.len() != count {
251 bail!(
252 "BUG: distribute_cpus_proportional: cell {} expected {} CPUs but got {}",
253 id,
254 count,
255 cpus_for_recipient.len()
256 );
257 }
258 if !cpus_for_recipient.is_empty() {
259 result.insert(id, cpus_for_recipient);
260 }
261 }
262
263 Ok(result)
264 }
265
266 fn reserve_requested_minimums(
273 &self,
274 recipients: &[CpuRecipient],
275 contention: &HashMap<usize, Vec<u32>>,
276 ) -> Result<HashMap<u32, Cpumask>> {
277 let by_id: HashMap<u32, &CpuRecipient> = recipients
278 .iter()
279 .map(|recipient| (recipient.id, recipient))
280 .collect();
281 let mut assignments = HashMap::new();
282 let mut reserved = Cpumask::new();
283
284 let mut requesters: Vec<&CpuRecipient> = recipients
285 .iter()
286 .filter(|recipient| recipient.minimum.requested > 0)
287 .collect();
288 requesters.sort_by_key(|recipient| recipient.id);
289
290 let mut partition_claimants: HashMap<usize, HashSet<u32>> = HashMap::new();
291 if let Some(cpu_to_partition) = self.cpu_to_partition {
292 for (cpu, claimants) in contention {
293 if let Some(&partition) = cpu_to_partition.get(cpu) {
294 partition_claimants
295 .entry(partition)
296 .or_default()
297 .extend(claimants);
298 }
299 }
300 }
301 let partition_pressure = |cpu: usize| {
302 self.cpu_to_partition
303 .and_then(|cpu_to_partition| cpu_to_partition.get(&cpu))
304 .and_then(|partition| partition_claimants.get(partition))
305 .map_or(0, HashSet::len)
306 };
307
308 let mut claimed_cpus: HashMap<u32, Vec<usize>> = HashMap::new();
309 for (&cpu, claimants) in contention {
310 for &id in claimants {
311 claimed_cpus.entry(id).or_default().push(cpu);
312 }
313 }
314
315 for recipient in requesters {
316 let mut recipient_reserved = Cpumask::new();
317 let mut unclaimed: Vec<usize> = self
318 .domain
319 .iter()
320 .filter(|&cpu| {
321 recipient.allowed.test_cpu(cpu)
322 && !reserved.test_cpu(cpu)
323 && !contention.contains_key(&cpu)
324 })
325 .collect();
326 unclaimed.sort_unstable();
327
328 let mut taken = 0usize;
329 for cpu in unclaimed.into_iter().take(recipient.minimum.requested) {
330 recipient_reserved.set_cpu(cpu).ok();
331 reserved.set_cpu(cpu).ok();
332 taken += 1;
333 }
334
335 let mut taken_from: HashMap<u32, usize> =
336 claimed_cpus.keys().map(|&id| (id, 0)).collect();
337 while taken < recipient.minimum.requested {
338 let remaining: HashMap<u32, usize> = claimed_cpus
342 .iter()
343 .map(|(&id, cpus)| {
344 (
345 id,
346 cpus.iter()
347 .filter(|&&cpu| {
348 !reserved.test_cpu(cpu) && contention[&cpu].len() == 1
349 })
350 .count(),
351 )
352 })
353 .collect();
354
355 let reservable = |cpu: usize| {
356 recipient.allowed.test_cpu(cpu)
357 && !reserved.test_cpu(cpu)
358 && contention.get(&cpu).is_some_and(|claimants| {
359 claimants.iter().all(|id| {
360 let protected = by_id
361 .get(id)
362 .map_or(0, |recipient| recipient.minimum.protected);
363 remaining.get(id).copied().unwrap_or(0) > protected
364 })
365 })
366 };
367
368 let donor = claimed_cpus
369 .iter()
370 .filter(|(_, cpus)| cpus.iter().any(|&cpu| reservable(cpu)))
371 .map(|(&id, _)| id)
372 .min_by(|&a, &b| {
373 taken_from[&a]
374 .cmp(&taken_from[&b])
375 .then(remaining[&b].cmp(&remaining[&a]))
376 .then(a.cmp(&b))
377 });
378 let Some(donor) = donor else {
379 break;
380 };
381
382 let cpu = claimed_cpus[&donor]
383 .iter()
384 .copied()
385 .filter(|&cpu| reservable(cpu))
386 .min_by(|&a, &b| {
387 let contested = |cpu: usize| contention.get(&cpu).map_or(0, Vec::len) > 1;
388 contested(a)
389 .cmp(&contested(b))
390 .then_with(|| partition_pressure(b).cmp(&partition_pressure(a)))
391 .then(a.cmp(&b))
392 })
393 .expect("donor has a reservable CPU");
394 recipient_reserved.set_cpu(cpu).ok();
395 reserved.set_cpu(cpu).ok();
396 *taken_from.entry(donor).or_insert(0) += 1;
397 taken += 1;
398 }
399
400 if recipient_reserved.weight() > 0 {
401 assignments.insert(recipient.id, recipient_reserved);
402 }
403 }
404
405 Ok(assignments)
406 }
407
408 fn compute_assignments(
422 &self,
423 recipients: &[CpuRecipient],
424 compute_borrowable: bool,
425 ) -> Result<Vec<CpuAssignment>> {
426 let domain = self.domain;
427
428 if recipients.is_empty() {
429 bail!("compute_cpu_assignments called with no recipients");
430 }
431
432 let mut seen_ids = HashSet::new();
433 for recipient in recipients {
434 if recipient.weight < 0.0 {
435 bail!(
436 "Recipient {} has negative weight {}",
437 recipient.id,
438 recipient.weight
439 );
440 }
441 if !seen_ids.insert(recipient.id) {
442 bail!("Duplicate recipient id {}", recipient.id);
443 }
444 if let Some(claimed) = &recipient.claimed {
445 let claimed_outside_allowed = claimed.and(&recipient.allowed.not()).and(domain);
446 if claimed_outside_allowed.weight() > 0 {
447 bail!(
448 "Recipient {} claims CPUs outside its allowed mask: {}",
449 recipient.id,
450 claimed_outside_allowed.to_cpulist()
451 );
452 }
453 }
454 }
455
456 let mut contention: HashMap<usize, Vec<u32>> = HashMap::new();
458 for recipient in recipients {
459 if let Some(ref claimed) = recipient.claimed {
460 for cpu in claimed.iter() {
461 if domain.test_cpu(cpu) {
462 contention.entry(cpu).or_default().push(recipient.id);
463 }
464 }
465 }
466 }
467
468 let mut recipient_cpus = self.reserve_requested_minimums(recipients, &contention)?;
469 let mut reserved = Cpumask::new();
470 for cpus in recipient_cpus.values() {
471 reserved = reserved.or(cpus);
472 }
473
474 let mut contested_cpus: Vec<usize> = Vec::new();
479 let mut unclaimed_cpus: Vec<usize> = Vec::new();
480
481 for cpu in domain.iter() {
482 if reserved.test_cpu(cpu) {
483 continue;
484 }
485 match contention.get(&cpu) {
486 None => unclaimed_cpus.push(cpu),
487 Some(claimants) if claimants.len() == 1 => {
488 let id = claimants[0];
489 recipient_cpus
490 .entry(id)
491 .or_insert_with(Cpumask::new)
492 .set_cpu(cpu)
493 .ok();
494 }
495 Some(_) => contested_cpus.push(cpu),
496 }
497 }
498
499 let mut assigned_count: HashMap<u32, usize> = recipient_cpus
505 .iter()
506 .map(|(&id, mask)| (id, mask.weight()))
507 .collect();
508 contested_cpus.sort_unstable();
509 let contested_choices: HashMap<u32, usize> = recipients
510 .iter()
511 .map(|recipient| {
512 let choices = contested_cpus
513 .iter()
514 .filter(|cpu| {
515 contention
516 .get(cpu)
517 .is_some_and(|claimants| claimants.contains(&recipient.id))
518 })
519 .count();
520 (recipient.id, choices)
521 })
522 .collect();
523 let mut weighted_contested_cpus = Vec::new();
524 for cpu in contested_cpus {
525 let claimants = contention
526 .get(&cpu)
527 .expect("contested CPU missing contention entry");
528 let protected_recipient = claimants
529 .iter()
530 .filter_map(|id| {
531 let recipient = recipients
532 .iter()
533 .find(|recipient| recipient.id == *id)
534 .expect("claimant missing recipient");
535 let assigned = assigned_count.get(id).copied().unwrap_or(0);
536 (assigned < recipient.minimum.protected)
537 .then(|| (*id, recipient.minimum.protected - assigned))
538 })
539 .min_by(|(a_id, a_deficit), (b_id, b_deficit)| {
542 contested_choices[a_id]
543 .cmp(&contested_choices[b_id])
544 .then(b_deficit.cmp(a_deficit))
545 .then(a_id.cmp(b_id))
546 })
547 .map(|(id, _)| id);
548
549 if let Some(id) = protected_recipient {
550 recipient_cpus
551 .entry(id)
552 .or_insert_with(Cpumask::new)
553 .set_cpu(cpu)
554 .ok();
555 *assigned_count.entry(id).or_insert(0) += 1;
556 } else {
557 weighted_contested_cpus.push(cpu);
558 }
559 }
560 let contested_cpus = weighted_contested_cpus;
561
562 let total_cpu_count = domain.weight();
564 let mut all_recipients_with_weights: Vec<(u32, f64)> =
565 recipients.iter().map(|r| (r.id, r.weight)).collect();
566 all_recipients_with_weights.sort_by_key(|(id, _)| *id);
567
568 let targets = Self::compute_targets(total_cpu_count, &all_recipients_with_weights)?;
569
570 let mut contested_groups: HashMap<Vec<u32>, Vec<usize>> = HashMap::new();
572 for cpu in contested_cpus {
573 if let Some(claimants) = contention.get(&cpu) {
574 let mut sorted_claimants = claimants.clone();
575 sorted_claimants.sort();
576 contested_groups
577 .entry(sorted_claimants)
578 .or_default()
579 .push(cpu);
580 }
581 }
582
583 let initial_deficit: HashMap<u32, f64> = targets
589 .iter()
590 .map(|(&id, &target)| {
591 let already = assigned_count.get(&id).copied().unwrap_or(0);
593 let deficit = if target > already {
594 (target - already) as f64
595 } else {
596 0.0
597 };
598 (id, deficit)
599 })
600 .collect();
601
602 for (claimants, cpus) in contested_groups {
603 let mut recipients_with_deficit: Vec<(u32, f64)> = Vec::new();
604 let mut all_zero = true;
605 for &id in &claimants {
606 let deficit = *initial_deficit.get(&id).ok_or_else(|| {
607 anyhow::anyhow!(
608 "BUG: recipient {} in contention map but missing from targets",
609 id
610 )
611 })?;
612 if deficit > 0.0 {
613 all_zero = false;
614 }
615 recipients_with_deficit.push((id, deficit));
616 }
617
618 if all_zero {
620 recipients_with_deficit = claimants.iter().map(|&id| (id, 1.0)).collect();
621 }
622
623 let distribution = Self::distribute_cpus_proportional(&cpus, &recipients_with_deficit)?;
624 for (id, assigned_cpus) in distribution {
625 let count = assigned_cpus.len();
626 for cpu in assigned_cpus {
627 recipient_cpus
628 .entry(id)
629 .or_insert_with(Cpumask::new)
630 .set_cpu(cpu)
631 .ok();
632 }
633 *assigned_count.entry(id).or_insert(0) += count;
634 }
635 }
636
637 if !unclaimed_cpus.is_empty() {
639 let mut unpinned_recipients: Vec<(u32, f64)> = Vec::new();
640 let mut all_zero = true;
641
642 for recipient in recipients {
643 if recipient.claimed.is_some() {
644 continue; }
646 let target = *targets.get(&recipient.id).ok_or_else(|| {
647 anyhow::anyhow!(
648 "BUG: recipient {} is unpinned but missing from targets",
649 recipient.id
650 )
651 })?;
652 let already = assigned_count.get(&recipient.id).copied().unwrap_or(0);
654 let deficit = if target > already {
655 (target - already) as f64
656 } else {
657 0.0
658 };
659 if deficit > 0.0 {
660 all_zero = false;
661 }
662 unpinned_recipients.push((recipient.id, deficit));
663 }
664 unpinned_recipients.sort_by_key(|(id, _)| *id);
665
666 if all_zero {
668 unpinned_recipients = unpinned_recipients
669 .iter()
670 .map(|(id, _)| (*id, 1.0))
671 .collect();
672 }
673
674 let distribution =
675 Self::distribute_cpus_proportional(&unclaimed_cpus, &unpinned_recipients)?;
676 for (id, assigned_cpus) in distribution {
677 let count = assigned_cpus.len();
678 for cpu in assigned_cpus {
679 recipient_cpus
680 .entry(id)
681 .or_insert_with(Cpumask::new)
682 .set_cpu(cpu)
683 .ok();
684 }
685 *assigned_count.entry(id).or_insert(0) += count;
686 }
687 }
688
689 for recipient in recipients {
692 let assigned = recipient_cpus.get(&recipient.id).map_or(0, Cpumask::weight);
693 let required = recipient.minimum.protected.max(1);
694 if assigned < required {
695 bail!(
696 "Recipient {} has {} CPUs assigned, below required minimum {} \
697 (nr_cpus={}, num_recipients={})",
698 recipient.id,
699 assigned,
700 required,
701 domain.weight(),
702 recipients.len()
703 );
704 }
705 }
706
707 let allowed_by_id: HashMap<u32, Cpumask> = recipients
708 .iter()
709 .map(|recipient| (recipient.id, recipient.allowed.clone()))
710 .collect();
711
712 Ok(recipient_cpus
714 .into_iter()
715 .map(|(id, primary)| {
716 let borrowable = if compute_borrowable {
717 let allowed = allowed_by_id
718 .get(&id)
719 .expect("recipient assignment missing allowed mask");
720 Some(domain.and(&primary.not()).and(allowed))
721 } else {
722 None
723 };
724 CpuAssignment {
725 id,
726 primary,
727 borrowable,
728 }
729 })
730 .collect())
731 }
732}
733
734fn read_dir_sorted(path: &Path) -> Result<Vec<DirEntry>> {
735 let readdir = std::fs::read_dir(path)
736 .with_context(|| format!("Failed to read directory: {}", path.display()))?;
737
738 let mut entries: Vec<_> = readdir
739 .map(|entry| {
740 entry.with_context(|| format!("Failed to read directory entry in: {}", path.display()))
741 })
742 .collect::<Result<Vec<_>>>()?;
743 entries.sort_by_key(|entry| entry.path());
744
745 Ok(entries)
746}
747
748pub struct CellManager {
750 cell_parent_path: PathBuf,
751 inotify: Inotify,
752 cells: HashMap<u64, CellInfo>,
754 cell_id_to_cgid: HashMap<u32, u64>,
756 free_cell_ids: BTreeSet<u32>,
758 next_cell_id: u32,
759 max_cells: u32,
760 all_cpus: Cpumask,
762 exclude_names: HashSet<String>,
764 cell0_min_cpus: usize,
767 cpu_to_llc: HashMap<usize, usize>,
771 enforced_holdout: AtomicBool,
776}
777
778impl CellManager {
779 pub fn new(
780 cell_parent_path: &str,
781 max_cells: u32,
782 all_cpus: Cpumask,
783 exclude: HashSet<String>,
784 cell0_min_cpus: usize,
785 cpu_to_llc: HashMap<usize, usize>,
786 ) -> Result<Self> {
787 let path = PathBuf::from(format!("/sys/fs/cgroup{}", cell_parent_path));
788 if !path.exists() {
789 bail!("Cell parent cgroup path does not exist: {}", path.display());
790 }
791 Self::new_with_path_opts(
792 path,
793 max_cells,
794 all_cpus,
795 exclude,
796 cell0_min_cpus,
797 cpu_to_llc,
798 )
799 }
800
801 #[cfg(test)]
805 fn new_with_path(
806 path: PathBuf,
807 max_cells: u32,
808 all_cpus: Cpumask,
809 exclude: HashSet<String>,
810 ) -> Result<Self> {
811 Self::new_with_path_opts(path, max_cells, all_cpus, exclude, 0, HashMap::new())
812 }
813
814 fn new_with_path_opts(
815 path: PathBuf,
816 max_cells: u32,
817 all_cpus: Cpumask,
818 exclude: HashSet<String>,
819 cell0_min_cpus: usize,
820 cpu_to_llc: HashMap<usize, usize>,
821 ) -> Result<Self> {
822 let inotify = Inotify::init().context("Failed to initialize inotify")?;
823 inotify
824 .watches()
825 .add(&path, WatchMask::CREATE | WatchMask::DELETE)
826 .context("Failed to add inotify watch")?;
827
828 let mut mgr = Self {
829 cell_parent_path: path.clone(),
830 inotify,
831 cells: HashMap::new(),
832 cell_id_to_cgid: HashMap::new(),
833 free_cell_ids: BTreeSet::new(),
834 next_cell_id: 1, max_cells,
836 all_cpus,
837 exclude_names: exclude,
838 cell0_min_cpus,
839 cpu_to_llc,
840 enforced_holdout: AtomicBool::new(false),
841 };
842
843 mgr.cells.insert(
846 0,
847 CellInfo {
848 cell_id: 0,
849 cgroup_path: None,
850 cgid: None,
851 cpuset: None,
852 },
853 );
854 mgr.cell_id_to_cgid.insert(0, 0);
855
856 mgr.scan_existing_children()
858 .context("Failed to scan existing child cgroups at startup")?;
859 Ok(mgr)
860 }
861
862 pub fn enforced_holdout(&self) -> bool {
865 self.enforced_holdout.load(Ordering::Relaxed)
866 }
867
868 fn should_exclude(&self, path: &Path) -> bool {
869 path.file_name()
870 .and_then(|n| n.to_str())
871 .map(|name| self.exclude_names.contains(name))
872 .unwrap_or(false)
873 }
874
875 fn scan_existing_children(&mut self) -> Result<Vec<(u64, u32)>> {
876 let mut assignments = Vec::new();
877
878 let entries = read_dir_sorted(&self.cell_parent_path)
879 .context("Failed to read cell parent directory")?;
880
881 for entry in entries {
882 let file_type = entry.file_type().with_context(|| {
883 format!("Failed to get file type for: {}", entry.path().display())
884 })?;
885 if file_type.is_dir() {
886 let path = entry.path();
887 if self.should_exclude(&path) {
888 continue;
889 }
890 let cgid = path
891 .metadata()
892 .with_context(|| {
893 format!("reading inode of cgroup directory {}", path.display())
894 })?
895 .ino();
896 let (cgid, cell_id) =
897 self.create_cell_for_cgroup(&path, cgid).with_context(|| {
898 format!("Failed to create cell for cgroup: {}", path.display())
899 })?;
900 assignments.push((cgid, cell_id));
901 }
902 }
903 Ok(assignments)
904 }
905
906 pub fn process_events(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
913 let mut buffer = [0; 1024];
914 let mut has_events = false;
915
916 loop {
918 match self.inotify.read_events(&mut buffer) {
919 Ok(events) => {
920 if events.into_iter().next().is_some() {
921 has_events = true;
922 } else {
923 break;
924 }
925 }
926 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
927 Err(e) => {
928 return Err(e).context("Failed to read inotify events");
929 }
930 }
931 }
932
933 if !has_events {
934 return Ok((Vec::new(), Vec::new()));
935 }
936
937 self.reconcile_cells()
939 }
940
941 fn reconcile_cells(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
944 let mut new_cells = Vec::new();
945
946 let mut current_entries: BTreeMap<PathBuf, u64> = BTreeMap::new();
950 let entries = read_dir_sorted(&self.cell_parent_path)?;
951 for entry in entries {
952 let file_type = entry.file_type().with_context(|| {
953 format!("Failed to get file type for: {}", entry.path().display())
954 })?;
955 if file_type.is_dir() {
956 let path = entry.path();
957 if self.should_exclude(&path) {
958 continue;
959 }
960
961 let metadata = match entry.metadata() {
962 Ok(metadata) => metadata,
963 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
964 continue;
967 }
968 Err(e) => {
969 return Err(e).with_context(|| {
970 format!("reading inode of cgroup directory {}", path.display())
971 });
972 }
973 };
974
975 current_entries.insert(path, metadata.ino());
976 }
977 }
978
979 let mut destroyed_cells: BTreeSet<u32> = BTreeSet::new();
981 self.cells.retain(|&cgid, info| {
982 if info.cell_id == 0 {
983 return true; }
985 let cgroup_path = info
987 .cgroup_path
988 .as_ref()
989 .expect("BUG: non-zero cell missing cgroup_path");
990 if current_entries.get(cgroup_path) == Some(&cgid) {
993 true
994 } else {
995 info!(
996 "Destroyed cell {} for cgroup {} (cgid={})",
997 info.cell_id,
998 cgroup_path.display(),
999 cgid
1000 );
1001 destroyed_cells.insert(info.cell_id);
1002 false
1003 }
1004 });
1005
1006 self.cell_id_to_cgid
1008 .retain(|cell_id, _| !destroyed_cells.contains(cell_id));
1009 self.free_cell_ids.extend(destroyed_cells.iter().copied());
1010
1011 for (path, cgid) in current_entries {
1013 if self.cells.contains_key(&cgid) {
1014 continue; }
1016 let (cgid, cell_id) = self
1017 .create_cell_for_cgroup(&path, cgid)
1018 .with_context(|| format!("Failed to create cell for cgroup: {}", path.display()))?;
1019 new_cells.push((cgid, cell_id));
1020 }
1021
1022 Ok((new_cells, destroyed_cells.into_iter().collect()))
1023 }
1024
1025 fn create_cell_for_cgroup(&mut self, path: &Path, cgid: u64) -> Result<(u64, u32)> {
1026 let cell_id = self.allocate_cell_id().context("allocating cell ID")?;
1027
1028 let cpuset = Self::read_cpuset(path)
1029 .with_context(|| format!("reading cpuset for cgroup {}", path.display()))?;
1030 if let Some(ref mask) = cpuset {
1031 debug!(
1032 "Cell {} has cpuset: {} (from {})",
1033 cell_id,
1034 mask.to_cpulist(),
1035 path.join("cpuset.cpus").display()
1036 );
1037 }
1038
1039 self.cells.insert(
1040 cgid,
1041 CellInfo {
1042 cell_id,
1043 cgroup_path: Some(path.to_path_buf()),
1044 cgid: Some(cgid),
1045 cpuset,
1046 },
1047 );
1048 self.cell_id_to_cgid.insert(cell_id, cgid);
1049
1050 info!(
1051 "Created cell {} for cgroup {} (cgid={})",
1052 cell_id,
1053 path.display(),
1054 cgid
1055 );
1056
1057 Ok((cgid, cell_id))
1058 }
1059
1060 fn read_cpuset(cgroup_path: &Path) -> Result<Option<Cpumask>> {
1062 let cpuset_path = cgroup_path.join("cpuset.cpus");
1063 match std::fs::read_to_string(&cpuset_path) {
1064 Ok(content) => {
1065 let content = content.trim();
1066 if content.is_empty() {
1067 Ok(None)
1068 } else {
1069 let mask = Cpumask::from_cpulist(content).with_context(|| {
1070 format!(
1071 "Failed to parse cpuset '{}' from {}",
1072 content,
1073 cpuset_path.display()
1074 )
1075 })?;
1076 Ok(Some(mask))
1077 }
1078 }
1079 Err(_) => Ok(None),
1081 }
1082 }
1083
1084 fn allocate_cell_id(&mut self) -> Result<u32> {
1085 if let Some(id) = self.free_cell_ids.pop_first() {
1087 return Ok(id);
1088 }
1089
1090 if self.next_cell_id >= self.max_cells {
1091 bail!("Cell ID space exhausted (max_cells={})", self.max_cells);
1092 }
1093
1094 let id = self.next_cell_id;
1095 self.next_cell_id += 1;
1096 Ok(id)
1097 }
1098
1099 pub fn compute_cpu_assignments(&self, compute_borrowable: bool) -> Result<Vec<CpuAssignment>> {
1111 self.compute_cpu_assignments_inner(None, compute_borrowable)
1113 }
1114
1115 pub fn compute_demand_cpu_assignments(
1123 &self,
1124 cell_demands: &HashMap<u32, f64>,
1125 compute_borrowable: bool,
1126 ) -> Result<Vec<CpuAssignment>> {
1127 self.compute_cpu_assignments_inner(Some(cell_demands), compute_borrowable)
1128 }
1129
1130 fn compute_cpu_assignments_inner(
1132 &self,
1133 cell_demands: Option<&HashMap<u32, f64>>,
1134 compute_borrowable: bool,
1135 ) -> Result<Vec<CpuAssignment>> {
1136 let recipients: Vec<CpuRecipient> = self
1137 .cells
1138 .values()
1139 .map(|info| {
1140 let weight = match cell_demands {
1141 Some(demands) => {
1142 let weight = *demands
1143 .get(&info.cell_id)
1144 .ok_or_else(|| {
1145 anyhow::anyhow!("Cell {} is missing from demands map", info.cell_id)
1146 })
1147 .context("building cell demand weights map")?;
1148 if weight < 0.0 {
1149 bail!(
1150 "Cell {} has negative demand weight {}",
1151 info.cell_id,
1152 weight
1153 );
1154 }
1155 weight
1156 }
1157 None => 1.0,
1158 };
1159 Ok(CpuRecipient {
1160 id: info.cell_id,
1161 weight,
1162 allowed: info.cpuset.clone().unwrap_or_else(|| self.all_cpus.clone()),
1163 claimed: info.cpuset.clone(),
1164 minimum: if info.cell_id == 0 {
1165 CpuMinimum {
1166 protected: 0,
1167 requested: self.cell0_min_cpus,
1168 }
1169 } else {
1170 CpuMinimum {
1171 protected: 1,
1172 requested: 0,
1173 }
1174 },
1175 })
1176 })
1177 .collect::<Result<Vec<_>>>()?;
1178
1179 let assignments = CpuManager::with_partitions(&self.all_cpus, &self.cpu_to_llc)
1180 .compute_assignments(&recipients, compute_borrowable)?;
1181
1182 if self.cell0_min_cpus > 0 {
1185 let mut workload_claims = Cpumask::new();
1186 for recipient in recipients.iter().filter(|recipient| recipient.id != 0) {
1187 if let Some(claimed) = &recipient.claimed {
1188 workload_claims = workload_claims.or(claimed);
1189 }
1190 }
1191 if assignments
1192 .iter()
1193 .find(|assignment| assignment.id == 0)
1194 .is_some_and(|assignment| assignment.primary.and(&workload_claims).weight() > 0)
1195 {
1196 self.enforced_holdout.store(true, Ordering::Relaxed);
1197 }
1198 }
1199
1200 Ok(assignments)
1201 }
1202
1203 pub fn get_cell_assignments(&self) -> Vec<(u64, u32)> {
1206 let mut assignments: Vec<_> = self
1207 .cells
1208 .values()
1209 .filter(|info| info.cell_id != 0)
1210 .map(|info| {
1211 (
1212 info.cgid.expect("BUG: non-zero cell missing cgid"),
1213 info.cell_id,
1214 )
1215 })
1216 .collect();
1217 assignments.sort_by_key(|(_cgid, cell_id)| *cell_id);
1218 assignments
1219 }
1220
1221 pub fn format_cell_config(&self, cpu_assignments: &[CpuAssignment]) -> String {
1224 let mut sorted: Vec<_> = cpu_assignments.iter().collect();
1225 sorted.sort_by_key(|a| a.id);
1226
1227 let mut parts = Vec::new();
1228 for assignment in sorted {
1229 let cpulist = assignment.primary.to_cpulist();
1230 if assignment.id == 0 {
1231 parts.push(format!("[0: {}]", cpulist));
1232 } else {
1233 let name = self
1235 .cells
1236 .values()
1237 .find(|info| info.cell_id == assignment.id)
1238 .and_then(|info| {
1239 info.cgroup_path
1240 .as_ref()
1241 .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
1242 })
1243 .unwrap_or_else(|| "?".to_string());
1244 parts.push(format!("[{}({}): {}]", assignment.id, name, cpulist));
1245 }
1246 }
1247 parts.join(" ")
1248 }
1249
1250 pub fn cgroup_path_for_cell(&self, cell_id: u32) -> String {
1254 self.cell_id_to_cgid
1255 .get(&cell_id)
1256 .and_then(|cgid| self.cells.get(cgid))
1257 .and_then(|info| info.cgroup_path.as_deref())
1258 .map(cgroup_root_relative)
1259 .unwrap_or_else(|| "/".to_string())
1260 }
1261
1262 pub fn refresh_cpusets(&mut self) -> Result<bool> {
1265 let mut changed = false;
1266 for info in self.cells.values_mut() {
1267 let Some(ref cgroup_path) = info.cgroup_path else {
1268 continue; };
1270 let new_cpuset = Self::read_cpuset(cgroup_path)
1271 .with_context(|| format!("reading cpuset for cgroup {}", cgroup_path.display()))?;
1272 if new_cpuset != info.cpuset {
1273 info!(
1274 "Cell {} cpuset changed: {:?} -> {:?} ({})",
1275 info.cell_id,
1276 info.cpuset.as_ref().map(|m| m.to_cpulist()),
1277 new_cpuset.as_ref().map(|m| m.to_cpulist()),
1278 cgroup_path.display(),
1279 );
1280 info.cpuset = new_cpuset;
1281 changed = true;
1282 }
1283 }
1284 Ok(changed)
1285 }
1286}
1287
1288impl AsFd for CellManager {
1289 fn as_fd(&self) -> BorrowedFd<'_> {
1290 self.inotify.as_fd()
1291 }
1292}
1293
1294#[cfg(test)]
1295impl CellManager {
1296 fn cell_count(&self) -> usize {
1299 self.cells.values().filter(|c| c.cell_id != 0).count()
1300 }
1301
1302 fn get_cell_ids(&self) -> Vec<u32> {
1305 self.cells
1306 .values()
1307 .filter(|c| c.cell_id != 0)
1308 .map(|c| c.cell_id)
1309 .collect()
1310 }
1311
1312 fn find_cell_by_name(&self, name: &str) -> Option<&CellInfo> {
1315 self.cells.values().filter(|c| c.cell_id != 0).find(|c| {
1316 c.cgroup_path
1317 .as_ref()
1318 .and_then(|p| p.file_name())
1319 .map(|n| n.to_str() == Some(name))
1320 .unwrap_or(false)
1321 })
1322 }
1323}
1324
1325#[cfg(test)]
1326#[allow(clippy::unwrap_used)]
1327mod tests {
1328 use super::*;
1329 use tempfile::TempDir;
1330
1331 fn cpumask_for_range(nr_cpus: usize) -> Cpumask {
1332 scx_utils::set_cpumask_test_width(nr_cpus);
1333 let mut mask = Cpumask::new();
1334 for cpu in 0..nr_cpus {
1335 mask.set_cpu(cpu).unwrap();
1336 }
1337 mask
1338 }
1339
1340 #[test]
1343 fn test_cgroup_path_for_cell() {
1344 let tmp = TempDir::new().unwrap();
1345 let child = tmp.path().join("container-a");
1346 std::fs::create_dir(&child).unwrap();
1347 let mgr = CellManager::new_with_path(
1348 tmp.path().to_path_buf(),
1349 256,
1350 cpumask_for_range(16),
1351 HashSet::new(),
1352 )
1353 .unwrap();
1354
1355 assert_eq!(mgr.cgroup_path_for_cell(0), "/");
1357
1358 let cell_id = mgr.find_cell_by_name("container-a").unwrap().cell_id;
1362 assert_eq!(
1363 mgr.cgroup_path_for_cell(cell_id),
1364 child.to_string_lossy().into_owned()
1365 );
1366
1367 assert_eq!(
1369 cgroup_root_relative(Path::new("/sys/fs/cgroup/test.slice/foobar")),
1370 "/test.slice/foobar"
1371 );
1372 }
1373
1374 fn cpumask_from_cpulist(nr_cpus: usize, cpulist: &str) -> Cpumask {
1375 scx_utils::set_cpumask_test_width(nr_cpus);
1376 Cpumask::from_cpulist(cpulist).unwrap()
1377 }
1378
1379 fn find_assignment(assignments: &[CpuAssignment], id: u32) -> &CpuAssignment {
1380 assignments.iter().find(|a| a.id == id).unwrap()
1381 }
1382
1383 #[test]
1386 fn test_scan_empty_directory() {
1387 let tmp = TempDir::new().unwrap();
1388 let mgr = CellManager::new_with_path(
1389 tmp.path().to_path_buf(),
1390 256,
1391 cpumask_for_range(16),
1392 HashSet::new(),
1393 )
1394 .unwrap();
1395
1396 assert_eq!(mgr.cell_count(), 0);
1397 }
1398
1399 #[test]
1400 fn test_scan_existing_subdirectories() {
1401 let tmp = TempDir::new().unwrap();
1402
1403 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1405 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1406
1407 let mgr = CellManager::new_with_path(
1408 tmp.path().to_path_buf(),
1409 256,
1410 cpumask_for_range(16),
1411 HashSet::new(),
1412 )
1413 .unwrap();
1414
1415 assert_eq!(mgr.cell_count(), 2);
1416
1417 let cell_ids = mgr.get_cell_ids();
1419 assert!(cell_ids.contains(&1));
1420 assert!(cell_ids.contains(&2));
1421 }
1422
1423 #[test]
1424 fn test_reconcile_detects_new_directories() {
1425 let tmp = TempDir::new().unwrap();
1426
1427 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1429 let mut mgr = CellManager::new_with_path(
1430 tmp.path().to_path_buf(),
1431 256,
1432 cpumask_for_range(16),
1433 HashSet::new(),
1434 )
1435 .unwrap();
1436 assert_eq!(mgr.cell_count(), 1);
1437
1438 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1440
1441 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1443 assert_eq!(new_cells.len(), 1);
1444 assert_eq!(destroyed_cells.len(), 0);
1445 assert_eq!(mgr.cell_count(), 2);
1446 }
1447
1448 #[test]
1449 fn test_reconcile_detects_removed_directories() {
1450 let tmp = TempDir::new().unwrap();
1451
1452 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1454 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1455 let mut mgr = CellManager::new_with_path(
1456 tmp.path().to_path_buf(),
1457 256,
1458 cpumask_for_range(16),
1459 HashSet::new(),
1460 )
1461 .unwrap();
1462 assert_eq!(mgr.cell_count(), 2);
1463
1464 std::fs::remove_dir(tmp.path().join("container-b")).unwrap();
1466
1467 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1469 assert_eq!(new_cells.len(), 0);
1470 assert_eq!(destroyed_cells.len(), 1);
1471 assert_eq!(mgr.cell_count(), 1);
1472 }
1473
1474 #[test]
1475 fn test_reconcile_replaces_reused_path_with_new_inode() {
1476 let tmp = TempDir::new().unwrap();
1477 let parked = TempDir::new().unwrap();
1478
1479 let original_path = tmp.path().join("container-a");
1480 std::fs::create_dir(&original_path).unwrap();
1481
1482 let mut mgr = CellManager::new_with_path(
1483 tmp.path().to_path_buf(),
1484 256,
1485 cpumask_for_range(16),
1486 HashSet::new(),
1487 )
1488 .unwrap();
1489 assert_eq!(mgr.cell_count(), 1);
1490
1491 let old_info = mgr.find_cell_by_name("container-a").unwrap();
1492 let old_cell_id = old_info.cell_id;
1493 let old_cgid = old_info.cgid.unwrap();
1494
1495 std::fs::rename(&original_path, parked.path().join("container-a-old")).unwrap();
1498 std::fs::create_dir(&original_path).unwrap();
1499
1500 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1501
1502 assert_eq!(new_cells.len(), 1);
1503 assert_eq!(destroyed_cells, vec![old_cell_id]);
1504 assert_eq!(mgr.cell_count(), 1);
1505
1506 let new_info = mgr.find_cell_by_name("container-a").unwrap();
1507 assert_eq!(new_info.cell_id, old_cell_id);
1508 assert_ne!(new_info.cgid.unwrap(), old_cgid);
1509 }
1510
1511 #[test]
1512 fn test_cell_id_reuse_after_destruction() {
1513 let tmp = TempDir::new().unwrap();
1514
1515 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1517 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1518 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1519
1520 let mut mgr = CellManager::new_with_path(
1521 tmp.path().to_path_buf(),
1522 256,
1523 cpumask_for_range(16),
1524 HashSet::new(),
1525 )
1526 .unwrap();
1527
1528 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1530 let cell2_id = cell2_info.cell_id;
1531
1532 std::fs::remove_dir(tmp.path().join("cell2")).unwrap();
1534 mgr.reconcile_cells().unwrap();
1535
1536 std::fs::create_dir(tmp.path().join("cell4")).unwrap();
1538 mgr.reconcile_cells().unwrap();
1539
1540 let cell4_info = mgr.find_cell_by_name("cell4").unwrap();
1541 assert_eq!(cell4_info.cell_id, cell2_id);
1542 }
1543
1544 #[test]
1547 fn test_cpu_assignments_no_cells() {
1548 let tmp = TempDir::new().unwrap();
1549 let mgr = CellManager::new_with_path(
1550 tmp.path().to_path_buf(),
1551 256,
1552 cpumask_for_range(16),
1553 HashSet::new(),
1554 )
1555 .unwrap();
1556
1557 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1558
1559 assert_eq!(assignments.len(), 1);
1561 assert_eq!(assignments[0].id, 0);
1562 assert_eq!(assignments[0].primary.weight(), 16);
1563 }
1564
1565 #[test]
1566 fn test_cpu_assignments_proportional() {
1567 let tmp = TempDir::new().unwrap();
1568 std::fs::create_dir(tmp.path().join("container")).unwrap();
1569
1570 let mgr = CellManager::new_with_path(
1571 tmp.path().to_path_buf(),
1572 256,
1573 cpumask_for_range(16),
1574 HashSet::new(),
1575 )
1576 .unwrap();
1577 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1578
1579 assert_eq!(assignments.len(), 2);
1581
1582 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
1583 let cell1 = assignments.iter().find(|a| a.id == 1).unwrap();
1584
1585 assert_eq!(cell0.primary.weight(), 8);
1586 assert_eq!(cell1.primary.weight(), 8);
1587 }
1588
1589 #[test]
1590 fn test_cpu_assignments_remainder_to_cell0() {
1591 let tmp = TempDir::new().unwrap();
1592 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1593 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1594
1595 let mgr = CellManager::new_with_path(
1596 tmp.path().to_path_buf(),
1597 256,
1598 cpumask_for_range(10),
1599 HashSet::new(),
1600 )
1601 .unwrap();
1602 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1603
1604 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
1606 assert_eq!(cell0.primary.weight(), 4); }
1608
1609 #[test]
1610 fn test_cpu_assignments_too_many_cells() {
1611 let tmp = TempDir::new().unwrap();
1612
1613 for i in 1..=5 {
1615 std::fs::create_dir(tmp.path().join(format!("cell{}", i))).unwrap();
1616 }
1617
1618 let mgr = CellManager::new_with_path(
1620 tmp.path().to_path_buf(),
1621 256,
1622 cpumask_for_range(4),
1623 HashSet::new(),
1624 )
1625 .unwrap();
1626 let result = mgr.compute_cpu_assignments(false);
1627
1628 assert!(result.is_err());
1629 let err_msg = format!("{:#}", result.unwrap_err());
1630 assert!(
1631 err_msg.contains("Not enough CPUs"),
1632 "Expected 'Not enough CPUs' error, got: {}",
1633 err_msg
1634 );
1635 }
1636
1637 #[test]
1638 fn test_cpu_assignments_with_cpusets() {
1639 let tmp = TempDir::new().unwrap();
1640
1641 let cell1_path = tmp.path().join("cell1");
1643 std::fs::create_dir(&cell1_path).unwrap();
1644 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1645
1646 let cell2_path = tmp.path().join("cell2");
1647 std::fs::create_dir(&cell2_path).unwrap();
1648 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1649
1650 let mgr = CellManager::new_with_path(
1651 tmp.path().to_path_buf(),
1652 256,
1653 cpumask_for_range(16),
1654 HashSet::new(),
1655 )
1656 .unwrap();
1657 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1658
1659 assert_eq!(assignments.len(), 3);
1661
1662 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1664 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1665
1666 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
1667 let cell1 = assignments
1668 .iter()
1669 .find(|a| a.id == cell1_info.cell_id)
1670 .unwrap();
1671 let cell2 = assignments
1672 .iter()
1673 .find(|a| a.id == cell2_info.cell_id)
1674 .unwrap();
1675
1676 assert_eq!(cell1.primary.weight(), 4);
1678 for cpu in 0..4 {
1679 assert!(cell1.primary.test_cpu(cpu));
1680 }
1681
1682 assert_eq!(cell2.primary.weight(), 4);
1684 for cpu in 8..12 {
1685 assert!(cell2.primary.test_cpu(cpu));
1686 }
1687
1688 assert_eq!(cell0.primary.weight(), 8);
1690 for cpu in 4..8 {
1691 assert!(cell0.primary.test_cpu(cpu));
1692 }
1693 for cpu in 12..16 {
1694 assert!(cell0.primary.test_cpu(cpu));
1695 }
1696 }
1697
1698 #[test]
1699 fn test_cpu_assignments_cpusets_cover_all_cpus() {
1700 let tmp = TempDir::new().unwrap();
1701
1702 let cell1_path = tmp.path().join("cell1");
1706 std::fs::create_dir(&cell1_path).unwrap();
1707 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1708
1709 let cell2_path = tmp.path().join("cell2");
1710 std::fs::create_dir(&cell2_path).unwrap();
1711 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1712
1713 let mgr = CellManager::new_with_path(
1714 tmp.path().to_path_buf(),
1715 256,
1716 cpumask_for_range(16),
1717 HashSet::new(),
1718 )
1719 .unwrap();
1720 let result = mgr.compute_cpu_assignments(false);
1721
1722 assert!(result.is_err());
1724 let err = result.unwrap_err();
1725 let err_msg = format!("{:#}", err);
1726 assert!(
1727 err_msg.contains("Recipient 0 has 0 CPUs assigned, below required minimum 1"),
1728 "Expected recipient minimum error, got: {}",
1729 err_msg
1730 );
1731 }
1732
1733 #[test]
1734 fn test_cpu_assignments_cpusets_cover_all_cpus_with_holdout() {
1735 let tmp = TempDir::new().unwrap();
1736
1737 let cell1_path = tmp.path().join("cell1");
1742 std::fs::create_dir(&cell1_path).unwrap();
1743 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1744
1745 let cell2_path = tmp.path().join("cell2");
1746 std::fs::create_dir(&cell2_path).unwrap();
1747 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1748
1749 let mgr = CellManager::new_with_path_opts(
1750 tmp.path().to_path_buf(),
1751 256,
1752 cpumask_for_range(16),
1753 HashSet::new(),
1754 1,
1755 HashMap::new(),
1756 )
1757 .unwrap();
1758 let assignments = mgr
1759 .compute_cpu_assignments(false)
1760 .expect("holdout should populate cell 0");
1761
1762 let cell0 = assignments
1763 .iter()
1764 .find(|a| a.id == 0)
1765 .expect("cell 0 present");
1766 assert_eq!(cell0.primary.weight(), 1);
1767
1768 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1769 assert_eq!(total, 16);
1770
1771 for assignment in &assignments {
1774 assert!(
1775 assignment.primary.weight() >= 1,
1776 "cell {} starved after holdout: {:?}",
1777 assignment.id,
1778 assignment.primary
1779 );
1780 }
1781 let donor_weights: Vec<usize> = assignments
1782 .iter()
1783 .filter(|a| a.id != 0)
1784 .map(|a| a.primary.weight())
1785 .collect();
1786 assert!(donor_weights.contains(&7));
1787 assert!(donor_weights.contains(&8));
1788 }
1789
1790 #[test]
1791 fn test_cpu_assignments_holdout_takes_from_largest_cell() {
1792 let tmp = TempDir::new().unwrap();
1793
1794 let cell1_path = tmp.path().join("cell1");
1800 std::fs::create_dir(&cell1_path).unwrap();
1801 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1802
1803 let cell2_path = tmp.path().join("cell2");
1804 std::fs::create_dir(&cell2_path).unwrap();
1805 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1806
1807 let cell3_path = tmp.path().join("cell3");
1808 std::fs::create_dir(&cell3_path).unwrap();
1809 std::fs::write(cell3_path.join("cpuset.cpus"), "12-15\n").unwrap();
1810
1811 let cpu_to_llc: HashMap<usize, usize> = (0..16usize).map(|cpu| (cpu, cpu / 8)).collect();
1812
1813 let mgr = CellManager::new_with_path_opts(
1814 tmp.path().to_path_buf(),
1815 256,
1816 cpumask_for_range(16),
1817 HashSet::new(),
1818 1,
1819 cpu_to_llc,
1820 )
1821 .unwrap();
1822 let assignments = mgr
1823 .compute_cpu_assignments(false)
1824 .expect("holdout should populate cell 0");
1825
1826 let cell0 = assignments
1827 .iter()
1828 .find(|a| a.id == 0)
1829 .expect("cell 0 present");
1830 assert_eq!(cell0.primary.weight(), 1);
1831 assert!(
1832 cell0.primary.test_cpu(0),
1833 "holdout should take CPU 0 (lowest CPU of the largest cell), got {:?}",
1834 cell0.primary
1835 );
1836 assert!(
1837 mgr.enforced_holdout(),
1838 "stealing a claimed CPU must set enforced_holdout"
1839 );
1840
1841 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1842 assert_eq!(total, 16);
1843
1844 let mut child_weights: Vec<usize> = assignments
1847 .iter()
1848 .filter(|a| a.id != 0)
1849 .map(|a| a.primary.weight())
1850 .collect();
1851 child_weights.sort_unstable();
1852 assert_eq!(child_weights, vec![4, 4, 7]);
1853
1854 for a in assignments.iter().filter(|a| a.id != 0) {
1856 assert!(
1857 !a.primary.test_cpu(0),
1858 "cell {} still holds held-out CPU 0",
1859 a.id
1860 );
1861 }
1862 }
1863
1864 #[test]
1865 fn test_cpu_assignments_holdout_steals_evenly_across_cells() {
1866 let tmp = TempDir::new().unwrap();
1867
1868 for (name, range) in [("cell1", "0-7"), ("cell2", "8-15"), ("cell3", "16-23")] {
1874 let p = tmp.path().join(name);
1875 std::fs::create_dir(&p).unwrap();
1876 std::fs::write(p.join("cpuset.cpus"), format!("{range}\n")).unwrap();
1877 }
1878
1879 let mgr = CellManager::new_with_path_opts(
1880 tmp.path().to_path_buf(),
1881 256,
1882 cpumask_for_range(24),
1883 HashSet::new(),
1884 6,
1885 HashMap::new(),
1886 )
1887 .unwrap();
1888 let assignments = mgr
1889 .compute_cpu_assignments(false)
1890 .expect("holdout should populate cell 0");
1891
1892 let cell0 = assignments
1893 .iter()
1894 .find(|a| a.id == 0)
1895 .expect("cell 0 present");
1896 assert_eq!(
1897 cell0.primary.weight(),
1898 6,
1899 "cell 0 holds the six reserved CPUs"
1900 );
1901
1902 let mut donor_weights: Vec<usize> = assignments
1905 .iter()
1906 .filter(|a| a.id != 0)
1907 .map(|a| a.primary.weight())
1908 .collect();
1909 donor_weights.sort_unstable();
1910 assert_eq!(
1911 donor_weights,
1912 vec![6, 6, 6],
1913 "holdout should steal evenly (two from each cell), got {donor_weights:?}"
1914 );
1915
1916 let count_in =
1918 |lo: usize, hi: usize| (lo..hi).filter(|&c| cell0.primary.test_cpu(c)).count();
1919 assert_eq!(count_in(0, 8), 2, "two CPUs taken from cell1");
1920 assert_eq!(count_in(8, 16), 2, "two CPUs taken from cell2");
1921 assert_eq!(count_in(16, 24), 2, "two CPUs taken from cell3");
1922 }
1923
1924 #[test]
1925 fn test_cpu_assignments_holdout_never_starves_a_child() {
1926 let tmp = TempDir::new().unwrap();
1927
1928 let cell1_path = tmp.path().join("cell1");
1935 std::fs::create_dir(&cell1_path).unwrap();
1936 std::fs::write(cell1_path.join("cpuset.cpus"), "0-1\n").unwrap();
1937
1938 let cell2_path = tmp.path().join("cell2");
1939 std::fs::create_dir(&cell2_path).unwrap();
1940 std::fs::write(cell2_path.join("cpuset.cpus"), "2-3\n").unwrap();
1941
1942 let mgr = CellManager::new_with_path_opts(
1943 tmp.path().to_path_buf(),
1944 256,
1945 cpumask_for_range(4),
1946 HashSet::new(),
1947 3,
1948 HashMap::new(),
1949 )
1950 .unwrap();
1951 let assignments = mgr
1952 .compute_cpu_assignments(false)
1953 .expect("holdout must not starve a child cell");
1954
1955 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1956 assert_eq!(total, 4);
1957
1958 for a in assignments.iter().filter(|a| a.id != 0) {
1960 assert!(
1961 a.primary.weight() >= 1,
1962 "cell {} starved by the holdout: {:?}",
1963 a.id,
1964 a.primary
1965 );
1966 }
1967
1968 let cell0 = assignments
1970 .iter()
1971 .find(|a| a.id == 0)
1972 .expect("cell 0 present");
1973 assert_eq!(
1974 cell0.primary.weight(),
1975 2,
1976 "holdout should cap at total - num_children, got {:?}",
1977 cell0.primary
1978 );
1979 }
1980
1981 #[test]
1982 fn test_cpu_assignments_holdout_overlapping_cpusets_no_starvation() {
1983 let tmp = TempDir::new().unwrap();
1984
1985 let cell1_path = tmp.path().join("cell1");
1994 std::fs::create_dir(&cell1_path).unwrap();
1995 std::fs::write(cell1_path.join("cpuset.cpus"), "0\n").unwrap();
1996
1997 let cell2_path = tmp.path().join("cell2");
1998 std::fs::create_dir(&cell2_path).unwrap();
1999 std::fs::write(cell2_path.join("cpuset.cpus"), "0-1\n").unwrap();
2000
2001 let mgr = CellManager::new_with_path_opts(
2002 tmp.path().to_path_buf(),
2003 256,
2004 cpumask_for_range(5),
2005 HashSet::new(),
2006 4,
2007 HashMap::new(),
2008 )
2009 .unwrap();
2010 let assignments = mgr
2011 .compute_cpu_assignments(false)
2012 .expect("holdout must not starve a cell that shares a contested CPU");
2013 assert!(
2017 !mgr.enforced_holdout(),
2018 "entering the steal loop without taking a claimed CPU must not set \
2019 enforced_holdout"
2020 );
2021
2022 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
2023 assert_eq!(total, 5);
2024
2025 for a in assignments.iter().filter(|a| a.id != 0) {
2027 assert!(
2028 a.primary.weight() >= 1,
2029 "cell {} starved by the holdout: {:?}",
2030 a.id,
2031 a.primary
2032 );
2033 }
2034
2035 let cell0 = assignments
2038 .iter()
2039 .find(|a| a.id == 0)
2040 .expect("cell 0 present");
2041 assert_eq!(
2042 cell0.primary.weight(),
2043 3,
2044 "holdout should cap at the 3 unclaimed CPUs, got {:?}",
2045 cell0.primary
2046 );
2047 }
2048
2049 #[test]
2050 fn test_cpu_assignments_single_cpuset() {
2051 let tmp = TempDir::new().unwrap();
2052
2053 let cell1_path = tmp.path().join("cell1");
2055 std::fs::create_dir(&cell1_path).unwrap();
2056 std::fs::write(cell1_path.join("cpuset.cpus"), "0,2,4,6\n").unwrap();
2057
2058 let mgr = CellManager::new_with_path(
2059 tmp.path().to_path_buf(),
2060 256,
2061 cpumask_for_range(8),
2062 HashSet::new(),
2063 )
2064 .unwrap();
2065 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2066
2067 assert_eq!(assignments.len(), 2);
2068
2069 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2070 let cell1 = assignments.iter().find(|a| a.id != 0).unwrap();
2071
2072 assert_eq!(cell1.primary.weight(), 4);
2074 for cpu in [0, 2, 4, 6] {
2075 assert!(cell1.primary.test_cpu(cpu));
2076 }
2077
2078 assert_eq!(cell0.primary.weight(), 4);
2080 for cpu in [1, 3, 5, 7] {
2081 assert!(cell0.primary.test_cpu(cpu));
2082 }
2083 }
2084
2085 #[test]
2086 fn test_cpuset_parsing_from_file() {
2087 let tmp = TempDir::new().unwrap();
2088
2089 let cell_path = tmp.path().join("cell1");
2091 std::fs::create_dir(&cell_path).unwrap();
2092 std::fs::write(cell_path.join("cpuset.cpus"), "0-3,8-11,16\n").unwrap();
2093
2094 let mgr = CellManager::new_with_path(
2095 tmp.path().to_path_buf(),
2096 256,
2097 cpumask_for_range(32),
2098 HashSet::new(),
2099 )
2100 .unwrap();
2101
2102 let cell_info = mgr.find_cell_by_name("cell1").unwrap();
2104 let cpuset = cell_info.cpuset.as_ref().unwrap();
2105
2106 assert_eq!(cpuset.weight(), 9); for cpu in 0..4 {
2108 assert!(cpuset.test_cpu(cpu));
2109 }
2110 for cpu in 8..12 {
2111 assert!(cpuset.test_cpu(cpu));
2112 }
2113 assert!(cpuset.test_cpu(16));
2114 }
2115
2116 #[test]
2117 fn test_cpu_assignments_mixed_cpuset_and_no_cpuset() {
2118 let tmp = TempDir::new().unwrap();
2119
2120 let cell1_path = tmp.path().join("cell1");
2122 std::fs::create_dir(&cell1_path).unwrap();
2123 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
2124
2125 let cell2_path = tmp.path().join("cell2");
2127 std::fs::create_dir(&cell2_path).unwrap();
2128
2129 let mgr = CellManager::new_with_path(
2130 tmp.path().to_path_buf(),
2131 256,
2132 cpumask_for_range(16),
2133 HashSet::new(),
2134 )
2135 .unwrap();
2136
2137 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2139 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2140 assert!(cell1_info.cpuset.is_some());
2141 assert!(cell2_info.cpuset.is_none());
2142
2143 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2144
2145 assert_eq!(assignments.len(), 3);
2151
2152 let cell1_assignment = assignments
2154 .iter()
2155 .find(|a| a.id == cell1_info.cell_id)
2156 .unwrap();
2157 assert_eq!(cell1_assignment.primary.weight(), 4);
2158
2159 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2161 assert_eq!(cell0.primary.weight(), 7);
2162
2163 let cell2_assignment = assignments
2165 .iter()
2166 .find(|a| a.id == cell2_info.cell_id)
2167 .unwrap();
2168 assert_eq!(cell2_assignment.primary.weight(), 5);
2169 }
2170
2171 #[test]
2174 fn test_cpu_assignments_partial_overlap() {
2175 let tmp = TempDir::new().unwrap();
2176
2177 let cell_a_path = tmp.path().join("cell_a");
2179 std::fs::create_dir(&cell_a_path).unwrap();
2180 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2181
2182 let cell_b_path = tmp.path().join("cell_b");
2183 std::fs::create_dir(&cell_b_path).unwrap();
2184 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
2185
2186 let mgr = CellManager::new_with_path(
2187 tmp.path().to_path_buf(),
2188 256,
2189 cpumask_for_range(16),
2190 HashSet::new(),
2191 )
2192 .unwrap();
2193 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2194
2195 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2196 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2197
2198 let cell_a = assignments
2199 .iter()
2200 .find(|a| a.id == cell_a_info.cell_id)
2201 .unwrap();
2202 let cell_b = assignments
2203 .iter()
2204 .find(|a| a.id == cell_b_info.cell_id)
2205 .unwrap();
2206 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2207
2208 assert_eq!(cell_a.primary.weight(), 6);
2212 assert_eq!(cell_b.primary.weight(), 6);
2213 assert_eq!(cell0.primary.weight(), 4);
2214
2215 for cpu in 0..4 {
2217 assert!(
2218 cell_a.primary.test_cpu(cpu),
2219 "CPU {} should be in cell_a",
2220 cpu
2221 );
2222 }
2223 for cpu in 8..12 {
2224 assert!(
2225 cell_b.primary.test_cpu(cpu),
2226 "CPU {} should be in cell_b",
2227 cpu
2228 );
2229 }
2230 for cpu in 12..16 {
2231 assert!(
2232 cell0.primary.test_cpu(cpu),
2233 "CPU {} should be in cell0",
2234 cpu
2235 );
2236 }
2237
2238 let cell_a_contested: Vec<_> = (4..8).filter(|&cpu| cell_a.primary.test_cpu(cpu)).collect();
2240 let cell_b_contested: Vec<_> = (4..8).filter(|&cpu| cell_b.primary.test_cpu(cpu)).collect();
2241 assert_eq!(cell_a_contested.len(), 2);
2242 assert_eq!(cell_b_contested.len(), 2);
2243
2244 for cpu in 0..16 {
2246 let mut count = 0;
2247 if cell_a.primary.test_cpu(cpu) {
2248 count += 1;
2249 }
2250 if cell_b.primary.test_cpu(cpu) {
2251 count += 1;
2252 }
2253 if cell0.primary.test_cpu(cpu) {
2254 count += 1;
2255 }
2256 assert!(count <= 1, "CPU {} is assigned to {} cells", cpu, count);
2257 }
2258 }
2259
2260 #[test]
2261 fn test_cpu_assignments_three_way_overlap() {
2262 let tmp = TempDir::new().unwrap();
2263
2264 let cell_a_path = tmp.path().join("cell_a");
2266 std::fs::create_dir(&cell_a_path).unwrap();
2267 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-5\n").unwrap();
2268
2269 let cell_b_path = tmp.path().join("cell_b");
2270 std::fs::create_dir(&cell_b_path).unwrap();
2271 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-5\n").unwrap();
2272
2273 let cell_c_path = tmp.path().join("cell_c");
2274 std::fs::create_dir(&cell_c_path).unwrap();
2275 std::fs::write(cell_c_path.join("cpuset.cpus"), "0-5\n").unwrap();
2276
2277 let mgr = CellManager::new_with_path(
2278 tmp.path().to_path_buf(),
2279 256,
2280 cpumask_for_range(12),
2281 HashSet::new(),
2282 )
2283 .unwrap();
2284 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2285
2286 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2287 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2288 let cell_c_info = mgr.find_cell_by_name("cell_c").unwrap();
2289
2290 let cell_a = assignments
2291 .iter()
2292 .find(|a| a.id == cell_a_info.cell_id)
2293 .unwrap();
2294 let cell_b = assignments
2295 .iter()
2296 .find(|a| a.id == cell_b_info.cell_id)
2297 .unwrap();
2298 let cell_c = assignments
2299 .iter()
2300 .find(|a| a.id == cell_c_info.cell_id)
2301 .unwrap();
2302 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2303
2304 assert_eq!(cell_a.primary.weight(), 2);
2306 assert_eq!(cell_b.primary.weight(), 2);
2307 assert_eq!(cell_c.primary.weight(), 2);
2308
2309 assert_eq!(cell0.primary.weight(), 6);
2311 for cpu in 6..12 {
2312 assert!(cell0.primary.test_cpu(cpu));
2313 }
2314
2315 let total_contested: usize = (0..6)
2317 .filter(|&cpu| {
2318 cell_a.primary.test_cpu(cpu)
2319 || cell_b.primary.test_cpu(cpu)
2320 || cell_c.primary.test_cpu(cpu)
2321 })
2322 .count();
2323 assert_eq!(total_contested, 6);
2324 }
2325
2326 #[test]
2327 fn test_cpu_assignments_odd_contested_count() {
2328 let tmp = TempDir::new().unwrap();
2329
2330 let cell_a_path = tmp.path().join("cell_a");
2332 std::fs::create_dir(&cell_a_path).unwrap();
2333 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-2\n").unwrap();
2334
2335 let cell_b_path = tmp.path().join("cell_b");
2336 std::fs::create_dir(&cell_b_path).unwrap();
2337 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-2\n").unwrap();
2338
2339 let mgr = CellManager::new_with_path(
2340 tmp.path().to_path_buf(),
2341 256,
2342 cpumask_for_range(8),
2343 HashSet::new(),
2344 )
2345 .unwrap();
2346 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2347
2348 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2349 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2350
2351 let cell_a = assignments
2352 .iter()
2353 .find(|a| a.id == cell_a_info.cell_id)
2354 .unwrap();
2355 let cell_b = assignments
2356 .iter()
2357 .find(|a| a.id == cell_b_info.cell_id)
2358 .unwrap();
2359
2360 let total = cell_a.primary.weight() + cell_b.primary.weight();
2363 assert_eq!(total, 3);
2364 assert!(cell_a.primary.weight() >= 1 && cell_a.primary.weight() <= 2);
2365 assert!(cell_b.primary.weight() >= 1 && cell_b.primary.weight() <= 2);
2366
2367 for cpu in 0..3 {
2369 let a_has = cell_a.primary.test_cpu(cpu);
2370 let b_has = cell_b.primary.test_cpu(cpu);
2371 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
2372 }
2373 }
2374
2375 #[test]
2376 fn test_cpu_assignments_complete_overlap() {
2377 let tmp = TempDir::new().unwrap();
2378
2379 let cell_a_path = tmp.path().join("cell_a");
2381 std::fs::create_dir(&cell_a_path).unwrap();
2382 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2383
2384 let cell_b_path = tmp.path().join("cell_b");
2385 std::fs::create_dir(&cell_b_path).unwrap();
2386 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-7\n").unwrap();
2387
2388 let mgr = CellManager::new_with_path(
2389 tmp.path().to_path_buf(),
2390 256,
2391 cpumask_for_range(16),
2392 HashSet::new(),
2393 )
2394 .unwrap();
2395 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2396
2397 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2398 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2399
2400 let cell_a = assignments
2401 .iter()
2402 .find(|a| a.id == cell_a_info.cell_id)
2403 .unwrap();
2404 let cell_b = assignments
2405 .iter()
2406 .find(|a| a.id == cell_b_info.cell_id)
2407 .unwrap();
2408 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2409
2410 assert_eq!(cell_a.primary.weight(), 4);
2412 assert_eq!(cell_b.primary.weight(), 4);
2413
2414 assert_eq!(cell0.primary.weight(), 8);
2416 for cpu in 8..16 {
2417 assert!(cell0.primary.test_cpu(cpu));
2418 }
2419
2420 for cpu in 0..8 {
2422 let a_has = cell_a.primary.test_cpu(cpu);
2423 let b_has = cell_b.primary.test_cpu(cpu);
2424 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
2425 }
2426 }
2427
2428 #[test]
2429 fn test_cpu_assignments_no_overlap() {
2430 let tmp = TempDir::new().unwrap();
2432
2433 let cell_a_path = tmp.path().join("cell_a");
2434 std::fs::create_dir(&cell_a_path).unwrap();
2435 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-3\n").unwrap();
2436
2437 let cell_b_path = tmp.path().join("cell_b");
2438 std::fs::create_dir(&cell_b_path).unwrap();
2439 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-7\n").unwrap();
2440
2441 let mgr = CellManager::new_with_path(
2442 tmp.path().to_path_buf(),
2443 256,
2444 cpumask_for_range(16),
2445 HashSet::new(),
2446 )
2447 .unwrap();
2448 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2449
2450 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2451 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2452
2453 let cell_a = assignments
2454 .iter()
2455 .find(|a| a.id == cell_a_info.cell_id)
2456 .unwrap();
2457 let cell_b = assignments
2458 .iter()
2459 .find(|a| a.id == cell_b_info.cell_id)
2460 .unwrap();
2461 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2462
2463 assert_eq!(cell_a.primary.weight(), 4);
2465 for cpu in 0..4 {
2466 assert!(cell_a.primary.test_cpu(cpu));
2467 }
2468
2469 assert_eq!(cell_b.primary.weight(), 4);
2470 for cpu in 4..8 {
2471 assert!(cell_b.primary.test_cpu(cpu));
2472 }
2473
2474 assert_eq!(cell0.primary.weight(), 8);
2476 for cpu in 8..16 {
2477 assert!(cell0.primary.test_cpu(cpu));
2478 }
2479 }
2480
2481 #[test]
2484 fn test_format_cell_config_only_cell0() {
2485 let tmp = TempDir::new().unwrap();
2486 let mgr = CellManager::new_with_path(
2487 tmp.path().to_path_buf(),
2488 256,
2489 cpumask_for_range(8),
2490 HashSet::new(),
2491 )
2492 .unwrap();
2493
2494 let mut mask = Cpumask::new();
2495 for cpu in 0..8 {
2496 mask.set_cpu(cpu).unwrap();
2497 }
2498
2499 let assignments = vec![CpuAssignment {
2500 id: 0,
2501 primary: mask,
2502 borrowable: None,
2503 }];
2504 let result = mgr.format_cell_config(&assignments);
2505
2506 assert_eq!(result, "[0: 0-7]");
2507 }
2508
2509 #[test]
2510 fn test_format_cell_config_with_cells() {
2511 let tmp = TempDir::new().unwrap();
2512 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2513
2514 let mgr = CellManager::new_with_path(
2515 tmp.path().to_path_buf(),
2516 256,
2517 cpumask_for_range(16),
2518 HashSet::new(),
2519 )
2520 .unwrap();
2521
2522 let mut mask0 = Cpumask::new();
2523 for cpu in 0..8 {
2524 mask0.set_cpu(cpu).unwrap();
2525 }
2526
2527 let mut mask1 = Cpumask::new();
2528 for cpu in 8..16 {
2529 mask1.set_cpu(cpu).unwrap();
2530 }
2531
2532 let assignments = vec![
2533 CpuAssignment {
2534 id: 0,
2535 primary: mask0,
2536 borrowable: None,
2537 },
2538 CpuAssignment {
2539 id: 1,
2540 primary: mask1,
2541 borrowable: None,
2542 },
2543 ];
2544 let result = mgr.format_cell_config(&assignments);
2545
2546 assert_eq!(result, "[0: 0-7] [1(container-a): 8-15]");
2547 }
2548
2549 #[test]
2552 fn test_cell_id_exhaustion() {
2553 let tmp = TempDir::new().unwrap();
2554
2555 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2558 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2559
2560 let mut mgr = CellManager::new_with_path(
2561 tmp.path().to_path_buf(),
2562 3,
2563 cpumask_for_range(16),
2564 HashSet::new(),
2565 )
2566 .unwrap();
2567 assert_eq!(mgr.cell_count(), 2); std::fs::create_dir(tmp.path().join("cell3")).unwrap();
2571 let result = mgr.reconcile_cells();
2572
2573 assert!(result.is_err());
2574 let err = result.unwrap_err();
2575 let err_chain = format!("{:#}", err);
2576 assert!(
2577 err_chain.contains("Cell ID space exhausted"),
2578 "Expected exhaustion error, got: {}",
2579 err_chain
2580 );
2581 }
2582
2583 #[test]
2584 fn test_cell_id_reuse_prevents_exhaustion() {
2585 let tmp = TempDir::new().unwrap();
2586
2587 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2589 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2590
2591 let mut mgr = CellManager::new_with_path(
2592 tmp.path().to_path_buf(),
2593 3,
2594 cpumask_for_range(16),
2595 HashSet::new(),
2596 )
2597 .unwrap();
2598 assert_eq!(mgr.cell_count(), 2);
2599
2600 std::fs::remove_dir(tmp.path().join("cell1")).unwrap();
2602 mgr.reconcile_cells().unwrap();
2603 assert_eq!(mgr.cell_count(), 1);
2604
2605 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
2607 let result = mgr.reconcile_cells();
2608 assert!(result.is_ok());
2609 assert_eq!(mgr.cell_count(), 2);
2610 }
2611
2612 #[test]
2615 fn test_scan_excludes_named_cgroups() {
2616 let tmp = TempDir::new().unwrap();
2617
2618 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2619 std::fs::create_dir(tmp.path().join("systemd-workaround.service")).unwrap();
2620 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
2621
2622 let exclude = HashSet::from(["systemd-workaround.service".to_string()]);
2623 let mgr = CellManager::new_with_path(
2624 tmp.path().to_path_buf(),
2625 256,
2626 cpumask_for_range(16),
2627 exclude,
2628 )
2629 .unwrap();
2630
2631 assert_eq!(mgr.cell_count(), 2);
2633 assert!(mgr.find_cell_by_name("container-a").is_some());
2634 assert!(mgr.find_cell_by_name("container-b").is_some());
2635 assert!(mgr
2636 .find_cell_by_name("systemd-workaround.service")
2637 .is_none());
2638 }
2639
2640 #[test]
2641 fn test_reconcile_excludes_named_cgroups() {
2642 let tmp = TempDir::new().unwrap();
2643
2644 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2645
2646 let exclude = HashSet::from(["ignored-service".to_string()]);
2647 let mut mgr = CellManager::new_with_path(
2648 tmp.path().to_path_buf(),
2649 256,
2650 cpumask_for_range(16),
2651 exclude,
2652 )
2653 .unwrap();
2654 assert_eq!(mgr.cell_count(), 1);
2655
2656 std::fs::create_dir(tmp.path().join("ignored-service")).unwrap();
2658 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
2659 assert_eq!(new_cells.len(), 0);
2660 assert_eq!(destroyed_cells.len(), 0);
2661 assert_eq!(mgr.cell_count(), 1);
2662
2663 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
2665 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
2666 assert_eq!(new_cells.len(), 1);
2667 assert_eq!(destroyed_cells.len(), 0);
2668 assert_eq!(mgr.cell_count(), 2);
2669 }
2670
2671 #[test]
2674 fn test_borrowable_cpumasks_basic() {
2675 let tmp = TempDir::new().unwrap();
2676
2677 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2679 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2680
2681 let mgr = CellManager::new_with_path(
2682 tmp.path().to_path_buf(),
2683 256,
2684 cpumask_for_range(16),
2685 HashSet::new(),
2686 )
2687 .unwrap();
2688 let assignments = mgr.compute_cpu_assignments(true).unwrap();
2689
2690 for assignment in &assignments {
2692 let borrow_mask = assignment.borrowable.as_ref().unwrap();
2693 let overlap = borrow_mask.and(&assignment.primary);
2695 assert_eq!(
2696 overlap.weight(),
2697 0,
2698 "Cell {} borrowable overlaps with primary",
2699 assignment.id
2700 );
2701 let union = borrow_mask.or(&assignment.primary);
2703 assert_eq!(
2704 union.weight(),
2705 16,
2706 "Cell {} union doesn't cover all CPUs",
2707 assignment.id
2708 );
2709 }
2710 }
2711
2712 #[test]
2713 fn test_borrowable_cpumasks_no_overlap() {
2714 let tmp = TempDir::new().unwrap();
2715
2716 let cell1_path = tmp.path().join("cell1");
2717 std::fs::create_dir(&cell1_path).unwrap();
2718 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
2719
2720 let cell2_path = tmp.path().join("cell2");
2721 std::fs::create_dir(&cell2_path).unwrap();
2722 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
2723
2724 let mgr = CellManager::new_with_path(
2725 tmp.path().to_path_buf(),
2726 256,
2727 cpumask_for_range(16),
2728 HashSet::new(),
2729 )
2730 .unwrap();
2731 let assignments = mgr.compute_cpu_assignments(true).unwrap();
2732
2733 for assignment in &assignments {
2735 let borrow_mask = assignment.borrowable.as_ref().unwrap();
2736 let overlap = borrow_mask.and(&assignment.primary);
2737 assert_eq!(
2738 overlap.weight(),
2739 0,
2740 "Cell {} borrowable overlaps with primary",
2741 assignment.id
2742 );
2743 }
2744 }
2745
2746 #[test]
2747 fn test_borrowable_cpumasks_respects_cpuset() {
2748 let tmp = TempDir::new().unwrap();
2749
2750 let cell1_path = tmp.path().join("cell1");
2752 std::fs::create_dir(&cell1_path).unwrap();
2753 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
2754
2755 let cell2_path = tmp.path().join("cell2");
2756 std::fs::create_dir(&cell2_path).unwrap();
2757 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
2758
2759 let mgr = CellManager::new_with_path(
2760 tmp.path().to_path_buf(),
2761 256,
2762 cpumask_for_range(32),
2763 HashSet::new(),
2764 )
2765 .unwrap();
2766 let assignments = mgr.compute_cpu_assignments(true).unwrap();
2767
2768 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2769 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2770
2771 let cell1_assignment = assignments
2775 .iter()
2776 .find(|a| a.id == cell1_info.cell_id)
2777 .unwrap();
2778 let cell1_borrow = cell1_assignment.borrowable.as_ref().unwrap();
2779 for cpu in 8..32 {
2781 assert!(
2782 !cell1_borrow.test_cpu(cpu),
2783 "Cell 1 borrowable should not include CPU {} (outside cpuset)",
2784 cpu
2785 );
2786 }
2787
2788 let cell2_assignment = assignments
2790 .iter()
2791 .find(|a| a.id == cell2_info.cell_id)
2792 .unwrap();
2793 let cell2_borrow = cell2_assignment.borrowable.as_ref().unwrap();
2794 for cpu in 0..8 {
2795 assert!(
2796 !cell2_borrow.test_cpu(cpu),
2797 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2798 cpu
2799 );
2800 }
2801 for cpu in 16..32 {
2802 assert!(
2803 !cell2_borrow.test_cpu(cpu),
2804 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2805 cpu
2806 );
2807 }
2808 }
2809
2810 #[test]
2813 fn test_demand_cpu_assignments_all_idle() {
2814 let tmp = TempDir::new().unwrap();
2815 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2816 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2817
2818 let mgr = CellManager::new_with_path(
2819 tmp.path().to_path_buf(),
2820 256,
2821 cpumask_for_range(12),
2822 HashSet::new(),
2823 )
2824 .unwrap();
2825
2826 let demands: HashMap<u32, f64> = [(0, 0.0), (1, 0.0), (2, 0.0)].into();
2828 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2829
2830 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2832 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2833 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2834 let c1 = assignments
2835 .iter()
2836 .find(|a| a.id == cell1_info.cell_id)
2837 .unwrap();
2838 let c2 = assignments
2839 .iter()
2840 .find(|a| a.id == cell2_info.cell_id)
2841 .unwrap();
2842
2843 assert_eq!(cell0.primary.weight(), 4);
2844 assert_eq!(c1.primary.weight(), 4);
2845 assert_eq!(c2.primary.weight(), 4);
2846 }
2847
2848 #[test]
2849 fn test_demand_cpu_assignments_uneven_demand() {
2850 let tmp = TempDir::new().unwrap();
2851 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2852 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2853
2854 let mgr = CellManager::new_with_path(
2855 tmp.path().to_path_buf(),
2856 256,
2857 cpumask_for_range(12),
2858 HashSet::new(),
2859 )
2860 .unwrap();
2861
2862 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2863 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2864
2865 let demands: HashMap<u32, f64> = [
2867 (0, 50.0),
2868 (cell1_info.cell_id, 100.0),
2869 (cell2_info.cell_id, 1.0),
2870 ]
2871 .into();
2872 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2873
2874 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2875 let c1 = assignments
2876 .iter()
2877 .find(|a| a.id == cell1_info.cell_id)
2878 .unwrap();
2879 let c2 = assignments
2880 .iter()
2881 .find(|a| a.id == cell2_info.cell_id)
2882 .unwrap();
2883
2884 assert!(
2886 c1.primary.weight() > c2.primary.weight(),
2887 "Busy cell ({}) should have more CPUs than idle cell ({})",
2888 c1.primary.weight(),
2889 c2.primary.weight()
2890 );
2891 assert!(c2.primary.weight() >= 1);
2893 assert!(cell0.primary.weight() >= 1);
2894 assert_eq!(
2896 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2897 12
2898 );
2899 }
2900
2901 #[test]
2902 fn test_demand_cpu_assignments_with_cpusets() {
2903 let tmp = TempDir::new().unwrap();
2904
2905 let cell_a_path = tmp.path().join("cell_a");
2907 std::fs::create_dir(&cell_a_path).unwrap();
2908 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2909
2910 let cell_b_path = tmp.path().join("cell_b");
2911 std::fs::create_dir(&cell_b_path).unwrap();
2912 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
2913
2914 let mgr = CellManager::new_with_path(
2915 tmp.path().to_path_buf(),
2916 256,
2917 cpumask_for_range(16),
2918 HashSet::new(),
2919 )
2920 .unwrap();
2921
2922 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2923 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2924
2925 let demands: HashMap<u32, f64> = [
2927 (0, 10.0),
2928 (cell_a_info.cell_id, 90.0),
2929 (cell_b_info.cell_id, 10.0),
2930 ]
2931 .into();
2932 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2933
2934 let cell_a = assignments
2935 .iter()
2936 .find(|a| a.id == cell_a_info.cell_id)
2937 .unwrap();
2938 let cell_b = assignments
2939 .iter()
2940 .find(|a| a.id == cell_b_info.cell_id)
2941 .unwrap();
2942
2943 assert!(
2947 cell_a.primary.weight() > cell_b.primary.weight(),
2948 "Cell A ({}) should have more CPUs than Cell B ({})",
2949 cell_a.primary.weight(),
2950 cell_b.primary.weight()
2951 );
2952
2953 assert!(cell_a.primary.weight() >= 4);
2955 assert!(cell_b.primary.weight() >= 4);
2956 }
2957
2958 #[test]
2959 fn test_demand_cpu_assignments_idle_cell_gets_floor() {
2960 let tmp = TempDir::new().unwrap();
2961 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2962 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2963
2964 let mgr = CellManager::new_with_path(
2965 tmp.path().to_path_buf(),
2966 256,
2967 cpumask_for_range(12),
2968 HashSet::new(),
2969 )
2970 .unwrap();
2971
2972 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2973 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2974
2975 let demands: HashMap<u32, f64> = [
2977 (0, 50.0),
2978 (cell1_info.cell_id, 100.0),
2979 (cell2_info.cell_id, 0.0),
2980 ]
2981 .into();
2982 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2983
2984 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
2985 let c1 = assignments
2986 .iter()
2987 .find(|a| a.id == cell1_info.cell_id)
2988 .unwrap();
2989 let c2 = assignments
2990 .iter()
2991 .find(|a| a.id == cell2_info.cell_id)
2992 .unwrap();
2993
2994 assert_eq!(
2997 c2.primary.weight(),
2998 1,
2999 "Idle cell should get minimum target of 1 CPU"
3000 );
3001 assert!(c1.primary.weight() > cell0.primary.weight());
3003 assert!(c1.primary.weight() > c2.primary.weight());
3004 assert_eq!(
3006 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
3007 12
3008 );
3009 }
3010
3011 #[test]
3012 fn test_demand_cpu_assignments_negative_weight_errors() {
3013 let tmp = TempDir::new().unwrap();
3014 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
3015
3016 let mgr = CellManager::new_with_path(
3017 tmp.path().to_path_buf(),
3018 256,
3019 cpumask_for_range(8),
3020 HashSet::new(),
3021 )
3022 .unwrap();
3023
3024 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
3025
3026 let demands: HashMap<u32, f64> = [(0, 50.0), (cell1_info.cell_id, -10.0)].into();
3027 let result = mgr.compute_demand_cpu_assignments(&demands, false);
3028 assert!(result.is_err(), "Negative weight should be rejected");
3029 assert!(
3030 result
3031 .unwrap_err()
3032 .to_string()
3033 .contains("negative demand weight"),
3034 "Error message should mention negative weight"
3035 );
3036 }
3037
3038 #[test]
3041 fn test_deficit_distribution_with_cpusets() {
3042 let tmp = TempDir::new().unwrap();
3046
3047 let cell_a_path = tmp.path().join("cell_a");
3049 std::fs::create_dir(&cell_a_path).unwrap();
3050 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
3051
3052 let cell_b_path = tmp.path().join("cell_b");
3054 std::fs::create_dir(&cell_b_path).unwrap();
3055 std::fs::write(cell_b_path.join("cpuset.cpus"), "6-11\n").unwrap();
3056
3057 let mgr = CellManager::new_with_path(
3059 tmp.path().to_path_buf(),
3060 256,
3061 cpumask_for_range(16),
3062 HashSet::new(),
3063 )
3064 .unwrap();
3065
3066 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
3067 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
3068
3069 let demands: HashMap<u32, f64> = [
3071 (0, 10.0),
3072 (cell_a_info.cell_id, 90.0),
3073 (cell_b_info.cell_id, 10.0),
3074 ]
3075 .into();
3076 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
3077
3078 let cell_a = assignments
3079 .iter()
3080 .find(|a| a.id == cell_a_info.cell_id)
3081 .unwrap();
3082 let cell_b = assignments
3083 .iter()
3084 .find(|a| a.id == cell_b_info.cell_id)
3085 .unwrap();
3086 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
3087
3088 assert_eq!(
3090 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
3091 16,
3092 );
3093
3094 assert!(
3096 cell_a.primary.weight() > cell_b.primary.weight(),
3097 "cell_a ({}) should have more CPUs than cell_b ({})",
3098 cell_a.primary.weight(),
3099 cell_b.primary.weight()
3100 );
3101
3102 assert!(cell_a.primary.weight() >= 1);
3104 assert!(cell_b.primary.weight() >= 1);
3105 assert!(cell0.primary.weight() >= 1);
3106 }
3107
3108 #[test]
3109 fn test_deficit_distribution_equal_weight_with_exclusive() {
3110 let tmp = TempDir::new().unwrap();
3113
3114 let cell1_path = tmp.path().join("cell1");
3116 std::fs::create_dir(&cell1_path).unwrap();
3117 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
3118
3119 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
3121
3122 let mgr = CellManager::new_with_path(
3124 tmp.path().to_path_buf(),
3125 256,
3126 cpumask_for_range(16),
3127 HashSet::new(),
3128 )
3129 .unwrap();
3130
3131 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
3132 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
3133
3134 let assignments = mgr.compute_cpu_assignments(false).unwrap();
3135
3136 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
3137 let cell1 = assignments
3138 .iter()
3139 .find(|a| a.id == cell1_info.cell_id)
3140 .unwrap();
3141 let cell2 = assignments
3142 .iter()
3143 .find(|a| a.id == cell2_info.cell_id)
3144 .unwrap();
3145
3146 assert_eq!(cell1.primary.weight(), 8); assert_eq!(
3156 cell0.primary.weight() + cell1.primary.weight() + cell2.primary.weight(),
3157 16,
3158 );
3159
3160 assert!(
3162 cell0.primary.weight() >= cell2.primary.weight(),
3163 "cell0 ({}) should have >= CPUs than cell2 ({})",
3164 cell0.primary.weight(),
3165 cell2.primary.weight()
3166 );
3167 }
3168
3169 #[test]
3170 fn test_deficit_all_cells_exceed_target() {
3171 let tmp = TempDir::new().unwrap();
3174
3175 let cell_a_path = tmp.path().join("cell_a");
3177 std::fs::create_dir(&cell_a_path).unwrap();
3178 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
3179
3180 let cell_b_path = tmp.path().join("cell_b");
3182 std::fs::create_dir(&cell_b_path).unwrap();
3183 std::fs::write(cell_b_path.join("cpuset.cpus"), "8-13\n").unwrap();
3184
3185 let mgr = CellManager::new_with_path(
3192 tmp.path().to_path_buf(),
3193 256,
3194 cpumask_for_range(20),
3195 HashSet::new(),
3196 )
3197 .unwrap();
3198
3199 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
3200 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
3201
3202 let assignments = mgr.compute_cpu_assignments(false).unwrap();
3203
3204 let cell_a = assignments
3205 .iter()
3206 .find(|a| a.id == cell_a_info.cell_id)
3207 .unwrap();
3208 let cell_b = assignments
3209 .iter()
3210 .find(|a| a.id == cell_b_info.cell_id)
3211 .unwrap();
3212 let cell0 = assignments.iter().find(|a| a.id == 0).unwrap();
3213
3214 let cell_a_contested: Vec<_> = (8..10)
3217 .filter(|&cpu| cell_a.primary.test_cpu(cpu))
3218 .collect();
3219 let cell_b_contested: Vec<_> = (8..10)
3220 .filter(|&cpu| cell_b.primary.test_cpu(cpu))
3221 .collect();
3222 assert_eq!(
3223 cell_a_contested.len(),
3224 0,
3225 "cell_a should get 0 contested CPUs (exceeded target)"
3226 );
3227 assert_eq!(
3228 cell_b_contested.len(),
3229 2,
3230 "cell_b should get all 2 contested CPUs (has deficit)"
3231 );
3232
3233 assert_eq!(
3235 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
3236 20,
3237 );
3238
3239 assert!(cell0.primary.weight() >= 1);
3241 assert!(cell_a.primary.weight() >= 1);
3242 assert!(cell_b.primary.weight() >= 1);
3243 }
3244
3245 #[test]
3248 fn test_distribute_proportional_basic() {
3249 let cpus: Vec<usize> = (0..8).collect();
3250 let recipients = vec![(0, 3.0), (1, 1.0)];
3251 let result = CpuManager::distribute_cpus_proportional(&cpus, &recipients).unwrap();
3252
3253 assert_eq!(result.get(&0).unwrap().len(), 6);
3255 assert_eq!(result.get(&1).unwrap().len(), 2);
3256 }
3257
3258 #[test]
3259 fn test_distribute_proportional_zero_weight_gets_nothing() {
3260 let cpus: Vec<usize> = (0..6).collect();
3261 let recipients = vec![(0, 1.0), (1, 0.0)];
3262 let result = CpuManager::distribute_cpus_proportional(&cpus, &recipients).unwrap();
3263
3264 assert_eq!(result.get(&0).unwrap().len(), 6);
3266 assert!(result.get(&1).is_none() || result.get(&1).unwrap().is_empty());
3267 }
3268
3269 #[test]
3270 fn test_distribute_proportional_all_zero_fallback() {
3271 let cpus: Vec<usize> = (0..6).collect();
3272 let recipients = vec![(0, 0.0), (1, 0.0)];
3273 let result = CpuManager::distribute_cpus_proportional(&cpus, &recipients).unwrap();
3274
3275 assert_eq!(result.get(&0).unwrap().len(), 3);
3277 assert_eq!(result.get(&1).unwrap().len(), 3);
3278 }
3279
3280 #[test]
3281 fn test_distribute_proportional_skewed_weights_floor_guarantee() {
3282 let cpus: Vec<usize> = (0..38).collect();
3285 let recipients = vec![(2, 1.0), (3, 100.0)];
3286 let result = CpuManager::distribute_cpus_proportional(&cpus, &recipients).unwrap();
3287
3288 let cell2_count = result.get(&2).map_or(0, |v| v.len());
3290 let cell3_count = result.get(&3).map_or(0, |v| v.len());
3291 assert!(
3292 cell2_count >= 1,
3293 "cell 2 must get at least 1 CPU, got {}",
3294 cell2_count,
3295 );
3296 assert!(
3297 cell3_count >= 1,
3298 "cell 3 must get at least 1 CPU, got {}",
3299 cell3_count,
3300 );
3301 assert_eq!(cell2_count + cell3_count, 38, "all CPUs must be assigned");
3302 }
3303
3304 #[test]
3305 fn test_distribute_proportional_overlapping_cpusets_no_starvation() {
3306 let cpus: Vec<usize> = (24..62).collect(); let recipients = vec![(2, 1.0), (3, 87.0)];
3311 let result = CpuManager::distribute_cpus_proportional(&cpus, &recipients).unwrap();
3312
3313 let cell2_count = result.get(&2).map_or(0, |v| v.len());
3314 let cell3_count = result.get(&3).map_or(0, |v| v.len());
3315 assert!(
3316 cell2_count >= 1,
3317 "cell 2 must not be starved, got {} CPUs",
3318 cell2_count,
3319 );
3320 assert!(
3321 cell3_count >= 1,
3322 "cell 3 must not be starved, got {} CPUs",
3323 cell3_count,
3324 );
3325 assert_eq!(cell2_count + cell3_count, 38);
3326 }
3327
3328 #[test]
3331 fn test_compute_targets_equal_weight() {
3332 let targets = CpuManager::compute_targets(12, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
3333 assert_eq!(*targets.get(&0).unwrap(), 4);
3334 assert_eq!(*targets.get(&1).unwrap(), 4);
3335 assert_eq!(*targets.get(&2).unwrap(), 4);
3336 }
3337
3338 #[test]
3339 fn test_compute_targets_with_remainder() {
3340 let targets = CpuManager::compute_targets(10, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
3341 let total: usize = targets.values().sum();
3343 assert_eq!(total, 10);
3344 for (_, &count) in &targets {
3345 assert!(count >= 3 && count <= 4);
3346 }
3347 }
3348
3349 #[test]
3352 fn test_cpu_manager_non_root_domain_unpinned_invariants() {
3353 let domain = cpumask_from_cpulist(32, "8-15");
3354 let recipients = vec![
3355 CpuRecipient {
3356 id: 10,
3357 weight: 1.0,
3358 allowed: domain.clone(),
3359 claimed: None,
3360 minimum: CpuMinimum::default(),
3361 },
3362 CpuRecipient {
3363 id: 11,
3364 weight: 1.0,
3365 allowed: domain.clone(),
3366 claimed: None,
3367 minimum: CpuMinimum::default(),
3368 },
3369 CpuRecipient {
3370 id: 12,
3371 weight: 2.0,
3372 allowed: domain.clone(),
3373 claimed: None,
3374 minimum: CpuMinimum::default(),
3375 },
3376 ];
3377
3378 let assignments = CpuManager::new(&domain)
3379 .compute_assignments(&recipients, false)
3380 .unwrap();
3381
3382 assert_eq!(assignments.len(), 3);
3383 assert_eq!(find_assignment(&assignments, 10).primary.weight(), 2);
3384 assert_eq!(find_assignment(&assignments, 11).primary.weight(), 2);
3385 assert_eq!(find_assignment(&assignments, 12).primary.weight(), 4);
3386
3387 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
3388 assert_eq!(total, domain.weight());
3389
3390 for assignment in &assignments {
3391 assert!(
3392 assignment.borrowable.is_none(),
3393 "borrowable should stay disabled in this case"
3394 );
3395 for cpu in assignment.primary.iter() {
3396 assert!(domain.test_cpu(cpu), "CPU {} escaped the domain", cpu);
3397 }
3398 }
3399
3400 for i in 0..assignments.len() {
3401 for j in (i + 1)..assignments.len() {
3402 let overlap = assignments[i].primary.and(&assignments[j].primary);
3403 assert_eq!(
3404 overlap.weight(),
3405 0,
3406 "Recipients {} and {} overlap",
3407 assignments[i].id,
3408 assignments[j].id
3409 );
3410 }
3411 }
3412
3413 let mut union = Cpumask::new();
3414 for assignment in &assignments {
3415 union = union.or(&assignment.primary);
3416 }
3417 assert_eq!(union, domain, "Primary masks should cover the whole domain");
3418 }
3419
3420 #[test]
3421 fn test_cpu_manager_non_contiguous_domain_borrowable_invariants() {
3422 let domain = cpumask_from_cpulist(32, "8-9,12,14-15,20-21,26");
3423 let allowed_1 = cpumask_from_cpulist(32, "0-1,8-9,12");
3424 let allowed_2 = cpumask_from_cpulist(32, "9,14-15,20-22");
3425 let recipients = vec![
3426 CpuRecipient {
3427 id: 1,
3428 weight: 1.0,
3429 allowed: allowed_1.clone(),
3430 claimed: Some(allowed_1.clone()),
3431 minimum: CpuMinimum::default(),
3432 },
3433 CpuRecipient {
3434 id: 2,
3435 weight: 2.0,
3436 allowed: allowed_2.clone(),
3437 claimed: Some(allowed_2.clone()),
3438 minimum: CpuMinimum::default(),
3439 },
3440 CpuRecipient {
3441 id: 3,
3442 weight: 1.0,
3443 allowed: domain.clone(),
3444 claimed: None,
3445 minimum: CpuMinimum::default(),
3446 },
3447 ];
3448
3449 let assignments = CpuManager::new(&domain)
3450 .compute_assignments(&recipients, true)
3451 .unwrap();
3452
3453 assert_eq!(assignments.len(), 3);
3454
3455 for assignment in &assignments {
3456 for cpu in assignment.primary.iter() {
3457 assert!(
3458 domain.test_cpu(cpu),
3459 "Primary CPU {} escaped the domain",
3460 cpu
3461 );
3462 }
3463 let borrowable = assignment.borrowable.as_ref().unwrap();
3464 for cpu in borrowable.iter() {
3465 assert!(
3466 domain.test_cpu(cpu),
3467 "Borrowable CPU {} escaped the domain",
3468 cpu
3469 );
3470 }
3471 assert_eq!(
3472 assignment.primary.and(borrowable).weight(),
3473 0,
3474 "Primary and borrowable overlap for recipient {}",
3475 assignment.id
3476 );
3477 }
3478
3479 let recipient_1 = find_assignment(&assignments, 1);
3480 let recipient_2 = find_assignment(&assignments, 2);
3481 let recipient_3 = find_assignment(&assignments, 3);
3482
3483 for cpu in recipient_1.primary.iter() {
3484 assert!(
3485 allowed_1.test_cpu(cpu),
3486 "Recipient 1 got CPU {} outside allowed",
3487 cpu
3488 );
3489 }
3490 for cpu in recipient_2.primary.iter() {
3491 assert!(
3492 allowed_2.test_cpu(cpu),
3493 "Recipient 2 got CPU {} outside allowed",
3494 cpu
3495 );
3496 }
3497
3498 let allowed_1_in_domain = domain.and(&allowed_1);
3499 let allowed_2_in_domain = domain.and(&allowed_2);
3500 let recipient_1_union = recipient_1
3501 .primary
3502 .or(recipient_1.borrowable.as_ref().unwrap());
3503 let recipient_2_union = recipient_2
3504 .primary
3505 .or(recipient_2.borrowable.as_ref().unwrap());
3506 assert_eq!(recipient_1_union, allowed_1_in_domain);
3507 assert_eq!(recipient_2_union, allowed_2_in_domain);
3508
3509 let recipient_3_union = recipient_3
3510 .primary
3511 .or(recipient_3.borrowable.as_ref().unwrap());
3512 assert_eq!(recipient_3_union, domain);
3513 assert_eq!(recipient_3.primary.weight(), 1);
3514 assert!(!recipient_3.primary.test_cpu(8));
3515 assert!(!recipient_3.primary.test_cpu(9));
3516 assert!(recipient_3.primary.test_cpu(26));
3517
3518 let mut union = Cpumask::new();
3519 for assignment in &assignments {
3520 union = union.or(&assignment.primary);
3521 }
3522 assert_eq!(union, domain, "Primary masks should cover the whole domain");
3523 }
3524
3525 #[test]
3528 fn test_symmetric_pairwise_overlap_produces_equal_cells() {
3529 let tmp = TempDir::new().unwrap();
3530
3531 let cpusets = [
3541 ("cell_a", "0-3,20-27"),
3542 ("cell_b", "4-7,20-21,28-33"),
3543 ("cell_c", "8-11,22-23,28-29,34-37"),
3544 ("cell_d", "12-15,24-25,30-31,34-35,38-39"),
3545 ("cell_e", "16-19,26-27,32-33,36-39"),
3546 ];
3547 for (name, cpus) in &cpusets {
3548 let p = tmp.path().join(name);
3549 std::fs::create_dir(&p).unwrap();
3550 std::fs::write(p.join("cpuset.cpus"), format!("{cpus}\n")).unwrap();
3551 }
3552
3553 let mgr = CellManager::new_with_path(
3554 tmp.path().to_path_buf(),
3555 256,
3556 cpumask_for_range(56),
3557 HashSet::new(),
3558 )
3559 .unwrap();
3560
3561 let assignments = mgr.compute_cpu_assignments(false).unwrap();
3562
3563 let workload: Vec<_> = assignments.iter().filter(|a| a.id != 0).collect();
3564 assert_eq!(workload.len(), 5, "Expected 5 workload cells");
3565
3566 let counts: Vec<(u32, usize)> = workload
3568 .iter()
3569 .map(|a| (a.id, a.primary.weight()))
3570 .collect();
3571 for &(cell_id, count) in &counts {
3572 assert_eq!(
3573 count, counts[0].1,
3574 "Cell {} has {} CPUs, expected {} — symmetric inputs \
3575 should produce equal cell sizes. All counts: {:?}",
3576 cell_id, count, counts[0].1, counts,
3577 );
3578 }
3579
3580 for i in 0..assignments.len() {
3582 for j in (i + 1)..assignments.len() {
3583 for cpu in 0..56 {
3584 assert!(
3585 !(assignments[i].primary.test_cpu(cpu)
3586 && assignments[j].primary.test_cpu(cpu)),
3587 "CPU {} assigned to both cell {} and cell {}",
3588 cpu,
3589 assignments[i].id,
3590 assignments[j].id,
3591 );
3592 }
3593 }
3594 }
3595
3596 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
3598 assert_eq!(total, 56, "All CPUs must be assigned");
3599 }
3600}