1use std::collections::{HashMap, HashSet};
13use std::os::unix::fs::MetadataExt;
14use std::os::unix::io::{AsFd, BorrowedFd};
15use std::path::{Path, PathBuf};
16
17use anyhow::{bail, Context, Result};
18use inotify::{Inotify, WatchMask};
19use scx_utils::Cpumask;
20use tracing::{debug, info};
21
22#[derive(Debug)]
24pub struct CellInfo {
25 pub cell_id: u32,
26 pub cgroup_path: Option<PathBuf>,
27 pub cgid: Option<u64>,
28 pub cpuset: Option<Cpumask>,
30}
31
32fn compute_targets(total_cpus: usize, cells: &[(u32, f64)]) -> Result<HashMap<u32, usize>> {
37 if cells.is_empty() {
38 bail!("compute_targets called with no cells");
39 }
40 if total_cpus < cells.len() {
41 bail!(
42 "Not enough CPUs ({}) for {} cells (need at least 1 each)",
43 total_cpus,
44 cells.len()
45 );
46 }
47
48 let total_weight: f64 = cells.iter().map(|(_, w)| w).sum();
49 let num_cells = cells.len();
50
51 if total_weight <= 0.0 {
52 let per = total_cpus / num_cells;
54 let remainder = total_cpus % num_cells;
55 return Ok(cells
56 .iter()
57 .enumerate()
58 .map(|(i, (cell_id, _))| {
59 let extra = if i < remainder { 1 } else { 0 };
60 (*cell_id, per + extra)
61 })
62 .collect());
63 }
64
65 let distributable = total_cpus - num_cells;
66 let mut assigned = num_cells;
67 let mut raw: Vec<(u32, f64, usize)> = cells
68 .iter()
69 .map(|(cell_id, w)| {
70 let frac = w / total_weight * distributable as f64;
71 let floored = frac.floor() as usize;
72 assigned += floored;
73 (*cell_id, frac - frac.floor(), 1 + floored)
74 })
75 .collect();
76
77 let mut remainder = total_cpus - assigned;
78 raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
79 for entry in raw.iter_mut() {
80 if remainder == 0 {
81 break;
82 }
83 entry.2 += 1;
84 remainder -= 1;
85 }
86
87 Ok(raw
88 .iter()
89 .map(|(cell_id, _, count)| (*cell_id, *count))
90 .collect())
91}
92
93fn distribute_cpus_proportional(
100 cpus: &[usize],
101 recipients: &[(u32, f64)],
102) -> Result<HashMap<u32, Vec<usize>>> {
103 if cpus.is_empty() {
104 bail!("distribute_cpus_proportional called with no CPUs");
105 }
106 if recipients.is_empty() {
107 bail!("distribute_cpus_proportional called with no recipients");
108 }
109
110 let total_weight: f64 = recipients.iter().map(|(_, w)| w).sum();
111 let n = cpus.len();
112
113 let mut allocs: Vec<(u32, usize)> = if total_weight <= 0.0 {
114 let per = n / recipients.len();
116 let remainder = n % recipients.len();
117 recipients
118 .iter()
119 .enumerate()
120 .map(|(i, (cell_id, _))| {
121 let extra = if i < remainder { 1 } else { 0 };
122 (*cell_id, per + extra)
123 })
124 .collect()
125 } else {
126 let mut assigned = 0usize;
128 let mut raw: Vec<(u32, f64, usize, bool)> = recipients
129 .iter()
130 .map(|(cell_id, w)| {
131 let frac = w / total_weight * n as f64;
132 let floored = frac.floor() as usize;
133 assigned += floored;
134 (*cell_id, frac - frac.floor(), floored, *w > 0.0)
135 })
136 .collect();
137
138 let mut remainder = n - assigned;
139 raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
140 for entry in raw.iter_mut() {
141 if remainder == 0 {
142 break;
143 }
144 entry.2 += 1;
145 remainder -= 1;
146 }
147
148 loop {
153 let Some(starved_idx) = raw
154 .iter()
155 .position(|(_, _, count, pos_weight)| *pos_weight && *count == 0)
156 else {
157 break;
158 };
159 let Some(donor_idx) = raw
160 .iter()
161 .enumerate()
162 .filter(|(_, (_, _, count, _))| *count > 1)
163 .max_by_key(|(_, (_, _, count, _))| *count)
164 .map(|(i, _)| i)
165 else {
166 break; };
168 raw[donor_idx].2 -= 1;
169 raw[starved_idx].2 += 1;
170 }
171
172 raw.iter()
173 .map(|(cell_id, _, count, _)| (*cell_id, *count))
174 .collect()
175 };
176
177 let mut cpu_iter = cpus.iter().copied();
179 allocs.sort_by_key(|(cell_id, _)| *cell_id);
180 let mut result: HashMap<u32, Vec<usize>> = HashMap::new();
181 for (cell_id, count) in allocs {
182 let cpus_for_cell: Vec<usize> = cpu_iter.by_ref().take(count).collect();
183 if cpus_for_cell.len() != count {
184 bail!(
185 "BUG: distribute_cpus_proportional: cell {} expected {} CPUs but got {}",
186 cell_id,
187 count,
188 cpus_for_cell.len()
189 );
190 }
191 if !cpus_for_cell.is_empty() {
192 result.insert(cell_id, cpus_for_cell);
193 }
194 }
195
196 Ok(result)
197}
198
199#[derive(Debug)]
201pub struct CpuAssignment {
202 pub cell_id: u32,
203 pub primary: Cpumask,
204 pub borrowable: Option<Cpumask>,
205}
206
207pub struct CellManager {
209 cell_parent_path: PathBuf,
210 inotify: Inotify,
211 cells: HashMap<u64, CellInfo>,
213 cell_id_to_cgid: HashMap<u32, u64>,
215 free_cell_ids: Vec<u32>,
217 next_cell_id: u32,
218 max_cells: u32,
219 all_cpus: Cpumask,
221 exclude_names: HashSet<String>,
223}
224
225impl CellManager {
226 pub fn new(
227 cell_parent_path: &str,
228 max_cells: u32,
229 all_cpus: Cpumask,
230 exclude: HashSet<String>,
231 ) -> Result<Self> {
232 let path = PathBuf::from(format!("/sys/fs/cgroup{}", cell_parent_path));
233 if !path.exists() {
234 bail!("Cell parent cgroup path does not exist: {}", path.display());
235 }
236 Self::new_with_path(path, max_cells, all_cpus, exclude)
237 }
238
239 fn new_with_path(
240 path: PathBuf,
241 max_cells: u32,
242 all_cpus: Cpumask,
243 exclude: HashSet<String>,
244 ) -> Result<Self> {
245 let inotify = Inotify::init().context("Failed to initialize inotify")?;
246 inotify
247 .watches()
248 .add(&path, WatchMask::CREATE | WatchMask::DELETE)
249 .context("Failed to add inotify watch")?;
250
251 let mut mgr = Self {
252 cell_parent_path: path.clone(),
253 inotify,
254 cells: HashMap::new(),
255 cell_id_to_cgid: HashMap::new(),
256 free_cell_ids: Vec::new(),
257 next_cell_id: 1, max_cells,
259 all_cpus,
260 exclude_names: exclude,
261 };
262
263 mgr.cells.insert(
266 0,
267 CellInfo {
268 cell_id: 0,
269 cgroup_path: None,
270 cgid: None,
271 cpuset: None,
272 },
273 );
274 mgr.cell_id_to_cgid.insert(0, 0);
275
276 mgr.scan_existing_children()
278 .context("Failed to scan existing child cgroups at startup")?;
279 Ok(mgr)
280 }
281
282 fn should_exclude(&self, path: &Path) -> bool {
283 path.file_name()
284 .and_then(|n| n.to_str())
285 .map(|name| self.exclude_names.contains(name))
286 .unwrap_or(false)
287 }
288
289 fn scan_existing_children(&mut self) -> Result<Vec<(u64, u32)>> {
290 let mut assignments = Vec::new();
291 let entries = std::fs::read_dir(&self.cell_parent_path).with_context(|| {
292 format!(
293 "Failed to read cell parent directory: {}",
294 self.cell_parent_path.display()
295 )
296 })?;
297 for entry in entries {
298 let entry = entry.with_context(|| {
299 format!(
300 "Failed to read directory entry in: {}",
301 self.cell_parent_path.display()
302 )
303 })?;
304 let file_type = entry.file_type().with_context(|| {
305 format!("Failed to get file type for: {}", entry.path().display())
306 })?;
307 if file_type.is_dir() {
308 let path = entry.path();
309 if self.should_exclude(&path) {
310 continue;
311 }
312 let cgid = path
313 .metadata()
314 .with_context(|| {
315 format!("reading inode of cgroup directory {}", path.display())
316 })?
317 .ino();
318 let (cgid, cell_id) =
319 self.create_cell_for_cgroup(&path, cgid).with_context(|| {
320 format!("Failed to create cell for cgroup: {}", path.display())
321 })?;
322 assignments.push((cgid, cell_id));
323 }
324 }
325 Ok(assignments)
326 }
327
328 pub fn process_events(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
335 let mut buffer = [0; 1024];
336 let mut has_events = false;
337
338 loop {
340 match self.inotify.read_events(&mut buffer) {
341 Ok(events) => {
342 if events.into_iter().next().is_some() {
343 has_events = true;
344 } else {
345 break;
346 }
347 }
348 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
349 Err(e) => {
350 return Err(e).context("Failed to read inotify events");
351 }
352 }
353 }
354
355 if !has_events {
356 return Ok((Vec::new(), Vec::new()));
357 }
358
359 self.reconcile_cells()
361 }
362
363 fn reconcile_cells(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
366 let mut new_cells = Vec::new();
367
368 let mut current_entries: HashMap<PathBuf, u64> = HashMap::new();
372 let entries = std::fs::read_dir(&self.cell_parent_path).with_context(|| {
373 format!(
374 "Failed to read cell parent directory: {}",
375 self.cell_parent_path.display()
376 )
377 })?;
378 for entry in entries {
379 let entry = entry.with_context(|| {
380 format!(
381 "Failed to read directory entry in: {}",
382 self.cell_parent_path.display()
383 )
384 })?;
385 let file_type = entry.file_type().with_context(|| {
386 format!("Failed to get file type for: {}", entry.path().display())
387 })?;
388 if file_type.is_dir() {
389 let path = entry.path();
390 if self.should_exclude(&path) {
391 continue;
392 }
393
394 let metadata = match entry.metadata() {
395 Ok(metadata) => metadata,
396 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
397 continue;
400 }
401 Err(e) => {
402 return Err(e).with_context(|| {
403 format!("reading inode of cgroup directory {}", path.display())
404 });
405 }
406 };
407
408 current_entries.insert(path, metadata.ino());
409 }
410 }
411
412 let mut destroyed_cells: HashSet<u32> = HashSet::new();
414 self.cells.retain(|&cgid, info| {
415 if info.cell_id == 0 {
416 return true; }
418 let cgroup_path = info
420 .cgroup_path
421 .as_ref()
422 .expect("BUG: non-zero cell missing cgroup_path");
423 if current_entries.get(cgroup_path) == Some(&cgid) {
426 true
427 } else {
428 info!(
429 "Destroyed cell {} for cgroup {} (cgid={})",
430 info.cell_id,
431 cgroup_path.display(),
432 cgid
433 );
434 destroyed_cells.insert(info.cell_id);
435 false
436 }
437 });
438
439 self.cell_id_to_cgid
441 .retain(|cell_id, _| !destroyed_cells.contains(cell_id));
442 self.free_cell_ids.extend(destroyed_cells.iter().copied());
443
444 for (path, cgid) in current_entries {
446 if self.cells.contains_key(&cgid) {
447 continue; }
449 let (cgid, cell_id) = self
450 .create_cell_for_cgroup(&path, cgid)
451 .with_context(|| format!("Failed to create cell for cgroup: {}", path.display()))?;
452 new_cells.push((cgid, cell_id));
453 }
454
455 Ok((new_cells, destroyed_cells.into_iter().collect()))
456 }
457
458 fn create_cell_for_cgroup(&mut self, path: &Path, cgid: u64) -> Result<(u64, u32)> {
459 let cell_id = self.allocate_cell_id().context("allocating cell ID")?;
460
461 let cpuset = Self::read_cpuset(path)
462 .with_context(|| format!("reading cpuset for cgroup {}", path.display()))?;
463 if let Some(ref mask) = cpuset {
464 debug!(
465 "Cell {} has cpuset: {} (from {})",
466 cell_id,
467 mask.to_cpulist(),
468 path.join("cpuset.cpus").display()
469 );
470 }
471
472 self.cells.insert(
473 cgid,
474 CellInfo {
475 cell_id,
476 cgroup_path: Some(path.to_path_buf()),
477 cgid: Some(cgid),
478 cpuset,
479 },
480 );
481 self.cell_id_to_cgid.insert(cell_id, cgid);
482
483 info!(
484 "Created cell {} for cgroup {} (cgid={})",
485 cell_id,
486 path.display(),
487 cgid
488 );
489
490 Ok((cgid, cell_id))
491 }
492
493 fn read_cpuset(cgroup_path: &Path) -> Result<Option<Cpumask>> {
495 let cpuset_path = cgroup_path.join("cpuset.cpus");
496 match std::fs::read_to_string(&cpuset_path) {
497 Ok(content) => {
498 let content = content.trim();
499 if content.is_empty() {
500 Ok(None)
501 } else {
502 let mask = Cpumask::from_cpulist(content).with_context(|| {
503 format!(
504 "Failed to parse cpuset '{}' from {}",
505 content,
506 cpuset_path.display()
507 )
508 })?;
509 Ok(Some(mask))
510 }
511 }
512 Err(_) => Ok(None),
514 }
515 }
516
517 fn allocate_cell_id(&mut self) -> Result<u32> {
518 if let Some(id) = self.free_cell_ids.pop() {
520 return Ok(id);
521 }
522
523 if self.next_cell_id >= self.max_cells {
524 bail!("Cell ID space exhausted (max_cells={})", self.max_cells);
525 }
526
527 let id = self.next_cell_id;
528 self.next_cell_id += 1;
529 Ok(id)
530 }
531
532 pub fn compute_cpu_assignments(&self, compute_borrowable: bool) -> Result<Vec<CpuAssignment>> {
544 self.compute_cpu_assignments_inner(None, compute_borrowable)
546 }
547
548 pub fn compute_demand_cpu_assignments(
556 &self,
557 cell_demands: &HashMap<u32, f64>,
558 compute_borrowable: bool,
559 ) -> Result<Vec<CpuAssignment>> {
560 self.compute_cpu_assignments_inner(Some(cell_demands), compute_borrowable)
561 }
562
563 fn compute_cpu_assignments_inner(
565 &self,
566 cell_demands: Option<&HashMap<u32, f64>>,
567 compute_borrowable: bool,
568 ) -> Result<Vec<CpuAssignment>> {
569 if let Some(demands) = cell_demands {
571 for (&cell_id, &weight) in demands {
572 if weight < 0.0 {
573 bail!("Cell {} has negative demand weight {}", cell_id, weight);
574 }
575 }
576 }
577
578 let mut contention: HashMap<usize, Vec<u32>> = HashMap::new();
580 for cell_info in self.cells.values() {
581 if let Some(ref cpuset) = cell_info.cpuset {
582 for cpu in cpuset.iter() {
583 contention.entry(cpu).or_default().push(cell_info.cell_id);
584 }
585 }
586 }
587
588 let mut cell_cpus: HashMap<u32, Cpumask> = HashMap::new();
593 let mut contested_cpus: Vec<usize> = Vec::new();
594 let mut unclaimed_cpus: Vec<usize> = Vec::new();
595
596 for cpu in self.all_cpus.iter() {
597 match contention.get(&cpu) {
598 None => unclaimed_cpus.push(cpu),
599 Some(claimants) if claimants.len() == 1 => {
600 let cell_id = claimants[0];
601 cell_cpus
602 .entry(cell_id)
603 .or_insert_with(Cpumask::new)
604 .set_cpu(cpu)
605 .ok();
606 }
607 Some(_) => contested_cpus.push(cpu),
608 }
609 }
610
611 let total_cpu_count = self.all_cpus.weight();
613 let mut all_cells_with_weights: Vec<(u32, f64)> = self
614 .cells
615 .values()
616 .map(|info| {
617 let weight = match cell_demands {
618 Some(demands) => *demands.get(&info.cell_id).ok_or_else(|| {
619 anyhow::anyhow!("Cell {} is missing from demands map", info.cell_id)
620 })?,
621 None => 1.0,
622 };
623 Ok((info.cell_id, weight))
624 })
625 .collect::<Result<Vec<_>>>()
626 .context("building cell demand weights map")?;
627 all_cells_with_weights.sort_by_key(|(cell_id, _)| *cell_id);
628
629 let targets = compute_targets(total_cpu_count, &all_cells_with_weights)
630 .context("computing per-cell CPU targets")?;
631
632 let mut assigned_count: HashMap<u32, usize> = HashMap::new();
634 for (cell_id, mask) in &cell_cpus {
635 assigned_count.insert(*cell_id, mask.weight());
636 }
637
638 let mut contested_groups: HashMap<Vec<u32>, Vec<usize>> = HashMap::new();
640 for cpu in contested_cpus {
641 if let Some(claimants) = contention.get(&cpu) {
642 let mut sorted_claimants = claimants.clone();
643 sorted_claimants.sort();
644 contested_groups
645 .entry(sorted_claimants)
646 .or_default()
647 .push(cpu);
648 }
649 }
650
651 let initial_deficit: HashMap<u32, f64> = targets
657 .iter()
658 .map(|(&cell_id, &target)| {
659 let already = assigned_count.get(&cell_id).copied().unwrap_or(0);
661 let deficit = if target > already {
662 (target - already) as f64
663 } else {
664 0.0
665 };
666 (cell_id, deficit)
667 })
668 .collect();
669
670 for (claimants, cpus) in contested_groups {
671 let mut recipients: Vec<(u32, f64)> = Vec::new();
672 let mut all_zero = true;
673 for &cell_id in &claimants {
674 let deficit = *initial_deficit.get(&cell_id).ok_or_else(|| {
675 anyhow::anyhow!(
676 "BUG: cell {} in contention map but missing from targets",
677 cell_id
678 )
679 })?;
680 if deficit > 0.0 {
681 all_zero = false;
682 }
683 recipients.push((cell_id, deficit));
684 }
685
686 if all_zero {
688 recipients = claimants.iter().map(|&cell_id| (cell_id, 1.0)).collect();
689 }
690
691 let distribution = distribute_cpus_proportional(&cpus, &recipients)
692 .context("distributing contested CPUs among claimants")?;
693 for (cell_id, assigned_cpus) in distribution {
694 let count = assigned_cpus.len();
695 for cpu in assigned_cpus {
696 cell_cpus
697 .entry(cell_id)
698 .or_insert_with(Cpumask::new)
699 .set_cpu(cpu)
700 .ok();
701 }
702 *assigned_count.entry(cell_id).or_insert(0) += count;
703 }
704 }
705
706 if !unclaimed_cpus.is_empty() {
708 let mut recipients: Vec<(u32, f64)> = Vec::new();
709 let mut all_zero = true;
710
711 for info in self.cells.values() {
712 if info.cpuset.is_some() {
713 continue; }
715 let target = *targets.get(&info.cell_id).ok_or_else(|| {
716 anyhow::anyhow!(
717 "BUG: cell {} is unpinned but missing from targets",
718 info.cell_id
719 )
720 })?;
721 let already = assigned_count.get(&info.cell_id).copied().unwrap_or(0);
723 let deficit = if target > already {
724 (target - already) as f64
725 } else {
726 0.0
727 };
728 if deficit > 0.0 {
729 all_zero = false;
730 }
731 recipients.push((info.cell_id, deficit));
732 }
733 recipients.sort_by_key(|(cell_id, _)| *cell_id);
734
735 if all_zero {
737 recipients = recipients
738 .iter()
739 .map(|(cell_id, _)| (*cell_id, 1.0))
740 .collect();
741 }
742
743 let distribution = distribute_cpus_proportional(&unclaimed_cpus, &recipients)
744 .context("distributing unclaimed CPUs among unpinned cells")?;
745 for (cell_id, assigned_cpus) in distribution {
746 let count = assigned_cpus.len();
747 for cpu in assigned_cpus {
748 cell_cpus
749 .entry(cell_id)
750 .or_insert_with(Cpumask::new)
751 .set_cpu(cpu)
752 .ok();
753 }
754 *assigned_count.entry(cell_id).or_insert(0) += count;
755 }
756 }
757
758 for info in self.cells.values() {
760 if !cell_cpus.contains_key(&info.cell_id)
761 || cell_cpus
762 .get(&info.cell_id)
763 .map_or(true, |m| m.weight() == 0)
764 {
765 bail!(
766 "Cell {} has no CPUs assigned (nr_cpus={}, num_cells={})",
767 info.cell_id,
768 self.all_cpus.weight(),
769 self.cells.len()
770 );
771 }
772 }
773
774 let assignments: Vec<CpuAssignment> = cell_cpus
776 .into_iter()
777 .map(|(cell_id, primary)| {
778 let borrowable = if compute_borrowable {
779 let mut borrow_mask = self.all_cpus.and(&primary.not());
780
781 if let Some(cell_info) = self.cells.values().find(|c| c.cell_id == cell_id) {
783 if let Some(ref cpuset) = cell_info.cpuset {
784 borrow_mask = borrow_mask.and(cpuset);
785 }
786 }
787
788 Some(borrow_mask)
789 } else {
790 None
791 };
792 CpuAssignment {
793 cell_id,
794 primary,
795 borrowable,
796 }
797 })
798 .collect();
799
800 Ok(assignments)
801 }
802
803 pub fn get_cell_assignments(&self) -> Vec<(u64, u32)> {
806 self.cells
807 .values()
808 .filter(|info| info.cell_id != 0)
809 .map(|info| {
810 (
811 info.cgid.expect("BUG: non-zero cell missing cgid"),
812 info.cell_id,
813 )
814 })
815 .collect()
816 }
817
818 pub fn format_cell_config(&self, cpu_assignments: &[CpuAssignment]) -> String {
821 let mut sorted: Vec<_> = cpu_assignments.iter().collect();
822 sorted.sort_by_key(|a| a.cell_id);
823
824 let mut parts = Vec::new();
825 for assignment in sorted {
826 let cpulist = assignment.primary.to_cpulist();
827 if assignment.cell_id == 0 {
828 parts.push(format!("[0: {}]", cpulist));
829 } else {
830 let name = self
832 .cells
833 .values()
834 .find(|info| info.cell_id == assignment.cell_id)
835 .and_then(|info| {
836 info.cgroup_path
837 .as_ref()
838 .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
839 })
840 .unwrap_or_else(|| "?".to_string());
841 parts.push(format!("[{}({}): {}]", assignment.cell_id, name, cpulist));
842 }
843 }
844 parts.join(" ")
845 }
846
847 pub fn refresh_cpusets(&mut self) -> Result<bool> {
850 let mut changed = false;
851 for info in self.cells.values_mut() {
852 let Some(ref cgroup_path) = info.cgroup_path else {
853 continue; };
855 let new_cpuset = Self::read_cpuset(cgroup_path)
856 .with_context(|| format!("reading cpuset for cgroup {}", cgroup_path.display()))?;
857 if new_cpuset != info.cpuset {
858 info!(
859 "Cell {} cpuset changed: {:?} -> {:?} ({})",
860 info.cell_id,
861 info.cpuset.as_ref().map(|m| m.to_cpulist()),
862 new_cpuset.as_ref().map(|m| m.to_cpulist()),
863 cgroup_path.display(),
864 );
865 info.cpuset = new_cpuset;
866 changed = true;
867 }
868 }
869 Ok(changed)
870 }
871}
872
873impl AsFd for CellManager {
874 fn as_fd(&self) -> BorrowedFd<'_> {
875 self.inotify.as_fd()
876 }
877}
878
879#[cfg(test)]
880impl CellManager {
881 fn cell_count(&self) -> usize {
884 self.cells.values().filter(|c| c.cell_id != 0).count()
885 }
886
887 fn get_cell_ids(&self) -> Vec<u32> {
890 self.cells
891 .values()
892 .filter(|c| c.cell_id != 0)
893 .map(|c| c.cell_id)
894 .collect()
895 }
896
897 fn find_cell_by_name(&self, name: &str) -> Option<&CellInfo> {
900 self.cells.values().filter(|c| c.cell_id != 0).find(|c| {
901 c.cgroup_path
902 .as_ref()
903 .and_then(|p| p.file_name())
904 .map(|n| n.to_str() == Some(name))
905 .unwrap_or(false)
906 })
907 }
908}
909
910#[cfg(test)]
911#[allow(clippy::unwrap_used)]
912mod tests {
913 use super::*;
914 use tempfile::TempDir;
915
916 fn cpumask_for_range(nr_cpus: usize) -> Cpumask {
917 scx_utils::set_cpumask_test_width(nr_cpus);
918 let mut mask = Cpumask::new();
919 for cpu in 0..nr_cpus {
920 mask.set_cpu(cpu).unwrap();
921 }
922 mask
923 }
924
925 #[test]
928 fn test_scan_empty_directory() {
929 let tmp = TempDir::new().unwrap();
930 let mgr = CellManager::new_with_path(
931 tmp.path().to_path_buf(),
932 256,
933 cpumask_for_range(16),
934 HashSet::new(),
935 )
936 .unwrap();
937
938 assert_eq!(mgr.cell_count(), 0);
939 }
940
941 #[test]
942 fn test_scan_existing_subdirectories() {
943 let tmp = TempDir::new().unwrap();
944
945 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
947 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
948
949 let mgr = CellManager::new_with_path(
950 tmp.path().to_path_buf(),
951 256,
952 cpumask_for_range(16),
953 HashSet::new(),
954 )
955 .unwrap();
956
957 assert_eq!(mgr.cell_count(), 2);
958
959 let cell_ids = mgr.get_cell_ids();
961 assert!(cell_ids.contains(&1));
962 assert!(cell_ids.contains(&2));
963 }
964
965 #[test]
966 fn test_reconcile_detects_new_directories() {
967 let tmp = TempDir::new().unwrap();
968
969 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
971 let mut mgr = CellManager::new_with_path(
972 tmp.path().to_path_buf(),
973 256,
974 cpumask_for_range(16),
975 HashSet::new(),
976 )
977 .unwrap();
978 assert_eq!(mgr.cell_count(), 1);
979
980 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
982
983 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
985 assert_eq!(new_cells.len(), 1);
986 assert_eq!(destroyed_cells.len(), 0);
987 assert_eq!(mgr.cell_count(), 2);
988 }
989
990 #[test]
991 fn test_reconcile_detects_removed_directories() {
992 let tmp = TempDir::new().unwrap();
993
994 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
996 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
997 let mut mgr = CellManager::new_with_path(
998 tmp.path().to_path_buf(),
999 256,
1000 cpumask_for_range(16),
1001 HashSet::new(),
1002 )
1003 .unwrap();
1004 assert_eq!(mgr.cell_count(), 2);
1005
1006 std::fs::remove_dir(tmp.path().join("container-b")).unwrap();
1008
1009 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1011 assert_eq!(new_cells.len(), 0);
1012 assert_eq!(destroyed_cells.len(), 1);
1013 assert_eq!(mgr.cell_count(), 1);
1014 }
1015
1016 #[test]
1017 fn test_reconcile_replaces_reused_path_with_new_inode() {
1018 let tmp = TempDir::new().unwrap();
1019 let parked = TempDir::new().unwrap();
1020
1021 let original_path = tmp.path().join("container-a");
1022 std::fs::create_dir(&original_path).unwrap();
1023
1024 let mut mgr = CellManager::new_with_path(
1025 tmp.path().to_path_buf(),
1026 256,
1027 cpumask_for_range(16),
1028 HashSet::new(),
1029 )
1030 .unwrap();
1031 assert_eq!(mgr.cell_count(), 1);
1032
1033 let old_info = mgr.find_cell_by_name("container-a").unwrap();
1034 let old_cell_id = old_info.cell_id;
1035 let old_cgid = old_info.cgid.unwrap();
1036
1037 std::fs::rename(&original_path, parked.path().join("container-a-old")).unwrap();
1040 std::fs::create_dir(&original_path).unwrap();
1041
1042 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1043
1044 assert_eq!(new_cells.len(), 1);
1045 assert_eq!(destroyed_cells, vec![old_cell_id]);
1046 assert_eq!(mgr.cell_count(), 1);
1047
1048 let new_info = mgr.find_cell_by_name("container-a").unwrap();
1049 assert_eq!(new_info.cell_id, old_cell_id);
1050 assert_ne!(new_info.cgid.unwrap(), old_cgid);
1051 }
1052
1053 #[test]
1054 fn test_cell_id_reuse_after_destruction() {
1055 let tmp = TempDir::new().unwrap();
1056
1057 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1059 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1060 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1061
1062 let mut mgr = CellManager::new_with_path(
1063 tmp.path().to_path_buf(),
1064 256,
1065 cpumask_for_range(16),
1066 HashSet::new(),
1067 )
1068 .unwrap();
1069
1070 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1072 let cell2_id = cell2_info.cell_id;
1073
1074 std::fs::remove_dir(tmp.path().join("cell2")).unwrap();
1076 mgr.reconcile_cells().unwrap();
1077
1078 std::fs::create_dir(tmp.path().join("cell4")).unwrap();
1080 mgr.reconcile_cells().unwrap();
1081
1082 let cell4_info = mgr.find_cell_by_name("cell4").unwrap();
1083 assert_eq!(cell4_info.cell_id, cell2_id);
1084 }
1085
1086 #[test]
1089 fn test_cpu_assignments_no_cells() {
1090 let tmp = TempDir::new().unwrap();
1091 let mgr = CellManager::new_with_path(
1092 tmp.path().to_path_buf(),
1093 256,
1094 cpumask_for_range(16),
1095 HashSet::new(),
1096 )
1097 .unwrap();
1098
1099 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1100
1101 assert_eq!(assignments.len(), 1);
1103 assert_eq!(assignments[0].cell_id, 0);
1104 assert_eq!(assignments[0].primary.weight(), 16);
1105 }
1106
1107 #[test]
1108 fn test_cpu_assignments_proportional() {
1109 let tmp = TempDir::new().unwrap();
1110 std::fs::create_dir(tmp.path().join("container")).unwrap();
1111
1112 let mgr = CellManager::new_with_path(
1113 tmp.path().to_path_buf(),
1114 256,
1115 cpumask_for_range(16),
1116 HashSet::new(),
1117 )
1118 .unwrap();
1119 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1120
1121 assert_eq!(assignments.len(), 2);
1123
1124 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1125 let cell1 = assignments.iter().find(|a| a.cell_id == 1).unwrap();
1126
1127 assert_eq!(cell0.primary.weight(), 8);
1128 assert_eq!(cell1.primary.weight(), 8);
1129 }
1130
1131 #[test]
1132 fn test_cpu_assignments_remainder_to_cell0() {
1133 let tmp = TempDir::new().unwrap();
1134 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1135 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1136
1137 let mgr = CellManager::new_with_path(
1138 tmp.path().to_path_buf(),
1139 256,
1140 cpumask_for_range(10),
1141 HashSet::new(),
1142 )
1143 .unwrap();
1144 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1145
1146 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1148 assert_eq!(cell0.primary.weight(), 4); }
1150
1151 #[test]
1152 fn test_cpu_assignments_too_many_cells() {
1153 let tmp = TempDir::new().unwrap();
1154
1155 for i in 1..=5 {
1157 std::fs::create_dir(tmp.path().join(format!("cell{}", i))).unwrap();
1158 }
1159
1160 let mgr = CellManager::new_with_path(
1162 tmp.path().to_path_buf(),
1163 256,
1164 cpumask_for_range(4),
1165 HashSet::new(),
1166 )
1167 .unwrap();
1168 let result = mgr.compute_cpu_assignments(false);
1169
1170 assert!(result.is_err());
1171 let err_msg = format!("{:#}", result.unwrap_err());
1172 assert!(
1173 err_msg.contains("Not enough CPUs"),
1174 "Expected 'Not enough CPUs' error, got: {}",
1175 err_msg
1176 );
1177 }
1178
1179 #[test]
1180 fn test_cpu_assignments_with_cpusets() {
1181 let tmp = TempDir::new().unwrap();
1182
1183 let cell1_path = tmp.path().join("cell1");
1185 std::fs::create_dir(&cell1_path).unwrap();
1186 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1187
1188 let cell2_path = tmp.path().join("cell2");
1189 std::fs::create_dir(&cell2_path).unwrap();
1190 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1191
1192 let mgr = CellManager::new_with_path(
1193 tmp.path().to_path_buf(),
1194 256,
1195 cpumask_for_range(16),
1196 HashSet::new(),
1197 )
1198 .unwrap();
1199 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1200
1201 assert_eq!(assignments.len(), 3);
1203
1204 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1206 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1207
1208 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1209 let cell1 = assignments
1210 .iter()
1211 .find(|a| a.cell_id == cell1_info.cell_id)
1212 .unwrap();
1213 let cell2 = assignments
1214 .iter()
1215 .find(|a| a.cell_id == cell2_info.cell_id)
1216 .unwrap();
1217
1218 assert_eq!(cell1.primary.weight(), 4);
1220 for cpu in 0..4 {
1221 assert!(cell1.primary.test_cpu(cpu));
1222 }
1223
1224 assert_eq!(cell2.primary.weight(), 4);
1226 for cpu in 8..12 {
1227 assert!(cell2.primary.test_cpu(cpu));
1228 }
1229
1230 assert_eq!(cell0.primary.weight(), 8);
1232 for cpu in 4..8 {
1233 assert!(cell0.primary.test_cpu(cpu));
1234 }
1235 for cpu in 12..16 {
1236 assert!(cell0.primary.test_cpu(cpu));
1237 }
1238 }
1239
1240 #[test]
1241 fn test_cpu_assignments_cpusets_cover_all_cpus() {
1242 let tmp = TempDir::new().unwrap();
1243
1244 let cell1_path = tmp.path().join("cell1");
1246 std::fs::create_dir(&cell1_path).unwrap();
1247 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1248
1249 let cell2_path = tmp.path().join("cell2");
1250 std::fs::create_dir(&cell2_path).unwrap();
1251 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1252
1253 let mgr = CellManager::new_with_path(
1254 tmp.path().to_path_buf(),
1255 256,
1256 cpumask_for_range(16),
1257 HashSet::new(),
1258 )
1259 .unwrap();
1260 let result = mgr.compute_cpu_assignments(false);
1261
1262 assert!(result.is_err());
1264 let err = result.unwrap_err();
1265 let err_msg = format!("{:#}", err);
1266 assert!(
1267 err_msg.contains("Cell 0 has no CPUs assigned"),
1268 "Expected 'Cell 0 has no CPUs assigned' error, got: {}",
1269 err_msg
1270 );
1271 }
1272
1273 #[test]
1274 fn test_cpu_assignments_single_cpuset() {
1275 let tmp = TempDir::new().unwrap();
1276
1277 let cell1_path = tmp.path().join("cell1");
1279 std::fs::create_dir(&cell1_path).unwrap();
1280 std::fs::write(cell1_path.join("cpuset.cpus"), "0,2,4,6\n").unwrap();
1281
1282 let mgr = CellManager::new_with_path(
1283 tmp.path().to_path_buf(),
1284 256,
1285 cpumask_for_range(8),
1286 HashSet::new(),
1287 )
1288 .unwrap();
1289 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1290
1291 assert_eq!(assignments.len(), 2);
1292
1293 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1294 let cell1 = assignments.iter().find(|a| a.cell_id != 0).unwrap();
1295
1296 assert_eq!(cell1.primary.weight(), 4);
1298 for cpu in [0, 2, 4, 6] {
1299 assert!(cell1.primary.test_cpu(cpu));
1300 }
1301
1302 assert_eq!(cell0.primary.weight(), 4);
1304 for cpu in [1, 3, 5, 7] {
1305 assert!(cell0.primary.test_cpu(cpu));
1306 }
1307 }
1308
1309 #[test]
1310 fn test_cpuset_parsing_from_file() {
1311 let tmp = TempDir::new().unwrap();
1312
1313 let cell_path = tmp.path().join("cell1");
1315 std::fs::create_dir(&cell_path).unwrap();
1316 std::fs::write(cell_path.join("cpuset.cpus"), "0-3,8-11,16\n").unwrap();
1317
1318 let mgr = CellManager::new_with_path(
1319 tmp.path().to_path_buf(),
1320 256,
1321 cpumask_for_range(32),
1322 HashSet::new(),
1323 )
1324 .unwrap();
1325
1326 let cell_info = mgr.find_cell_by_name("cell1").unwrap();
1328 let cpuset = cell_info.cpuset.as_ref().unwrap();
1329
1330 assert_eq!(cpuset.weight(), 9); for cpu in 0..4 {
1332 assert!(cpuset.test_cpu(cpu));
1333 }
1334 for cpu in 8..12 {
1335 assert!(cpuset.test_cpu(cpu));
1336 }
1337 assert!(cpuset.test_cpu(16));
1338 }
1339
1340 #[test]
1341 fn test_cpu_assignments_mixed_cpuset_and_no_cpuset() {
1342 let tmp = TempDir::new().unwrap();
1343
1344 let cell1_path = tmp.path().join("cell1");
1346 std::fs::create_dir(&cell1_path).unwrap();
1347 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1348
1349 let cell2_path = tmp.path().join("cell2");
1351 std::fs::create_dir(&cell2_path).unwrap();
1352
1353 let mgr = CellManager::new_with_path(
1354 tmp.path().to_path_buf(),
1355 256,
1356 cpumask_for_range(16),
1357 HashSet::new(),
1358 )
1359 .unwrap();
1360
1361 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1363 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1364 assert!(cell1_info.cpuset.is_some());
1365 assert!(cell2_info.cpuset.is_none());
1366
1367 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1368
1369 assert_eq!(assignments.len(), 3);
1375
1376 let cell1_assignment = assignments
1378 .iter()
1379 .find(|a| a.cell_id == cell1_info.cell_id)
1380 .unwrap();
1381 assert_eq!(cell1_assignment.primary.weight(), 4);
1382
1383 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1385 assert_eq!(cell0.primary.weight(), 7);
1386
1387 let cell2_assignment = assignments
1389 .iter()
1390 .find(|a| a.cell_id == cell2_info.cell_id)
1391 .unwrap();
1392 assert_eq!(cell2_assignment.primary.weight(), 5);
1393 }
1394
1395 #[test]
1398 fn test_cpu_assignments_partial_overlap() {
1399 let tmp = TempDir::new().unwrap();
1400
1401 let cell_a_path = tmp.path().join("cell_a");
1403 std::fs::create_dir(&cell_a_path).unwrap();
1404 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
1405
1406 let cell_b_path = tmp.path().join("cell_b");
1407 std::fs::create_dir(&cell_b_path).unwrap();
1408 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
1409
1410 let mgr = CellManager::new_with_path(
1411 tmp.path().to_path_buf(),
1412 256,
1413 cpumask_for_range(16),
1414 HashSet::new(),
1415 )
1416 .unwrap();
1417 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1418
1419 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1420 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1421
1422 let cell_a = assignments
1423 .iter()
1424 .find(|a| a.cell_id == cell_a_info.cell_id)
1425 .unwrap();
1426 let cell_b = assignments
1427 .iter()
1428 .find(|a| a.cell_id == cell_b_info.cell_id)
1429 .unwrap();
1430 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1431
1432 assert_eq!(cell_a.primary.weight(), 6);
1436 assert_eq!(cell_b.primary.weight(), 6);
1437 assert_eq!(cell0.primary.weight(), 4);
1438
1439 for cpu in 0..4 {
1441 assert!(
1442 cell_a.primary.test_cpu(cpu),
1443 "CPU {} should be in cell_a",
1444 cpu
1445 );
1446 }
1447 for cpu in 8..12 {
1448 assert!(
1449 cell_b.primary.test_cpu(cpu),
1450 "CPU {} should be in cell_b",
1451 cpu
1452 );
1453 }
1454 for cpu in 12..16 {
1455 assert!(
1456 cell0.primary.test_cpu(cpu),
1457 "CPU {} should be in cell0",
1458 cpu
1459 );
1460 }
1461
1462 let cell_a_contested: Vec<_> = (4..8).filter(|&cpu| cell_a.primary.test_cpu(cpu)).collect();
1464 let cell_b_contested: Vec<_> = (4..8).filter(|&cpu| cell_b.primary.test_cpu(cpu)).collect();
1465 assert_eq!(cell_a_contested.len(), 2);
1466 assert_eq!(cell_b_contested.len(), 2);
1467
1468 for cpu in 0..16 {
1470 let mut count = 0;
1471 if cell_a.primary.test_cpu(cpu) {
1472 count += 1;
1473 }
1474 if cell_b.primary.test_cpu(cpu) {
1475 count += 1;
1476 }
1477 if cell0.primary.test_cpu(cpu) {
1478 count += 1;
1479 }
1480 assert!(count <= 1, "CPU {} is assigned to {} cells", cpu, count);
1481 }
1482 }
1483
1484 #[test]
1485 fn test_cpu_assignments_three_way_overlap() {
1486 let tmp = TempDir::new().unwrap();
1487
1488 let cell_a_path = tmp.path().join("cell_a");
1490 std::fs::create_dir(&cell_a_path).unwrap();
1491 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-5\n").unwrap();
1492
1493 let cell_b_path = tmp.path().join("cell_b");
1494 std::fs::create_dir(&cell_b_path).unwrap();
1495 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-5\n").unwrap();
1496
1497 let cell_c_path = tmp.path().join("cell_c");
1498 std::fs::create_dir(&cell_c_path).unwrap();
1499 std::fs::write(cell_c_path.join("cpuset.cpus"), "0-5\n").unwrap();
1500
1501 let mgr = CellManager::new_with_path(
1502 tmp.path().to_path_buf(),
1503 256,
1504 cpumask_for_range(12),
1505 HashSet::new(),
1506 )
1507 .unwrap();
1508 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1509
1510 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1511 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1512 let cell_c_info = mgr.find_cell_by_name("cell_c").unwrap();
1513
1514 let cell_a = assignments
1515 .iter()
1516 .find(|a| a.cell_id == cell_a_info.cell_id)
1517 .unwrap();
1518 let cell_b = assignments
1519 .iter()
1520 .find(|a| a.cell_id == cell_b_info.cell_id)
1521 .unwrap();
1522 let cell_c = assignments
1523 .iter()
1524 .find(|a| a.cell_id == cell_c_info.cell_id)
1525 .unwrap();
1526 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1527
1528 assert_eq!(cell_a.primary.weight(), 2);
1530 assert_eq!(cell_b.primary.weight(), 2);
1531 assert_eq!(cell_c.primary.weight(), 2);
1532
1533 assert_eq!(cell0.primary.weight(), 6);
1535 for cpu in 6..12 {
1536 assert!(cell0.primary.test_cpu(cpu));
1537 }
1538
1539 let total_contested: usize = (0..6)
1541 .filter(|&cpu| {
1542 cell_a.primary.test_cpu(cpu)
1543 || cell_b.primary.test_cpu(cpu)
1544 || cell_c.primary.test_cpu(cpu)
1545 })
1546 .count();
1547 assert_eq!(total_contested, 6);
1548 }
1549
1550 #[test]
1551 fn test_cpu_assignments_odd_contested_count() {
1552 let tmp = TempDir::new().unwrap();
1553
1554 let cell_a_path = tmp.path().join("cell_a");
1556 std::fs::create_dir(&cell_a_path).unwrap();
1557 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-2\n").unwrap();
1558
1559 let cell_b_path = tmp.path().join("cell_b");
1560 std::fs::create_dir(&cell_b_path).unwrap();
1561 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-2\n").unwrap();
1562
1563 let mgr = CellManager::new_with_path(
1564 tmp.path().to_path_buf(),
1565 256,
1566 cpumask_for_range(8),
1567 HashSet::new(),
1568 )
1569 .unwrap();
1570 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1571
1572 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1573 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1574
1575 let cell_a = assignments
1576 .iter()
1577 .find(|a| a.cell_id == cell_a_info.cell_id)
1578 .unwrap();
1579 let cell_b = assignments
1580 .iter()
1581 .find(|a| a.cell_id == cell_b_info.cell_id)
1582 .unwrap();
1583
1584 let total = cell_a.primary.weight() + cell_b.primary.weight();
1587 assert_eq!(total, 3);
1588 assert!(cell_a.primary.weight() >= 1 && cell_a.primary.weight() <= 2);
1589 assert!(cell_b.primary.weight() >= 1 && cell_b.primary.weight() <= 2);
1590
1591 for cpu in 0..3 {
1593 let a_has = cell_a.primary.test_cpu(cpu);
1594 let b_has = cell_b.primary.test_cpu(cpu);
1595 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
1596 }
1597 }
1598
1599 #[test]
1600 fn test_cpu_assignments_complete_overlap() {
1601 let tmp = TempDir::new().unwrap();
1602
1603 let cell_a_path = tmp.path().join("cell_a");
1605 std::fs::create_dir(&cell_a_path).unwrap();
1606 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
1607
1608 let cell_b_path = tmp.path().join("cell_b");
1609 std::fs::create_dir(&cell_b_path).unwrap();
1610 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-7\n").unwrap();
1611
1612 let mgr = CellManager::new_with_path(
1613 tmp.path().to_path_buf(),
1614 256,
1615 cpumask_for_range(16),
1616 HashSet::new(),
1617 )
1618 .unwrap();
1619 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1620
1621 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1622 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1623
1624 let cell_a = assignments
1625 .iter()
1626 .find(|a| a.cell_id == cell_a_info.cell_id)
1627 .unwrap();
1628 let cell_b = assignments
1629 .iter()
1630 .find(|a| a.cell_id == cell_b_info.cell_id)
1631 .unwrap();
1632 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1633
1634 assert_eq!(cell_a.primary.weight(), 4);
1636 assert_eq!(cell_b.primary.weight(), 4);
1637
1638 assert_eq!(cell0.primary.weight(), 8);
1640 for cpu in 8..16 {
1641 assert!(cell0.primary.test_cpu(cpu));
1642 }
1643
1644 for cpu in 0..8 {
1646 let a_has = cell_a.primary.test_cpu(cpu);
1647 let b_has = cell_b.primary.test_cpu(cpu);
1648 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
1649 }
1650 }
1651
1652 #[test]
1653 fn test_cpu_assignments_no_overlap() {
1654 let tmp = TempDir::new().unwrap();
1656
1657 let cell_a_path = tmp.path().join("cell_a");
1658 std::fs::create_dir(&cell_a_path).unwrap();
1659 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-3\n").unwrap();
1660
1661 let cell_b_path = tmp.path().join("cell_b");
1662 std::fs::create_dir(&cell_b_path).unwrap();
1663 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-7\n").unwrap();
1664
1665 let mgr = CellManager::new_with_path(
1666 tmp.path().to_path_buf(),
1667 256,
1668 cpumask_for_range(16),
1669 HashSet::new(),
1670 )
1671 .unwrap();
1672 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1673
1674 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1675 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1676
1677 let cell_a = assignments
1678 .iter()
1679 .find(|a| a.cell_id == cell_a_info.cell_id)
1680 .unwrap();
1681 let cell_b = assignments
1682 .iter()
1683 .find(|a| a.cell_id == cell_b_info.cell_id)
1684 .unwrap();
1685 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1686
1687 assert_eq!(cell_a.primary.weight(), 4);
1689 for cpu in 0..4 {
1690 assert!(cell_a.primary.test_cpu(cpu));
1691 }
1692
1693 assert_eq!(cell_b.primary.weight(), 4);
1694 for cpu in 4..8 {
1695 assert!(cell_b.primary.test_cpu(cpu));
1696 }
1697
1698 assert_eq!(cell0.primary.weight(), 8);
1700 for cpu in 8..16 {
1701 assert!(cell0.primary.test_cpu(cpu));
1702 }
1703 }
1704
1705 #[test]
1708 fn test_format_cell_config_only_cell0() {
1709 let tmp = TempDir::new().unwrap();
1710 let mgr = CellManager::new_with_path(
1711 tmp.path().to_path_buf(),
1712 256,
1713 cpumask_for_range(8),
1714 HashSet::new(),
1715 )
1716 .unwrap();
1717
1718 let mut mask = Cpumask::new();
1719 for cpu in 0..8 {
1720 mask.set_cpu(cpu).unwrap();
1721 }
1722
1723 let assignments = vec![CpuAssignment {
1724 cell_id: 0,
1725 primary: mask,
1726 borrowable: None,
1727 }];
1728 let result = mgr.format_cell_config(&assignments);
1729
1730 assert_eq!(result, "[0: 0-7]");
1731 }
1732
1733 #[test]
1734 fn test_format_cell_config_with_cells() {
1735 let tmp = TempDir::new().unwrap();
1736 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1737
1738 let mgr = CellManager::new_with_path(
1739 tmp.path().to_path_buf(),
1740 256,
1741 cpumask_for_range(16),
1742 HashSet::new(),
1743 )
1744 .unwrap();
1745
1746 let mut mask0 = Cpumask::new();
1747 for cpu in 0..8 {
1748 mask0.set_cpu(cpu).unwrap();
1749 }
1750
1751 let mut mask1 = Cpumask::new();
1752 for cpu in 8..16 {
1753 mask1.set_cpu(cpu).unwrap();
1754 }
1755
1756 let assignments = vec![
1757 CpuAssignment {
1758 cell_id: 0,
1759 primary: mask0,
1760 borrowable: None,
1761 },
1762 CpuAssignment {
1763 cell_id: 1,
1764 primary: mask1,
1765 borrowable: None,
1766 },
1767 ];
1768 let result = mgr.format_cell_config(&assignments);
1769
1770 assert_eq!(result, "[0: 0-7] [1(container-a): 8-15]");
1771 }
1772
1773 #[test]
1776 fn test_cell_id_exhaustion() {
1777 let tmp = TempDir::new().unwrap();
1778
1779 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1782 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1783
1784 let mut mgr = CellManager::new_with_path(
1785 tmp.path().to_path_buf(),
1786 3,
1787 cpumask_for_range(16),
1788 HashSet::new(),
1789 )
1790 .unwrap();
1791 assert_eq!(mgr.cell_count(), 2); std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1795 let result = mgr.reconcile_cells();
1796
1797 assert!(result.is_err());
1798 let err = result.unwrap_err();
1799 let err_chain = format!("{:#}", err);
1800 assert!(
1801 err_chain.contains("Cell ID space exhausted"),
1802 "Expected exhaustion error, got: {}",
1803 err_chain
1804 );
1805 }
1806
1807 #[test]
1808 fn test_cell_id_reuse_prevents_exhaustion() {
1809 let tmp = TempDir::new().unwrap();
1810
1811 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1813 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1814
1815 let mut mgr = CellManager::new_with_path(
1816 tmp.path().to_path_buf(),
1817 3,
1818 cpumask_for_range(16),
1819 HashSet::new(),
1820 )
1821 .unwrap();
1822 assert_eq!(mgr.cell_count(), 2);
1823
1824 std::fs::remove_dir(tmp.path().join("cell1")).unwrap();
1826 mgr.reconcile_cells().unwrap();
1827 assert_eq!(mgr.cell_count(), 1);
1828
1829 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1831 let result = mgr.reconcile_cells();
1832 assert!(result.is_ok());
1833 assert_eq!(mgr.cell_count(), 2);
1834 }
1835
1836 #[test]
1839 fn test_scan_excludes_named_cgroups() {
1840 let tmp = TempDir::new().unwrap();
1841
1842 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1843 std::fs::create_dir(tmp.path().join("systemd-workaround.service")).unwrap();
1844 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1845
1846 let exclude = HashSet::from(["systemd-workaround.service".to_string()]);
1847 let mgr = CellManager::new_with_path(
1848 tmp.path().to_path_buf(),
1849 256,
1850 cpumask_for_range(16),
1851 exclude,
1852 )
1853 .unwrap();
1854
1855 assert_eq!(mgr.cell_count(), 2);
1857 assert!(mgr.find_cell_by_name("container-a").is_some());
1858 assert!(mgr.find_cell_by_name("container-b").is_some());
1859 assert!(mgr
1860 .find_cell_by_name("systemd-workaround.service")
1861 .is_none());
1862 }
1863
1864 #[test]
1865 fn test_reconcile_excludes_named_cgroups() {
1866 let tmp = TempDir::new().unwrap();
1867
1868 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1869
1870 let exclude = HashSet::from(["ignored-service".to_string()]);
1871 let mut mgr = CellManager::new_with_path(
1872 tmp.path().to_path_buf(),
1873 256,
1874 cpumask_for_range(16),
1875 exclude,
1876 )
1877 .unwrap();
1878 assert_eq!(mgr.cell_count(), 1);
1879
1880 std::fs::create_dir(tmp.path().join("ignored-service")).unwrap();
1882 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1883 assert_eq!(new_cells.len(), 0);
1884 assert_eq!(destroyed_cells.len(), 0);
1885 assert_eq!(mgr.cell_count(), 1);
1886
1887 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1889 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1890 assert_eq!(new_cells.len(), 1);
1891 assert_eq!(destroyed_cells.len(), 0);
1892 assert_eq!(mgr.cell_count(), 2);
1893 }
1894
1895 #[test]
1898 fn test_borrowable_cpumasks_basic() {
1899 let tmp = TempDir::new().unwrap();
1900
1901 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1903 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1904
1905 let mgr = CellManager::new_with_path(
1906 tmp.path().to_path_buf(),
1907 256,
1908 cpumask_for_range(16),
1909 HashSet::new(),
1910 )
1911 .unwrap();
1912 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1913
1914 for assignment in &assignments {
1916 let borrow_mask = assignment.borrowable.as_ref().unwrap();
1917 let overlap = borrow_mask.and(&assignment.primary);
1919 assert_eq!(
1920 overlap.weight(),
1921 0,
1922 "Cell {} borrowable overlaps with primary",
1923 assignment.cell_id
1924 );
1925 let union = borrow_mask.or(&assignment.primary);
1927 assert_eq!(
1928 union.weight(),
1929 16,
1930 "Cell {} union doesn't cover all CPUs",
1931 assignment.cell_id
1932 );
1933 }
1934 }
1935
1936 #[test]
1937 fn test_borrowable_cpumasks_no_overlap() {
1938 let tmp = TempDir::new().unwrap();
1939
1940 let cell1_path = tmp.path().join("cell1");
1941 std::fs::create_dir(&cell1_path).unwrap();
1942 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1943
1944 let cell2_path = tmp.path().join("cell2");
1945 std::fs::create_dir(&cell2_path).unwrap();
1946 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1947
1948 let mgr = CellManager::new_with_path(
1949 tmp.path().to_path_buf(),
1950 256,
1951 cpumask_for_range(16),
1952 HashSet::new(),
1953 )
1954 .unwrap();
1955 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1956
1957 for assignment in &assignments {
1959 let borrow_mask = assignment.borrowable.as_ref().unwrap();
1960 let overlap = borrow_mask.and(&assignment.primary);
1961 assert_eq!(
1962 overlap.weight(),
1963 0,
1964 "Cell {} borrowable overlaps with primary",
1965 assignment.cell_id
1966 );
1967 }
1968 }
1969
1970 #[test]
1971 fn test_borrowable_cpumasks_respects_cpuset() {
1972 let tmp = TempDir::new().unwrap();
1973
1974 let cell1_path = tmp.path().join("cell1");
1976 std::fs::create_dir(&cell1_path).unwrap();
1977 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1978
1979 let cell2_path = tmp.path().join("cell2");
1980 std::fs::create_dir(&cell2_path).unwrap();
1981 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1982
1983 let mgr = CellManager::new_with_path(
1984 tmp.path().to_path_buf(),
1985 256,
1986 cpumask_for_range(32),
1987 HashSet::new(),
1988 )
1989 .unwrap();
1990 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1991
1992 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1993 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1994
1995 let cell1_assignment = assignments
1999 .iter()
2000 .find(|a| a.cell_id == cell1_info.cell_id)
2001 .unwrap();
2002 let cell1_borrow = cell1_assignment.borrowable.as_ref().unwrap();
2003 for cpu in 8..32 {
2005 assert!(
2006 !cell1_borrow.test_cpu(cpu),
2007 "Cell 1 borrowable should not include CPU {} (outside cpuset)",
2008 cpu
2009 );
2010 }
2011
2012 let cell2_assignment = assignments
2014 .iter()
2015 .find(|a| a.cell_id == cell2_info.cell_id)
2016 .unwrap();
2017 let cell2_borrow = cell2_assignment.borrowable.as_ref().unwrap();
2018 for cpu in 0..8 {
2019 assert!(
2020 !cell2_borrow.test_cpu(cpu),
2021 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2022 cpu
2023 );
2024 }
2025 for cpu in 16..32 {
2026 assert!(
2027 !cell2_borrow.test_cpu(cpu),
2028 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2029 cpu
2030 );
2031 }
2032 }
2033
2034 #[test]
2037 fn test_demand_cpu_assignments_all_idle() {
2038 let tmp = TempDir::new().unwrap();
2039 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2040 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2041
2042 let mgr = CellManager::new_with_path(
2043 tmp.path().to_path_buf(),
2044 256,
2045 cpumask_for_range(12),
2046 HashSet::new(),
2047 )
2048 .unwrap();
2049
2050 let demands: HashMap<u32, f64> = [(0, 0.0), (1, 0.0), (2, 0.0)].into();
2052 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2053
2054 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2056 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2057 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2058 let c1 = assignments
2059 .iter()
2060 .find(|a| a.cell_id == cell1_info.cell_id)
2061 .unwrap();
2062 let c2 = assignments
2063 .iter()
2064 .find(|a| a.cell_id == cell2_info.cell_id)
2065 .unwrap();
2066
2067 assert_eq!(cell0.primary.weight(), 4);
2068 assert_eq!(c1.primary.weight(), 4);
2069 assert_eq!(c2.primary.weight(), 4);
2070 }
2071
2072 #[test]
2073 fn test_demand_cpu_assignments_uneven_demand() {
2074 let tmp = TempDir::new().unwrap();
2075 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2076 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2077
2078 let mgr = CellManager::new_with_path(
2079 tmp.path().to_path_buf(),
2080 256,
2081 cpumask_for_range(12),
2082 HashSet::new(),
2083 )
2084 .unwrap();
2085
2086 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2087 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2088
2089 let demands: HashMap<u32, f64> = [
2091 (0, 50.0),
2092 (cell1_info.cell_id, 100.0),
2093 (cell2_info.cell_id, 1.0),
2094 ]
2095 .into();
2096 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2097
2098 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2099 let c1 = assignments
2100 .iter()
2101 .find(|a| a.cell_id == cell1_info.cell_id)
2102 .unwrap();
2103 let c2 = assignments
2104 .iter()
2105 .find(|a| a.cell_id == cell2_info.cell_id)
2106 .unwrap();
2107
2108 assert!(
2110 c1.primary.weight() > c2.primary.weight(),
2111 "Busy cell ({}) should have more CPUs than idle cell ({})",
2112 c1.primary.weight(),
2113 c2.primary.weight()
2114 );
2115 assert!(c2.primary.weight() >= 1);
2117 assert!(cell0.primary.weight() >= 1);
2118 assert_eq!(
2120 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2121 12
2122 );
2123 }
2124
2125 #[test]
2126 fn test_demand_cpu_assignments_with_cpusets() {
2127 let tmp = TempDir::new().unwrap();
2128
2129 let cell_a_path = tmp.path().join("cell_a");
2131 std::fs::create_dir(&cell_a_path).unwrap();
2132 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2133
2134 let cell_b_path = tmp.path().join("cell_b");
2135 std::fs::create_dir(&cell_b_path).unwrap();
2136 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
2137
2138 let mgr = CellManager::new_with_path(
2139 tmp.path().to_path_buf(),
2140 256,
2141 cpumask_for_range(16),
2142 HashSet::new(),
2143 )
2144 .unwrap();
2145
2146 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2147 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2148
2149 let demands: HashMap<u32, f64> = [
2151 (0, 10.0),
2152 (cell_a_info.cell_id, 90.0),
2153 (cell_b_info.cell_id, 10.0),
2154 ]
2155 .into();
2156 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2157
2158 let cell_a = assignments
2159 .iter()
2160 .find(|a| a.cell_id == cell_a_info.cell_id)
2161 .unwrap();
2162 let cell_b = assignments
2163 .iter()
2164 .find(|a| a.cell_id == cell_b_info.cell_id)
2165 .unwrap();
2166
2167 assert!(
2171 cell_a.primary.weight() > cell_b.primary.weight(),
2172 "Cell A ({}) should have more CPUs than Cell B ({})",
2173 cell_a.primary.weight(),
2174 cell_b.primary.weight()
2175 );
2176
2177 assert!(cell_a.primary.weight() >= 4);
2179 assert!(cell_b.primary.weight() >= 4);
2180 }
2181
2182 #[test]
2183 fn test_demand_cpu_assignments_idle_cell_gets_floor() {
2184 let tmp = TempDir::new().unwrap();
2185 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2186 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2187
2188 let mgr = CellManager::new_with_path(
2189 tmp.path().to_path_buf(),
2190 256,
2191 cpumask_for_range(12),
2192 HashSet::new(),
2193 )
2194 .unwrap();
2195
2196 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2197 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2198
2199 let demands: HashMap<u32, f64> = [
2201 (0, 50.0),
2202 (cell1_info.cell_id, 100.0),
2203 (cell2_info.cell_id, 0.0),
2204 ]
2205 .into();
2206 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2207
2208 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2209 let c1 = assignments
2210 .iter()
2211 .find(|a| a.cell_id == cell1_info.cell_id)
2212 .unwrap();
2213 let c2 = assignments
2214 .iter()
2215 .find(|a| a.cell_id == cell2_info.cell_id)
2216 .unwrap();
2217
2218 assert_eq!(
2221 c2.primary.weight(),
2222 1,
2223 "Idle cell should get minimum target of 1 CPU"
2224 );
2225 assert!(c1.primary.weight() > cell0.primary.weight());
2227 assert!(c1.primary.weight() > c2.primary.weight());
2228 assert_eq!(
2230 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2231 12
2232 );
2233 }
2234
2235 #[test]
2236 fn test_demand_cpu_assignments_negative_weight_errors() {
2237 let tmp = TempDir::new().unwrap();
2238 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2239
2240 let mgr = CellManager::new_with_path(
2241 tmp.path().to_path_buf(),
2242 256,
2243 cpumask_for_range(8),
2244 HashSet::new(),
2245 )
2246 .unwrap();
2247
2248 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2249
2250 let demands: HashMap<u32, f64> = [(0, 50.0), (cell1_info.cell_id, -10.0)].into();
2251 let result = mgr.compute_demand_cpu_assignments(&demands, false);
2252 assert!(result.is_err(), "Negative weight should be rejected");
2253 assert!(
2254 result
2255 .unwrap_err()
2256 .to_string()
2257 .contains("negative demand weight"),
2258 "Error message should mention negative weight"
2259 );
2260 }
2261
2262 #[test]
2265 fn test_deficit_distribution_with_cpusets() {
2266 let tmp = TempDir::new().unwrap();
2270
2271 let cell_a_path = tmp.path().join("cell_a");
2273 std::fs::create_dir(&cell_a_path).unwrap();
2274 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2275
2276 let cell_b_path = tmp.path().join("cell_b");
2278 std::fs::create_dir(&cell_b_path).unwrap();
2279 std::fs::write(cell_b_path.join("cpuset.cpus"), "6-11\n").unwrap();
2280
2281 let mgr = CellManager::new_with_path(
2283 tmp.path().to_path_buf(),
2284 256,
2285 cpumask_for_range(16),
2286 HashSet::new(),
2287 )
2288 .unwrap();
2289
2290 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2291 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2292
2293 let demands: HashMap<u32, f64> = [
2295 (0, 10.0),
2296 (cell_a_info.cell_id, 90.0),
2297 (cell_b_info.cell_id, 10.0),
2298 ]
2299 .into();
2300 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2301
2302 let cell_a = assignments
2303 .iter()
2304 .find(|a| a.cell_id == cell_a_info.cell_id)
2305 .unwrap();
2306 let cell_b = assignments
2307 .iter()
2308 .find(|a| a.cell_id == cell_b_info.cell_id)
2309 .unwrap();
2310 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2311
2312 assert_eq!(
2314 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2315 16,
2316 );
2317
2318 assert!(
2320 cell_a.primary.weight() > cell_b.primary.weight(),
2321 "cell_a ({}) should have more CPUs than cell_b ({})",
2322 cell_a.primary.weight(),
2323 cell_b.primary.weight()
2324 );
2325
2326 assert!(cell_a.primary.weight() >= 1);
2328 assert!(cell_b.primary.weight() >= 1);
2329 assert!(cell0.primary.weight() >= 1);
2330 }
2331
2332 #[test]
2333 fn test_deficit_distribution_equal_weight_with_exclusive() {
2334 let tmp = TempDir::new().unwrap();
2337
2338 let cell1_path = tmp.path().join("cell1");
2340 std::fs::create_dir(&cell1_path).unwrap();
2341 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
2342
2343 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2345
2346 let mgr = CellManager::new_with_path(
2348 tmp.path().to_path_buf(),
2349 256,
2350 cpumask_for_range(16),
2351 HashSet::new(),
2352 )
2353 .unwrap();
2354
2355 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2356 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2357
2358 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2359
2360 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2361 let cell1 = assignments
2362 .iter()
2363 .find(|a| a.cell_id == cell1_info.cell_id)
2364 .unwrap();
2365 let cell2 = assignments
2366 .iter()
2367 .find(|a| a.cell_id == cell2_info.cell_id)
2368 .unwrap();
2369
2370 assert_eq!(cell1.primary.weight(), 8); assert_eq!(
2380 cell0.primary.weight() + cell1.primary.weight() + cell2.primary.weight(),
2381 16,
2382 );
2383
2384 assert!(
2386 cell0.primary.weight() >= cell2.primary.weight(),
2387 "cell0 ({}) should have >= CPUs than cell2 ({})",
2388 cell0.primary.weight(),
2389 cell2.primary.weight()
2390 );
2391 }
2392
2393 #[test]
2394 fn test_deficit_all_cells_exceed_target() {
2395 let tmp = TempDir::new().unwrap();
2398
2399 let cell_a_path = tmp.path().join("cell_a");
2401 std::fs::create_dir(&cell_a_path).unwrap();
2402 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2403
2404 let cell_b_path = tmp.path().join("cell_b");
2406 std::fs::create_dir(&cell_b_path).unwrap();
2407 std::fs::write(cell_b_path.join("cpuset.cpus"), "8-13\n").unwrap();
2408
2409 let mgr = CellManager::new_with_path(
2416 tmp.path().to_path_buf(),
2417 256,
2418 cpumask_for_range(20),
2419 HashSet::new(),
2420 )
2421 .unwrap();
2422
2423 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2424 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2425
2426 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2427
2428 let cell_a = assignments
2429 .iter()
2430 .find(|a| a.cell_id == cell_a_info.cell_id)
2431 .unwrap();
2432 let cell_b = assignments
2433 .iter()
2434 .find(|a| a.cell_id == cell_b_info.cell_id)
2435 .unwrap();
2436 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2437
2438 let cell_a_contested: Vec<_> = (8..10)
2441 .filter(|&cpu| cell_a.primary.test_cpu(cpu))
2442 .collect();
2443 let cell_b_contested: Vec<_> = (8..10)
2444 .filter(|&cpu| cell_b.primary.test_cpu(cpu))
2445 .collect();
2446 assert_eq!(
2447 cell_a_contested.len(),
2448 0,
2449 "cell_a should get 0 contested CPUs (exceeded target)"
2450 );
2451 assert_eq!(
2452 cell_b_contested.len(),
2453 2,
2454 "cell_b should get all 2 contested CPUs (has deficit)"
2455 );
2456
2457 assert_eq!(
2459 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2460 20,
2461 );
2462
2463 assert!(cell0.primary.weight() >= 1);
2465 assert!(cell_a.primary.weight() >= 1);
2466 assert!(cell_b.primary.weight() >= 1);
2467 }
2468
2469 #[test]
2472 fn test_distribute_proportional_basic() {
2473 let cpus: Vec<usize> = (0..8).collect();
2474 let recipients = vec![(0, 3.0), (1, 1.0)];
2475 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2476
2477 assert_eq!(result.get(&0).unwrap().len(), 6);
2479 assert_eq!(result.get(&1).unwrap().len(), 2);
2480 }
2481
2482 #[test]
2483 fn test_distribute_proportional_zero_weight_gets_nothing() {
2484 let cpus: Vec<usize> = (0..6).collect();
2485 let recipients = vec![(0, 1.0), (1, 0.0)];
2486 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2487
2488 assert_eq!(result.get(&0).unwrap().len(), 6);
2490 assert!(result.get(&1).is_none() || result.get(&1).unwrap().is_empty());
2491 }
2492
2493 #[test]
2494 fn test_distribute_proportional_all_zero_fallback() {
2495 let cpus: Vec<usize> = (0..6).collect();
2496 let recipients = vec![(0, 0.0), (1, 0.0)];
2497 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2498
2499 assert_eq!(result.get(&0).unwrap().len(), 3);
2501 assert_eq!(result.get(&1).unwrap().len(), 3);
2502 }
2503
2504 #[test]
2505 fn test_distribute_proportional_skewed_weights_floor_guarantee() {
2506 let cpus: Vec<usize> = (0..38).collect();
2509 let recipients = vec![(2, 1.0), (3, 100.0)];
2510 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2511
2512 let cell2_count = result.get(&2).map_or(0, |v| v.len());
2514 let cell3_count = result.get(&3).map_or(0, |v| v.len());
2515 assert!(
2516 cell2_count >= 1,
2517 "cell 2 must get at least 1 CPU, got {}",
2518 cell2_count,
2519 );
2520 assert!(
2521 cell3_count >= 1,
2522 "cell 3 must get at least 1 CPU, got {}",
2523 cell3_count,
2524 );
2525 assert_eq!(cell2_count + cell3_count, 38, "all CPUs must be assigned");
2526 }
2527
2528 #[test]
2529 fn test_distribute_proportional_overlapping_cpusets_no_starvation() {
2530 let cpus: Vec<usize> = (24..62).collect(); let recipients = vec![(2, 1.0), (3, 87.0)];
2535 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2536
2537 let cell2_count = result.get(&2).map_or(0, |v| v.len());
2538 let cell3_count = result.get(&3).map_or(0, |v| v.len());
2539 assert!(
2540 cell2_count >= 1,
2541 "cell 2 must not be starved, got {} CPUs",
2542 cell2_count,
2543 );
2544 assert!(
2545 cell3_count >= 1,
2546 "cell 3 must not be starved, got {} CPUs",
2547 cell3_count,
2548 );
2549 assert_eq!(cell2_count + cell3_count, 38);
2550 }
2551
2552 #[test]
2555 fn test_compute_targets_equal_weight() {
2556 let targets = compute_targets(12, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
2557 assert_eq!(*targets.get(&0).unwrap(), 4);
2558 assert_eq!(*targets.get(&1).unwrap(), 4);
2559 assert_eq!(*targets.get(&2).unwrap(), 4);
2560 }
2561
2562 #[test]
2563 fn test_compute_targets_with_remainder() {
2564 let targets = compute_targets(10, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
2565 let total: usize = targets.values().sum();
2567 assert_eq!(total, 10);
2568 for (_, &count) in &targets {
2569 assert!(count >= 3 && count <= 4);
2570 }
2571 }
2572
2573 #[test]
2576 fn test_symmetric_pairwise_overlap_produces_equal_cells() {
2577 let tmp = TempDir::new().unwrap();
2578
2579 let cpusets = [
2589 ("cell_a", "0-3,20-27"),
2590 ("cell_b", "4-7,20-21,28-33"),
2591 ("cell_c", "8-11,22-23,28-29,34-37"),
2592 ("cell_d", "12-15,24-25,30-31,34-35,38-39"),
2593 ("cell_e", "16-19,26-27,32-33,36-39"),
2594 ];
2595 for (name, cpus) in &cpusets {
2596 let p = tmp.path().join(name);
2597 std::fs::create_dir(&p).unwrap();
2598 std::fs::write(p.join("cpuset.cpus"), format!("{cpus}\n")).unwrap();
2599 }
2600
2601 let mgr = CellManager::new_with_path(
2602 tmp.path().to_path_buf(),
2603 256,
2604 cpumask_for_range(56),
2605 HashSet::new(),
2606 )
2607 .unwrap();
2608
2609 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2610
2611 let workload: Vec<_> = assignments.iter().filter(|a| a.cell_id != 0).collect();
2612 assert_eq!(workload.len(), 5, "Expected 5 workload cells");
2613
2614 let counts: Vec<(u32, usize)> = workload
2616 .iter()
2617 .map(|a| (a.cell_id, a.primary.weight()))
2618 .collect();
2619 for &(cell_id, count) in &counts {
2620 assert_eq!(
2621 count, counts[0].1,
2622 "Cell {} has {} CPUs, expected {} — symmetric inputs \
2623 should produce equal cell sizes. All counts: {:?}",
2624 cell_id, count, counts[0].1, counts,
2625 );
2626 }
2627
2628 for i in 0..assignments.len() {
2630 for j in (i + 1)..assignments.len() {
2631 for cpu in 0..56 {
2632 assert!(
2633 !(assignments[i].primary.test_cpu(cpu)
2634 && assignments[j].primary.test_cpu(cpu)),
2635 "CPU {} assigned to both cell {} and cell {}",
2636 cpu,
2637 assignments[i].cell_id,
2638 assignments[j].cell_id,
2639 );
2640 }
2641 }
2642 }
2643
2644 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
2646 assert_eq!(total, 56, "All CPUs must be assigned");
2647 }
2648}