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.metadata()?.ino();
313 let (cgid, cell_id) =
314 self.create_cell_for_cgroup(&path, cgid).with_context(|| {
315 format!("Failed to create cell for cgroup: {}", path.display())
316 })?;
317 assignments.push((cgid, cell_id));
318 }
319 }
320 Ok(assignments)
321 }
322
323 pub fn process_events(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
330 let mut buffer = [0; 1024];
331 let mut has_events = false;
332
333 loop {
335 match self.inotify.read_events(&mut buffer) {
336 Ok(events) => {
337 if events.into_iter().next().is_some() {
338 has_events = true;
339 } else {
340 break;
341 }
342 }
343 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
344 Err(e) => {
345 return Err(e).context("Failed to read inotify events");
346 }
347 }
348 }
349
350 if !has_events {
351 return Ok((Vec::new(), Vec::new()));
352 }
353
354 self.reconcile_cells()
356 }
357
358 fn reconcile_cells(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
361 let mut new_cells = Vec::new();
362
363 let mut current_paths: HashSet<PathBuf> = HashSet::new();
365 let entries = std::fs::read_dir(&self.cell_parent_path).with_context(|| {
366 format!(
367 "Failed to read cell parent directory: {}",
368 self.cell_parent_path.display()
369 )
370 })?;
371 for entry in entries {
372 let entry = entry.with_context(|| {
373 format!(
374 "Failed to read directory entry in: {}",
375 self.cell_parent_path.display()
376 )
377 })?;
378 let file_type = entry.file_type().with_context(|| {
379 format!("Failed to get file type for: {}", entry.path().display())
380 })?;
381 if file_type.is_dir() {
382 let path = entry.path();
383 if !self.should_exclude(&path) {
384 current_paths.insert(path);
385 }
386 }
387 }
388
389 let mut destroyed_cells: HashSet<u32> = HashSet::new();
391 self.cells.retain(|&cgid, info| {
392 if info.cell_id == 0 {
393 return true; }
395 let cgroup_path = info
397 .cgroup_path
398 .as_ref()
399 .expect("BUG: non-zero cell missing cgroup_path");
400 if current_paths.contains(cgroup_path) {
401 true
402 } else {
403 info!(
404 "Destroyed cell {} for cgroup {} (cgid={})",
405 info.cell_id,
406 cgroup_path.display(),
407 cgid
408 );
409 destroyed_cells.insert(info.cell_id);
410 false
411 }
412 });
413
414 self.cell_id_to_cgid
416 .retain(|cell_id, _| !destroyed_cells.contains(cell_id));
417 self.free_cell_ids.extend(destroyed_cells.iter().copied());
418
419 for path in current_paths {
421 let cgid = path.metadata()?.ino();
422 if self.cells.contains_key(&cgid) {
423 continue; }
425 let (cgid, cell_id) = self
426 .create_cell_for_cgroup(&path, cgid)
427 .with_context(|| format!("Failed to create cell for cgroup: {}", path.display()))?;
428 new_cells.push((cgid, cell_id));
429 }
430
431 Ok((new_cells, destroyed_cells.into_iter().collect()))
432 }
433
434 fn create_cell_for_cgroup(&mut self, path: &Path, cgid: u64) -> Result<(u64, u32)> {
435 let cell_id = self.allocate_cell_id()?;
436
437 let cpuset = Self::read_cpuset(path)?;
438 if let Some(ref mask) = cpuset {
439 debug!(
440 "Cell {} has cpuset: {} (from {})",
441 cell_id,
442 mask.to_cpulist(),
443 path.join("cpuset.cpus").display()
444 );
445 }
446
447 self.cells.insert(
448 cgid,
449 CellInfo {
450 cell_id,
451 cgroup_path: Some(path.to_path_buf()),
452 cgid: Some(cgid),
453 cpuset,
454 },
455 );
456 self.cell_id_to_cgid.insert(cell_id, cgid);
457
458 info!(
459 "Created cell {} for cgroup {} (cgid={})",
460 cell_id,
461 path.display(),
462 cgid
463 );
464
465 Ok((cgid, cell_id))
466 }
467
468 fn read_cpuset(cgroup_path: &Path) -> Result<Option<Cpumask>> {
470 let cpuset_path = cgroup_path.join("cpuset.cpus");
471 match std::fs::read_to_string(&cpuset_path) {
472 Ok(content) => {
473 let content = content.trim();
474 if content.is_empty() {
475 Ok(None)
476 } else {
477 let mask = Cpumask::from_cpulist(content).with_context(|| {
478 format!(
479 "Failed to parse cpuset '{}' from {}",
480 content,
481 cpuset_path.display()
482 )
483 })?;
484 Ok(Some(mask))
485 }
486 }
487 Err(_) => Ok(None),
489 }
490 }
491
492 fn allocate_cell_id(&mut self) -> Result<u32> {
493 if let Some(id) = self.free_cell_ids.pop() {
495 return Ok(id);
496 }
497
498 if self.next_cell_id >= self.max_cells {
499 bail!("Cell ID space exhausted (max_cells={})", self.max_cells);
500 }
501
502 let id = self.next_cell_id;
503 self.next_cell_id += 1;
504 Ok(id)
505 }
506
507 pub fn compute_cpu_assignments(&self, compute_borrowable: bool) -> Result<Vec<CpuAssignment>> {
519 self.compute_cpu_assignments_inner(None, compute_borrowable)
521 }
522
523 pub fn compute_demand_cpu_assignments(
531 &self,
532 cell_demands: &HashMap<u32, f64>,
533 compute_borrowable: bool,
534 ) -> Result<Vec<CpuAssignment>> {
535 self.compute_cpu_assignments_inner(Some(cell_demands), compute_borrowable)
536 }
537
538 fn compute_cpu_assignments_inner(
540 &self,
541 cell_demands: Option<&HashMap<u32, f64>>,
542 compute_borrowable: bool,
543 ) -> Result<Vec<CpuAssignment>> {
544 if let Some(demands) = cell_demands {
546 for (&cell_id, &weight) in demands {
547 if weight < 0.0 {
548 bail!("Cell {} has negative demand weight {}", cell_id, weight);
549 }
550 }
551 }
552
553 let mut contention: HashMap<usize, Vec<u32>> = HashMap::new();
555 for cell_info in self.cells.values() {
556 if let Some(ref cpuset) = cell_info.cpuset {
557 for cpu in cpuset.iter() {
558 contention.entry(cpu).or_default().push(cell_info.cell_id);
559 }
560 }
561 }
562
563 let mut cell_cpus: HashMap<u32, Cpumask> = HashMap::new();
568 let mut contested_cpus: Vec<usize> = Vec::new();
569 let mut unclaimed_cpus: Vec<usize> = Vec::new();
570
571 for cpu in self.all_cpus.iter() {
572 match contention.get(&cpu) {
573 None => unclaimed_cpus.push(cpu),
574 Some(claimants) if claimants.len() == 1 => {
575 let cell_id = claimants[0];
576 cell_cpus
577 .entry(cell_id)
578 .or_insert_with(Cpumask::new)
579 .set_cpu(cpu)
580 .ok();
581 }
582 Some(_) => contested_cpus.push(cpu),
583 }
584 }
585
586 let total_cpu_count = self.all_cpus.weight();
588 let mut all_cells_with_weights: Vec<(u32, f64)> = self
589 .cells
590 .values()
591 .map(|info| {
592 let weight = match cell_demands {
593 Some(demands) => *demands.get(&info.cell_id).ok_or_else(|| {
594 anyhow::anyhow!("Cell {} is missing from demands map", info.cell_id)
595 })?,
596 None => 1.0,
597 };
598 Ok((info.cell_id, weight))
599 })
600 .collect::<Result<Vec<_>>>()?;
601 all_cells_with_weights.sort_by_key(|(cell_id, _)| *cell_id);
602
603 let targets = compute_targets(total_cpu_count, &all_cells_with_weights)?;
604
605 let mut assigned_count: HashMap<u32, usize> = HashMap::new();
607 for (cell_id, mask) in &cell_cpus {
608 assigned_count.insert(*cell_id, mask.weight());
609 }
610
611 let mut contested_groups: HashMap<Vec<u32>, Vec<usize>> = HashMap::new();
613 for cpu in contested_cpus {
614 if let Some(claimants) = contention.get(&cpu) {
615 let mut sorted_claimants = claimants.clone();
616 sorted_claimants.sort();
617 contested_groups
618 .entry(sorted_claimants)
619 .or_default()
620 .push(cpu);
621 }
622 }
623
624 let initial_deficit: HashMap<u32, f64> = targets
630 .iter()
631 .map(|(&cell_id, &target)| {
632 let already = assigned_count.get(&cell_id).copied().unwrap_or(0);
634 let deficit = if target > already {
635 (target - already) as f64
636 } else {
637 0.0
638 };
639 (cell_id, deficit)
640 })
641 .collect();
642
643 for (claimants, cpus) in contested_groups {
644 let mut recipients: Vec<(u32, f64)> = Vec::new();
645 let mut all_zero = true;
646 for &cell_id in &claimants {
647 let deficit = *initial_deficit.get(&cell_id).ok_or_else(|| {
648 anyhow::anyhow!(
649 "BUG: cell {} in contention map but missing from targets",
650 cell_id
651 )
652 })?;
653 if deficit > 0.0 {
654 all_zero = false;
655 }
656 recipients.push((cell_id, deficit));
657 }
658
659 if all_zero {
661 recipients = claimants.iter().map(|&cell_id| (cell_id, 1.0)).collect();
662 }
663
664 let distribution = distribute_cpus_proportional(&cpus, &recipients)?;
665 for (cell_id, assigned_cpus) in distribution {
666 let count = assigned_cpus.len();
667 for cpu in assigned_cpus {
668 cell_cpus
669 .entry(cell_id)
670 .or_insert_with(Cpumask::new)
671 .set_cpu(cpu)
672 .ok();
673 }
674 *assigned_count.entry(cell_id).or_insert(0) += count;
675 }
676 }
677
678 if !unclaimed_cpus.is_empty() {
680 let mut recipients: Vec<(u32, f64)> = Vec::new();
681 let mut all_zero = true;
682
683 for info in self.cells.values() {
684 if info.cpuset.is_some() {
685 continue; }
687 let target = *targets.get(&info.cell_id).ok_or_else(|| {
688 anyhow::anyhow!(
689 "BUG: cell {} is unpinned but missing from targets",
690 info.cell_id
691 )
692 })?;
693 let already = assigned_count.get(&info.cell_id).copied().unwrap_or(0);
695 let deficit = if target > already {
696 (target - already) as f64
697 } else {
698 0.0
699 };
700 if deficit > 0.0 {
701 all_zero = false;
702 }
703 recipients.push((info.cell_id, deficit));
704 }
705 recipients.sort_by_key(|(cell_id, _)| *cell_id);
706
707 if all_zero {
709 recipients = recipients
710 .iter()
711 .map(|(cell_id, _)| (*cell_id, 1.0))
712 .collect();
713 }
714
715 let distribution = distribute_cpus_proportional(&unclaimed_cpus, &recipients)?;
716 for (cell_id, assigned_cpus) in distribution {
717 let count = assigned_cpus.len();
718 for cpu in assigned_cpus {
719 cell_cpus
720 .entry(cell_id)
721 .or_insert_with(Cpumask::new)
722 .set_cpu(cpu)
723 .ok();
724 }
725 *assigned_count.entry(cell_id).or_insert(0) += count;
726 }
727 }
728
729 for info in self.cells.values() {
731 if !cell_cpus.contains_key(&info.cell_id)
732 || cell_cpus
733 .get(&info.cell_id)
734 .map_or(true, |m| m.weight() == 0)
735 {
736 bail!(
737 "Cell {} has no CPUs assigned (nr_cpus={}, num_cells={})",
738 info.cell_id,
739 self.all_cpus.weight(),
740 self.cells.len()
741 );
742 }
743 }
744
745 let assignments: Vec<CpuAssignment> = cell_cpus
747 .into_iter()
748 .map(|(cell_id, primary)| {
749 let borrowable = if compute_borrowable {
750 let mut borrow_mask = self.all_cpus.and(&primary.not());
751
752 if let Some(cell_info) = self.cells.values().find(|c| c.cell_id == cell_id) {
754 if let Some(ref cpuset) = cell_info.cpuset {
755 borrow_mask = borrow_mask.and(cpuset);
756 }
757 }
758
759 Some(borrow_mask)
760 } else {
761 None
762 };
763 CpuAssignment {
764 cell_id,
765 primary,
766 borrowable,
767 }
768 })
769 .collect();
770
771 Ok(assignments)
772 }
773
774 pub fn get_cell_assignments(&self) -> Vec<(u64, u32)> {
777 self.cells
778 .values()
779 .filter(|info| info.cell_id != 0)
780 .map(|info| {
781 (
782 info.cgid.expect("BUG: non-zero cell missing cgid"),
783 info.cell_id,
784 )
785 })
786 .collect()
787 }
788
789 pub fn format_cell_config(&self, cpu_assignments: &[CpuAssignment]) -> String {
792 let mut sorted: Vec<_> = cpu_assignments.iter().collect();
793 sorted.sort_by_key(|a| a.cell_id);
794
795 let mut parts = Vec::new();
796 for assignment in sorted {
797 let cpulist = assignment.primary.to_cpulist();
798 if assignment.cell_id == 0 {
799 parts.push(format!("[0: {}]", cpulist));
800 } else {
801 let name = self
803 .cells
804 .values()
805 .find(|info| info.cell_id == assignment.cell_id)
806 .and_then(|info| {
807 info.cgroup_path
808 .as_ref()
809 .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
810 })
811 .unwrap_or_else(|| "?".to_string());
812 parts.push(format!("[{}({}): {}]", assignment.cell_id, name, cpulist));
813 }
814 }
815 parts.join(" ")
816 }
817
818 pub fn refresh_cpusets(&mut self) -> Result<bool> {
821 let mut changed = false;
822 for info in self.cells.values_mut() {
823 let Some(ref cgroup_path) = info.cgroup_path else {
824 continue; };
826 let new_cpuset = Self::read_cpuset(cgroup_path)?;
827 if new_cpuset != info.cpuset {
828 info!(
829 "Cell {} cpuset changed: {:?} -> {:?} ({})",
830 info.cell_id,
831 info.cpuset.as_ref().map(|m| m.to_cpulist()),
832 new_cpuset.as_ref().map(|m| m.to_cpulist()),
833 cgroup_path.display(),
834 );
835 info.cpuset = new_cpuset;
836 changed = true;
837 }
838 }
839 Ok(changed)
840 }
841}
842
843impl AsFd for CellManager {
844 fn as_fd(&self) -> BorrowedFd<'_> {
845 self.inotify.as_fd()
846 }
847}
848
849#[cfg(test)]
850impl CellManager {
851 fn cell_count(&self) -> usize {
854 self.cells.values().filter(|c| c.cell_id != 0).count()
855 }
856
857 fn get_cell_ids(&self) -> Vec<u32> {
860 self.cells
861 .values()
862 .filter(|c| c.cell_id != 0)
863 .map(|c| c.cell_id)
864 .collect()
865 }
866
867 fn find_cell_by_name(&self, name: &str) -> Option<&CellInfo> {
870 self.cells.values().filter(|c| c.cell_id != 0).find(|c| {
871 c.cgroup_path
872 .as_ref()
873 .and_then(|p| p.file_name())
874 .map(|n| n.to_str() == Some(name))
875 .unwrap_or(false)
876 })
877 }
878}
879
880#[cfg(test)]
881mod tests {
882 use super::*;
883 use tempfile::TempDir;
884
885 fn cpumask_for_range(nr_cpus: usize) -> Cpumask {
886 scx_utils::set_cpumask_test_width(nr_cpus);
887 let mut mask = Cpumask::new();
888 for cpu in 0..nr_cpus {
889 mask.set_cpu(cpu).unwrap();
890 }
891 mask
892 }
893
894 #[test]
897 fn test_scan_empty_directory() {
898 let tmp = TempDir::new().unwrap();
899 let mgr = CellManager::new_with_path(
900 tmp.path().to_path_buf(),
901 256,
902 cpumask_for_range(16),
903 HashSet::new(),
904 )
905 .unwrap();
906
907 assert_eq!(mgr.cell_count(), 0);
908 }
909
910 #[test]
911 fn test_scan_existing_subdirectories() {
912 let tmp = TempDir::new().unwrap();
913
914 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
916 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
917
918 let mgr = CellManager::new_with_path(
919 tmp.path().to_path_buf(),
920 256,
921 cpumask_for_range(16),
922 HashSet::new(),
923 )
924 .unwrap();
925
926 assert_eq!(mgr.cell_count(), 2);
927
928 let cell_ids = mgr.get_cell_ids();
930 assert!(cell_ids.contains(&1));
931 assert!(cell_ids.contains(&2));
932 }
933
934 #[test]
935 fn test_reconcile_detects_new_directories() {
936 let tmp = TempDir::new().unwrap();
937
938 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
940 let mut mgr = CellManager::new_with_path(
941 tmp.path().to_path_buf(),
942 256,
943 cpumask_for_range(16),
944 HashSet::new(),
945 )
946 .unwrap();
947 assert_eq!(mgr.cell_count(), 1);
948
949 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
951
952 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
954 assert_eq!(new_cells.len(), 1);
955 assert_eq!(destroyed_cells.len(), 0);
956 assert_eq!(mgr.cell_count(), 2);
957 }
958
959 #[test]
960 fn test_reconcile_detects_removed_directories() {
961 let tmp = TempDir::new().unwrap();
962
963 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
965 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
966 let mut mgr = CellManager::new_with_path(
967 tmp.path().to_path_buf(),
968 256,
969 cpumask_for_range(16),
970 HashSet::new(),
971 )
972 .unwrap();
973 assert_eq!(mgr.cell_count(), 2);
974
975 std::fs::remove_dir(tmp.path().join("container-b")).unwrap();
977
978 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
980 assert_eq!(new_cells.len(), 0);
981 assert_eq!(destroyed_cells.len(), 1);
982 assert_eq!(mgr.cell_count(), 1);
983 }
984
985 #[test]
986 fn test_cell_id_reuse_after_destruction() {
987 let tmp = TempDir::new().unwrap();
988
989 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
991 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
992 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
993
994 let mut mgr = CellManager::new_with_path(
995 tmp.path().to_path_buf(),
996 256,
997 cpumask_for_range(16),
998 HashSet::new(),
999 )
1000 .unwrap();
1001
1002 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1004 let cell2_id = cell2_info.cell_id;
1005
1006 std::fs::remove_dir(tmp.path().join("cell2")).unwrap();
1008 mgr.reconcile_cells().unwrap();
1009
1010 std::fs::create_dir(tmp.path().join("cell4")).unwrap();
1012 mgr.reconcile_cells().unwrap();
1013
1014 let cell4_info = mgr.find_cell_by_name("cell4").unwrap();
1015 assert_eq!(cell4_info.cell_id, cell2_id);
1016 }
1017
1018 #[test]
1021 fn test_cpu_assignments_no_cells() {
1022 let tmp = TempDir::new().unwrap();
1023 let mgr = CellManager::new_with_path(
1024 tmp.path().to_path_buf(),
1025 256,
1026 cpumask_for_range(16),
1027 HashSet::new(),
1028 )
1029 .unwrap();
1030
1031 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1032
1033 assert_eq!(assignments.len(), 1);
1035 assert_eq!(assignments[0].cell_id, 0);
1036 assert_eq!(assignments[0].primary.weight(), 16);
1037 }
1038
1039 #[test]
1040 fn test_cpu_assignments_proportional() {
1041 let tmp = TempDir::new().unwrap();
1042 std::fs::create_dir(tmp.path().join("container")).unwrap();
1043
1044 let mgr = CellManager::new_with_path(
1045 tmp.path().to_path_buf(),
1046 256,
1047 cpumask_for_range(16),
1048 HashSet::new(),
1049 )
1050 .unwrap();
1051 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1052
1053 assert_eq!(assignments.len(), 2);
1055
1056 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1057 let cell1 = assignments.iter().find(|a| a.cell_id == 1).unwrap();
1058
1059 assert_eq!(cell0.primary.weight(), 8);
1060 assert_eq!(cell1.primary.weight(), 8);
1061 }
1062
1063 #[test]
1064 fn test_cpu_assignments_remainder_to_cell0() {
1065 let tmp = TempDir::new().unwrap();
1066 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1067 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1068
1069 let mgr = CellManager::new_with_path(
1070 tmp.path().to_path_buf(),
1071 256,
1072 cpumask_for_range(10),
1073 HashSet::new(),
1074 )
1075 .unwrap();
1076 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1077
1078 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1080 assert_eq!(cell0.primary.weight(), 4); }
1082
1083 #[test]
1084 fn test_cpu_assignments_too_many_cells() {
1085 let tmp = TempDir::new().unwrap();
1086
1087 for i in 1..=5 {
1089 std::fs::create_dir(tmp.path().join(format!("cell{}", i))).unwrap();
1090 }
1091
1092 let mgr = CellManager::new_with_path(
1094 tmp.path().to_path_buf(),
1095 256,
1096 cpumask_for_range(4),
1097 HashSet::new(),
1098 )
1099 .unwrap();
1100 let result = mgr.compute_cpu_assignments(false);
1101
1102 assert!(result.is_err());
1103 let err_msg = format!("{:#}", result.unwrap_err());
1104 assert!(
1105 err_msg.contains("Not enough CPUs"),
1106 "Expected 'Not enough CPUs' error, got: {}",
1107 err_msg
1108 );
1109 }
1110
1111 #[test]
1112 fn test_cpu_assignments_with_cpusets() {
1113 let tmp = TempDir::new().unwrap();
1114
1115 let cell1_path = tmp.path().join("cell1");
1117 std::fs::create_dir(&cell1_path).unwrap();
1118 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1119
1120 let cell2_path = tmp.path().join("cell2");
1121 std::fs::create_dir(&cell2_path).unwrap();
1122 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1123
1124 let mgr = CellManager::new_with_path(
1125 tmp.path().to_path_buf(),
1126 256,
1127 cpumask_for_range(16),
1128 HashSet::new(),
1129 )
1130 .unwrap();
1131 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1132
1133 assert_eq!(assignments.len(), 3);
1135
1136 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1138 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1139
1140 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1141 let cell1 = assignments
1142 .iter()
1143 .find(|a| a.cell_id == cell1_info.cell_id)
1144 .unwrap();
1145 let cell2 = assignments
1146 .iter()
1147 .find(|a| a.cell_id == cell2_info.cell_id)
1148 .unwrap();
1149
1150 assert_eq!(cell1.primary.weight(), 4);
1152 for cpu in 0..4 {
1153 assert!(cell1.primary.test_cpu(cpu));
1154 }
1155
1156 assert_eq!(cell2.primary.weight(), 4);
1158 for cpu in 8..12 {
1159 assert!(cell2.primary.test_cpu(cpu));
1160 }
1161
1162 assert_eq!(cell0.primary.weight(), 8);
1164 for cpu in 4..8 {
1165 assert!(cell0.primary.test_cpu(cpu));
1166 }
1167 for cpu in 12..16 {
1168 assert!(cell0.primary.test_cpu(cpu));
1169 }
1170 }
1171
1172 #[test]
1173 fn test_cpu_assignments_cpusets_cover_all_cpus() {
1174 let tmp = TempDir::new().unwrap();
1175
1176 let cell1_path = tmp.path().join("cell1");
1178 std::fs::create_dir(&cell1_path).unwrap();
1179 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1180
1181 let cell2_path = tmp.path().join("cell2");
1182 std::fs::create_dir(&cell2_path).unwrap();
1183 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1184
1185 let mgr = CellManager::new_with_path(
1186 tmp.path().to_path_buf(),
1187 256,
1188 cpumask_for_range(16),
1189 HashSet::new(),
1190 )
1191 .unwrap();
1192 let result = mgr.compute_cpu_assignments(false);
1193
1194 assert!(result.is_err());
1196 let err = result.unwrap_err();
1197 let err_msg = format!("{:#}", err);
1198 assert!(
1199 err_msg.contains("Cell 0 has no CPUs assigned"),
1200 "Expected 'Cell 0 has no CPUs assigned' error, got: {}",
1201 err_msg
1202 );
1203 }
1204
1205 #[test]
1206 fn test_cpu_assignments_single_cpuset() {
1207 let tmp = TempDir::new().unwrap();
1208
1209 let cell1_path = tmp.path().join("cell1");
1211 std::fs::create_dir(&cell1_path).unwrap();
1212 std::fs::write(cell1_path.join("cpuset.cpus"), "0,2,4,6\n").unwrap();
1213
1214 let mgr = CellManager::new_with_path(
1215 tmp.path().to_path_buf(),
1216 256,
1217 cpumask_for_range(8),
1218 HashSet::new(),
1219 )
1220 .unwrap();
1221 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1222
1223 assert_eq!(assignments.len(), 2);
1224
1225 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1226 let cell1 = assignments.iter().find(|a| a.cell_id != 0).unwrap();
1227
1228 assert_eq!(cell1.primary.weight(), 4);
1230 for cpu in [0, 2, 4, 6] {
1231 assert!(cell1.primary.test_cpu(cpu));
1232 }
1233
1234 assert_eq!(cell0.primary.weight(), 4);
1236 for cpu in [1, 3, 5, 7] {
1237 assert!(cell0.primary.test_cpu(cpu));
1238 }
1239 }
1240
1241 #[test]
1242 fn test_cpuset_parsing_from_file() {
1243 let tmp = TempDir::new().unwrap();
1244
1245 let cell_path = tmp.path().join("cell1");
1247 std::fs::create_dir(&cell_path).unwrap();
1248 std::fs::write(cell_path.join("cpuset.cpus"), "0-3,8-11,16\n").unwrap();
1249
1250 let mgr = CellManager::new_with_path(
1251 tmp.path().to_path_buf(),
1252 256,
1253 cpumask_for_range(32),
1254 HashSet::new(),
1255 )
1256 .unwrap();
1257
1258 let cell_info = mgr.find_cell_by_name("cell1").unwrap();
1260 let cpuset = cell_info.cpuset.as_ref().unwrap();
1261
1262 assert_eq!(cpuset.weight(), 9); for cpu in 0..4 {
1264 assert!(cpuset.test_cpu(cpu));
1265 }
1266 for cpu in 8..12 {
1267 assert!(cpuset.test_cpu(cpu));
1268 }
1269 assert!(cpuset.test_cpu(16));
1270 }
1271
1272 #[test]
1273 fn test_cpu_assignments_mixed_cpuset_and_no_cpuset() {
1274 let tmp = TempDir::new().unwrap();
1275
1276 let cell1_path = tmp.path().join("cell1");
1278 std::fs::create_dir(&cell1_path).unwrap();
1279 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1280
1281 let cell2_path = tmp.path().join("cell2");
1283 std::fs::create_dir(&cell2_path).unwrap();
1284
1285 let mgr = CellManager::new_with_path(
1286 tmp.path().to_path_buf(),
1287 256,
1288 cpumask_for_range(16),
1289 HashSet::new(),
1290 )
1291 .unwrap();
1292
1293 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1295 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1296 assert!(cell1_info.cpuset.is_some());
1297 assert!(cell2_info.cpuset.is_none());
1298
1299 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1300
1301 assert_eq!(assignments.len(), 3);
1307
1308 let cell1_assignment = assignments
1310 .iter()
1311 .find(|a| a.cell_id == cell1_info.cell_id)
1312 .unwrap();
1313 assert_eq!(cell1_assignment.primary.weight(), 4);
1314
1315 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1317 assert_eq!(cell0.primary.weight(), 7);
1318
1319 let cell2_assignment = assignments
1321 .iter()
1322 .find(|a| a.cell_id == cell2_info.cell_id)
1323 .unwrap();
1324 assert_eq!(cell2_assignment.primary.weight(), 5);
1325 }
1326
1327 #[test]
1330 fn test_cpu_assignments_partial_overlap() {
1331 let tmp = TempDir::new().unwrap();
1332
1333 let cell_a_path = tmp.path().join("cell_a");
1335 std::fs::create_dir(&cell_a_path).unwrap();
1336 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
1337
1338 let cell_b_path = tmp.path().join("cell_b");
1339 std::fs::create_dir(&cell_b_path).unwrap();
1340 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
1341
1342 let mgr = CellManager::new_with_path(
1343 tmp.path().to_path_buf(),
1344 256,
1345 cpumask_for_range(16),
1346 HashSet::new(),
1347 )
1348 .unwrap();
1349 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1350
1351 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1352 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1353
1354 let cell_a = assignments
1355 .iter()
1356 .find(|a| a.cell_id == cell_a_info.cell_id)
1357 .unwrap();
1358 let cell_b = assignments
1359 .iter()
1360 .find(|a| a.cell_id == cell_b_info.cell_id)
1361 .unwrap();
1362 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1363
1364 assert_eq!(cell_a.primary.weight(), 6);
1368 assert_eq!(cell_b.primary.weight(), 6);
1369 assert_eq!(cell0.primary.weight(), 4);
1370
1371 for cpu in 0..4 {
1373 assert!(
1374 cell_a.primary.test_cpu(cpu),
1375 "CPU {} should be in cell_a",
1376 cpu
1377 );
1378 }
1379 for cpu in 8..12 {
1380 assert!(
1381 cell_b.primary.test_cpu(cpu),
1382 "CPU {} should be in cell_b",
1383 cpu
1384 );
1385 }
1386 for cpu in 12..16 {
1387 assert!(
1388 cell0.primary.test_cpu(cpu),
1389 "CPU {} should be in cell0",
1390 cpu
1391 );
1392 }
1393
1394 let cell_a_contested: Vec<_> = (4..8).filter(|&cpu| cell_a.primary.test_cpu(cpu)).collect();
1396 let cell_b_contested: Vec<_> = (4..8).filter(|&cpu| cell_b.primary.test_cpu(cpu)).collect();
1397 assert_eq!(cell_a_contested.len(), 2);
1398 assert_eq!(cell_b_contested.len(), 2);
1399
1400 for cpu in 0..16 {
1402 let mut count = 0;
1403 if cell_a.primary.test_cpu(cpu) {
1404 count += 1;
1405 }
1406 if cell_b.primary.test_cpu(cpu) {
1407 count += 1;
1408 }
1409 if cell0.primary.test_cpu(cpu) {
1410 count += 1;
1411 }
1412 assert!(count <= 1, "CPU {} is assigned to {} cells", cpu, count);
1413 }
1414 }
1415
1416 #[test]
1417 fn test_cpu_assignments_three_way_overlap() {
1418 let tmp = TempDir::new().unwrap();
1419
1420 let cell_a_path = tmp.path().join("cell_a");
1422 std::fs::create_dir(&cell_a_path).unwrap();
1423 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-5\n").unwrap();
1424
1425 let cell_b_path = tmp.path().join("cell_b");
1426 std::fs::create_dir(&cell_b_path).unwrap();
1427 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-5\n").unwrap();
1428
1429 let cell_c_path = tmp.path().join("cell_c");
1430 std::fs::create_dir(&cell_c_path).unwrap();
1431 std::fs::write(cell_c_path.join("cpuset.cpus"), "0-5\n").unwrap();
1432
1433 let mgr = CellManager::new_with_path(
1434 tmp.path().to_path_buf(),
1435 256,
1436 cpumask_for_range(12),
1437 HashSet::new(),
1438 )
1439 .unwrap();
1440 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1441
1442 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1443 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1444 let cell_c_info = mgr.find_cell_by_name("cell_c").unwrap();
1445
1446 let cell_a = assignments
1447 .iter()
1448 .find(|a| a.cell_id == cell_a_info.cell_id)
1449 .unwrap();
1450 let cell_b = assignments
1451 .iter()
1452 .find(|a| a.cell_id == cell_b_info.cell_id)
1453 .unwrap();
1454 let cell_c = assignments
1455 .iter()
1456 .find(|a| a.cell_id == cell_c_info.cell_id)
1457 .unwrap();
1458 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1459
1460 assert_eq!(cell_a.primary.weight(), 2);
1462 assert_eq!(cell_b.primary.weight(), 2);
1463 assert_eq!(cell_c.primary.weight(), 2);
1464
1465 assert_eq!(cell0.primary.weight(), 6);
1467 for cpu in 6..12 {
1468 assert!(cell0.primary.test_cpu(cpu));
1469 }
1470
1471 let total_contested: usize = (0..6)
1473 .filter(|&cpu| {
1474 cell_a.primary.test_cpu(cpu)
1475 || cell_b.primary.test_cpu(cpu)
1476 || cell_c.primary.test_cpu(cpu)
1477 })
1478 .count();
1479 assert_eq!(total_contested, 6);
1480 }
1481
1482 #[test]
1483 fn test_cpu_assignments_odd_contested_count() {
1484 let tmp = TempDir::new().unwrap();
1485
1486 let cell_a_path = tmp.path().join("cell_a");
1488 std::fs::create_dir(&cell_a_path).unwrap();
1489 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-2\n").unwrap();
1490
1491 let cell_b_path = tmp.path().join("cell_b");
1492 std::fs::create_dir(&cell_b_path).unwrap();
1493 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-2\n").unwrap();
1494
1495 let mgr = CellManager::new_with_path(
1496 tmp.path().to_path_buf(),
1497 256,
1498 cpumask_for_range(8),
1499 HashSet::new(),
1500 )
1501 .unwrap();
1502 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1503
1504 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1505 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1506
1507 let cell_a = assignments
1508 .iter()
1509 .find(|a| a.cell_id == cell_a_info.cell_id)
1510 .unwrap();
1511 let cell_b = assignments
1512 .iter()
1513 .find(|a| a.cell_id == cell_b_info.cell_id)
1514 .unwrap();
1515
1516 let total = cell_a.primary.weight() + cell_b.primary.weight();
1519 assert_eq!(total, 3);
1520 assert!(cell_a.primary.weight() >= 1 && cell_a.primary.weight() <= 2);
1521 assert!(cell_b.primary.weight() >= 1 && cell_b.primary.weight() <= 2);
1522
1523 for cpu in 0..3 {
1525 let a_has = cell_a.primary.test_cpu(cpu);
1526 let b_has = cell_b.primary.test_cpu(cpu);
1527 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
1528 }
1529 }
1530
1531 #[test]
1532 fn test_cpu_assignments_complete_overlap() {
1533 let tmp = TempDir::new().unwrap();
1534
1535 let cell_a_path = tmp.path().join("cell_a");
1537 std::fs::create_dir(&cell_a_path).unwrap();
1538 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
1539
1540 let cell_b_path = tmp.path().join("cell_b");
1541 std::fs::create_dir(&cell_b_path).unwrap();
1542 std::fs::write(cell_b_path.join("cpuset.cpus"), "0-7\n").unwrap();
1543
1544 let mgr = CellManager::new_with_path(
1545 tmp.path().to_path_buf(),
1546 256,
1547 cpumask_for_range(16),
1548 HashSet::new(),
1549 )
1550 .unwrap();
1551 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1552
1553 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1554 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1555
1556 let cell_a = assignments
1557 .iter()
1558 .find(|a| a.cell_id == cell_a_info.cell_id)
1559 .unwrap();
1560 let cell_b = assignments
1561 .iter()
1562 .find(|a| a.cell_id == cell_b_info.cell_id)
1563 .unwrap();
1564 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1565
1566 assert_eq!(cell_a.primary.weight(), 4);
1568 assert_eq!(cell_b.primary.weight(), 4);
1569
1570 assert_eq!(cell0.primary.weight(), 8);
1572 for cpu in 8..16 {
1573 assert!(cell0.primary.test_cpu(cpu));
1574 }
1575
1576 for cpu in 0..8 {
1578 let a_has = cell_a.primary.test_cpu(cpu);
1579 let b_has = cell_b.primary.test_cpu(cpu);
1580 assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
1581 }
1582 }
1583
1584 #[test]
1585 fn test_cpu_assignments_no_overlap() {
1586 let tmp = TempDir::new().unwrap();
1588
1589 let cell_a_path = tmp.path().join("cell_a");
1590 std::fs::create_dir(&cell_a_path).unwrap();
1591 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-3\n").unwrap();
1592
1593 let cell_b_path = tmp.path().join("cell_b");
1594 std::fs::create_dir(&cell_b_path).unwrap();
1595 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-7\n").unwrap();
1596
1597 let mgr = CellManager::new_with_path(
1598 tmp.path().to_path_buf(),
1599 256,
1600 cpumask_for_range(16),
1601 HashSet::new(),
1602 )
1603 .unwrap();
1604 let assignments = mgr.compute_cpu_assignments(false).unwrap();
1605
1606 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1607 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1608
1609 let cell_a = assignments
1610 .iter()
1611 .find(|a| a.cell_id == cell_a_info.cell_id)
1612 .unwrap();
1613 let cell_b = assignments
1614 .iter()
1615 .find(|a| a.cell_id == cell_b_info.cell_id)
1616 .unwrap();
1617 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1618
1619 assert_eq!(cell_a.primary.weight(), 4);
1621 for cpu in 0..4 {
1622 assert!(cell_a.primary.test_cpu(cpu));
1623 }
1624
1625 assert_eq!(cell_b.primary.weight(), 4);
1626 for cpu in 4..8 {
1627 assert!(cell_b.primary.test_cpu(cpu));
1628 }
1629
1630 assert_eq!(cell0.primary.weight(), 8);
1632 for cpu in 8..16 {
1633 assert!(cell0.primary.test_cpu(cpu));
1634 }
1635 }
1636
1637 #[test]
1640 fn test_format_cell_config_only_cell0() {
1641 let tmp = TempDir::new().unwrap();
1642 let mgr = CellManager::new_with_path(
1643 tmp.path().to_path_buf(),
1644 256,
1645 cpumask_for_range(8),
1646 HashSet::new(),
1647 )
1648 .unwrap();
1649
1650 let mut mask = Cpumask::new();
1651 for cpu in 0..8 {
1652 mask.set_cpu(cpu).unwrap();
1653 }
1654
1655 let assignments = vec![CpuAssignment {
1656 cell_id: 0,
1657 primary: mask,
1658 borrowable: None,
1659 }];
1660 let result = mgr.format_cell_config(&assignments);
1661
1662 assert_eq!(result, "[0: 0-7]");
1663 }
1664
1665 #[test]
1666 fn test_format_cell_config_with_cells() {
1667 let tmp = TempDir::new().unwrap();
1668 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1669
1670 let mgr = CellManager::new_with_path(
1671 tmp.path().to_path_buf(),
1672 256,
1673 cpumask_for_range(16),
1674 HashSet::new(),
1675 )
1676 .unwrap();
1677
1678 let mut mask0 = Cpumask::new();
1679 for cpu in 0..8 {
1680 mask0.set_cpu(cpu).unwrap();
1681 }
1682
1683 let mut mask1 = Cpumask::new();
1684 for cpu in 8..16 {
1685 mask1.set_cpu(cpu).unwrap();
1686 }
1687
1688 let assignments = vec![
1689 CpuAssignment {
1690 cell_id: 0,
1691 primary: mask0,
1692 borrowable: None,
1693 },
1694 CpuAssignment {
1695 cell_id: 1,
1696 primary: mask1,
1697 borrowable: None,
1698 },
1699 ];
1700 let result = mgr.format_cell_config(&assignments);
1701
1702 assert_eq!(result, "[0: 0-7] [1(container-a): 8-15]");
1703 }
1704
1705 #[test]
1708 fn test_cell_id_exhaustion() {
1709 let tmp = TempDir::new().unwrap();
1710
1711 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1714 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1715
1716 let mut mgr = CellManager::new_with_path(
1717 tmp.path().to_path_buf(),
1718 3,
1719 cpumask_for_range(16),
1720 HashSet::new(),
1721 )
1722 .unwrap();
1723 assert_eq!(mgr.cell_count(), 2); std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1727 let result = mgr.reconcile_cells();
1728
1729 assert!(result.is_err());
1730 let err = result.unwrap_err();
1731 let err_chain = format!("{:#}", err);
1732 assert!(
1733 err_chain.contains("Cell ID space exhausted"),
1734 "Expected exhaustion error, got: {}",
1735 err_chain
1736 );
1737 }
1738
1739 #[test]
1740 fn test_cell_id_reuse_prevents_exhaustion() {
1741 let tmp = TempDir::new().unwrap();
1742
1743 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1745 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1746
1747 let mut mgr = CellManager::new_with_path(
1748 tmp.path().to_path_buf(),
1749 3,
1750 cpumask_for_range(16),
1751 HashSet::new(),
1752 )
1753 .unwrap();
1754 assert_eq!(mgr.cell_count(), 2);
1755
1756 std::fs::remove_dir(tmp.path().join("cell1")).unwrap();
1758 mgr.reconcile_cells().unwrap();
1759 assert_eq!(mgr.cell_count(), 1);
1760
1761 std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1763 let result = mgr.reconcile_cells();
1764 assert!(result.is_ok());
1765 assert_eq!(mgr.cell_count(), 2);
1766 }
1767
1768 #[test]
1771 fn test_scan_excludes_named_cgroups() {
1772 let tmp = TempDir::new().unwrap();
1773
1774 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1775 std::fs::create_dir(tmp.path().join("systemd-workaround.service")).unwrap();
1776 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1777
1778 let exclude = HashSet::from(["systemd-workaround.service".to_string()]);
1779 let mgr = CellManager::new_with_path(
1780 tmp.path().to_path_buf(),
1781 256,
1782 cpumask_for_range(16),
1783 exclude,
1784 )
1785 .unwrap();
1786
1787 assert_eq!(mgr.cell_count(), 2);
1789 assert!(mgr.find_cell_by_name("container-a").is_some());
1790 assert!(mgr.find_cell_by_name("container-b").is_some());
1791 assert!(mgr
1792 .find_cell_by_name("systemd-workaround.service")
1793 .is_none());
1794 }
1795
1796 #[test]
1797 fn test_reconcile_excludes_named_cgroups() {
1798 let tmp = TempDir::new().unwrap();
1799
1800 std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1801
1802 let exclude = HashSet::from(["ignored-service".to_string()]);
1803 let mut mgr = CellManager::new_with_path(
1804 tmp.path().to_path_buf(),
1805 256,
1806 cpumask_for_range(16),
1807 exclude,
1808 )
1809 .unwrap();
1810 assert_eq!(mgr.cell_count(), 1);
1811
1812 std::fs::create_dir(tmp.path().join("ignored-service")).unwrap();
1814 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1815 assert_eq!(new_cells.len(), 0);
1816 assert_eq!(destroyed_cells.len(), 0);
1817 assert_eq!(mgr.cell_count(), 1);
1818
1819 std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1821 let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1822 assert_eq!(new_cells.len(), 1);
1823 assert_eq!(destroyed_cells.len(), 0);
1824 assert_eq!(mgr.cell_count(), 2);
1825 }
1826
1827 #[test]
1830 fn test_borrowable_cpumasks_basic() {
1831 let tmp = TempDir::new().unwrap();
1832
1833 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1835 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1836
1837 let mgr = CellManager::new_with_path(
1838 tmp.path().to_path_buf(),
1839 256,
1840 cpumask_for_range(16),
1841 HashSet::new(),
1842 )
1843 .unwrap();
1844 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1845
1846 for assignment in &assignments {
1848 let borrow_mask = assignment.borrowable.as_ref().unwrap();
1849 let overlap = borrow_mask.and(&assignment.primary);
1851 assert_eq!(
1852 overlap.weight(),
1853 0,
1854 "Cell {} borrowable overlaps with primary",
1855 assignment.cell_id
1856 );
1857 let union = borrow_mask.or(&assignment.primary);
1859 assert_eq!(
1860 union.weight(),
1861 16,
1862 "Cell {} union doesn't cover all CPUs",
1863 assignment.cell_id
1864 );
1865 }
1866 }
1867
1868 #[test]
1869 fn test_borrowable_cpumasks_no_overlap() {
1870 let tmp = TempDir::new().unwrap();
1871
1872 let cell1_path = tmp.path().join("cell1");
1873 std::fs::create_dir(&cell1_path).unwrap();
1874 std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1875
1876 let cell2_path = tmp.path().join("cell2");
1877 std::fs::create_dir(&cell2_path).unwrap();
1878 std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1879
1880 let mgr = CellManager::new_with_path(
1881 tmp.path().to_path_buf(),
1882 256,
1883 cpumask_for_range(16),
1884 HashSet::new(),
1885 )
1886 .unwrap();
1887 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1888
1889 for assignment in &assignments {
1891 let borrow_mask = assignment.borrowable.as_ref().unwrap();
1892 let overlap = borrow_mask.and(&assignment.primary);
1893 assert_eq!(
1894 overlap.weight(),
1895 0,
1896 "Cell {} borrowable overlaps with primary",
1897 assignment.cell_id
1898 );
1899 }
1900 }
1901
1902 #[test]
1903 fn test_borrowable_cpumasks_respects_cpuset() {
1904 let tmp = TempDir::new().unwrap();
1905
1906 let cell1_path = tmp.path().join("cell1");
1908 std::fs::create_dir(&cell1_path).unwrap();
1909 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1910
1911 let cell2_path = tmp.path().join("cell2");
1912 std::fs::create_dir(&cell2_path).unwrap();
1913 std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1914
1915 let mgr = CellManager::new_with_path(
1916 tmp.path().to_path_buf(),
1917 256,
1918 cpumask_for_range(32),
1919 HashSet::new(),
1920 )
1921 .unwrap();
1922 let assignments = mgr.compute_cpu_assignments(true).unwrap();
1923
1924 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1925 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1926
1927 let cell1_assignment = assignments
1931 .iter()
1932 .find(|a| a.cell_id == cell1_info.cell_id)
1933 .unwrap();
1934 let cell1_borrow = cell1_assignment.borrowable.as_ref().unwrap();
1935 for cpu in 8..32 {
1937 assert!(
1938 !cell1_borrow.test_cpu(cpu),
1939 "Cell 1 borrowable should not include CPU {} (outside cpuset)",
1940 cpu
1941 );
1942 }
1943
1944 let cell2_assignment = assignments
1946 .iter()
1947 .find(|a| a.cell_id == cell2_info.cell_id)
1948 .unwrap();
1949 let cell2_borrow = cell2_assignment.borrowable.as_ref().unwrap();
1950 for cpu in 0..8 {
1951 assert!(
1952 !cell2_borrow.test_cpu(cpu),
1953 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
1954 cpu
1955 );
1956 }
1957 for cpu in 16..32 {
1958 assert!(
1959 !cell2_borrow.test_cpu(cpu),
1960 "Cell 2 borrowable should not include CPU {} (outside cpuset)",
1961 cpu
1962 );
1963 }
1964 }
1965
1966 #[test]
1969 fn test_demand_cpu_assignments_all_idle() {
1970 let tmp = TempDir::new().unwrap();
1971 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1972 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1973
1974 let mgr = CellManager::new_with_path(
1975 tmp.path().to_path_buf(),
1976 256,
1977 cpumask_for_range(12),
1978 HashSet::new(),
1979 )
1980 .unwrap();
1981
1982 let demands: HashMap<u32, f64> = [(0, 0.0), (1, 0.0), (2, 0.0)].into();
1984 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
1985
1986 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1988 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1989 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1990 let c1 = assignments
1991 .iter()
1992 .find(|a| a.cell_id == cell1_info.cell_id)
1993 .unwrap();
1994 let c2 = assignments
1995 .iter()
1996 .find(|a| a.cell_id == cell2_info.cell_id)
1997 .unwrap();
1998
1999 assert_eq!(cell0.primary.weight(), 4);
2000 assert_eq!(c1.primary.weight(), 4);
2001 assert_eq!(c2.primary.weight(), 4);
2002 }
2003
2004 #[test]
2005 fn test_demand_cpu_assignments_uneven_demand() {
2006 let tmp = TempDir::new().unwrap();
2007 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2008 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2009
2010 let mgr = CellManager::new_with_path(
2011 tmp.path().to_path_buf(),
2012 256,
2013 cpumask_for_range(12),
2014 HashSet::new(),
2015 )
2016 .unwrap();
2017
2018 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2019 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2020
2021 let demands: HashMap<u32, f64> = [
2023 (0, 50.0),
2024 (cell1_info.cell_id, 100.0),
2025 (cell2_info.cell_id, 1.0),
2026 ]
2027 .into();
2028 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2029
2030 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2031 let c1 = assignments
2032 .iter()
2033 .find(|a| a.cell_id == cell1_info.cell_id)
2034 .unwrap();
2035 let c2 = assignments
2036 .iter()
2037 .find(|a| a.cell_id == cell2_info.cell_id)
2038 .unwrap();
2039
2040 assert!(
2042 c1.primary.weight() > c2.primary.weight(),
2043 "Busy cell ({}) should have more CPUs than idle cell ({})",
2044 c1.primary.weight(),
2045 c2.primary.weight()
2046 );
2047 assert!(c2.primary.weight() >= 1);
2049 assert!(cell0.primary.weight() >= 1);
2050 assert_eq!(
2052 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2053 12
2054 );
2055 }
2056
2057 #[test]
2058 fn test_demand_cpu_assignments_with_cpusets() {
2059 let tmp = TempDir::new().unwrap();
2060
2061 let cell_a_path = tmp.path().join("cell_a");
2063 std::fs::create_dir(&cell_a_path).unwrap();
2064 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2065
2066 let cell_b_path = tmp.path().join("cell_b");
2067 std::fs::create_dir(&cell_b_path).unwrap();
2068 std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
2069
2070 let mgr = CellManager::new_with_path(
2071 tmp.path().to_path_buf(),
2072 256,
2073 cpumask_for_range(16),
2074 HashSet::new(),
2075 )
2076 .unwrap();
2077
2078 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2079 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2080
2081 let demands: HashMap<u32, f64> = [
2083 (0, 10.0),
2084 (cell_a_info.cell_id, 90.0),
2085 (cell_b_info.cell_id, 10.0),
2086 ]
2087 .into();
2088 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2089
2090 let cell_a = assignments
2091 .iter()
2092 .find(|a| a.cell_id == cell_a_info.cell_id)
2093 .unwrap();
2094 let cell_b = assignments
2095 .iter()
2096 .find(|a| a.cell_id == cell_b_info.cell_id)
2097 .unwrap();
2098
2099 assert!(
2103 cell_a.primary.weight() > cell_b.primary.weight(),
2104 "Cell A ({}) should have more CPUs than Cell B ({})",
2105 cell_a.primary.weight(),
2106 cell_b.primary.weight()
2107 );
2108
2109 assert!(cell_a.primary.weight() >= 4);
2111 assert!(cell_b.primary.weight() >= 4);
2112 }
2113
2114 #[test]
2115 fn test_demand_cpu_assignments_idle_cell_gets_floor() {
2116 let tmp = TempDir::new().unwrap();
2117 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2118 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2119
2120 let mgr = CellManager::new_with_path(
2121 tmp.path().to_path_buf(),
2122 256,
2123 cpumask_for_range(12),
2124 HashSet::new(),
2125 )
2126 .unwrap();
2127
2128 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2129 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2130
2131 let demands: HashMap<u32, f64> = [
2133 (0, 50.0),
2134 (cell1_info.cell_id, 100.0),
2135 (cell2_info.cell_id, 0.0),
2136 ]
2137 .into();
2138 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2139
2140 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2141 let c1 = assignments
2142 .iter()
2143 .find(|a| a.cell_id == cell1_info.cell_id)
2144 .unwrap();
2145 let c2 = assignments
2146 .iter()
2147 .find(|a| a.cell_id == cell2_info.cell_id)
2148 .unwrap();
2149
2150 assert_eq!(
2153 c2.primary.weight(),
2154 1,
2155 "Idle cell should get minimum target of 1 CPU"
2156 );
2157 assert!(c1.primary.weight() > cell0.primary.weight());
2159 assert!(c1.primary.weight() > c2.primary.weight());
2160 assert_eq!(
2162 cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2163 12
2164 );
2165 }
2166
2167 #[test]
2168 fn test_demand_cpu_assignments_negative_weight_errors() {
2169 let tmp = TempDir::new().unwrap();
2170 std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2171
2172 let mgr = CellManager::new_with_path(
2173 tmp.path().to_path_buf(),
2174 256,
2175 cpumask_for_range(8),
2176 HashSet::new(),
2177 )
2178 .unwrap();
2179
2180 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2181
2182 let demands: HashMap<u32, f64> = [(0, 50.0), (cell1_info.cell_id, -10.0)].into();
2183 let result = mgr.compute_demand_cpu_assignments(&demands, false);
2184 assert!(result.is_err(), "Negative weight should be rejected");
2185 assert!(
2186 result
2187 .unwrap_err()
2188 .to_string()
2189 .contains("negative demand weight"),
2190 "Error message should mention negative weight"
2191 );
2192 }
2193
2194 #[test]
2197 fn test_deficit_distribution_with_cpusets() {
2198 let tmp = TempDir::new().unwrap();
2202
2203 let cell_a_path = tmp.path().join("cell_a");
2205 std::fs::create_dir(&cell_a_path).unwrap();
2206 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2207
2208 let cell_b_path = tmp.path().join("cell_b");
2210 std::fs::create_dir(&cell_b_path).unwrap();
2211 std::fs::write(cell_b_path.join("cpuset.cpus"), "6-11\n").unwrap();
2212
2213 let mgr = CellManager::new_with_path(
2215 tmp.path().to_path_buf(),
2216 256,
2217 cpumask_for_range(16),
2218 HashSet::new(),
2219 )
2220 .unwrap();
2221
2222 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2223 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2224
2225 let demands: HashMap<u32, f64> = [
2227 (0, 10.0),
2228 (cell_a_info.cell_id, 90.0),
2229 (cell_b_info.cell_id, 10.0),
2230 ]
2231 .into();
2232 let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2233
2234 let cell_a = assignments
2235 .iter()
2236 .find(|a| a.cell_id == cell_a_info.cell_id)
2237 .unwrap();
2238 let cell_b = assignments
2239 .iter()
2240 .find(|a| a.cell_id == cell_b_info.cell_id)
2241 .unwrap();
2242 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2243
2244 assert_eq!(
2246 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2247 16,
2248 );
2249
2250 assert!(
2252 cell_a.primary.weight() > cell_b.primary.weight(),
2253 "cell_a ({}) should have more CPUs than cell_b ({})",
2254 cell_a.primary.weight(),
2255 cell_b.primary.weight()
2256 );
2257
2258 assert!(cell_a.primary.weight() >= 1);
2260 assert!(cell_b.primary.weight() >= 1);
2261 assert!(cell0.primary.weight() >= 1);
2262 }
2263
2264 #[test]
2265 fn test_deficit_distribution_equal_weight_with_exclusive() {
2266 let tmp = TempDir::new().unwrap();
2269
2270 let cell1_path = tmp.path().join("cell1");
2272 std::fs::create_dir(&cell1_path).unwrap();
2273 std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
2274
2275 std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2277
2278 let mgr = CellManager::new_with_path(
2280 tmp.path().to_path_buf(),
2281 256,
2282 cpumask_for_range(16),
2283 HashSet::new(),
2284 )
2285 .unwrap();
2286
2287 let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2288 let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2289
2290 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2291
2292 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2293 let cell1 = assignments
2294 .iter()
2295 .find(|a| a.cell_id == cell1_info.cell_id)
2296 .unwrap();
2297 let cell2 = assignments
2298 .iter()
2299 .find(|a| a.cell_id == cell2_info.cell_id)
2300 .unwrap();
2301
2302 assert_eq!(cell1.primary.weight(), 8); assert_eq!(
2312 cell0.primary.weight() + cell1.primary.weight() + cell2.primary.weight(),
2313 16,
2314 );
2315
2316 assert!(
2318 cell0.primary.weight() >= cell2.primary.weight(),
2319 "cell0 ({}) should have >= CPUs than cell2 ({})",
2320 cell0.primary.weight(),
2321 cell2.primary.weight()
2322 );
2323 }
2324
2325 #[test]
2326 fn test_deficit_all_cells_exceed_target() {
2327 let tmp = TempDir::new().unwrap();
2330
2331 let cell_a_path = tmp.path().join("cell_a");
2333 std::fs::create_dir(&cell_a_path).unwrap();
2334 std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2335
2336 let cell_b_path = tmp.path().join("cell_b");
2338 std::fs::create_dir(&cell_b_path).unwrap();
2339 std::fs::write(cell_b_path.join("cpuset.cpus"), "8-13\n").unwrap();
2340
2341 let mgr = CellManager::new_with_path(
2348 tmp.path().to_path_buf(),
2349 256,
2350 cpumask_for_range(20),
2351 HashSet::new(),
2352 )
2353 .unwrap();
2354
2355 let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2356 let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2357
2358 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2359
2360 let cell_a = assignments
2361 .iter()
2362 .find(|a| a.cell_id == cell_a_info.cell_id)
2363 .unwrap();
2364 let cell_b = assignments
2365 .iter()
2366 .find(|a| a.cell_id == cell_b_info.cell_id)
2367 .unwrap();
2368 let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2369
2370 let cell_a_contested: Vec<_> = (8..10)
2373 .filter(|&cpu| cell_a.primary.test_cpu(cpu))
2374 .collect();
2375 let cell_b_contested: Vec<_> = (8..10)
2376 .filter(|&cpu| cell_b.primary.test_cpu(cpu))
2377 .collect();
2378 assert_eq!(
2379 cell_a_contested.len(),
2380 0,
2381 "cell_a should get 0 contested CPUs (exceeded target)"
2382 );
2383 assert_eq!(
2384 cell_b_contested.len(),
2385 2,
2386 "cell_b should get all 2 contested CPUs (has deficit)"
2387 );
2388
2389 assert_eq!(
2391 cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2392 20,
2393 );
2394
2395 assert!(cell0.primary.weight() >= 1);
2397 assert!(cell_a.primary.weight() >= 1);
2398 assert!(cell_b.primary.weight() >= 1);
2399 }
2400
2401 #[test]
2404 fn test_distribute_proportional_basic() {
2405 let cpus: Vec<usize> = (0..8).collect();
2406 let recipients = vec![(0, 3.0), (1, 1.0)];
2407 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2408
2409 assert_eq!(result.get(&0).unwrap().len(), 6);
2411 assert_eq!(result.get(&1).unwrap().len(), 2);
2412 }
2413
2414 #[test]
2415 fn test_distribute_proportional_zero_weight_gets_nothing() {
2416 let cpus: Vec<usize> = (0..6).collect();
2417 let recipients = vec![(0, 1.0), (1, 0.0)];
2418 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2419
2420 assert_eq!(result.get(&0).unwrap().len(), 6);
2422 assert!(result.get(&1).is_none() || result.get(&1).unwrap().is_empty());
2423 }
2424
2425 #[test]
2426 fn test_distribute_proportional_all_zero_fallback() {
2427 let cpus: Vec<usize> = (0..6).collect();
2428 let recipients = vec![(0, 0.0), (1, 0.0)];
2429 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2430
2431 assert_eq!(result.get(&0).unwrap().len(), 3);
2433 assert_eq!(result.get(&1).unwrap().len(), 3);
2434 }
2435
2436 #[test]
2437 fn test_distribute_proportional_skewed_weights_floor_guarantee() {
2438 let cpus: Vec<usize> = (0..38).collect();
2441 let recipients = vec![(2, 1.0), (3, 100.0)];
2442 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2443
2444 let cell2_count = result.get(&2).map_or(0, |v| v.len());
2446 let cell3_count = result.get(&3).map_or(0, |v| v.len());
2447 assert!(
2448 cell2_count >= 1,
2449 "cell 2 must get at least 1 CPU, got {}",
2450 cell2_count,
2451 );
2452 assert!(
2453 cell3_count >= 1,
2454 "cell 3 must get at least 1 CPU, got {}",
2455 cell3_count,
2456 );
2457 assert_eq!(cell2_count + cell3_count, 38, "all CPUs must be assigned");
2458 }
2459
2460 #[test]
2461 fn test_distribute_proportional_overlapping_cpusets_no_starvation() {
2462 let cpus: Vec<usize> = (24..62).collect(); let recipients = vec![(2, 1.0), (3, 87.0)];
2467 let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
2468
2469 let cell2_count = result.get(&2).map_or(0, |v| v.len());
2470 let cell3_count = result.get(&3).map_or(0, |v| v.len());
2471 assert!(
2472 cell2_count >= 1,
2473 "cell 2 must not be starved, got {} CPUs",
2474 cell2_count,
2475 );
2476 assert!(
2477 cell3_count >= 1,
2478 "cell 3 must not be starved, got {} CPUs",
2479 cell3_count,
2480 );
2481 assert_eq!(cell2_count + cell3_count, 38);
2482 }
2483
2484 #[test]
2487 fn test_compute_targets_equal_weight() {
2488 let targets = compute_targets(12, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
2489 assert_eq!(*targets.get(&0).unwrap(), 4);
2490 assert_eq!(*targets.get(&1).unwrap(), 4);
2491 assert_eq!(*targets.get(&2).unwrap(), 4);
2492 }
2493
2494 #[test]
2495 fn test_compute_targets_with_remainder() {
2496 let targets = compute_targets(10, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
2497 let total: usize = targets.values().sum();
2499 assert_eq!(total, 10);
2500 for (_, &count) in &targets {
2501 assert!(count >= 3 && count <= 4);
2502 }
2503 }
2504
2505 #[test]
2508 fn test_symmetric_pairwise_overlap_produces_equal_cells() {
2509 let tmp = TempDir::new().unwrap();
2510
2511 let cpusets = [
2521 ("cell_a", "0-3,20-27"),
2522 ("cell_b", "4-7,20-21,28-33"),
2523 ("cell_c", "8-11,22-23,28-29,34-37"),
2524 ("cell_d", "12-15,24-25,30-31,34-35,38-39"),
2525 ("cell_e", "16-19,26-27,32-33,36-39"),
2526 ];
2527 for (name, cpus) in &cpusets {
2528 let p = tmp.path().join(name);
2529 std::fs::create_dir(&p).unwrap();
2530 std::fs::write(p.join("cpuset.cpus"), format!("{cpus}\n")).unwrap();
2531 }
2532
2533 let mgr = CellManager::new_with_path(
2534 tmp.path().to_path_buf(),
2535 256,
2536 cpumask_for_range(56),
2537 HashSet::new(),
2538 )
2539 .unwrap();
2540
2541 let assignments = mgr.compute_cpu_assignments(false).unwrap();
2542
2543 let workload: Vec<_> = assignments.iter().filter(|a| a.cell_id != 0).collect();
2544 assert_eq!(workload.len(), 5, "Expected 5 workload cells");
2545
2546 let counts: Vec<(u32, usize)> = workload
2548 .iter()
2549 .map(|a| (a.cell_id, a.primary.weight()))
2550 .collect();
2551 for &(cell_id, count) in &counts {
2552 assert_eq!(
2553 count, counts[0].1,
2554 "Cell {} has {} CPUs, expected {} — symmetric inputs \
2555 should produce equal cell sizes. All counts: {:?}",
2556 cell_id, count, counts[0].1, counts,
2557 );
2558 }
2559
2560 for i in 0..assignments.len() {
2562 for j in (i + 1)..assignments.len() {
2563 for cpu in 0..56 {
2564 assert!(
2565 !(assignments[i].primary.test_cpu(cpu)
2566 && assignments[j].primary.test_cpu(cpu)),
2567 "CPU {} assigned to both cell {} and cell {}",
2568 cpu,
2569 assignments[i].cell_id,
2570 assignments[j].cell_id,
2571 );
2572 }
2573 }
2574 }
2575
2576 let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
2578 assert_eq!(total, 56, "All CPUs must be assigned");
2579 }
2580}