Skip to main content

scx_mitosis/
cell_manager.rs

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