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