Skip to main content

scx_mitosis/
cell_manager.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5
6//! Cell manager for userspace-driven cell creation.
7//!
8//! This module implements the `--cell-parent-cgroup` mode where cells are created
9//! for direct child cgroups of a specified parent. Uses inotify to watch for
10//! cgroup creation/destruction and manages cell ID allocation.
11
12use 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/// Information about a cell created for a cgroup
23#[derive(Debug)]
24pub struct CellInfo {
25    pub cell_id: u32,
26    pub cgroup_path: Option<PathBuf>,
27    pub cgid: Option<u64>,
28    /// Optional cpuset mask if the cgroup has cpuset.cpus configured
29    pub cpuset: Option<Cpumask>,
30}
31
32/// Compute the global target CPU count for each cell.
33///
34/// Each cell gets a floor of 1 CPU, with the remainder distributed proportionally
35/// by weight. Returns only counts, not actual CPU assignments.
36fn 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        // Equal division fallback
53        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
93/// Distribute CPUs among recipients proportionally by weight.
94///
95/// Recipients with weight 0 receive 0 CPUs. If all weights are 0, falls back to
96/// equal division. As a post-processing step, if any positive-weight recipient got
97/// 0 CPUs, 1 CPU is stolen from the recipient with the highest allocation (that has
98/// > 1) to prevent starvation.
99fn 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        // Equal division fallback
115        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        // Standard proportional distribution
127        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        // Post-processing floor guarantee: if any positive-weight recipient got 0
149        // CPUs, steal 1 from the recipient with the highest allocation (> 1).
150        // This prevents death spirals where cells with low but non-zero demand
151        // get starved of CPUs entirely.
152        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; // No donor with > 1 CPU available
167            };
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    // Assign actual CPU numbers
178    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/// Result of CPU assignment computation, containing both primary and optional borrowable masks.
200#[derive(Debug)]
201pub struct CpuAssignment {
202    pub cell_id: u32,
203    pub primary: Cpumask,
204    pub borrowable: Option<Cpumask>,
205}
206
207/// Manages cells for direct child cgroups of a specified parent
208pub struct CellManager {
209    cell_parent_path: PathBuf,
210    inotify: Inotify,
211    /// Maps cgroup ID to cell info
212    cells: HashMap<u64, CellInfo>,
213    /// Maps cell ID to cgroup ID (for reverse lookup)
214    cell_id_to_cgid: HashMap<u32, u64>,
215    /// Freed cell IDs available for reuse
216    free_cell_ids: Vec<u32>,
217    next_cell_id: u32,
218    max_cells: u32,
219    /// Cpumask of all CPUs in the system (from topology)
220    all_cpus: Cpumask,
221    /// Cgroup directory names to exclude from cell creation
222    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, // Cell 0 is reserved for root
258            max_cells,
259            all_cpus,
260            exclude_names: exclude,
261        };
262
263        // Insert cell 0 as a permanent entry. cgid 0 is a safe sentinel —
264        // real cgroup inode numbers are always > 0.
265        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        // Scan for existing children at startup
277        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    /// Process pending inotify events. Returns list of (cgid, cell_id) for new cells
324    /// and list of cell_ids that were destroyed.
325    ///
326    /// Rather than processing individual events, we simply check if any events occurred
327    /// and then rescan the directory to reconcile state. This is simpler and handles
328    /// edge cases like inotify queue overflow gracefully.
329    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        // Drain all pending events
334        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        // Rescan directory and reconcile with our tracked state
355        self.reconcile_cells()
356    }
357
358    /// Reconcile our tracked cells with the actual cgroup directory contents.
359    /// Returns (new_cells, destroyed_cells).
360    fn reconcile_cells(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
361        let mut new_cells = Vec::new();
362
363        // Find current cgroups on disk
364        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        // Remove cells for cgroups that no longer exist
390        let mut destroyed_cells: HashSet<u32> = HashSet::new();
391        self.cells.retain(|&cgid, info| {
392            if info.cell_id == 0 {
393                return true; // Cell 0 is permanent
394            }
395            // Non-zero cells always have a cgroup_path
396            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        // Update tracking structures for destroyed cells
415        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        // Find new cgroups that we don't have cells for
420        for path in current_paths {
421            let cgid = path.metadata()?.ino();
422            if self.cells.contains_key(&cgid) {
423                continue; // Already have a cell for this cgroup
424            }
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    /// Read cpuset.cpus from a cgroup path. Returns None if empty or unavailable.
469    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            // File doesn't exist - cpuset controller is not enabled for this cgroup
488            Err(_) => Ok(None),
489        }
490    }
491
492    fn allocate_cell_id(&mut self) -> Result<u32> {
493        // Prefer reusing freed IDs to keep cell ID space compact
494        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    /// Compute CPU assignments for all cells.
508    ///
509    /// When cpusets overlap, contested CPUs are divided proportionally among claimants.
510    /// Unclaimed CPUs go to cell 0 and any unpinned cells (cells without cpusets).
511    ///
512    /// If `compute_borrowable` is true, each assignment includes a borrowable cpumask
513    /// (all system CPUs minus the cell's own, intersected with cpuset if present).
514    /// Without demand data, borrowable masks are uncapped.
515    ///
516    /// Returns a Vec of CpuAssignment, or an error if any cell would
517    /// receive zero CPUs (which indicates too many cells for available CPUs).
518    pub fn compute_cpu_assignments(&self, compute_borrowable: bool) -> Result<Vec<CpuAssignment>> {
519        // Use equal weights for all cells (no demand data)
520        self.compute_cpu_assignments_inner(None, compute_borrowable)
521    }
522
523    /// Compute CPU assignments weighted by per-cell demand.
524    ///
525    /// `cell_demands` maps cell_id -> smoothed_util_pct. All active cells must be
526    /// present in the map; missing entries or negative weights are errors.
527    ///
528    /// If `compute_borrowable` is true, each assignment includes a borrowable cpumask
529    /// (all system CPUs minus the cell's own, intersected with cpuset if present).
530    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    /// Internal implementation shared by equal-weight and demand-weighted assignment.
539    fn compute_cpu_assignments_inner(
540        &self,
541        cell_demands: Option<&HashMap<u32, f64>>,
542        compute_borrowable: bool,
543    ) -> Result<Vec<CpuAssignment>> {
544        // Validate that all demand weights are non-negative
545        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        // Phase 1: Build contention map - for each CPU, track which cells claim it
554        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        // Phase 2: Categorize CPUs and build initial assignments
564        // - Exclusive: claimed by exactly 1 cell -> assigned directly
565        // - Contested: claimed by 2+ cells -> distributed by weight
566        // - Unclaimed: no cpuset claims it -> shared among cell 0 + unpinned cells
567        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        // Compute global targets and initialize running count of CPUs assigned per cell
587        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        // Seed assigned_count from exclusive assignments
606        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        // Phase 3: Distribute contested CPUs among claimants using deficit weights
612        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        // Freeze deficit weights before processing any group. Each group is an
625        // independent allocation decision, so a cell's weight should not depend
626        // on HashMap iteration order (i.e., which other groups were processed
627        // first). Using the initial deficit (target - exclusive_count) as weight
628        // for all groups makes the result deterministic.
629        let initial_deficit: HashMap<u32, f64> = targets
630            .iter()
631            .map(|(&cell_id, &target)| {
632                // Cells with no exclusive CPUs have no entry yet; 0 is correct.
633                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 claimants already meet/exceed their target, fall back to equal weights
660            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        // Phase 4: Distribute unclaimed CPUs among unpinned cells (including cell 0) using deficit weights
679        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; // pinned cells don't receive unclaimed CPUs
686                }
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                // Cells with no exclusive CPUs have no entry yet; 0 is correct.
694                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 recipients already meet/exceed their target, fall back to equal weights
708            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        // Phase 5: Verify all cells have at least one CPU assigned
730        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        // Phase 6: Build CpuAssignment results, optionally computing borrowable masks
746        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 this cell has a cpuset, restrict borrowable to it
753                    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    /// Returns all cell assignments as (cgid, cell_id) pairs.
775    /// Used to configure BPF with cgroup-to-cell mappings.
776    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    /// Format the cell configuration as a compact string for logging.
790    /// Example output: "[0: 0-7] [1(container-a): 8-15] [2(container-b): 16-23]"
791    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                // Find cgroup name for this cell
802                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    /// Re-read cpuset.cpus for all cells and update stored cpusets.
819    /// Returns true if any cell's cpuset changed.
820    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; // cell 0 has no cgroup
825            };
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    /// Returns the number of cells created for cgroups.
852    /// Does not include cell 0 (the implicit root cell).
853    fn cell_count(&self) -> usize {
854        self.cells.values().filter(|c| c.cell_id != 0).count()
855    }
856
857    /// Get all cell IDs for cells created for cgroups.
858    /// Does not include cell 0 (the implicit root cell).
859    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    /// Find a cell by cgroup directory name.
868    /// Only searches cells created for cgroups, not cell 0.
869    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    // ==================== Cell scanning and creation tests ====================
895
896    #[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        // Create some "cgroup" directories
915        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        // Verify cells were assigned IDs 1 and 2
929        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        // Start with one directory
939        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        // Add another directory
950        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
951
952        // Reconcile should detect it
953        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        // Start with two directories
964        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        // Remove one directory
976        std::fs::remove_dir(tmp.path().join("container-b")).unwrap();
977
978        // Reconcile should detect it
979        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        // Create directories
990        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        // Find cell2's ID
1003        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1004        let cell2_id = cell2_info.cell_id;
1005
1006        // Remove cell2
1007        std::fs::remove_dir(tmp.path().join("cell2")).unwrap();
1008        mgr.reconcile_cells().unwrap();
1009
1010        // Add a new directory - should reuse cell2's ID
1011        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    // ==================== compute_cpu_assignments tests ====================
1019
1020    #[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        // Only cell 0 with all CPUs
1034        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        // 16 CPUs / 2 cells = 8 each
1054        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        // 10 CPUs / 3 cells = 3 each + 1 remainder to cell 0
1079        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1080        assert_eq!(cell0.primary.weight(), 4); // 3 + 1 remainder
1081    }
1082
1083    #[test]
1084    fn test_cpu_assignments_too_many_cells() {
1085        let tmp = TempDir::new().unwrap();
1086
1087        // Create more cells than CPUs
1088        for i in 1..=5 {
1089            std::fs::create_dir(tmp.path().join(format!("cell{}", i))).unwrap();
1090        }
1091
1092        // Only 4 CPUs but 6 cells (cell 0 + 5 user cells)
1093        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        // Create cgroup directories with cpuset files
1116        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        // Should have 3 assignments: cell1, cell2, and cell0
1134        assert_eq!(assignments.len(), 3);
1135
1136        // Find each cell's assignment using find_cell_by_name
1137        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        // cell1 gets CPUs 0-3
1151        assert_eq!(cell1.primary.weight(), 4);
1152        for cpu in 0..4 {
1153            assert!(cell1.primary.test_cpu(cpu));
1154        }
1155
1156        // cell2 gets CPUs 8-11
1157        assert_eq!(cell2.primary.weight(), 4);
1158        for cpu in 8..12 {
1159            assert!(cell2.primary.test_cpu(cpu));
1160        }
1161
1162        // cell0 gets remaining CPUs: 4-7, 12-15
1163        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        // Create cgroups that cover all CPUs - cell 0 gets nothing, which is an error
1177        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        // Should error because cell 0 has no CPUs
1195        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        // Only one cell with a cpuset
1210        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        // cell1 gets even CPUs
1229        assert_eq!(cell1.primary.weight(), 4);
1230        for cpu in [0, 2, 4, 6] {
1231            assert!(cell1.primary.test_cpu(cpu));
1232        }
1233
1234        // cell0 gets odd CPUs
1235        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        // Test various cpuset formats
1246        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        // Find the cell and verify its cpuset was parsed correctly
1259        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); // 4 + 4 + 1
1263        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        // cell1 has a cpuset
1277        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        // cell2 has NO cpuset (no cpuset.cpus file)
1282        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        // Verify cell1 has cpuset, cell2 doesn't
1294        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        // cell1 (pinned) gets its cpuset: 0-3 (4 CPUs)
1302        // Targets (equal weight, 3 cells, 16 CPUs): cell0=6, cell1=5, cell2=5
1303        // cell1 has 4 exclusive. Deficit = 1 (but can't participate in unclaimed)
1304        // 12 unclaimed CPUs split by deficit: cell0 deficit=6, cell2 deficit=5
1305        // Result: cell0=7, cell1=4, cell2=5
1306        assert_eq!(assignments.len(), 3);
1307
1308        // cell1 gets its cpuset (0-3)
1309        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        // cell0 gets 7 CPUs (deficit-proportional share of unclaimed)
1316        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1317        assert_eq!(cell0.primary.weight(), 7);
1318
1319        // cell2 (unpinned) gets 5 CPUs
1320        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    // ==================== Overlapping cpuset tests ====================
1328
1329    #[test]
1330    fn test_cpu_assignments_partial_overlap() {
1331        let tmp = TempDir::new().unwrap();
1332
1333        // Cell A (cpuset 0-7) and Cell B (cpuset 4-11) - overlap on 4-7
1334        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        // Cell A gets exclusive 0-3 (4 CPUs) + half of contested 4-7 (2 CPUs) = 6 CPUs
1365        // Cell B gets half of contested 4-7 (2 CPUs) + exclusive 8-11 (4 CPUs) = 6 CPUs
1366        // Cell 0 gets unclaimed 12-15 (4 CPUs)
1367        assert_eq!(cell_a.primary.weight(), 6);
1368        assert_eq!(cell_b.primary.weight(), 6);
1369        assert_eq!(cell0.primary.weight(), 4);
1370
1371        // Verify exclusive CPUs went to correct cells
1372        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        // Verify contested CPUs 4-7 are split - each cell gets exactly 2
1395        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        // No CPU should be assigned to multiple cells
1401        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        // All three cells claim CPUs 0-5
1421        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        // 6 contested CPUs / 3 cells = 2 each
1461        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        // Cell 0 gets unclaimed 6-11 (6 CPUs)
1466        assert_eq!(cell0.primary.weight(), 6);
1467        for cpu in 6..12 {
1468            assert!(cell0.primary.test_cpu(cpu));
1469        }
1470
1471        // Verify total contested CPUs assigned = 6 (no duplicates)
1472        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        // Two cells contesting 3 CPUs (odd number - can't split evenly)
1487        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        // 3 CPUs / 2 cells = 1 each + 1 remainder
1517        // One cell gets 2, the other gets 1
1518        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        // No overlap in assignments
1524        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        // Two cells with identical cpusets
1536        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        // 8 contested CPUs / 2 cells = 4 each
1567        assert_eq!(cell_a.primary.weight(), 4);
1568        assert_eq!(cell_b.primary.weight(), 4);
1569
1570        // Cell 0 gets unclaimed 8-15 (8 CPUs)
1571        assert_eq!(cell0.primary.weight(), 8);
1572        for cpu in 8..16 {
1573            assert!(cell0.primary.test_cpu(cpu));
1574        }
1575
1576        // Verify no overlap between cell_a and cell_b
1577        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        // This verifies existing non-overlapping behavior still works
1587        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        // No overlap - each cell gets its exact cpuset
1620        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        // Cell 0 gets remaining 8-15
1631        assert_eq!(cell0.primary.weight(), 8);
1632        for cpu in 8..16 {
1633            assert!(cell0.primary.test_cpu(cpu));
1634        }
1635    }
1636
1637    // ==================== format_cell_config tests ====================
1638
1639    #[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    // ==================== Cell ID exhaustion tests ====================
1706
1707    #[test]
1708    fn test_cell_id_exhaustion() {
1709        let tmp = TempDir::new().unwrap();
1710
1711        // Create a manager with max_cells=3 (can allocate cell IDs 1 and 2)
1712        // Cell 0 is reserved, so we can create 2 cells before exhaustion
1713        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); // cell1 + cell2
1724
1725        // Adding a third cell should fail due to exhaustion
1726        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        // Create a manager with max_cells=3
1744        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        // Remove cell1 to free up its ID
1757        std::fs::remove_dir(tmp.path().join("cell1")).unwrap();
1758        mgr.reconcile_cells().unwrap();
1759        assert_eq!(mgr.cell_count(), 1);
1760
1761        // Now adding cell3 should succeed by reusing the freed ID
1762        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    // ==================== Exclusion tests ====================
1769
1770    #[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        // Only 2 cells — the excluded cgroup is not a cell
1788        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        // Add an excluded cgroup — should not become a cell
1813        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        // Add a non-excluded cgroup — should become a cell
1820        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    // ==================== Borrowable cpumask tests ====================
1828
1829    #[test]
1830    fn test_borrowable_cpumasks_basic() {
1831        let tmp = TempDir::new().unwrap();
1832
1833        // Create 2 cells without cpusets
1834        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        // Each cell should be able to borrow CPUs from other cells
1847        for assignment in &assignments {
1848            let borrow_mask = assignment.borrowable.as_ref().unwrap();
1849            // borrowable should have no overlap with primary
1850            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            // borrowable + primary should cover all CPUs
1858            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        // Verify no cell's borrowable mask overlaps with its own primary
1890        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        // Cell 1 has cpuset 0-7, Cell 2 has cpuset 8-15
1907        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        // Cell 1's borrowable should be restricted to its cpuset (0-7),
1928        // minus its own CPUs. Since cell1 gets some of 0-7 as primary,
1929        // the borrowable within 0-7 is whatever it doesn't own.
1930        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        // Cell 1's borrowable should NOT include CPUs outside its cpuset (0-7)
1936        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        // Cell 2's borrowable should be restricted to its cpuset (8-15)
1945        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    // ==================== compute_demand_cpu_assignments tests ====================
1967
1968    #[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        // All cells idle (0 demand) -> falls back to equal division
1983        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        // 12 / 3 = 4 each
1987        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        // cell1 is very busy (100%), cell2 is idle (1%), cell0 is moderate (50%)
2022        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        // Busy cell should get more CPUs than idle cell
2041        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        // Each cell should have at least 1 CPU (floor guarantee)
2048        assert!(c2.primary.weight() >= 1);
2049        assert!(cell0.primary.weight() >= 1);
2050        // Total should be 12
2051        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        // Two cells with overlapping cpusets
2062        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        // Cell A is much busier than Cell B
2082        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        // Cell A should get more of the contested CPUs 4-7
2100        // Exclusive: A gets 0-3, B gets 8-11
2101        // Contested 4-7: A should get more due to higher demand
2102        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        // Both should have at least their exclusive CPUs
2110        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        // cell1 is very busy, cell2 is completely idle (weight 0)
2132        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        // Idle cell gets its minimum target (1 CPU) via deficit-based distribution:
2151        // compute_targets assigns a target of 1, giving it a deficit of 1.
2152        assert_eq!(
2153            c2.primary.weight(),
2154            1,
2155            "Idle cell should get minimum target of 1 CPU"
2156        );
2157        // Busy cell should get the most CPUs
2158        assert!(c1.primary.weight() > cell0.primary.weight());
2159        assert!(c1.primary.weight() > c2.primary.weight());
2160        // Total should be 12
2161        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    // ==================== Deficit distribution tests ====================
2195
2196    #[test]
2197    fn test_deficit_distribution_with_cpusets() {
2198        // Two pinned cells with overlapping cpusets + cell0.
2199        // One cell has high demand weight -> gets more contested CPUs.
2200        // Cell that exceeds target from exclusive gets 0 contested.
2201        let tmp = TempDir::new().unwrap();
2202
2203        // cell_a: cpuset 0-9 (10 CPUs)
2204        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        // cell_b: cpuset 6-11 (6 CPUs), overlaps with cell_a on 6-9
2209        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        // 16 CPUs total
2214        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        // cell_a has high demand (90), cell_b has low demand (10), cell0 moderate (10)
2226        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        // Verify total = 16
2245        assert_eq!(
2246            cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2247            16,
2248        );
2249
2250        // cell_a should get the most CPUs (high demand)
2251        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        // Every cell should have at least 1 CPU
2259        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        // Equal weights, one cell has cpuset covering half the CPUs.
2267        // The deficit adjustment should give the other cell(s) more unclaimed CPUs.
2268        let tmp = TempDir::new().unwrap();
2269
2270        // cell1 has cpuset 0-7 (8 CPUs exclusive, no overlap)
2271        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        // cell2 has no cpuset (unpinned)
2276        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2277
2278        // 16 CPUs total, 3 cells (cell0, cell1, cell2), equal weight
2279        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        // Targets (equal weight, 3 cells, 16 CPUs): cell0=6, cell1=5, cell2=5
2303        // cell1 has 8 exclusive (exceeds target of 5), deficit = 0
2304        // 8 unclaimed CPUs split by deficit: cell0 deficit=6, cell2 deficit=5
2305        // cell0 gets ceil(6/11*8) ~= 4, cell2 gets floor(5/11*8) ~= 4
2306        // But exact split: 6/11*8 = 4.36, 5/11*8 = 3.63
2307        // floor: cell0=4, cell2=3 -> assigned=7, remainder=1 -> cell0 gets it
2308        // Result: cell0=5, cell1=8, cell2=3
2309
2310        assert_eq!(cell1.primary.weight(), 8); // Gets all its exclusive CPUs
2311        assert_eq!(
2312            cell0.primary.weight() + cell1.primary.weight() + cell2.primary.weight(),
2313            16,
2314        );
2315
2316        // cell0 should get more unclaimed CPUs than cell2 (higher deficit)
2317        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        // Scenario where all claimants in a contested group already exceed their
2328        // global target from exclusive CPUs alone. Verify fallback to equal distribution.
2329        let tmp = TempDir::new().unwrap();
2330
2331        // cell_a: cpuset 0-9 (10 CPUs, overlaps with cell_b on 8-9)
2332        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        // cell_b: cpuset 8-13 (6 CPUs, overlaps with cell_a on 8-9)
2337        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        // 20 CPUs total, 3 cells, equal weight
2342        // Targets: cell0=7, cell_a=7, cell_b=6 (or similar)
2343        // cell_a exclusive: 0-7 (8 CPUs) -> exceeds target of 7
2344        // cell_b exclusive: 10-13 (4 CPUs) -> below target of 6
2345        // Contested: 8-9 (2 CPUs) -> cell_a deficit=0, cell_b deficit=2
2346        // cell_b should get both contested CPUs
2347        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        // cell_a exceeded its target from exclusive alone, so it gets 0 contested CPUs.
2371        // cell_b has deficit, so it gets all 2 contested CPUs.
2372        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        // Total should be 20
2390        assert_eq!(
2391            cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2392            20,
2393        );
2394
2395        // Every cell should have at least 1 CPU
2396        assert!(cell0.primary.weight() >= 1);
2397        assert!(cell_a.primary.weight() >= 1);
2398        assert!(cell_b.primary.weight() >= 1);
2399    }
2400
2401    // ==================== distribute_cpus_proportional tests ====================
2402
2403    #[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        // 3/4 * 8 = 6 for cell 0, 1/4 * 8 = 2 for cell 1
2410        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        // Cell 0 gets all, cell 1 gets nothing
2421        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        // Falls back to equal division
2432        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        // Heavily skewed weights: 1.0 vs 100.0 with 38 CPUs (mirrors the bug scenario).
2439        // Without floor guarantee, cell 0 would get floor(1/101 * 38) = 0.
2440        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        // Both cells must get at least 1 CPU
2445        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        // Simulates two cells sharing an identical cpuset (all CPUs contested).
2463        // Cell A has deficit=1 (low demand), Cell B has deficit=87 (high demand).
2464        // Without floor guarantee, cell A gets 0 → death spiral.
2465        let cpus: Vec<usize> = (24..62).collect(); // 38 CPUs (24-61)
2466        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    // ==================== compute_targets tests ====================
2485
2486    #[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        // 10 / 3 = 3 each + 1 remainder
2498        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    /// Symmetric pairwise overlaps must produce equal cell sizes regardless
2506    /// of HashMap iteration order.
2507    #[test]
2508    fn test_symmetric_pairwise_overlap_produces_equal_cells() {
2509        let tmp = TempDir::new().unwrap();
2510
2511        // 5 cells on 56 CPUs. Every pair shares exactly 2 CPUs.
2512        // Each cell: 4 exclusive + 8 contested (2 per pair) = 12 in cpuset.
2513        // Exclusive: 0-19, contested: 20-39, unclaimed: 40-55 (cell 0).
2514        //
2515        // Shared pairs:
2516        //   AB: 20-21, AC: 22-23, AD: 24-25, AE: 26-27
2517        //   BC: 28-29, BD: 30-31, BE: 32-33
2518        //   CD: 34-35, CE: 36-37
2519        //   DE: 38-39
2520        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        // All workload cells must have equal CPU counts.
2547        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        // Verify no overlap between any pair of cells.
2561        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        // Verify all 56 CPUs are assigned.
2577        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
2578        assert_eq!(total, 56, "All CPUs must be assigned");
2579    }
2580}