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};
16use std::sync::atomic::{AtomicBool, Ordering};
17
18use anyhow::{bail, Context, Result};
19use inotify::{Inotify, WatchMask};
20use scx_utils::Cpumask;
21use tracing::{debug, info};
22
23/// Strip the cgroup mount prefix from a stored cell path, yielding the
24/// root-relative cgroup path (e.g. `/sys/fs/cgroup/a/b` -> `/a/b`).
25fn cgroup_root_relative(path: &Path) -> String {
26    path.strip_prefix("/sys/fs/cgroup")
27        .map(|rel| format!("/{}", rel.to_string_lossy()))
28        .unwrap_or_else(|_| path.to_string_lossy().into_owned())
29}
30
31/// Information about a cell created for a cgroup
32#[derive(Debug)]
33pub struct CellInfo {
34    pub cell_id: u32,
35    pub cgroup_path: Option<PathBuf>,
36    pub cgid: Option<u64>,
37    /// Optional cpuset mask if the cgroup has cpuset.cpus configured
38    pub cpuset: Option<Cpumask>,
39}
40
41/// Compute the global target CPU count for each cell.
42///
43/// Each cell gets a floor of 1 CPU, with the remainder distributed proportionally
44/// by weight. Returns only counts, not actual CPU assignments.
45fn compute_targets(total_cpus: usize, cells: &[(u32, f64)]) -> Result<HashMap<u32, usize>> {
46    if cells.is_empty() {
47        bail!("compute_targets called with no cells");
48    }
49    if total_cpus < cells.len() {
50        bail!(
51            "Not enough CPUs ({}) for {} cells (need at least 1 each)",
52            total_cpus,
53            cells.len()
54        );
55    }
56
57    let total_weight: f64 = cells.iter().map(|(_, w)| w).sum();
58    let num_cells = cells.len();
59
60    if total_weight <= 0.0 {
61        // Equal division fallback
62        let per = total_cpus / num_cells;
63        let remainder = total_cpus % num_cells;
64        return Ok(cells
65            .iter()
66            .enumerate()
67            .map(|(i, (cell_id, _))| {
68                let extra = if i < remainder { 1 } else { 0 };
69                (*cell_id, per + extra)
70            })
71            .collect());
72    }
73
74    let distributable = total_cpus - num_cells;
75    let mut assigned = num_cells;
76    let mut raw: Vec<(u32, f64, usize)> = cells
77        .iter()
78        .map(|(cell_id, w)| {
79            let frac = w / total_weight * distributable as f64;
80            let floored = frac.floor() as usize;
81            assigned += floored;
82            (*cell_id, frac - frac.floor(), 1 + floored)
83        })
84        .collect();
85
86    let mut remainder = total_cpus - assigned;
87    raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
88    for entry in raw.iter_mut() {
89        if remainder == 0 {
90            break;
91        }
92        entry.2 += 1;
93        remainder -= 1;
94    }
95
96    Ok(raw
97        .iter()
98        .map(|(cell_id, _, count)| (*cell_id, *count))
99        .collect())
100}
101
102/// Distribute CPUs among recipients proportionally by weight.
103///
104/// Recipients with weight 0 receive 0 CPUs. If all weights are 0, falls back to
105/// equal division. As a post-processing step, if any positive-weight recipient got
106/// 0 CPUs, 1 CPU is stolen from the recipient with the highest allocation (that has
107/// > 1) to prevent starvation.
108fn distribute_cpus_proportional(
109    cpus: &[usize],
110    recipients: &[(u32, f64)],
111) -> Result<HashMap<u32, Vec<usize>>> {
112    if cpus.is_empty() {
113        bail!("distribute_cpus_proportional called with no CPUs");
114    }
115    if recipients.is_empty() {
116        bail!("distribute_cpus_proportional called with no recipients");
117    }
118
119    let total_weight: f64 = recipients.iter().map(|(_, w)| w).sum();
120    let n = cpus.len();
121
122    let mut allocs: Vec<(u32, usize)> = if total_weight <= 0.0 {
123        // Equal division fallback
124        let per = n / recipients.len();
125        let remainder = n % recipients.len();
126        recipients
127            .iter()
128            .enumerate()
129            .map(|(i, (cell_id, _))| {
130                let extra = if i < remainder { 1 } else { 0 };
131                (*cell_id, per + extra)
132            })
133            .collect()
134    } else {
135        // Standard proportional distribution
136        let mut assigned = 0usize;
137        let mut raw: Vec<(u32, f64, usize, bool)> = recipients
138            .iter()
139            .map(|(cell_id, w)| {
140                let frac = w / total_weight * n as f64;
141                let floored = frac.floor() as usize;
142                assigned += floored;
143                (*cell_id, frac - frac.floor(), floored, *w > 0.0)
144            })
145            .collect();
146
147        let mut remainder = n - assigned;
148        raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
149        for entry in raw.iter_mut() {
150            if remainder == 0 {
151                break;
152            }
153            entry.2 += 1;
154            remainder -= 1;
155        }
156
157        // Post-processing floor guarantee: if any positive-weight recipient got 0
158        // CPUs, steal 1 from the recipient with the highest allocation (> 1).
159        // This prevents death spirals where cells with low but non-zero demand
160        // get starved of CPUs entirely.
161        loop {
162            let Some(starved_idx) = raw
163                .iter()
164                .position(|(_, _, count, pos_weight)| *pos_weight && *count == 0)
165            else {
166                break;
167            };
168            let Some(donor_idx) = raw
169                .iter()
170                .enumerate()
171                .filter(|(_, (_, _, count, _))| *count > 1)
172                .max_by_key(|(_, (_, _, count, _))| *count)
173                .map(|(i, _)| i)
174            else {
175                break; // No donor with > 1 CPU available
176            };
177            raw[donor_idx].2 -= 1;
178            raw[starved_idx].2 += 1;
179        }
180
181        raw.iter()
182            .map(|(cell_id, _, count, _)| (*cell_id, *count))
183            .collect()
184    };
185
186    // Assign actual CPU numbers
187    let mut cpu_iter = cpus.iter().copied();
188    allocs.sort_by_key(|(cell_id, _)| *cell_id);
189    let mut result: HashMap<u32, Vec<usize>> = HashMap::new();
190    for (cell_id, count) in allocs {
191        let cpus_for_cell: Vec<usize> = cpu_iter.by_ref().take(count).collect();
192        if cpus_for_cell.len() != count {
193            bail!(
194                "BUG: distribute_cpus_proportional: cell {} expected {} CPUs but got {}",
195                cell_id,
196                count,
197                cpus_for_cell.len()
198            );
199        }
200        if !cpus_for_cell.is_empty() {
201            result.insert(cell_id, cpus_for_cell);
202        }
203    }
204
205    Ok(result)
206}
207
208/// Result of CPU assignment computation, containing both primary and optional borrowable masks.
209#[derive(Debug)]
210pub struct CpuAssignment {
211    pub cell_id: u32,
212    pub primary: Cpumask,
213    pub borrowable: Option<Cpumask>,
214}
215
216/// Manages cells for direct child cgroups of a specified parent
217pub struct CellManager {
218    cell_parent_path: PathBuf,
219    inotify: Inotify,
220    /// Maps cgroup ID to cell info
221    cells: HashMap<u64, CellInfo>,
222    /// Maps cell ID to cgroup ID (for reverse lookup)
223    cell_id_to_cgid: HashMap<u32, u64>,
224    /// Freed cell IDs available for reuse
225    free_cell_ids: Vec<u32>,
226    next_cell_id: u32,
227    max_cells: u32,
228    /// Cpumask of all CPUs in the system (from topology)
229    all_cpus: Cpumask,
230    /// Cgroup directory names to exclude from cell creation
231    exclude_names: HashSet<String>,
232    /// Number of CPUs to hold out for cell 0 before child cpusets are
233    /// applied. 0 disables the holdout.
234    cell0_min_cpus: usize,
235    /// CPU -> LLC id. Steers holdout selection toward the LLC split across the
236    /// most cells once unclaimed CPUs run out; empty falls back to lowest CPU
237    /// number.
238    cpu_to_llc: HashMap<usize, usize>,
239    /// Set true (and never reset) the first time the holdout has to take a CPU
240    /// already claimed by a workload cell rather than only unclaimed CPUs.
241    /// Surfaced via stats so a host with unexpected performance can be checked
242    /// for whether the holdout took CPUs away from a workload.
243    enforced_holdout: AtomicBool,
244}
245
246impl CellManager {
247    pub fn new(
248        cell_parent_path: &str,
249        max_cells: u32,
250        all_cpus: Cpumask,
251        exclude: HashSet<String>,
252        cell0_min_cpus: usize,
253        cpu_to_llc: HashMap<usize, usize>,
254    ) -> Result<Self> {
255        let path = PathBuf::from(format!("/sys/fs/cgroup{}", cell_parent_path));
256        if !path.exists() {
257            bail!("Cell parent cgroup path does not exist: {}", path.display());
258        }
259        Self::new_with_path_opts(
260            path,
261            max_cells,
262            all_cpus,
263            exclude,
264            cell0_min_cpus,
265            cpu_to_llc,
266        )
267    }
268
269    /// Test-only 4-argument constructor with the holdout defaulted off
270    /// (`cell0_min_cpus = 0`, empty `cpu_to_llc`). Tests that exercise the
271    /// holdout call [`Self::new_with_path_opts`] directly with those arguments.
272    #[cfg(test)]
273    fn new_with_path(
274        path: PathBuf,
275        max_cells: u32,
276        all_cpus: Cpumask,
277        exclude: HashSet<String>,
278    ) -> Result<Self> {
279        Self::new_with_path_opts(path, max_cells, all_cpus, exclude, 0, HashMap::new())
280    }
281
282    fn new_with_path_opts(
283        path: PathBuf,
284        max_cells: u32,
285        all_cpus: Cpumask,
286        exclude: HashSet<String>,
287        cell0_min_cpus: usize,
288        cpu_to_llc: HashMap<usize, usize>,
289    ) -> Result<Self> {
290        let inotify = Inotify::init().context("Failed to initialize inotify")?;
291        inotify
292            .watches()
293            .add(&path, WatchMask::CREATE | WatchMask::DELETE)
294            .context("Failed to add inotify watch")?;
295
296        let mut mgr = Self {
297            cell_parent_path: path.clone(),
298            inotify,
299            cells: HashMap::new(),
300            cell_id_to_cgid: HashMap::new(),
301            free_cell_ids: Vec::new(),
302            next_cell_id: 1, // Cell 0 is reserved for root
303            max_cells,
304            all_cpus,
305            exclude_names: exclude,
306            cell0_min_cpus,
307            cpu_to_llc,
308            enforced_holdout: AtomicBool::new(false),
309        };
310
311        // Insert cell 0 as a permanent entry. cgid 0 is a safe sentinel —
312        // real cgroup inode numbers are always > 0.
313        mgr.cells.insert(
314            0,
315            CellInfo {
316                cell_id: 0,
317                cgroup_path: None,
318                cgid: None,
319                cpuset: None,
320            },
321        );
322        mgr.cell_id_to_cgid.insert(0, 0);
323
324        // Scan for existing children at startup
325        mgr.scan_existing_children()
326            .context("Failed to scan existing child cgroups at startup")?;
327        Ok(mgr)
328    }
329
330    /// True once the holdout has had to take a CPU already claimed by a workload
331    /// cell (sticky from scheduler start). Surfaced via stats.
332    pub fn enforced_holdout(&self) -> bool {
333        self.enforced_holdout.load(Ordering::Relaxed)
334    }
335
336    fn should_exclude(&self, path: &Path) -> bool {
337        path.file_name()
338            .and_then(|n| n.to_str())
339            .map(|name| self.exclude_names.contains(name))
340            .unwrap_or(false)
341    }
342
343    fn scan_existing_children(&mut self) -> Result<Vec<(u64, u32)>> {
344        let mut assignments = Vec::new();
345        let entries = std::fs::read_dir(&self.cell_parent_path).with_context(|| {
346            format!(
347                "Failed to read cell parent directory: {}",
348                self.cell_parent_path.display()
349            )
350        })?;
351        for entry in entries {
352            let entry = entry.with_context(|| {
353                format!(
354                    "Failed to read directory entry in: {}",
355                    self.cell_parent_path.display()
356                )
357            })?;
358            let file_type = entry.file_type().with_context(|| {
359                format!("Failed to get file type for: {}", entry.path().display())
360            })?;
361            if file_type.is_dir() {
362                let path = entry.path();
363                if self.should_exclude(&path) {
364                    continue;
365                }
366                let cgid = path
367                    .metadata()
368                    .with_context(|| {
369                        format!("reading inode of cgroup directory {}", path.display())
370                    })?
371                    .ino();
372                let (cgid, cell_id) =
373                    self.create_cell_for_cgroup(&path, cgid).with_context(|| {
374                        format!("Failed to create cell for cgroup: {}", path.display())
375                    })?;
376                assignments.push((cgid, cell_id));
377            }
378        }
379        Ok(assignments)
380    }
381
382    /// Process pending inotify events. Returns list of (cgid, cell_id) for new cells
383    /// and list of cell_ids that were destroyed.
384    ///
385    /// Rather than processing individual events, we simply check if any events occurred
386    /// and then rescan the directory to reconcile state. This is simpler and handles
387    /// edge cases like inotify queue overflow gracefully.
388    pub fn process_events(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
389        let mut buffer = [0; 1024];
390        let mut has_events = false;
391
392        // Drain all pending events
393        loop {
394            match self.inotify.read_events(&mut buffer) {
395                Ok(events) => {
396                    if events.into_iter().next().is_some() {
397                        has_events = true;
398                    } else {
399                        break;
400                    }
401                }
402                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
403                Err(e) => {
404                    return Err(e).context("Failed to read inotify events");
405                }
406            }
407        }
408
409        if !has_events {
410            return Ok((Vec::new(), Vec::new()));
411        }
412
413        // Rescan directory and reconcile with our tracked state
414        self.reconcile_cells()
415    }
416
417    /// Reconcile our tracked cells with the actual cgroup directory contents.
418    /// Returns (new_cells, destroyed_cells).
419    fn reconcile_cells(&mut self) -> Result<(Vec<(u64, u32)>, Vec<u32>)> {
420        let mut new_cells = Vec::new();
421
422        // Snapshot current child cgroups by path and inode.
423        // Reconcile by identity, not path alone, so path reuse doesn't keep
424        // the old cell and create a second one for the new inode.
425        let mut current_entries: HashMap<PathBuf, u64> = HashMap::new();
426        let entries = std::fs::read_dir(&self.cell_parent_path).with_context(|| {
427            format!(
428                "Failed to read cell parent directory: {}",
429                self.cell_parent_path.display()
430            )
431        })?;
432        for entry in entries {
433            let entry = entry.with_context(|| {
434                format!(
435                    "Failed to read directory entry in: {}",
436                    self.cell_parent_path.display()
437                )
438            })?;
439            let file_type = entry.file_type().with_context(|| {
440                format!("Failed to get file type for: {}", entry.path().display())
441            })?;
442            if file_type.is_dir() {
443                let path = entry.path();
444                if self.should_exclude(&path) {
445                    continue;
446                }
447
448                let metadata = match entry.metadata() {
449                    Ok(metadata) => metadata,
450                    Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
451                        // Directory disappeared after readdir(); retry on the
452                        // next reconcile instead of failing this scan.
453                        continue;
454                    }
455                    Err(e) => {
456                        return Err(e).with_context(|| {
457                            format!("reading inode of cgroup directory {}", path.display())
458                        });
459                    }
460                };
461
462                current_entries.insert(path, metadata.ino());
463            }
464        }
465
466        // Remove cells for cgroups that no longer exist
467        let mut destroyed_cells: HashSet<u32> = HashSet::new();
468        self.cells.retain(|&cgid, info| {
469            if info.cell_id == 0 {
470                return true; // Cell 0 is permanent
471            }
472            // Non-zero cells always have a cgroup_path
473            let cgroup_path = info
474                .cgroup_path
475                .as_ref()
476                .expect("BUG: non-zero cell missing cgroup_path");
477            // Same path with a different inode means the old cgroup was
478            // replaced and this tracked cell must be dropped.
479            if current_entries.get(cgroup_path) == Some(&cgid) {
480                true
481            } else {
482                info!(
483                    "Destroyed cell {} for cgroup {} (cgid={})",
484                    info.cell_id,
485                    cgroup_path.display(),
486                    cgid
487                );
488                destroyed_cells.insert(info.cell_id);
489                false
490            }
491        });
492
493        // Update tracking structures for destroyed cells
494        self.cell_id_to_cgid
495            .retain(|cell_id, _| !destroyed_cells.contains(cell_id));
496        self.free_cell_ids.extend(destroyed_cells.iter().copied());
497
498        // Find new cgroups that we don't have cells for
499        for (path, cgid) in current_entries {
500            if self.cells.contains_key(&cgid) {
501                continue; // Already have a cell for this cgroup
502            }
503            let (cgid, cell_id) = self
504                .create_cell_for_cgroup(&path, cgid)
505                .with_context(|| format!("Failed to create cell for cgroup: {}", path.display()))?;
506            new_cells.push((cgid, cell_id));
507        }
508
509        Ok((new_cells, destroyed_cells.into_iter().collect()))
510    }
511
512    fn create_cell_for_cgroup(&mut self, path: &Path, cgid: u64) -> Result<(u64, u32)> {
513        let cell_id = self.allocate_cell_id().context("allocating cell ID")?;
514
515        let cpuset = Self::read_cpuset(path)
516            .with_context(|| format!("reading cpuset for cgroup {}", path.display()))?;
517        if let Some(ref mask) = cpuset {
518            debug!(
519                "Cell {} has cpuset: {} (from {})",
520                cell_id,
521                mask.to_cpulist(),
522                path.join("cpuset.cpus").display()
523            );
524        }
525
526        self.cells.insert(
527            cgid,
528            CellInfo {
529                cell_id,
530                cgroup_path: Some(path.to_path_buf()),
531                cgid: Some(cgid),
532                cpuset,
533            },
534        );
535        self.cell_id_to_cgid.insert(cell_id, cgid);
536
537        info!(
538            "Created cell {} for cgroup {} (cgid={})",
539            cell_id,
540            path.display(),
541            cgid
542        );
543
544        Ok((cgid, cell_id))
545    }
546
547    /// Read cpuset.cpus from a cgroup path. Returns None if empty or unavailable.
548    fn read_cpuset(cgroup_path: &Path) -> Result<Option<Cpumask>> {
549        let cpuset_path = cgroup_path.join("cpuset.cpus");
550        match std::fs::read_to_string(&cpuset_path) {
551            Ok(content) => {
552                let content = content.trim();
553                if content.is_empty() {
554                    Ok(None)
555                } else {
556                    let mask = Cpumask::from_cpulist(content).with_context(|| {
557                        format!(
558                            "Failed to parse cpuset '{}' from {}",
559                            content,
560                            cpuset_path.display()
561                        )
562                    })?;
563                    Ok(Some(mask))
564                }
565            }
566            // File doesn't exist - cpuset controller is not enabled for this cgroup
567            Err(_) => Ok(None),
568        }
569    }
570
571    fn allocate_cell_id(&mut self) -> Result<u32> {
572        // Prefer reusing freed IDs to keep cell ID space compact
573        if let Some(id) = self.free_cell_ids.pop() {
574            return Ok(id);
575        }
576
577        if self.next_cell_id >= self.max_cells {
578            bail!("Cell ID space exhausted (max_cells={})", self.max_cells);
579        }
580
581        let id = self.next_cell_id;
582        self.next_cell_id += 1;
583        Ok(id)
584    }
585
586    /// Compute CPU assignments for all cells.
587    ///
588    /// When cpusets overlap, contested CPUs are divided proportionally among claimants.
589    /// Unclaimed CPUs go to cell 0 and any unpinned cells (cells without cpusets).
590    ///
591    /// If `compute_borrowable` is true, each assignment includes a borrowable cpumask
592    /// (all system CPUs minus the cell's own, intersected with cpuset if present).
593    /// Without demand data, borrowable masks are uncapped.
594    ///
595    /// Returns a Vec of CpuAssignment, or an error if any cell would
596    /// receive zero CPUs (which indicates too many cells for available CPUs).
597    pub fn compute_cpu_assignments(&self, compute_borrowable: bool) -> Result<Vec<CpuAssignment>> {
598        // Use equal weights for all cells (no demand data)
599        self.compute_cpu_assignments_inner(None, compute_borrowable)
600    }
601
602    /// Compute CPU assignments weighted by per-cell demand.
603    ///
604    /// `cell_demands` maps cell_id -> smoothed_util_pct. All active cells must be
605    /// present in the map; missing entries or negative weights are errors.
606    ///
607    /// If `compute_borrowable` is true, each assignment includes a borrowable cpumask
608    /// (all system CPUs minus the cell's own, intersected with cpuset if present).
609    pub fn compute_demand_cpu_assignments(
610        &self,
611        cell_demands: &HashMap<u32, f64>,
612        compute_borrowable: bool,
613    ) -> Result<Vec<CpuAssignment>> {
614        self.compute_cpu_assignments_inner(Some(cell_demands), compute_borrowable)
615    }
616
617    /// Internal implementation shared by equal-weight and demand-weighted assignment.
618    fn compute_cpu_assignments_inner(
619        &self,
620        cell_demands: Option<&HashMap<u32, f64>>,
621        compute_borrowable: bool,
622    ) -> Result<Vec<CpuAssignment>> {
623        // Validate that all demand weights are non-negative
624        if let Some(demands) = cell_demands {
625            for (&cell_id, &weight) in demands {
626                if weight < 0.0 {
627                    bail!("Cell {} has negative demand weight {}", cell_id, weight);
628                }
629            }
630        }
631
632        // Phase 1: Build contention map - for each CPU, track which cells claim it
633        let mut contention: HashMap<usize, Vec<u32>> = HashMap::new();
634        for cell_info in self.cells.values() {
635            if let Some(ref cpuset) = cell_info.cpuset {
636                for cpu in cpuset.iter() {
637                    contention.entry(cpu).or_default().push(cell_info.cell_id);
638                }
639            }
640        }
641
642        // Holdout: try to reserve `cell0_min_cpus` CPUs for cell 0 before the
643        // cpuset-based assignment below, but reserve fewer if reserving more
644        // would mean taking a workload cell's last CPU. Prefer CPUs no child
645        // cpuset claims; when those run out, steal from cells — largest cell
646        // first, but never a cell's last CPU — so no child is starved. Reserved
647        // CPUs are skipped by the phases below and pre-seeded onto cell 0.
648        let mut holdout = Cpumask::new();
649        if self.cell0_min_cpus > 0 {
650            let mut unclaimed: Vec<usize> = self
651                .all_cpus
652                .iter()
653                .filter(|cpu| !contention.contains_key(cpu))
654                .collect();
655            unclaimed.sort_unstable();
656
657            // Distinct cells claiming a CPU in each LLC.
658            let mut llc_cells: HashMap<usize, HashSet<u32>> = HashMap::new();
659            for (cpu, claimants) in &contention {
660                if let Some(&llc) = self.cpu_to_llc.get(cpu) {
661                    llc_cells.entry(llc).or_default().extend(claimants);
662                }
663            }
664            let partition = |cpu: &usize| -> usize {
665                self.cpu_to_llc
666                    .get(cpu)
667                    .and_then(|llc| llc_cells.get(llc))
668                    .map_or(0, |cells| cells.len())
669            };
670
671            // Take unclaimed CPUs first — they cost no cell anything.
672            let mut taken = 0usize;
673            for cpu in unclaimed.into_iter().take(self.cell0_min_cpus) {
674                holdout.set_cpu(cpu).ok();
675                taken += 1;
676            }
677
678            // Steal the remainder from cells round-robin: each cell yields one CPU
679            // in turn, largest cell first, so the reservation is spread across
680            // cells rather than draining one. A CPU is reservable only if every
681            // cell claiming it keeps at least one CPU afterward -- otherwise the
682            // reservation would zero that cell and re-trigger the Phase 5 "no CPUs
683            // assigned" bail this holdout exists to prevent (mirrors
684            // distribute_cpus_proportional's anti-starvation floor). Within the
685            // donor, prefer an exclusive CPU, then the LLC split across the most
686            // cells, then the lowest CPU. When no CPU is reservable, stop -- cell 0
687            // receives fewer than cell0_min_cpus rather than starving a child.
688            let mut cell_claimed: HashMap<u32, Vec<usize>> = HashMap::new();
689            for (&cpu, claimants) in &contention {
690                for &cell_id in claimants {
691                    cell_claimed.entry(cell_id).or_default().push(cpu);
692                }
693            }
694            // CPUs yielded to the holdout per cell, for round-robin selection.
695            let mut taken_from: HashMap<u32, usize> =
696                cell_claimed.keys().map(|&id| (id, 0usize)).collect();
697            while taken < self.cell0_min_cpus {
698                // Unreserved EXCLUSIVE CPUs per cell, recomputed as the holdout
699                // grows. Only exclusive CPUs count toward survival: a contested CPU
700                // may be awarded to a co-claimant in Phase 3, so it cannot be what
701                // keeps a cell alive -- counting it would let the steal take a
702                // cell's last exclusive CPU and zero the cell (the Phase-5 bail).
703                let remaining: HashMap<u32, usize> = cell_claimed
704                    .iter()
705                    .map(|(&cell_id, cpus)| {
706                        (
707                            cell_id,
708                            cpus.iter()
709                                .filter(|&&c| !holdout.test_cpu(c) && contention[&c].len() == 1)
710                                .count(),
711                        )
712                    })
713                    .collect();
714                // Reservable iff every cell claiming the CPU keeps >=1 exclusive CPU.
715                let reservable = |cpu: usize| {
716                    contention
717                        .get(&cpu)
718                        .is_some_and(|claimants| claimants.iter().all(|c| remaining[c] >= 2))
719                };
720                // Donor = cell that has yielded the fewest CPUs so far
721                // (round-robin); ties break toward the larger cell, then the
722                // lowest cell id. Only cells with a reservable CPU are eligible.
723                let donor = cell_claimed
724                    .iter()
725                    .filter(|(_, cpus)| {
726                        cpus.iter()
727                            .any(|&cpu| !holdout.test_cpu(cpu) && reservable(cpu))
728                    })
729                    .map(|(&cell_id, _)| cell_id)
730                    .min_by(|&a, &b| {
731                        taken_from[&a]
732                            .cmp(&taken_from[&b])
733                            .then(remaining[&b].cmp(&remaining[&a]))
734                            .then(a.cmp(&b))
735                    });
736                let Some(donor) = donor else {
737                    break; // no CPU can be reserved without starving a child cell
738                };
739                let cpu = cell_claimed[&donor]
740                    .iter()
741                    .copied()
742                    .filter(|&cpu| !holdout.test_cpu(cpu) && reservable(cpu))
743                    .min_by(|&a, &b| {
744                        let contested =
745                            |cpu: usize| contention.get(&cpu).map_or(0, |v| v.len()) > 1;
746                        contested(a)
747                            .cmp(&contested(b))
748                            .then_with(|| partition(&b).cmp(&partition(&a)))
749                            .then(a.cmp(&b))
750                    })
751                    .expect("donor has a reservable CPU");
752                holdout.set_cpu(cpu).ok();
753                // Taking a CPU a workload cell already claims -- record that the
754                // holdout put its hand on the scales (queryable via stats).
755                self.enforced_holdout.store(true, Ordering::Relaxed);
756                *taken_from.entry(donor).or_insert(0) += 1;
757                taken += 1;
758            }
759        }
760
761        // Phase 2: Categorize CPUs and build initial assignments
762        // - Exclusive: claimed by exactly 1 cell -> assigned directly
763        // - Contested: claimed by 2+ cells -> distributed by weight
764        // - Unclaimed: no cpuset claims it -> shared among cell 0 + unpinned cells
765        let mut cell_cpus: HashMap<u32, Cpumask> = HashMap::new();
766        if holdout.weight() > 0 {
767            cell_cpus.insert(0, holdout.clone());
768        }
769        let mut contested_cpus: Vec<usize> = Vec::new();
770        let mut unclaimed_cpus: Vec<usize> = Vec::new();
771
772        for cpu in self.all_cpus.iter() {
773            if holdout.test_cpu(cpu) {
774                continue; // reserved for cell 0 by the holdout above
775            }
776            match contention.get(&cpu) {
777                None => unclaimed_cpus.push(cpu),
778                Some(claimants) if claimants.len() == 1 => {
779                    let cell_id = claimants[0];
780                    cell_cpus
781                        .entry(cell_id)
782                        .or_insert_with(Cpumask::new)
783                        .set_cpu(cpu)
784                        .ok();
785                }
786                Some(_) => contested_cpus.push(cpu),
787            }
788        }
789
790        // Compute global targets and initialize running count of CPUs assigned per cell
791        let total_cpu_count = self.all_cpus.weight();
792        let mut all_cells_with_weights: Vec<(u32, f64)> = self
793            .cells
794            .values()
795            .map(|info| {
796                let weight = match cell_demands {
797                    Some(demands) => *demands.get(&info.cell_id).ok_or_else(|| {
798                        anyhow::anyhow!("Cell {} is missing from demands map", info.cell_id)
799                    })?,
800                    None => 1.0,
801                };
802                Ok((info.cell_id, weight))
803            })
804            .collect::<Result<Vec<_>>>()
805            .context("building cell demand weights map")?;
806        all_cells_with_weights.sort_by_key(|(cell_id, _)| *cell_id);
807
808        let targets = compute_targets(total_cpu_count, &all_cells_with_weights)
809            .context("computing per-cell CPU targets")?;
810
811        // Seed assigned_count from exclusive assignments
812        let mut assigned_count: HashMap<u32, usize> = HashMap::new();
813        for (cell_id, mask) in &cell_cpus {
814            assigned_count.insert(*cell_id, mask.weight());
815        }
816
817        // Phase 3: Distribute contested CPUs among claimants using deficit weights
818        let mut contested_groups: HashMap<Vec<u32>, Vec<usize>> = HashMap::new();
819        for cpu in contested_cpus {
820            if let Some(claimants) = contention.get(&cpu) {
821                let mut sorted_claimants = claimants.clone();
822                sorted_claimants.sort();
823                contested_groups
824                    .entry(sorted_claimants)
825                    .or_default()
826                    .push(cpu);
827            }
828        }
829
830        // Freeze deficit weights before processing any group. Each group is an
831        // independent allocation decision, so a cell's weight should not depend
832        // on HashMap iteration order (i.e., which other groups were processed
833        // first). Using the initial deficit (target - exclusive_count) as weight
834        // for all groups makes the result deterministic.
835        let initial_deficit: HashMap<u32, f64> = targets
836            .iter()
837            .map(|(&cell_id, &target)| {
838                // Cells with no exclusive CPUs have no entry yet; 0 is correct.
839                let already = assigned_count.get(&cell_id).copied().unwrap_or(0);
840                let deficit = if target > already {
841                    (target - already) as f64
842                } else {
843                    0.0
844                };
845                (cell_id, deficit)
846            })
847            .collect();
848
849        for (claimants, cpus) in contested_groups {
850            let mut recipients: Vec<(u32, f64)> = Vec::new();
851            let mut all_zero = true;
852            for &cell_id in &claimants {
853                let deficit = *initial_deficit.get(&cell_id).ok_or_else(|| {
854                    anyhow::anyhow!(
855                        "BUG: cell {} in contention map but missing from targets",
856                        cell_id
857                    )
858                })?;
859                if deficit > 0.0 {
860                    all_zero = false;
861                }
862                recipients.push((cell_id, deficit));
863            }
864
865            // If all claimants already meet/exceed their target, fall back to equal weights
866            if all_zero {
867                recipients = claimants.iter().map(|&cell_id| (cell_id, 1.0)).collect();
868            }
869
870            let distribution = distribute_cpus_proportional(&cpus, &recipients)
871                .context("distributing contested CPUs among claimants")?;
872            for (cell_id, assigned_cpus) in distribution {
873                let count = assigned_cpus.len();
874                for cpu in assigned_cpus {
875                    cell_cpus
876                        .entry(cell_id)
877                        .or_insert_with(Cpumask::new)
878                        .set_cpu(cpu)
879                        .ok();
880                }
881                *assigned_count.entry(cell_id).or_insert(0) += count;
882            }
883        }
884
885        // Phase 4: Distribute unclaimed CPUs among unpinned cells (including cell 0) using deficit weights
886        if !unclaimed_cpus.is_empty() {
887            let mut recipients: Vec<(u32, f64)> = Vec::new();
888            let mut all_zero = true;
889
890            for info in self.cells.values() {
891                if info.cpuset.is_some() {
892                    continue; // pinned cells don't receive unclaimed CPUs
893                }
894                let target = *targets.get(&info.cell_id).ok_or_else(|| {
895                    anyhow::anyhow!(
896                        "BUG: cell {} is unpinned but missing from targets",
897                        info.cell_id
898                    )
899                })?;
900                // Cells with no exclusive CPUs have no entry yet; 0 is correct.
901                let already = assigned_count.get(&info.cell_id).copied().unwrap_or(0);
902                let deficit = if target > already {
903                    (target - already) as f64
904                } else {
905                    0.0
906                };
907                if deficit > 0.0 {
908                    all_zero = false;
909                }
910                recipients.push((info.cell_id, deficit));
911            }
912            recipients.sort_by_key(|(cell_id, _)| *cell_id);
913
914            // If all recipients already meet/exceed their target, fall back to equal weights
915            if all_zero {
916                recipients = recipients
917                    .iter()
918                    .map(|(cell_id, _)| (*cell_id, 1.0))
919                    .collect();
920            }
921
922            let distribution = distribute_cpus_proportional(&unclaimed_cpus, &recipients)
923                .context("distributing unclaimed CPUs among unpinned cells")?;
924            for (cell_id, assigned_cpus) in distribution {
925                let count = assigned_cpus.len();
926                for cpu in assigned_cpus {
927                    cell_cpus
928                        .entry(cell_id)
929                        .or_insert_with(Cpumask::new)
930                        .set_cpu(cpu)
931                        .ok();
932                }
933                *assigned_count.entry(cell_id).or_insert(0) += count;
934            }
935        }
936
937        // Phase 5: Verify all cells have at least one CPU assigned
938        for info in self.cells.values() {
939            if !cell_cpus.contains_key(&info.cell_id)
940                || cell_cpus
941                    .get(&info.cell_id)
942                    .map_or(true, |m| m.weight() == 0)
943            {
944                bail!(
945                    "Cell {} has no CPUs assigned (nr_cpus={}, num_cells={})",
946                    info.cell_id,
947                    self.all_cpus.weight(),
948                    self.cells.len()
949                );
950            }
951        }
952
953        // Phase 6: Build CpuAssignment results, optionally computing borrowable masks
954        let assignments: Vec<CpuAssignment> = cell_cpus
955            .into_iter()
956            .map(|(cell_id, primary)| {
957                let borrowable = if compute_borrowable {
958                    let mut borrow_mask = self.all_cpus.and(&primary.not());
959
960                    // If this cell has a cpuset, restrict borrowable to it
961                    if let Some(cell_info) = self.cells.values().find(|c| c.cell_id == cell_id) {
962                        if let Some(ref cpuset) = cell_info.cpuset {
963                            borrow_mask = borrow_mask.and(cpuset);
964                        }
965                    }
966
967                    Some(borrow_mask)
968                } else {
969                    None
970                };
971                CpuAssignment {
972                    cell_id,
973                    primary,
974                    borrowable,
975                }
976            })
977            .collect();
978
979        Ok(assignments)
980    }
981
982    /// Returns all cell assignments as (cgid, cell_id) pairs.
983    /// Used to configure BPF with cgroup-to-cell mappings.
984    pub fn get_cell_assignments(&self) -> Vec<(u64, u32)> {
985        self.cells
986            .values()
987            .filter(|info| info.cell_id != 0)
988            .map(|info| {
989                (
990                    info.cgid.expect("BUG: non-zero cell missing cgid"),
991                    info.cell_id,
992                )
993            })
994            .collect()
995    }
996
997    /// Format the cell configuration as a compact string for logging.
998    /// Example output: "[0: 0-7] [1(container-a): 8-15] [2(container-b): 16-23]"
999    pub fn format_cell_config(&self, cpu_assignments: &[CpuAssignment]) -> String {
1000        let mut sorted: Vec<_> = cpu_assignments.iter().collect();
1001        sorted.sort_by_key(|a| a.cell_id);
1002
1003        let mut parts = Vec::new();
1004        for assignment in sorted {
1005            let cpulist = assignment.primary.to_cpulist();
1006            if assignment.cell_id == 0 {
1007                parts.push(format!("[0: {}]", cpulist));
1008            } else {
1009                // Find cgroup name for this cell
1010                let name = self
1011                    .cells
1012                    .values()
1013                    .find(|info| info.cell_id == assignment.cell_id)
1014                    .and_then(|info| {
1015                        info.cgroup_path
1016                            .as_ref()
1017                            .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
1018                    })
1019                    .unwrap_or_else(|| "?".to_string());
1020                parts.push(format!("[{}({}): {}]", assignment.cell_id, name, cpulist));
1021            }
1022        }
1023        parts.join(" ")
1024    }
1025
1026    /// Return the cgroup path for a cell, relative to the cgroup root as it
1027    /// appears in `/proc/<pid>/cgroup` (e.g. `/test.slice/foobar`), so consumers
1028    /// can map a cell id to its cgroup. Cell 0 (the root cell) is `/`.
1029    pub fn cgroup_path_for_cell(&self, cell_id: u32) -> String {
1030        self.cell_id_to_cgid
1031            .get(&cell_id)
1032            .and_then(|cgid| self.cells.get(cgid))
1033            .and_then(|info| info.cgroup_path.as_deref())
1034            .map(cgroup_root_relative)
1035            .unwrap_or_else(|| "/".to_string())
1036    }
1037
1038    /// Re-read cpuset.cpus for all cells and update stored cpusets.
1039    /// Returns true if any cell's cpuset changed.
1040    pub fn refresh_cpusets(&mut self) -> Result<bool> {
1041        let mut changed = false;
1042        for info in self.cells.values_mut() {
1043            let Some(ref cgroup_path) = info.cgroup_path else {
1044                continue; // cell 0 has no cgroup
1045            };
1046            let new_cpuset = Self::read_cpuset(cgroup_path)
1047                .with_context(|| format!("reading cpuset for cgroup {}", cgroup_path.display()))?;
1048            if new_cpuset != info.cpuset {
1049                info!(
1050                    "Cell {} cpuset changed: {:?} -> {:?} ({})",
1051                    info.cell_id,
1052                    info.cpuset.as_ref().map(|m| m.to_cpulist()),
1053                    new_cpuset.as_ref().map(|m| m.to_cpulist()),
1054                    cgroup_path.display(),
1055                );
1056                info.cpuset = new_cpuset;
1057                changed = true;
1058            }
1059        }
1060        Ok(changed)
1061    }
1062}
1063
1064impl AsFd for CellManager {
1065    fn as_fd(&self) -> BorrowedFd<'_> {
1066        self.inotify.as_fd()
1067    }
1068}
1069
1070#[cfg(test)]
1071impl CellManager {
1072    /// Returns the number of cells created for cgroups.
1073    /// Does not include cell 0 (the implicit root cell).
1074    fn cell_count(&self) -> usize {
1075        self.cells.values().filter(|c| c.cell_id != 0).count()
1076    }
1077
1078    /// Get all cell IDs for cells created for cgroups.
1079    /// Does not include cell 0 (the implicit root cell).
1080    fn get_cell_ids(&self) -> Vec<u32> {
1081        self.cells
1082            .values()
1083            .filter(|c| c.cell_id != 0)
1084            .map(|c| c.cell_id)
1085            .collect()
1086    }
1087
1088    /// Find a cell by cgroup directory name.
1089    /// Only searches cells created for cgroups, not cell 0.
1090    fn find_cell_by_name(&self, name: &str) -> Option<&CellInfo> {
1091        self.cells.values().filter(|c| c.cell_id != 0).find(|c| {
1092            c.cgroup_path
1093                .as_ref()
1094                .and_then(|p| p.file_name())
1095                .map(|n| n.to_str() == Some(name))
1096                .unwrap_or(false)
1097        })
1098    }
1099}
1100
1101#[cfg(test)]
1102#[allow(clippy::unwrap_used)]
1103mod tests {
1104    use super::*;
1105    use tempfile::TempDir;
1106
1107    fn cpumask_for_range(nr_cpus: usize) -> Cpumask {
1108        scx_utils::set_cpumask_test_width(nr_cpus);
1109        let mut mask = Cpumask::new();
1110        for cpu in 0..nr_cpus {
1111            mask.set_cpu(cpu).unwrap();
1112        }
1113        mask
1114    }
1115
1116    // ==================== cgroup path mapping tests ====================
1117
1118    #[test]
1119    fn test_cgroup_path_for_cell() {
1120        let tmp = TempDir::new().unwrap();
1121        let child = tmp.path().join("container-a");
1122        std::fs::create_dir(&child).unwrap();
1123        let mgr = CellManager::new_with_path(
1124            tmp.path().to_path_buf(),
1125            256,
1126            cpumask_for_range(16),
1127            HashSet::new(),
1128        )
1129        .unwrap();
1130
1131        // Cell 0 is the root cell -> "/".
1132        assert_eq!(mgr.cgroup_path_for_cell(0), "/");
1133
1134        // A created cell maps to its cgroup. The TempDir is not under the cgroup
1135        // mount, so the path is returned as-is (the strip falls back to the
1136        // absolute path); the mount-strip itself is covered below.
1137        let cell_id = mgr.find_cell_by_name("container-a").unwrap().cell_id;
1138        assert_eq!(
1139            mgr.cgroup_path_for_cell(cell_id),
1140            child.to_string_lossy().into_owned()
1141        );
1142
1143        // A real (mount-absolute) cgroup path is reported relative to the root.
1144        assert_eq!(
1145            cgroup_root_relative(Path::new("/sys/fs/cgroup/test.slice/foobar")),
1146            "/test.slice/foobar"
1147        );
1148    }
1149
1150    // ==================== Cell scanning and creation tests ====================
1151
1152    #[test]
1153    fn test_scan_empty_directory() {
1154        let tmp = TempDir::new().unwrap();
1155        let mgr = CellManager::new_with_path(
1156            tmp.path().to_path_buf(),
1157            256,
1158            cpumask_for_range(16),
1159            HashSet::new(),
1160        )
1161        .unwrap();
1162
1163        assert_eq!(mgr.cell_count(), 0);
1164    }
1165
1166    #[test]
1167    fn test_scan_existing_subdirectories() {
1168        let tmp = TempDir::new().unwrap();
1169
1170        // Create some "cgroup" directories
1171        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1172        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1173
1174        let mgr = CellManager::new_with_path(
1175            tmp.path().to_path_buf(),
1176            256,
1177            cpumask_for_range(16),
1178            HashSet::new(),
1179        )
1180        .unwrap();
1181
1182        assert_eq!(mgr.cell_count(), 2);
1183
1184        // Verify cells were assigned IDs 1 and 2
1185        let cell_ids = mgr.get_cell_ids();
1186        assert!(cell_ids.contains(&1));
1187        assert!(cell_ids.contains(&2));
1188    }
1189
1190    #[test]
1191    fn test_reconcile_detects_new_directories() {
1192        let tmp = TempDir::new().unwrap();
1193
1194        // Start with one directory
1195        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1196        let mut mgr = CellManager::new_with_path(
1197            tmp.path().to_path_buf(),
1198            256,
1199            cpumask_for_range(16),
1200            HashSet::new(),
1201        )
1202        .unwrap();
1203        assert_eq!(mgr.cell_count(), 1);
1204
1205        // Add another directory
1206        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1207
1208        // Reconcile should detect it
1209        let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1210        assert_eq!(new_cells.len(), 1);
1211        assert_eq!(destroyed_cells.len(), 0);
1212        assert_eq!(mgr.cell_count(), 2);
1213    }
1214
1215    #[test]
1216    fn test_reconcile_detects_removed_directories() {
1217        let tmp = TempDir::new().unwrap();
1218
1219        // Start with two directories
1220        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
1221        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
1222        let mut mgr = CellManager::new_with_path(
1223            tmp.path().to_path_buf(),
1224            256,
1225            cpumask_for_range(16),
1226            HashSet::new(),
1227        )
1228        .unwrap();
1229        assert_eq!(mgr.cell_count(), 2);
1230
1231        // Remove one directory
1232        std::fs::remove_dir(tmp.path().join("container-b")).unwrap();
1233
1234        // Reconcile should detect it
1235        let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1236        assert_eq!(new_cells.len(), 0);
1237        assert_eq!(destroyed_cells.len(), 1);
1238        assert_eq!(mgr.cell_count(), 1);
1239    }
1240
1241    #[test]
1242    fn test_reconcile_replaces_reused_path_with_new_inode() {
1243        let tmp = TempDir::new().unwrap();
1244        let parked = TempDir::new().unwrap();
1245
1246        let original_path = tmp.path().join("container-a");
1247        std::fs::create_dir(&original_path).unwrap();
1248
1249        let mut mgr = CellManager::new_with_path(
1250            tmp.path().to_path_buf(),
1251            256,
1252            cpumask_for_range(16),
1253            HashSet::new(),
1254        )
1255        .unwrap();
1256        assert_eq!(mgr.cell_count(), 1);
1257
1258        let old_info = mgr.find_cell_by_name("container-a").unwrap();
1259        let old_cell_id = old_info.cell_id;
1260        let old_cgid = old_info.cgid.unwrap();
1261
1262        // Move the original cgroup out of the watched directory so the inode
1263        // stays alive while a new cgroup is created at the same path.
1264        std::fs::rename(&original_path, parked.path().join("container-a-old")).unwrap();
1265        std::fs::create_dir(&original_path).unwrap();
1266
1267        let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
1268
1269        assert_eq!(new_cells.len(), 1);
1270        assert_eq!(destroyed_cells, vec![old_cell_id]);
1271        assert_eq!(mgr.cell_count(), 1);
1272
1273        let new_info = mgr.find_cell_by_name("container-a").unwrap();
1274        assert_eq!(new_info.cell_id, old_cell_id);
1275        assert_ne!(new_info.cgid.unwrap(), old_cgid);
1276    }
1277
1278    #[test]
1279    fn test_cell_id_reuse_after_destruction() {
1280        let tmp = TempDir::new().unwrap();
1281
1282        // Create directories
1283        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1284        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1285        std::fs::create_dir(tmp.path().join("cell3")).unwrap();
1286
1287        let mut mgr = CellManager::new_with_path(
1288            tmp.path().to_path_buf(),
1289            256,
1290            cpumask_for_range(16),
1291            HashSet::new(),
1292        )
1293        .unwrap();
1294
1295        // Find cell2's ID
1296        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1297        let cell2_id = cell2_info.cell_id;
1298
1299        // Remove cell2
1300        std::fs::remove_dir(tmp.path().join("cell2")).unwrap();
1301        mgr.reconcile_cells().unwrap();
1302
1303        // Add a new directory - should reuse cell2's ID
1304        std::fs::create_dir(tmp.path().join("cell4")).unwrap();
1305        mgr.reconcile_cells().unwrap();
1306
1307        let cell4_info = mgr.find_cell_by_name("cell4").unwrap();
1308        assert_eq!(cell4_info.cell_id, cell2_id);
1309    }
1310
1311    // ==================== compute_cpu_assignments tests ====================
1312
1313    #[test]
1314    fn test_cpu_assignments_no_cells() {
1315        let tmp = TempDir::new().unwrap();
1316        let mgr = CellManager::new_with_path(
1317            tmp.path().to_path_buf(),
1318            256,
1319            cpumask_for_range(16),
1320            HashSet::new(),
1321        )
1322        .unwrap();
1323
1324        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1325
1326        // Only cell 0 with all CPUs
1327        assert_eq!(assignments.len(), 1);
1328        assert_eq!(assignments[0].cell_id, 0);
1329        assert_eq!(assignments[0].primary.weight(), 16);
1330    }
1331
1332    #[test]
1333    fn test_cpu_assignments_proportional() {
1334        let tmp = TempDir::new().unwrap();
1335        std::fs::create_dir(tmp.path().join("container")).unwrap();
1336
1337        let mgr = CellManager::new_with_path(
1338            tmp.path().to_path_buf(),
1339            256,
1340            cpumask_for_range(16),
1341            HashSet::new(),
1342        )
1343        .unwrap();
1344        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1345
1346        // 16 CPUs / 2 cells = 8 each
1347        assert_eq!(assignments.len(), 2);
1348
1349        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1350        let cell1 = assignments.iter().find(|a| a.cell_id == 1).unwrap();
1351
1352        assert_eq!(cell0.primary.weight(), 8);
1353        assert_eq!(cell1.primary.weight(), 8);
1354    }
1355
1356    #[test]
1357    fn test_cpu_assignments_remainder_to_cell0() {
1358        let tmp = TempDir::new().unwrap();
1359        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
1360        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
1361
1362        let mgr = CellManager::new_with_path(
1363            tmp.path().to_path_buf(),
1364            256,
1365            cpumask_for_range(10),
1366            HashSet::new(),
1367        )
1368        .unwrap();
1369        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1370
1371        // 10 CPUs / 3 cells = 3 each + 1 remainder to cell 0
1372        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1373        assert_eq!(cell0.primary.weight(), 4); // 3 + 1 remainder
1374    }
1375
1376    #[test]
1377    fn test_cpu_assignments_too_many_cells() {
1378        let tmp = TempDir::new().unwrap();
1379
1380        // Create more cells than CPUs
1381        for i in 1..=5 {
1382            std::fs::create_dir(tmp.path().join(format!("cell{}", i))).unwrap();
1383        }
1384
1385        // Only 4 CPUs but 6 cells (cell 0 + 5 user cells)
1386        let mgr = CellManager::new_with_path(
1387            tmp.path().to_path_buf(),
1388            256,
1389            cpumask_for_range(4),
1390            HashSet::new(),
1391        )
1392        .unwrap();
1393        let result = mgr.compute_cpu_assignments(false);
1394
1395        assert!(result.is_err());
1396        let err_msg = format!("{:#}", result.unwrap_err());
1397        assert!(
1398            err_msg.contains("Not enough CPUs"),
1399            "Expected 'Not enough CPUs' error, got: {}",
1400            err_msg
1401        );
1402    }
1403
1404    #[test]
1405    fn test_cpu_assignments_with_cpusets() {
1406        let tmp = TempDir::new().unwrap();
1407
1408        // Create cgroup directories with cpuset files
1409        let cell1_path = tmp.path().join("cell1");
1410        std::fs::create_dir(&cell1_path).unwrap();
1411        std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1412
1413        let cell2_path = tmp.path().join("cell2");
1414        std::fs::create_dir(&cell2_path).unwrap();
1415        std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1416
1417        let mgr = CellManager::new_with_path(
1418            tmp.path().to_path_buf(),
1419            256,
1420            cpumask_for_range(16),
1421            HashSet::new(),
1422        )
1423        .unwrap();
1424        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1425
1426        // Should have 3 assignments: cell1, cell2, and cell0
1427        assert_eq!(assignments.len(), 3);
1428
1429        // Find each cell's assignment using find_cell_by_name
1430        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1431        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1432
1433        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1434        let cell1 = assignments
1435            .iter()
1436            .find(|a| a.cell_id == cell1_info.cell_id)
1437            .unwrap();
1438        let cell2 = assignments
1439            .iter()
1440            .find(|a| a.cell_id == cell2_info.cell_id)
1441            .unwrap();
1442
1443        // cell1 gets CPUs 0-3
1444        assert_eq!(cell1.primary.weight(), 4);
1445        for cpu in 0..4 {
1446            assert!(cell1.primary.test_cpu(cpu));
1447        }
1448
1449        // cell2 gets CPUs 8-11
1450        assert_eq!(cell2.primary.weight(), 4);
1451        for cpu in 8..12 {
1452            assert!(cell2.primary.test_cpu(cpu));
1453        }
1454
1455        // cell0 gets remaining CPUs: 4-7, 12-15
1456        assert_eq!(cell0.primary.weight(), 8);
1457        for cpu in 4..8 {
1458            assert!(cell0.primary.test_cpu(cpu));
1459        }
1460        for cpu in 12..16 {
1461            assert!(cell0.primary.test_cpu(cpu));
1462        }
1463    }
1464
1465    #[test]
1466    fn test_cpu_assignments_cpusets_cover_all_cpus() {
1467        let tmp = TempDir::new().unwrap();
1468
1469        // Cell 1 and cell 2's cpusets together claim every CPU. Cell 0
1470        // (the catch-all for unclaimed CPUs) would otherwise get nothing,
1471        // which is an error when cell0_min_cpus is 0 (the default).
1472        let cell1_path = tmp.path().join("cell1");
1473        std::fs::create_dir(&cell1_path).unwrap();
1474        std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1475
1476        let cell2_path = tmp.path().join("cell2");
1477        std::fs::create_dir(&cell2_path).unwrap();
1478        std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1479
1480        let mgr = CellManager::new_with_path(
1481            tmp.path().to_path_buf(),
1482            256,
1483            cpumask_for_range(16),
1484            HashSet::new(),
1485        )
1486        .unwrap();
1487        let result = mgr.compute_cpu_assignments(false);
1488
1489        // Should error because cell 0 has no CPUs
1490        assert!(result.is_err());
1491        let err = result.unwrap_err();
1492        let err_msg = format!("{:#}", err);
1493        assert!(
1494            err_msg.contains("Cell 0 has no CPUs assigned"),
1495            "Expected 'Cell 0 has no CPUs assigned' error, got: {}",
1496            err_msg
1497        );
1498    }
1499
1500    #[test]
1501    fn test_cpu_assignments_cpusets_cover_all_cpus_with_holdout() {
1502        let tmp = TempDir::new().unwrap();
1503
1504        // Same as test_cpu_assignments_cpusets_cover_all_cpus, but with
1505        // cell0_min_cpus=1: the holdout reserves one CPU for cell 0 before
1506        // assignment. Every CPU is claimed, so the holdout falls back to the
1507        // lowest-numbered claimed CPU (CPU 0, from cell 1).
1508        let cell1_path = tmp.path().join("cell1");
1509        std::fs::create_dir(&cell1_path).unwrap();
1510        std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1511
1512        let cell2_path = tmp.path().join("cell2");
1513        std::fs::create_dir(&cell2_path).unwrap();
1514        std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
1515
1516        let mgr = CellManager::new_with_path_opts(
1517            tmp.path().to_path_buf(),
1518            256,
1519            cpumask_for_range(16),
1520            HashSet::new(),
1521            1,
1522            HashMap::new(),
1523        )
1524        .unwrap();
1525        let assignments = mgr
1526            .compute_cpu_assignments(false)
1527            .expect("holdout should populate cell 0");
1528
1529        let cell0 = assignments
1530            .iter()
1531            .find(|a| a.cell_id == 0)
1532            .expect("cell 0 present");
1533        assert_eq!(cell0.primary.weight(), 1);
1534
1535        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1536        assert_eq!(total, 16);
1537
1538        // Every cell has at least one CPU; cell 1 yielded its lowest CPU to
1539        // the holdout.
1540        for assignment in &assignments {
1541            assert!(
1542                assignment.primary.weight() >= 1,
1543                "cell {} starved after holdout: {:?}",
1544                assignment.cell_id,
1545                assignment.primary
1546            );
1547        }
1548        let donor_weights: Vec<usize> = assignments
1549            .iter()
1550            .filter(|a| a.cell_id != 0)
1551            .map(|a| a.primary.weight())
1552            .collect();
1553        assert!(donor_weights.contains(&7));
1554        assert!(donor_weights.contains(&8));
1555    }
1556
1557    #[test]
1558    fn test_cpu_assignments_holdout_takes_from_largest_cell() {
1559        let tmp = TempDir::new().unwrap();
1560
1561        // 16 CPUs. cell1 claims 8 (0-7), cell2 and cell3 four each. With cpusets
1562        // covering every CPU and cell0_min_cpus=1, the holdout steals from the
1563        // largest cell first -> cell1, its lowest CPU -> CPU 0. (cpu_to_llc is
1564        // set but does not change the pick here: it is only the inner tie-break
1565        // among a donor's CPUs, and cell1's all share one LLC.)
1566        let cell1_path = tmp.path().join("cell1");
1567        std::fs::create_dir(&cell1_path).unwrap();
1568        std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
1569
1570        let cell2_path = tmp.path().join("cell2");
1571        std::fs::create_dir(&cell2_path).unwrap();
1572        std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
1573
1574        let cell3_path = tmp.path().join("cell3");
1575        std::fs::create_dir(&cell3_path).unwrap();
1576        std::fs::write(cell3_path.join("cpuset.cpus"), "12-15\n").unwrap();
1577
1578        let cpu_to_llc: HashMap<usize, usize> = (0..16usize).map(|cpu| (cpu, cpu / 8)).collect();
1579
1580        let mgr = CellManager::new_with_path_opts(
1581            tmp.path().to_path_buf(),
1582            256,
1583            cpumask_for_range(16),
1584            HashSet::new(),
1585            1,
1586            cpu_to_llc,
1587        )
1588        .unwrap();
1589        let assignments = mgr
1590            .compute_cpu_assignments(false)
1591            .expect("holdout should populate cell 0");
1592
1593        let cell0 = assignments
1594            .iter()
1595            .find(|a| a.cell_id == 0)
1596            .expect("cell 0 present");
1597        assert_eq!(cell0.primary.weight(), 1);
1598        assert!(
1599            cell0.primary.test_cpu(0),
1600            "holdout should take CPU 0 (lowest CPU of the largest cell), got {:?}",
1601            cell0.primary
1602        );
1603        assert!(
1604            mgr.enforced_holdout(),
1605            "stealing a claimed CPU must set enforced_holdout"
1606        );
1607
1608        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1609        assert_eq!(total, 16);
1610
1611        // The largest cell yields exactly one CPU to the holdout (8 -> 7); the
1612        // two smaller cells are untouched.
1613        let mut child_weights: Vec<usize> = assignments
1614            .iter()
1615            .filter(|a| a.cell_id != 0)
1616            .map(|a| a.primary.weight())
1617            .collect();
1618        child_weights.sort_unstable();
1619        assert_eq!(child_weights, vec![4, 4, 7]);
1620
1621        // No child cell still holds the held-out CPU.
1622        for a in assignments.iter().filter(|a| a.cell_id != 0) {
1623            assert!(
1624                !a.primary.test_cpu(0),
1625                "cell {} still holds held-out CPU 0",
1626                a.cell_id
1627            );
1628        }
1629    }
1630
1631    #[test]
1632    fn test_cpu_assignments_holdout_steals_evenly_across_cells() {
1633        let tmp = TempDir::new().unwrap();
1634
1635        // Three equal cells with disjoint cpusets covering all 24 CPUs, one LLC
1636        // (no LLC steering), cell0_min_cpus=6. Every CPU is claimed, so the
1637        // holdout must steal six. An even steal takes two from each cell, leaving
1638        // each at six — not six from a single cell (which a global lowest-CPU
1639        // prefix would do, draining cell1 to two while cell2/cell3 keep eight).
1640        for (name, range) in [("cell1", "0-7"), ("cell2", "8-15"), ("cell3", "16-23")] {
1641            let p = tmp.path().join(name);
1642            std::fs::create_dir(&p).unwrap();
1643            std::fs::write(p.join("cpuset.cpus"), format!("{range}\n")).unwrap();
1644        }
1645
1646        let mgr = CellManager::new_with_path_opts(
1647            tmp.path().to_path_buf(),
1648            256,
1649            cpumask_for_range(24),
1650            HashSet::new(),
1651            6,
1652            HashMap::new(),
1653        )
1654        .unwrap();
1655        let assignments = mgr
1656            .compute_cpu_assignments(false)
1657            .expect("holdout should populate cell 0");
1658
1659        let cell0 = assignments
1660            .iter()
1661            .find(|a| a.cell_id == 0)
1662            .expect("cell 0 present");
1663        assert_eq!(
1664            cell0.primary.weight(),
1665            6,
1666            "cell 0 holds the six reserved CPUs"
1667        );
1668
1669        // The steal is spread evenly: each of the three equal cells yields
1670        // exactly two, ending at six.
1671        let mut donor_weights: Vec<usize> = assignments
1672            .iter()
1673            .filter(|a| a.cell_id != 0)
1674            .map(|a| a.primary.weight())
1675            .collect();
1676        donor_weights.sort_unstable();
1677        assert_eq!(
1678            donor_weights,
1679            vec![6, 6, 6],
1680            "holdout should steal evenly (two from each cell), got {donor_weights:?}"
1681        );
1682
1683        // Concretely, cell 0 holds two CPUs from each cell's range.
1684        let count_in =
1685            |lo: usize, hi: usize| (lo..hi).filter(|&c| cell0.primary.test_cpu(c)).count();
1686        assert_eq!(count_in(0, 8), 2, "two CPUs taken from cell1");
1687        assert_eq!(count_in(8, 16), 2, "two CPUs taken from cell2");
1688        assert_eq!(count_in(16, 24), 2, "two CPUs taken from cell3");
1689    }
1690
1691    #[test]
1692    fn test_cpu_assignments_holdout_never_starves_a_child() {
1693        let tmp = TempDir::new().unwrap();
1694
1695        // 4 CPUs fully claimed by two 2-CPU children, cell0_min_cpus=3 -- more
1696        // than the 2 CPUs that can be reserved while each child keeps one. Without
1697        // a per-donor floor the steal drains a child to zero and re-triggers the
1698        // Phase 5 "no CPUs assigned" bail this holdout exists to prevent. The floor
1699        // caps the reservation so every child keeps >=1 CPU and cell 0 receives as
1700        // many as it safely can (2, not the requested 3).
1701        let cell1_path = tmp.path().join("cell1");
1702        std::fs::create_dir(&cell1_path).unwrap();
1703        std::fs::write(cell1_path.join("cpuset.cpus"), "0-1\n").unwrap();
1704
1705        let cell2_path = tmp.path().join("cell2");
1706        std::fs::create_dir(&cell2_path).unwrap();
1707        std::fs::write(cell2_path.join("cpuset.cpus"), "2-3\n").unwrap();
1708
1709        let mgr = CellManager::new_with_path_opts(
1710            tmp.path().to_path_buf(),
1711            256,
1712            cpumask_for_range(4),
1713            HashSet::new(),
1714            3,
1715            HashMap::new(),
1716        )
1717        .unwrap();
1718        let assignments = mgr
1719            .compute_cpu_assignments(false)
1720            .expect("holdout must not starve a child cell");
1721
1722        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1723        assert_eq!(total, 4);
1724
1725        // Every child keeps at least one CPU -- no Phase 5 bail.
1726        for a in assignments.iter().filter(|a| a.cell_id != 0) {
1727            assert!(
1728                a.primary.weight() >= 1,
1729                "cell {} starved by the holdout: {:?}",
1730                a.cell_id,
1731                a.primary
1732            );
1733        }
1734
1735        // Cell 0 received the floor-capped count (2), not the requested 3.
1736        let cell0 = assignments
1737            .iter()
1738            .find(|a| a.cell_id == 0)
1739            .expect("cell 0 present");
1740        assert_eq!(
1741            cell0.primary.weight(),
1742            2,
1743            "holdout should cap at total - num_children, got {:?}",
1744            cell0.primary
1745        );
1746    }
1747
1748    #[test]
1749    fn test_cpu_assignments_holdout_overlapping_cpusets_no_starvation() {
1750        let tmp = TempDir::new().unwrap();
1751
1752        // Overlapping cpusets: cell1 claims only CPU 0; cell2 claims CPUs 0-1, so
1753        // CPU 0 is contested. 5 CPUs total (2,3,4 unclaimed). With cell0_min_cpus=4
1754        // the holdout must dip past the unclaimed CPUs into a claimed one. It must
1755        // NOT take cell2's exclusive CPU 1: that would leave cell2 holding only the
1756        // contested CPU 0, which Phase 3 awards to cell1, zeroing cell2 -> the
1757        // Phase-5 bail. The exclusive-only floor refuses to steal a cell's last
1758        // exclusive CPU, so cell 0 gets only the 3 unclaimed CPUs (capped below 4)
1759        // and every cell keeps >=1.
1760        let cell1_path = tmp.path().join("cell1");
1761        std::fs::create_dir(&cell1_path).unwrap();
1762        std::fs::write(cell1_path.join("cpuset.cpus"), "0\n").unwrap();
1763
1764        let cell2_path = tmp.path().join("cell2");
1765        std::fs::create_dir(&cell2_path).unwrap();
1766        std::fs::write(cell2_path.join("cpuset.cpus"), "0-1\n").unwrap();
1767
1768        let mgr = CellManager::new_with_path_opts(
1769            tmp.path().to_path_buf(),
1770            256,
1771            cpumask_for_range(5),
1772            HashSet::new(),
1773            4,
1774            HashMap::new(),
1775        )
1776        .unwrap();
1777        let assignments = mgr
1778            .compute_cpu_assignments(false)
1779            .expect("holdout must not starve a cell that shares a contested CPU");
1780        // The steal loop runs (3 unclaimed < 4 requested) but no claimed CPU is
1781        // safely reservable, so it breaks without taking one -- enforced_holdout
1782        // tracks an actual steal, not loop entry.
1783        assert!(
1784            !mgr.enforced_holdout(),
1785            "entering the steal loop without taking a claimed CPU must not set \
1786             enforced_holdout"
1787        );
1788
1789        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
1790        assert_eq!(total, 5);
1791
1792        // Every child keeps at least one CPU -- no Phase 5 bail.
1793        for a in assignments.iter().filter(|a| a.cell_id != 0) {
1794            assert!(
1795                a.primary.weight() >= 1,
1796                "cell {} starved by the holdout: {:?}",
1797                a.cell_id,
1798                a.primary
1799            );
1800        }
1801
1802        // The holdout could not safely steal a claimed CPU, so cell 0 holds only
1803        // the 3 unclaimed CPUs -- capped below the requested 4.
1804        let cell0 = assignments
1805            .iter()
1806            .find(|a| a.cell_id == 0)
1807            .expect("cell 0 present");
1808        assert_eq!(
1809            cell0.primary.weight(),
1810            3,
1811            "holdout should cap at the 3 unclaimed CPUs, got {:?}",
1812            cell0.primary
1813        );
1814    }
1815
1816    #[test]
1817    fn test_cpu_assignments_single_cpuset() {
1818        let tmp = TempDir::new().unwrap();
1819
1820        // Only one cell with a cpuset
1821        let cell1_path = tmp.path().join("cell1");
1822        std::fs::create_dir(&cell1_path).unwrap();
1823        std::fs::write(cell1_path.join("cpuset.cpus"), "0,2,4,6\n").unwrap();
1824
1825        let mgr = CellManager::new_with_path(
1826            tmp.path().to_path_buf(),
1827            256,
1828            cpumask_for_range(8),
1829            HashSet::new(),
1830        )
1831        .unwrap();
1832        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1833
1834        assert_eq!(assignments.len(), 2);
1835
1836        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1837        let cell1 = assignments.iter().find(|a| a.cell_id != 0).unwrap();
1838
1839        // cell1 gets even CPUs
1840        assert_eq!(cell1.primary.weight(), 4);
1841        for cpu in [0, 2, 4, 6] {
1842            assert!(cell1.primary.test_cpu(cpu));
1843        }
1844
1845        // cell0 gets odd CPUs
1846        assert_eq!(cell0.primary.weight(), 4);
1847        for cpu in [1, 3, 5, 7] {
1848            assert!(cell0.primary.test_cpu(cpu));
1849        }
1850    }
1851
1852    #[test]
1853    fn test_cpuset_parsing_from_file() {
1854        let tmp = TempDir::new().unwrap();
1855
1856        // Test various cpuset formats
1857        let cell_path = tmp.path().join("cell1");
1858        std::fs::create_dir(&cell_path).unwrap();
1859        std::fs::write(cell_path.join("cpuset.cpus"), "0-3,8-11,16\n").unwrap();
1860
1861        let mgr = CellManager::new_with_path(
1862            tmp.path().to_path_buf(),
1863            256,
1864            cpumask_for_range(32),
1865            HashSet::new(),
1866        )
1867        .unwrap();
1868
1869        // Find the cell and verify its cpuset was parsed correctly
1870        let cell_info = mgr.find_cell_by_name("cell1").unwrap();
1871        let cpuset = cell_info.cpuset.as_ref().unwrap();
1872
1873        assert_eq!(cpuset.weight(), 9); // 4 + 4 + 1
1874        for cpu in 0..4 {
1875            assert!(cpuset.test_cpu(cpu));
1876        }
1877        for cpu in 8..12 {
1878            assert!(cpuset.test_cpu(cpu));
1879        }
1880        assert!(cpuset.test_cpu(16));
1881    }
1882
1883    #[test]
1884    fn test_cpu_assignments_mixed_cpuset_and_no_cpuset() {
1885        let tmp = TempDir::new().unwrap();
1886
1887        // cell1 has a cpuset
1888        let cell1_path = tmp.path().join("cell1");
1889        std::fs::create_dir(&cell1_path).unwrap();
1890        std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
1891
1892        // cell2 has NO cpuset (no cpuset.cpus file)
1893        let cell2_path = tmp.path().join("cell2");
1894        std::fs::create_dir(&cell2_path).unwrap();
1895
1896        let mgr = CellManager::new_with_path(
1897            tmp.path().to_path_buf(),
1898            256,
1899            cpumask_for_range(16),
1900            HashSet::new(),
1901        )
1902        .unwrap();
1903
1904        // Verify cell1 has cpuset, cell2 doesn't
1905        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
1906        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
1907        assert!(cell1_info.cpuset.is_some());
1908        assert!(cell2_info.cpuset.is_none());
1909
1910        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1911
1912        // cell1 (pinned) gets its cpuset: 0-3 (4 CPUs)
1913        // Targets (equal weight, 3 cells, 16 CPUs): cell0=6, cell1=5, cell2=5
1914        // cell1 has 4 exclusive. Deficit = 1 (but can't participate in unclaimed)
1915        // 12 unclaimed CPUs split by deficit: cell0 deficit=6, cell2 deficit=5
1916        // Result: cell0=7, cell1=4, cell2=5
1917        assert_eq!(assignments.len(), 3);
1918
1919        // cell1 gets its cpuset (0-3)
1920        let cell1_assignment = assignments
1921            .iter()
1922            .find(|a| a.cell_id == cell1_info.cell_id)
1923            .unwrap();
1924        assert_eq!(cell1_assignment.primary.weight(), 4);
1925
1926        // cell0 gets 7 CPUs (deficit-proportional share of unclaimed)
1927        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1928        assert_eq!(cell0.primary.weight(), 7);
1929
1930        // cell2 (unpinned) gets 5 CPUs
1931        let cell2_assignment = assignments
1932            .iter()
1933            .find(|a| a.cell_id == cell2_info.cell_id)
1934            .unwrap();
1935        assert_eq!(cell2_assignment.primary.weight(), 5);
1936    }
1937
1938    // ==================== Overlapping cpuset tests ====================
1939
1940    #[test]
1941    fn test_cpu_assignments_partial_overlap() {
1942        let tmp = TempDir::new().unwrap();
1943
1944        // Cell A (cpuset 0-7) and Cell B (cpuset 4-11) - overlap on 4-7
1945        let cell_a_path = tmp.path().join("cell_a");
1946        std::fs::create_dir(&cell_a_path).unwrap();
1947        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
1948
1949        let cell_b_path = tmp.path().join("cell_b");
1950        std::fs::create_dir(&cell_b_path).unwrap();
1951        std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
1952
1953        let mgr = CellManager::new_with_path(
1954            tmp.path().to_path_buf(),
1955            256,
1956            cpumask_for_range(16),
1957            HashSet::new(),
1958        )
1959        .unwrap();
1960        let assignments = mgr.compute_cpu_assignments(false).unwrap();
1961
1962        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
1963        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
1964
1965        let cell_a = assignments
1966            .iter()
1967            .find(|a| a.cell_id == cell_a_info.cell_id)
1968            .unwrap();
1969        let cell_b = assignments
1970            .iter()
1971            .find(|a| a.cell_id == cell_b_info.cell_id)
1972            .unwrap();
1973        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
1974
1975        // Cell A gets exclusive 0-3 (4 CPUs) + half of contested 4-7 (2 CPUs) = 6 CPUs
1976        // Cell B gets half of contested 4-7 (2 CPUs) + exclusive 8-11 (4 CPUs) = 6 CPUs
1977        // Cell 0 gets unclaimed 12-15 (4 CPUs)
1978        assert_eq!(cell_a.primary.weight(), 6);
1979        assert_eq!(cell_b.primary.weight(), 6);
1980        assert_eq!(cell0.primary.weight(), 4);
1981
1982        // Verify exclusive CPUs went to correct cells
1983        for cpu in 0..4 {
1984            assert!(
1985                cell_a.primary.test_cpu(cpu),
1986                "CPU {} should be in cell_a",
1987                cpu
1988            );
1989        }
1990        for cpu in 8..12 {
1991            assert!(
1992                cell_b.primary.test_cpu(cpu),
1993                "CPU {} should be in cell_b",
1994                cpu
1995            );
1996        }
1997        for cpu in 12..16 {
1998            assert!(
1999                cell0.primary.test_cpu(cpu),
2000                "CPU {} should be in cell0",
2001                cpu
2002            );
2003        }
2004
2005        // Verify contested CPUs 4-7 are split - each cell gets exactly 2
2006        let cell_a_contested: Vec<_> = (4..8).filter(|&cpu| cell_a.primary.test_cpu(cpu)).collect();
2007        let cell_b_contested: Vec<_> = (4..8).filter(|&cpu| cell_b.primary.test_cpu(cpu)).collect();
2008        assert_eq!(cell_a_contested.len(), 2);
2009        assert_eq!(cell_b_contested.len(), 2);
2010
2011        // No CPU should be assigned to multiple cells
2012        for cpu in 0..16 {
2013            let mut count = 0;
2014            if cell_a.primary.test_cpu(cpu) {
2015                count += 1;
2016            }
2017            if cell_b.primary.test_cpu(cpu) {
2018                count += 1;
2019            }
2020            if cell0.primary.test_cpu(cpu) {
2021                count += 1;
2022            }
2023            assert!(count <= 1, "CPU {} is assigned to {} cells", cpu, count);
2024        }
2025    }
2026
2027    #[test]
2028    fn test_cpu_assignments_three_way_overlap() {
2029        let tmp = TempDir::new().unwrap();
2030
2031        // All three cells claim CPUs 0-5
2032        let cell_a_path = tmp.path().join("cell_a");
2033        std::fs::create_dir(&cell_a_path).unwrap();
2034        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-5\n").unwrap();
2035
2036        let cell_b_path = tmp.path().join("cell_b");
2037        std::fs::create_dir(&cell_b_path).unwrap();
2038        std::fs::write(cell_b_path.join("cpuset.cpus"), "0-5\n").unwrap();
2039
2040        let cell_c_path = tmp.path().join("cell_c");
2041        std::fs::create_dir(&cell_c_path).unwrap();
2042        std::fs::write(cell_c_path.join("cpuset.cpus"), "0-5\n").unwrap();
2043
2044        let mgr = CellManager::new_with_path(
2045            tmp.path().to_path_buf(),
2046            256,
2047            cpumask_for_range(12),
2048            HashSet::new(),
2049        )
2050        .unwrap();
2051        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2052
2053        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2054        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2055        let cell_c_info = mgr.find_cell_by_name("cell_c").unwrap();
2056
2057        let cell_a = assignments
2058            .iter()
2059            .find(|a| a.cell_id == cell_a_info.cell_id)
2060            .unwrap();
2061        let cell_b = assignments
2062            .iter()
2063            .find(|a| a.cell_id == cell_b_info.cell_id)
2064            .unwrap();
2065        let cell_c = assignments
2066            .iter()
2067            .find(|a| a.cell_id == cell_c_info.cell_id)
2068            .unwrap();
2069        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2070
2071        // 6 contested CPUs / 3 cells = 2 each
2072        assert_eq!(cell_a.primary.weight(), 2);
2073        assert_eq!(cell_b.primary.weight(), 2);
2074        assert_eq!(cell_c.primary.weight(), 2);
2075
2076        // Cell 0 gets unclaimed 6-11 (6 CPUs)
2077        assert_eq!(cell0.primary.weight(), 6);
2078        for cpu in 6..12 {
2079            assert!(cell0.primary.test_cpu(cpu));
2080        }
2081
2082        // Verify total contested CPUs assigned = 6 (no duplicates)
2083        let total_contested: usize = (0..6)
2084            .filter(|&cpu| {
2085                cell_a.primary.test_cpu(cpu)
2086                    || cell_b.primary.test_cpu(cpu)
2087                    || cell_c.primary.test_cpu(cpu)
2088            })
2089            .count();
2090        assert_eq!(total_contested, 6);
2091    }
2092
2093    #[test]
2094    fn test_cpu_assignments_odd_contested_count() {
2095        let tmp = TempDir::new().unwrap();
2096
2097        // Two cells contesting 3 CPUs (odd number - can't split evenly)
2098        let cell_a_path = tmp.path().join("cell_a");
2099        std::fs::create_dir(&cell_a_path).unwrap();
2100        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-2\n").unwrap();
2101
2102        let cell_b_path = tmp.path().join("cell_b");
2103        std::fs::create_dir(&cell_b_path).unwrap();
2104        std::fs::write(cell_b_path.join("cpuset.cpus"), "0-2\n").unwrap();
2105
2106        let mgr = CellManager::new_with_path(
2107            tmp.path().to_path_buf(),
2108            256,
2109            cpumask_for_range(8),
2110            HashSet::new(),
2111        )
2112        .unwrap();
2113        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2114
2115        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2116        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2117
2118        let cell_a = assignments
2119            .iter()
2120            .find(|a| a.cell_id == cell_a_info.cell_id)
2121            .unwrap();
2122        let cell_b = assignments
2123            .iter()
2124            .find(|a| a.cell_id == cell_b_info.cell_id)
2125            .unwrap();
2126
2127        // 3 CPUs / 2 cells = 1 each + 1 remainder
2128        // One cell gets 2, the other gets 1
2129        let total = cell_a.primary.weight() + cell_b.primary.weight();
2130        assert_eq!(total, 3);
2131        assert!(cell_a.primary.weight() >= 1 && cell_a.primary.weight() <= 2);
2132        assert!(cell_b.primary.weight() >= 1 && cell_b.primary.weight() <= 2);
2133
2134        // No overlap in assignments
2135        for cpu in 0..3 {
2136            let a_has = cell_a.primary.test_cpu(cpu);
2137            let b_has = cell_b.primary.test_cpu(cpu);
2138            assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
2139        }
2140    }
2141
2142    #[test]
2143    fn test_cpu_assignments_complete_overlap() {
2144        let tmp = TempDir::new().unwrap();
2145
2146        // Two cells with identical cpusets
2147        let cell_a_path = tmp.path().join("cell_a");
2148        std::fs::create_dir(&cell_a_path).unwrap();
2149        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2150
2151        let cell_b_path = tmp.path().join("cell_b");
2152        std::fs::create_dir(&cell_b_path).unwrap();
2153        std::fs::write(cell_b_path.join("cpuset.cpus"), "0-7\n").unwrap();
2154
2155        let mgr = CellManager::new_with_path(
2156            tmp.path().to_path_buf(),
2157            256,
2158            cpumask_for_range(16),
2159            HashSet::new(),
2160        )
2161        .unwrap();
2162        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2163
2164        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2165        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2166
2167        let cell_a = assignments
2168            .iter()
2169            .find(|a| a.cell_id == cell_a_info.cell_id)
2170            .unwrap();
2171        let cell_b = assignments
2172            .iter()
2173            .find(|a| a.cell_id == cell_b_info.cell_id)
2174            .unwrap();
2175        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2176
2177        // 8 contested CPUs / 2 cells = 4 each
2178        assert_eq!(cell_a.primary.weight(), 4);
2179        assert_eq!(cell_b.primary.weight(), 4);
2180
2181        // Cell 0 gets unclaimed 8-15 (8 CPUs)
2182        assert_eq!(cell0.primary.weight(), 8);
2183        for cpu in 8..16 {
2184            assert!(cell0.primary.test_cpu(cpu));
2185        }
2186
2187        // Verify no overlap between cell_a and cell_b
2188        for cpu in 0..8 {
2189            let a_has = cell_a.primary.test_cpu(cpu);
2190            let b_has = cell_b.primary.test_cpu(cpu);
2191            assert!(!(a_has && b_has), "CPU {} assigned to both cells", cpu);
2192        }
2193    }
2194
2195    #[test]
2196    fn test_cpu_assignments_no_overlap() {
2197        // This verifies existing non-overlapping behavior still works
2198        let tmp = TempDir::new().unwrap();
2199
2200        let cell_a_path = tmp.path().join("cell_a");
2201        std::fs::create_dir(&cell_a_path).unwrap();
2202        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-3\n").unwrap();
2203
2204        let cell_b_path = tmp.path().join("cell_b");
2205        std::fs::create_dir(&cell_b_path).unwrap();
2206        std::fs::write(cell_b_path.join("cpuset.cpus"), "4-7\n").unwrap();
2207
2208        let mgr = CellManager::new_with_path(
2209            tmp.path().to_path_buf(),
2210            256,
2211            cpumask_for_range(16),
2212            HashSet::new(),
2213        )
2214        .unwrap();
2215        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2216
2217        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2218        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2219
2220        let cell_a = assignments
2221            .iter()
2222            .find(|a| a.cell_id == cell_a_info.cell_id)
2223            .unwrap();
2224        let cell_b = assignments
2225            .iter()
2226            .find(|a| a.cell_id == cell_b_info.cell_id)
2227            .unwrap();
2228        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2229
2230        // No overlap - each cell gets its exact cpuset
2231        assert_eq!(cell_a.primary.weight(), 4);
2232        for cpu in 0..4 {
2233            assert!(cell_a.primary.test_cpu(cpu));
2234        }
2235
2236        assert_eq!(cell_b.primary.weight(), 4);
2237        for cpu in 4..8 {
2238            assert!(cell_b.primary.test_cpu(cpu));
2239        }
2240
2241        // Cell 0 gets remaining 8-15
2242        assert_eq!(cell0.primary.weight(), 8);
2243        for cpu in 8..16 {
2244            assert!(cell0.primary.test_cpu(cpu));
2245        }
2246    }
2247
2248    // ==================== format_cell_config tests ====================
2249
2250    #[test]
2251    fn test_format_cell_config_only_cell0() {
2252        let tmp = TempDir::new().unwrap();
2253        let mgr = CellManager::new_with_path(
2254            tmp.path().to_path_buf(),
2255            256,
2256            cpumask_for_range(8),
2257            HashSet::new(),
2258        )
2259        .unwrap();
2260
2261        let mut mask = Cpumask::new();
2262        for cpu in 0..8 {
2263            mask.set_cpu(cpu).unwrap();
2264        }
2265
2266        let assignments = vec![CpuAssignment {
2267            cell_id: 0,
2268            primary: mask,
2269            borrowable: None,
2270        }];
2271        let result = mgr.format_cell_config(&assignments);
2272
2273        assert_eq!(result, "[0: 0-7]");
2274    }
2275
2276    #[test]
2277    fn test_format_cell_config_with_cells() {
2278        let tmp = TempDir::new().unwrap();
2279        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2280
2281        let mgr = CellManager::new_with_path(
2282            tmp.path().to_path_buf(),
2283            256,
2284            cpumask_for_range(16),
2285            HashSet::new(),
2286        )
2287        .unwrap();
2288
2289        let mut mask0 = Cpumask::new();
2290        for cpu in 0..8 {
2291            mask0.set_cpu(cpu).unwrap();
2292        }
2293
2294        let mut mask1 = Cpumask::new();
2295        for cpu in 8..16 {
2296            mask1.set_cpu(cpu).unwrap();
2297        }
2298
2299        let assignments = vec![
2300            CpuAssignment {
2301                cell_id: 0,
2302                primary: mask0,
2303                borrowable: None,
2304            },
2305            CpuAssignment {
2306                cell_id: 1,
2307                primary: mask1,
2308                borrowable: None,
2309            },
2310        ];
2311        let result = mgr.format_cell_config(&assignments);
2312
2313        assert_eq!(result, "[0: 0-7] [1(container-a): 8-15]");
2314    }
2315
2316    // ==================== Cell ID exhaustion tests ====================
2317
2318    #[test]
2319    fn test_cell_id_exhaustion() {
2320        let tmp = TempDir::new().unwrap();
2321
2322        // Create a manager with max_cells=3 (can allocate cell IDs 1 and 2)
2323        // Cell 0 is reserved, so we can create 2 cells before exhaustion
2324        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2325        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2326
2327        let mut mgr = CellManager::new_with_path(
2328            tmp.path().to_path_buf(),
2329            3,
2330            cpumask_for_range(16),
2331            HashSet::new(),
2332        )
2333        .unwrap();
2334        assert_eq!(mgr.cell_count(), 2); // cell1 + cell2
2335
2336        // Adding a third cell should fail due to exhaustion
2337        std::fs::create_dir(tmp.path().join("cell3")).unwrap();
2338        let result = mgr.reconcile_cells();
2339
2340        assert!(result.is_err());
2341        let err = result.unwrap_err();
2342        let err_chain = format!("{:#}", err);
2343        assert!(
2344            err_chain.contains("Cell ID space exhausted"),
2345            "Expected exhaustion error, got: {}",
2346            err_chain
2347        );
2348    }
2349
2350    #[test]
2351    fn test_cell_id_reuse_prevents_exhaustion() {
2352        let tmp = TempDir::new().unwrap();
2353
2354        // Create a manager with max_cells=3
2355        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2356        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2357
2358        let mut mgr = CellManager::new_with_path(
2359            tmp.path().to_path_buf(),
2360            3,
2361            cpumask_for_range(16),
2362            HashSet::new(),
2363        )
2364        .unwrap();
2365        assert_eq!(mgr.cell_count(), 2);
2366
2367        // Remove cell1 to free up its ID
2368        std::fs::remove_dir(tmp.path().join("cell1")).unwrap();
2369        mgr.reconcile_cells().unwrap();
2370        assert_eq!(mgr.cell_count(), 1);
2371
2372        // Now adding cell3 should succeed by reusing the freed ID
2373        std::fs::create_dir(tmp.path().join("cell3")).unwrap();
2374        let result = mgr.reconcile_cells();
2375        assert!(result.is_ok());
2376        assert_eq!(mgr.cell_count(), 2);
2377    }
2378
2379    // ==================== Exclusion tests ====================
2380
2381    #[test]
2382    fn test_scan_excludes_named_cgroups() {
2383        let tmp = TempDir::new().unwrap();
2384
2385        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2386        std::fs::create_dir(tmp.path().join("systemd-workaround.service")).unwrap();
2387        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
2388
2389        let exclude = HashSet::from(["systemd-workaround.service".to_string()]);
2390        let mgr = CellManager::new_with_path(
2391            tmp.path().to_path_buf(),
2392            256,
2393            cpumask_for_range(16),
2394            exclude,
2395        )
2396        .unwrap();
2397
2398        // Only 2 cells — the excluded cgroup is not a cell
2399        assert_eq!(mgr.cell_count(), 2);
2400        assert!(mgr.find_cell_by_name("container-a").is_some());
2401        assert!(mgr.find_cell_by_name("container-b").is_some());
2402        assert!(mgr
2403            .find_cell_by_name("systemd-workaround.service")
2404            .is_none());
2405    }
2406
2407    #[test]
2408    fn test_reconcile_excludes_named_cgroups() {
2409        let tmp = TempDir::new().unwrap();
2410
2411        std::fs::create_dir(tmp.path().join("container-a")).unwrap();
2412
2413        let exclude = HashSet::from(["ignored-service".to_string()]);
2414        let mut mgr = CellManager::new_with_path(
2415            tmp.path().to_path_buf(),
2416            256,
2417            cpumask_for_range(16),
2418            exclude,
2419        )
2420        .unwrap();
2421        assert_eq!(mgr.cell_count(), 1);
2422
2423        // Add an excluded cgroup — should not become a cell
2424        std::fs::create_dir(tmp.path().join("ignored-service")).unwrap();
2425        let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
2426        assert_eq!(new_cells.len(), 0);
2427        assert_eq!(destroyed_cells.len(), 0);
2428        assert_eq!(mgr.cell_count(), 1);
2429
2430        // Add a non-excluded cgroup — should become a cell
2431        std::fs::create_dir(tmp.path().join("container-b")).unwrap();
2432        let (new_cells, destroyed_cells) = mgr.reconcile_cells().unwrap();
2433        assert_eq!(new_cells.len(), 1);
2434        assert_eq!(destroyed_cells.len(), 0);
2435        assert_eq!(mgr.cell_count(), 2);
2436    }
2437
2438    // ==================== Borrowable cpumask tests ====================
2439
2440    #[test]
2441    fn test_borrowable_cpumasks_basic() {
2442        let tmp = TempDir::new().unwrap();
2443
2444        // Create 2 cells without cpusets
2445        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2446        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2447
2448        let mgr = CellManager::new_with_path(
2449            tmp.path().to_path_buf(),
2450            256,
2451            cpumask_for_range(16),
2452            HashSet::new(),
2453        )
2454        .unwrap();
2455        let assignments = mgr.compute_cpu_assignments(true).unwrap();
2456
2457        // Each cell should be able to borrow CPUs from other cells
2458        for assignment in &assignments {
2459            let borrow_mask = assignment.borrowable.as_ref().unwrap();
2460            // borrowable should have no overlap with primary
2461            let overlap = borrow_mask.and(&assignment.primary);
2462            assert_eq!(
2463                overlap.weight(),
2464                0,
2465                "Cell {} borrowable overlaps with primary",
2466                assignment.cell_id
2467            );
2468            // borrowable + primary should cover all CPUs
2469            let union = borrow_mask.or(&assignment.primary);
2470            assert_eq!(
2471                union.weight(),
2472                16,
2473                "Cell {} union doesn't cover all CPUs",
2474                assignment.cell_id
2475            );
2476        }
2477    }
2478
2479    #[test]
2480    fn test_borrowable_cpumasks_no_overlap() {
2481        let tmp = TempDir::new().unwrap();
2482
2483        let cell1_path = tmp.path().join("cell1");
2484        std::fs::create_dir(&cell1_path).unwrap();
2485        std::fs::write(cell1_path.join("cpuset.cpus"), "0-3\n").unwrap();
2486
2487        let cell2_path = tmp.path().join("cell2");
2488        std::fs::create_dir(&cell2_path).unwrap();
2489        std::fs::write(cell2_path.join("cpuset.cpus"), "8-11\n").unwrap();
2490
2491        let mgr = CellManager::new_with_path(
2492            tmp.path().to_path_buf(),
2493            256,
2494            cpumask_for_range(16),
2495            HashSet::new(),
2496        )
2497        .unwrap();
2498        let assignments = mgr.compute_cpu_assignments(true).unwrap();
2499
2500        // Verify no cell's borrowable mask overlaps with its own primary
2501        for assignment in &assignments {
2502            let borrow_mask = assignment.borrowable.as_ref().unwrap();
2503            let overlap = borrow_mask.and(&assignment.primary);
2504            assert_eq!(
2505                overlap.weight(),
2506                0,
2507                "Cell {} borrowable overlaps with primary",
2508                assignment.cell_id
2509            );
2510        }
2511    }
2512
2513    #[test]
2514    fn test_borrowable_cpumasks_respects_cpuset() {
2515        let tmp = TempDir::new().unwrap();
2516
2517        // Cell 1 has cpuset 0-7, Cell 2 has cpuset 8-15
2518        let cell1_path = tmp.path().join("cell1");
2519        std::fs::create_dir(&cell1_path).unwrap();
2520        std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
2521
2522        let cell2_path = tmp.path().join("cell2");
2523        std::fs::create_dir(&cell2_path).unwrap();
2524        std::fs::write(cell2_path.join("cpuset.cpus"), "8-15\n").unwrap();
2525
2526        let mgr = CellManager::new_with_path(
2527            tmp.path().to_path_buf(),
2528            256,
2529            cpumask_for_range(32),
2530            HashSet::new(),
2531        )
2532        .unwrap();
2533        let assignments = mgr.compute_cpu_assignments(true).unwrap();
2534
2535        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2536        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2537
2538        // Cell 1's borrowable should be restricted to its cpuset (0-7),
2539        // minus its own CPUs. Since cell1 gets some of 0-7 as primary,
2540        // the borrowable within 0-7 is whatever it doesn't own.
2541        let cell1_assignment = assignments
2542            .iter()
2543            .find(|a| a.cell_id == cell1_info.cell_id)
2544            .unwrap();
2545        let cell1_borrow = cell1_assignment.borrowable.as_ref().unwrap();
2546        // Cell 1's borrowable should NOT include CPUs outside its cpuset (0-7)
2547        for cpu in 8..32 {
2548            assert!(
2549                !cell1_borrow.test_cpu(cpu),
2550                "Cell 1 borrowable should not include CPU {} (outside cpuset)",
2551                cpu
2552            );
2553        }
2554
2555        // Cell 2's borrowable should be restricted to its cpuset (8-15)
2556        let cell2_assignment = assignments
2557            .iter()
2558            .find(|a| a.cell_id == cell2_info.cell_id)
2559            .unwrap();
2560        let cell2_borrow = cell2_assignment.borrowable.as_ref().unwrap();
2561        for cpu in 0..8 {
2562            assert!(
2563                !cell2_borrow.test_cpu(cpu),
2564                "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2565                cpu
2566            );
2567        }
2568        for cpu in 16..32 {
2569            assert!(
2570                !cell2_borrow.test_cpu(cpu),
2571                "Cell 2 borrowable should not include CPU {} (outside cpuset)",
2572                cpu
2573            );
2574        }
2575    }
2576
2577    // ==================== compute_demand_cpu_assignments tests ====================
2578
2579    #[test]
2580    fn test_demand_cpu_assignments_all_idle() {
2581        let tmp = TempDir::new().unwrap();
2582        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2583        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2584
2585        let mgr = CellManager::new_with_path(
2586            tmp.path().to_path_buf(),
2587            256,
2588            cpumask_for_range(12),
2589            HashSet::new(),
2590        )
2591        .unwrap();
2592
2593        // All cells idle (0 demand) -> falls back to equal division
2594        let demands: HashMap<u32, f64> = [(0, 0.0), (1, 0.0), (2, 0.0)].into();
2595        let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2596
2597        // 12 / 3 = 4 each
2598        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2599        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2600        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2601        let c1 = assignments
2602            .iter()
2603            .find(|a| a.cell_id == cell1_info.cell_id)
2604            .unwrap();
2605        let c2 = assignments
2606            .iter()
2607            .find(|a| a.cell_id == cell2_info.cell_id)
2608            .unwrap();
2609
2610        assert_eq!(cell0.primary.weight(), 4);
2611        assert_eq!(c1.primary.weight(), 4);
2612        assert_eq!(c2.primary.weight(), 4);
2613    }
2614
2615    #[test]
2616    fn test_demand_cpu_assignments_uneven_demand() {
2617        let tmp = TempDir::new().unwrap();
2618        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2619        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2620
2621        let mgr = CellManager::new_with_path(
2622            tmp.path().to_path_buf(),
2623            256,
2624            cpumask_for_range(12),
2625            HashSet::new(),
2626        )
2627        .unwrap();
2628
2629        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2630        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2631
2632        // cell1 is very busy (100%), cell2 is idle (1%), cell0 is moderate (50%)
2633        let demands: HashMap<u32, f64> = [
2634            (0, 50.0),
2635            (cell1_info.cell_id, 100.0),
2636            (cell2_info.cell_id, 1.0),
2637        ]
2638        .into();
2639        let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2640
2641        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2642        let c1 = assignments
2643            .iter()
2644            .find(|a| a.cell_id == cell1_info.cell_id)
2645            .unwrap();
2646        let c2 = assignments
2647            .iter()
2648            .find(|a| a.cell_id == cell2_info.cell_id)
2649            .unwrap();
2650
2651        // Busy cell should get more CPUs than idle cell
2652        assert!(
2653            c1.primary.weight() > c2.primary.weight(),
2654            "Busy cell ({}) should have more CPUs than idle cell ({})",
2655            c1.primary.weight(),
2656            c2.primary.weight()
2657        );
2658        // Each cell should have at least 1 CPU (floor guarantee)
2659        assert!(c2.primary.weight() >= 1);
2660        assert!(cell0.primary.weight() >= 1);
2661        // Total should be 12
2662        assert_eq!(
2663            cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2664            12
2665        );
2666    }
2667
2668    #[test]
2669    fn test_demand_cpu_assignments_with_cpusets() {
2670        let tmp = TempDir::new().unwrap();
2671
2672        // Two cells with overlapping cpusets
2673        let cell_a_path = tmp.path().join("cell_a");
2674        std::fs::create_dir(&cell_a_path).unwrap();
2675        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-7\n").unwrap();
2676
2677        let cell_b_path = tmp.path().join("cell_b");
2678        std::fs::create_dir(&cell_b_path).unwrap();
2679        std::fs::write(cell_b_path.join("cpuset.cpus"), "4-11\n").unwrap();
2680
2681        let mgr = CellManager::new_with_path(
2682            tmp.path().to_path_buf(),
2683            256,
2684            cpumask_for_range(16),
2685            HashSet::new(),
2686        )
2687        .unwrap();
2688
2689        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2690        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2691
2692        // Cell A is much busier than Cell B
2693        let demands: HashMap<u32, f64> = [
2694            (0, 10.0),
2695            (cell_a_info.cell_id, 90.0),
2696            (cell_b_info.cell_id, 10.0),
2697        ]
2698        .into();
2699        let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2700
2701        let cell_a = assignments
2702            .iter()
2703            .find(|a| a.cell_id == cell_a_info.cell_id)
2704            .unwrap();
2705        let cell_b = assignments
2706            .iter()
2707            .find(|a| a.cell_id == cell_b_info.cell_id)
2708            .unwrap();
2709
2710        // Cell A should get more of the contested CPUs 4-7
2711        // Exclusive: A gets 0-3, B gets 8-11
2712        // Contested 4-7: A should get more due to higher demand
2713        assert!(
2714            cell_a.primary.weight() > cell_b.primary.weight(),
2715            "Cell A ({}) should have more CPUs than Cell B ({})",
2716            cell_a.primary.weight(),
2717            cell_b.primary.weight()
2718        );
2719
2720        // Both should have at least their exclusive CPUs
2721        assert!(cell_a.primary.weight() >= 4);
2722        assert!(cell_b.primary.weight() >= 4);
2723    }
2724
2725    #[test]
2726    fn test_demand_cpu_assignments_idle_cell_gets_floor() {
2727        let tmp = TempDir::new().unwrap();
2728        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2729        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2730
2731        let mgr = CellManager::new_with_path(
2732            tmp.path().to_path_buf(),
2733            256,
2734            cpumask_for_range(12),
2735            HashSet::new(),
2736        )
2737        .unwrap();
2738
2739        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2740        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2741
2742        // cell1 is very busy, cell2 is completely idle (weight 0)
2743        let demands: HashMap<u32, f64> = [
2744            (0, 50.0),
2745            (cell1_info.cell_id, 100.0),
2746            (cell2_info.cell_id, 0.0),
2747        ]
2748        .into();
2749        let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2750
2751        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2752        let c1 = assignments
2753            .iter()
2754            .find(|a| a.cell_id == cell1_info.cell_id)
2755            .unwrap();
2756        let c2 = assignments
2757            .iter()
2758            .find(|a| a.cell_id == cell2_info.cell_id)
2759            .unwrap();
2760
2761        // Idle cell gets its minimum target (1 CPU) via deficit-based distribution:
2762        // compute_targets assigns a target of 1, giving it a deficit of 1.
2763        assert_eq!(
2764            c2.primary.weight(),
2765            1,
2766            "Idle cell should get minimum target of 1 CPU"
2767        );
2768        // Busy cell should get the most CPUs
2769        assert!(c1.primary.weight() > cell0.primary.weight());
2770        assert!(c1.primary.weight() > c2.primary.weight());
2771        // Total should be 12
2772        assert_eq!(
2773            cell0.primary.weight() + c1.primary.weight() + c2.primary.weight(),
2774            12
2775        );
2776    }
2777
2778    #[test]
2779    fn test_demand_cpu_assignments_negative_weight_errors() {
2780        let tmp = TempDir::new().unwrap();
2781        std::fs::create_dir(tmp.path().join("cell1")).unwrap();
2782
2783        let mgr = CellManager::new_with_path(
2784            tmp.path().to_path_buf(),
2785            256,
2786            cpumask_for_range(8),
2787            HashSet::new(),
2788        )
2789        .unwrap();
2790
2791        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2792
2793        let demands: HashMap<u32, f64> = [(0, 50.0), (cell1_info.cell_id, -10.0)].into();
2794        let result = mgr.compute_demand_cpu_assignments(&demands, false);
2795        assert!(result.is_err(), "Negative weight should be rejected");
2796        assert!(
2797            result
2798                .unwrap_err()
2799                .to_string()
2800                .contains("negative demand weight"),
2801            "Error message should mention negative weight"
2802        );
2803    }
2804
2805    // ==================== Deficit distribution tests ====================
2806
2807    #[test]
2808    fn test_deficit_distribution_with_cpusets() {
2809        // Two pinned cells with overlapping cpusets + cell0.
2810        // One cell has high demand weight -> gets more contested CPUs.
2811        // Cell that exceeds target from exclusive gets 0 contested.
2812        let tmp = TempDir::new().unwrap();
2813
2814        // cell_a: cpuset 0-9 (10 CPUs)
2815        let cell_a_path = tmp.path().join("cell_a");
2816        std::fs::create_dir(&cell_a_path).unwrap();
2817        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2818
2819        // cell_b: cpuset 6-11 (6 CPUs), overlaps with cell_a on 6-9
2820        let cell_b_path = tmp.path().join("cell_b");
2821        std::fs::create_dir(&cell_b_path).unwrap();
2822        std::fs::write(cell_b_path.join("cpuset.cpus"), "6-11\n").unwrap();
2823
2824        // 16 CPUs total
2825        let mgr = CellManager::new_with_path(
2826            tmp.path().to_path_buf(),
2827            256,
2828            cpumask_for_range(16),
2829            HashSet::new(),
2830        )
2831        .unwrap();
2832
2833        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2834        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2835
2836        // cell_a has high demand (90), cell_b has low demand (10), cell0 moderate (10)
2837        let demands: HashMap<u32, f64> = [
2838            (0, 10.0),
2839            (cell_a_info.cell_id, 90.0),
2840            (cell_b_info.cell_id, 10.0),
2841        ]
2842        .into();
2843        let assignments = mgr.compute_demand_cpu_assignments(&demands, false).unwrap();
2844
2845        let cell_a = assignments
2846            .iter()
2847            .find(|a| a.cell_id == cell_a_info.cell_id)
2848            .unwrap();
2849        let cell_b = assignments
2850            .iter()
2851            .find(|a| a.cell_id == cell_b_info.cell_id)
2852            .unwrap();
2853        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2854
2855        // Verify total = 16
2856        assert_eq!(
2857            cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
2858            16,
2859        );
2860
2861        // cell_a should get the most CPUs (high demand)
2862        assert!(
2863            cell_a.primary.weight() > cell_b.primary.weight(),
2864            "cell_a ({}) should have more CPUs than cell_b ({})",
2865            cell_a.primary.weight(),
2866            cell_b.primary.weight()
2867        );
2868
2869        // Every cell should have at least 1 CPU
2870        assert!(cell_a.primary.weight() >= 1);
2871        assert!(cell_b.primary.weight() >= 1);
2872        assert!(cell0.primary.weight() >= 1);
2873    }
2874
2875    #[test]
2876    fn test_deficit_distribution_equal_weight_with_exclusive() {
2877        // Equal weights, one cell has cpuset covering half the CPUs.
2878        // The deficit adjustment should give the other cell(s) more unclaimed CPUs.
2879        let tmp = TempDir::new().unwrap();
2880
2881        // cell1 has cpuset 0-7 (8 CPUs exclusive, no overlap)
2882        let cell1_path = tmp.path().join("cell1");
2883        std::fs::create_dir(&cell1_path).unwrap();
2884        std::fs::write(cell1_path.join("cpuset.cpus"), "0-7\n").unwrap();
2885
2886        // cell2 has no cpuset (unpinned)
2887        std::fs::create_dir(tmp.path().join("cell2")).unwrap();
2888
2889        // 16 CPUs total, 3 cells (cell0, cell1, cell2), equal weight
2890        let mgr = CellManager::new_with_path(
2891            tmp.path().to_path_buf(),
2892            256,
2893            cpumask_for_range(16),
2894            HashSet::new(),
2895        )
2896        .unwrap();
2897
2898        let cell1_info = mgr.find_cell_by_name("cell1").unwrap();
2899        let cell2_info = mgr.find_cell_by_name("cell2").unwrap();
2900
2901        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2902
2903        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2904        let cell1 = assignments
2905            .iter()
2906            .find(|a| a.cell_id == cell1_info.cell_id)
2907            .unwrap();
2908        let cell2 = assignments
2909            .iter()
2910            .find(|a| a.cell_id == cell2_info.cell_id)
2911            .unwrap();
2912
2913        // Targets (equal weight, 3 cells, 16 CPUs): cell0=6, cell1=5, cell2=5
2914        // cell1 has 8 exclusive (exceeds target of 5), deficit = 0
2915        // 8 unclaimed CPUs split by deficit: cell0 deficit=6, cell2 deficit=5
2916        // cell0 gets ceil(6/11*8) ~= 4, cell2 gets floor(5/11*8) ~= 4
2917        // But exact split: 6/11*8 = 4.36, 5/11*8 = 3.63
2918        // floor: cell0=4, cell2=3 -> assigned=7, remainder=1 -> cell0 gets it
2919        // Result: cell0=5, cell1=8, cell2=3
2920
2921        assert_eq!(cell1.primary.weight(), 8); // Gets all its exclusive CPUs
2922        assert_eq!(
2923            cell0.primary.weight() + cell1.primary.weight() + cell2.primary.weight(),
2924            16,
2925        );
2926
2927        // cell0 should get more unclaimed CPUs than cell2 (higher deficit)
2928        assert!(
2929            cell0.primary.weight() >= cell2.primary.weight(),
2930            "cell0 ({}) should have >= CPUs than cell2 ({})",
2931            cell0.primary.weight(),
2932            cell2.primary.weight()
2933        );
2934    }
2935
2936    #[test]
2937    fn test_deficit_all_cells_exceed_target() {
2938        // Scenario where all claimants in a contested group already exceed their
2939        // global target from exclusive CPUs alone. Verify fallback to equal distribution.
2940        let tmp = TempDir::new().unwrap();
2941
2942        // cell_a: cpuset 0-9 (10 CPUs, overlaps with cell_b on 8-9)
2943        let cell_a_path = tmp.path().join("cell_a");
2944        std::fs::create_dir(&cell_a_path).unwrap();
2945        std::fs::write(cell_a_path.join("cpuset.cpus"), "0-9\n").unwrap();
2946
2947        // cell_b: cpuset 8-13 (6 CPUs, overlaps with cell_a on 8-9)
2948        let cell_b_path = tmp.path().join("cell_b");
2949        std::fs::create_dir(&cell_b_path).unwrap();
2950        std::fs::write(cell_b_path.join("cpuset.cpus"), "8-13\n").unwrap();
2951
2952        // 20 CPUs total, 3 cells, equal weight
2953        // Targets: cell0=7, cell_a=7, cell_b=6 (or similar)
2954        // cell_a exclusive: 0-7 (8 CPUs) -> exceeds target of 7
2955        // cell_b exclusive: 10-13 (4 CPUs) -> below target of 6
2956        // Contested: 8-9 (2 CPUs) -> cell_a deficit=0, cell_b deficit=2
2957        // cell_b should get both contested CPUs
2958        let mgr = CellManager::new_with_path(
2959            tmp.path().to_path_buf(),
2960            256,
2961            cpumask_for_range(20),
2962            HashSet::new(),
2963        )
2964        .unwrap();
2965
2966        let cell_a_info = mgr.find_cell_by_name("cell_a").unwrap();
2967        let cell_b_info = mgr.find_cell_by_name("cell_b").unwrap();
2968
2969        let assignments = mgr.compute_cpu_assignments(false).unwrap();
2970
2971        let cell_a = assignments
2972            .iter()
2973            .find(|a| a.cell_id == cell_a_info.cell_id)
2974            .unwrap();
2975        let cell_b = assignments
2976            .iter()
2977            .find(|a| a.cell_id == cell_b_info.cell_id)
2978            .unwrap();
2979        let cell0 = assignments.iter().find(|a| a.cell_id == 0).unwrap();
2980
2981        // cell_a exceeded its target from exclusive alone, so it gets 0 contested CPUs.
2982        // cell_b has deficit, so it gets all 2 contested CPUs.
2983        let cell_a_contested: Vec<_> = (8..10)
2984            .filter(|&cpu| cell_a.primary.test_cpu(cpu))
2985            .collect();
2986        let cell_b_contested: Vec<_> = (8..10)
2987            .filter(|&cpu| cell_b.primary.test_cpu(cpu))
2988            .collect();
2989        assert_eq!(
2990            cell_a_contested.len(),
2991            0,
2992            "cell_a should get 0 contested CPUs (exceeded target)"
2993        );
2994        assert_eq!(
2995            cell_b_contested.len(),
2996            2,
2997            "cell_b should get all 2 contested CPUs (has deficit)"
2998        );
2999
3000        // Total should be 20
3001        assert_eq!(
3002            cell_a.primary.weight() + cell_b.primary.weight() + cell0.primary.weight(),
3003            20,
3004        );
3005
3006        // Every cell should have at least 1 CPU
3007        assert!(cell0.primary.weight() >= 1);
3008        assert!(cell_a.primary.weight() >= 1);
3009        assert!(cell_b.primary.weight() >= 1);
3010    }
3011
3012    // ==================== distribute_cpus_proportional tests ====================
3013
3014    #[test]
3015    fn test_distribute_proportional_basic() {
3016        let cpus: Vec<usize> = (0..8).collect();
3017        let recipients = vec![(0, 3.0), (1, 1.0)];
3018        let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
3019
3020        // 3/4 * 8 = 6 for cell 0, 1/4 * 8 = 2 for cell 1
3021        assert_eq!(result.get(&0).unwrap().len(), 6);
3022        assert_eq!(result.get(&1).unwrap().len(), 2);
3023    }
3024
3025    #[test]
3026    fn test_distribute_proportional_zero_weight_gets_nothing() {
3027        let cpus: Vec<usize> = (0..6).collect();
3028        let recipients = vec![(0, 1.0), (1, 0.0)];
3029        let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
3030
3031        // Cell 0 gets all, cell 1 gets nothing
3032        assert_eq!(result.get(&0).unwrap().len(), 6);
3033        assert!(result.get(&1).is_none() || result.get(&1).unwrap().is_empty());
3034    }
3035
3036    #[test]
3037    fn test_distribute_proportional_all_zero_fallback() {
3038        let cpus: Vec<usize> = (0..6).collect();
3039        let recipients = vec![(0, 0.0), (1, 0.0)];
3040        let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
3041
3042        // Falls back to equal division
3043        assert_eq!(result.get(&0).unwrap().len(), 3);
3044        assert_eq!(result.get(&1).unwrap().len(), 3);
3045    }
3046
3047    #[test]
3048    fn test_distribute_proportional_skewed_weights_floor_guarantee() {
3049        // Heavily skewed weights: 1.0 vs 100.0 with 38 CPUs (mirrors the bug scenario).
3050        // Without floor guarantee, cell 0 would get floor(1/101 * 38) = 0.
3051        let cpus: Vec<usize> = (0..38).collect();
3052        let recipients = vec![(2, 1.0), (3, 100.0)];
3053        let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
3054
3055        // Both cells must get at least 1 CPU
3056        let cell2_count = result.get(&2).map_or(0, |v| v.len());
3057        let cell3_count = result.get(&3).map_or(0, |v| v.len());
3058        assert!(
3059            cell2_count >= 1,
3060            "cell 2 must get at least 1 CPU, got {}",
3061            cell2_count,
3062        );
3063        assert!(
3064            cell3_count >= 1,
3065            "cell 3 must get at least 1 CPU, got {}",
3066            cell3_count,
3067        );
3068        assert_eq!(cell2_count + cell3_count, 38, "all CPUs must be assigned");
3069    }
3070
3071    #[test]
3072    fn test_distribute_proportional_overlapping_cpusets_no_starvation() {
3073        // Simulates two cells sharing an identical cpuset (all CPUs contested).
3074        // Cell A has deficit=1 (low demand), Cell B has deficit=87 (high demand).
3075        // Without floor guarantee, cell A gets 0 → death spiral.
3076        let cpus: Vec<usize> = (24..62).collect(); // 38 CPUs (24-61)
3077        let recipients = vec![(2, 1.0), (3, 87.0)];
3078        let result = distribute_cpus_proportional(&cpus, &recipients).unwrap();
3079
3080        let cell2_count = result.get(&2).map_or(0, |v| v.len());
3081        let cell3_count = result.get(&3).map_or(0, |v| v.len());
3082        assert!(
3083            cell2_count >= 1,
3084            "cell 2 must not be starved, got {} CPUs",
3085            cell2_count,
3086        );
3087        assert!(
3088            cell3_count >= 1,
3089            "cell 3 must not be starved, got {} CPUs",
3090            cell3_count,
3091        );
3092        assert_eq!(cell2_count + cell3_count, 38);
3093    }
3094
3095    // ==================== compute_targets tests ====================
3096
3097    #[test]
3098    fn test_compute_targets_equal_weight() {
3099        let targets = compute_targets(12, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
3100        assert_eq!(*targets.get(&0).unwrap(), 4);
3101        assert_eq!(*targets.get(&1).unwrap(), 4);
3102        assert_eq!(*targets.get(&2).unwrap(), 4);
3103    }
3104
3105    #[test]
3106    fn test_compute_targets_with_remainder() {
3107        let targets = compute_targets(10, &[(0, 1.0), (1, 1.0), (2, 1.0)]).unwrap();
3108        // 10 / 3 = 3 each + 1 remainder
3109        let total: usize = targets.values().sum();
3110        assert_eq!(total, 10);
3111        for (_, &count) in &targets {
3112            assert!(count >= 3 && count <= 4);
3113        }
3114    }
3115
3116    /// Symmetric pairwise overlaps must produce equal cell sizes regardless
3117    /// of HashMap iteration order.
3118    #[test]
3119    fn test_symmetric_pairwise_overlap_produces_equal_cells() {
3120        let tmp = TempDir::new().unwrap();
3121
3122        // 5 cells on 56 CPUs. Every pair shares exactly 2 CPUs.
3123        // Each cell: 4 exclusive + 8 contested (2 per pair) = 12 in cpuset.
3124        // Exclusive: 0-19, contested: 20-39, unclaimed: 40-55 (cell 0).
3125        //
3126        // Shared pairs:
3127        //   AB: 20-21, AC: 22-23, AD: 24-25, AE: 26-27
3128        //   BC: 28-29, BD: 30-31, BE: 32-33
3129        //   CD: 34-35, CE: 36-37
3130        //   DE: 38-39
3131        let cpusets = [
3132            ("cell_a", "0-3,20-27"),
3133            ("cell_b", "4-7,20-21,28-33"),
3134            ("cell_c", "8-11,22-23,28-29,34-37"),
3135            ("cell_d", "12-15,24-25,30-31,34-35,38-39"),
3136            ("cell_e", "16-19,26-27,32-33,36-39"),
3137        ];
3138        for (name, cpus) in &cpusets {
3139            let p = tmp.path().join(name);
3140            std::fs::create_dir(&p).unwrap();
3141            std::fs::write(p.join("cpuset.cpus"), format!("{cpus}\n")).unwrap();
3142        }
3143
3144        let mgr = CellManager::new_with_path(
3145            tmp.path().to_path_buf(),
3146            256,
3147            cpumask_for_range(56),
3148            HashSet::new(),
3149        )
3150        .unwrap();
3151
3152        let assignments = mgr.compute_cpu_assignments(false).unwrap();
3153
3154        let workload: Vec<_> = assignments.iter().filter(|a| a.cell_id != 0).collect();
3155        assert_eq!(workload.len(), 5, "Expected 5 workload cells");
3156
3157        // All workload cells must have equal CPU counts.
3158        let counts: Vec<(u32, usize)> = workload
3159            .iter()
3160            .map(|a| (a.cell_id, a.primary.weight()))
3161            .collect();
3162        for &(cell_id, count) in &counts {
3163            assert_eq!(
3164                count, counts[0].1,
3165                "Cell {} has {} CPUs, expected {} — symmetric inputs \
3166                 should produce equal cell sizes. All counts: {:?}",
3167                cell_id, count, counts[0].1, counts,
3168            );
3169        }
3170
3171        // Verify no overlap between any pair of cells.
3172        for i in 0..assignments.len() {
3173            for j in (i + 1)..assignments.len() {
3174                for cpu in 0..56 {
3175                    assert!(
3176                        !(assignments[i].primary.test_cpu(cpu)
3177                            && assignments[j].primary.test_cpu(cpu)),
3178                        "CPU {} assigned to both cell {} and cell {}",
3179                        cpu,
3180                        assignments[i].cell_id,
3181                        assignments[j].cell_id,
3182                    );
3183                }
3184            }
3185        }
3186
3187        // Verify all 56 CPUs are assigned.
3188        let total: usize = assignments.iter().map(|a| a.primary.weight()).sum();
3189        assert_eq!(total, 56, "All CPUs must be assigned");
3190    }
3191}