Skip to main content

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