Skip to main content

scx_mitosis/
cell_manager.rs

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