// Copyright 2026 Trace Machina, Inc. All rights reserved.

//

// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    See LICENSE file for details

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

//! What this worker can see of its own CPU and memory, and what it should

//! advertise from that.

//!

//! Every deployment so far typed the advertisement by hand next to the pod's

//! limits, and the two drifted: chinchilla advertised 16 cores and 60 GiB

//! inside 15-core, 56 GiB limits. The limit is the number the kernel

//! enforces, so the advertisement is derived from it here.

use core::hash::BuildHasher;

use std::collections::HashMap;

use std::path::{Path, PathBuf};

use nativelink_config::cas_server::{

    CapacityConfig, CpuUnit, MemoryEnforcement, ResourceEnforcementConfig, WorkerProperty,

};

use nativelink_error::{Code, Error, make_err};

use tracing::{info, warn};

/// CPU and memory as the cgroup (or, failing a limit, the host) reports them.

#[derive(Debug, Clone, Copy, PartialEq, Eq)]

pub struct ObservedCapacity {

    pub cpu_millicores: u64,

    pub memory_kb: u64,

}

/// `cpu.max` is `"<quota> <period>"` in microseconds, or `"max <period>"` when

/// unlimited; unlimited means the host's cores.

pub fn parse_cpu_max(cpu_max: &str, host_cpus: u64) -> Option<u64> {

    let mut fields = cpu_max.split_whitespace();

    let quota = fields.next()
?0
;

    if quota == "max" {

        return Some(host_cpus.saturating_mul(1000));

    }

    let 
quota2
: 
u642
 = quota.parse().ok()
?1
;

    let period: u64 = fields.next()
?0
.parse().ok()
?0
;

    if period == 0 {

        return None;

    }

    Some(quota.saturating_mul(1000) / period)

}

/// `memory.max` is bytes, or `"max"` when unlimited; unlimited means the

/// host's memory.

pub fn parse_memory_max(memory_max: &str, host_memory_kb: u64) -> Option<u64> {

    let value = memory_max.trim();

    if value == "max" {

        return Some(host_memory_kb);

    }

    value.parse::<u64>().ok().map(|bytes| 
bytes1
 / 1024)

}

/// The number after `key` on the first line that starts with it, as

/// `/proc/meminfo` (`MemTotal:  65536000 kB`) and a cgroup's `memory.stat`

/// (`inactive_file 18672467968`) both lay them out. A key without its own

/// colon must be followed by whitespace, so `inactive_file` does not match

/// `inactive_file_extra`.

fn keyed_value(text: &str, key: &str) -> Option<u64> {

    
text.lines()69.0k
.
find_map69.0k
(|line| {

        let 
rest69.0k
 = line.strip_prefix(key)
?137k
;

        if !key.ends_with(':') && 
!rest.starts_with(char::is_whitespace)16
 {

            return None;

        }

        rest.split_whitespace().next()
?0
.parse().ok()

    })

}

/// `MemTotal` from `/proc/meminfo`, in KiB.

pub fn parse_meminfo_total_kb(meminfo: &str) -> Option<u64> {

    keyed_value(meminfo, "MemTotal:")

}

/// The headroom the advertisement is divided by: the block's own figure

/// when set, else the enforcement headroom while memory enforcement is on,

/// so the two numbers cannot drift apart, else nothing.

pub fn memory_headroom_percent(

    config: &CapacityConfig,

    enforcement: Option<&ResourceEnforcementConfig>,

) -> u64 {

    config.memory_headroom_percent.unwrap_or_else(|| 
{3

        enforcement

            .filter(|enforcement| 
enforcement.memory2
 == 
MemoryEnforcement::Soft2
)

            .map_or(0, |enforcement| enforcement.memory_headroom_percent)

    })

}

/// The observed capacity less the worker's own share, memory divided by

/// `memory_headroom_percent`. Returns `(cpu, memory_kb)` with the CPU on

/// the scale `cpu_unit` names: whole cores rounded down, so a 14-core pod

/// keeping one core back advertises `13`, not `13000`, to a scheduler

/// whose actions ask for `cpu_count=1`.

pub const fn advertised(

    observed: ObservedCapacity,

    config: &CapacityConfig,

    memory_headroom_percent: u64,

) -> (u64, u64) {

    let cpu_millicores = observed

        .cpu_millicores

        .saturating_sub(config.overhead_cpu_millicores);

    let cpu = match config.cpu_unit {

        CpuUnit::Cores => cpu_millicores / 1000,

        CpuUnit::Millicores => cpu_millicores,

    };

    let memory = observed

        .memory_kb

        .saturating_sub(config.overhead_memory_kb)

        .saturating_mul(100)

        / (100 + memory_headroom_percent);

    (cpu, memory)

}

/// The largest memory reservation the scheduler can match to this worker:

/// what it advertises under `property`, from the cgroup when the capacity

/// block sets that property and the cgroup was read, else from the static

/// platform properties. `None` when the worker advertises no such figure.

pub fn advertised_memory_kb<S: BuildHasher>(

    capacity: Option<&CapacityConfig>,

    memory_headroom_percent: u64,

    observed: Option<ObservedCapacity>,

    properties: &HashMap<String, WorkerProperty, S>,

    property: &str,

) -> Option<u64> {

    if let (Some(
config3
), Some(
observed3
)) = (capacity, observed)

        && config.memory_property_name == property

    {

        return Some(advertised(observed, config, memory_headroom_percent).1);

    }

    match properties.get(property)
?1
 {

        WorkerProperty::Values(values) => values.iter().filter_map(|v| v.parse().ok()).max(),

        WorkerProperty::QueryCmd(_) => None,

    }

}

/// What to do when the cgroup cannot be read: the configured properties

/// stand if they carry both numbers, since an operator who kept them has a

/// worker that still takes work; a worker with neither would register and

/// then satisfy no action that asks for CPU or memory, idling for good on

/// one warning line, so that is refused.

pub fn without_cgroup<S: BuildHasher>(

    config: &CapacityConfig,

    properties: &HashMap<String, WorkerProperty, S>,

) -> Result<(), Error> {

    let missing: Vec<&str> = [&config.cpu_property_name, &config.memory_property_name]

        .into_iter()

        .
filter3
(|name| !properties.contains_key(name.as_str()))

        .map(String::as_str)

        .collect();

    if missing.is_empty() {

        warn!(

            "capacity is configured but no cgroup v2 limit applies to the worker (see the line above for the cgroup looked at); advertising the configured platform_properties instead"

        );

        return Ok(());

    }

    Err(make_err!(

        Code::FailedPrecondition,

        "capacity is configured but no cgroup v2 limit applies to the worker (not Linux, cgroup v1, no permission, or no limit at any level; see the line above), and platform_properties carries no {}: set them, or drop the capacity block",

        missing.join(" or ")

    ))

}

/// The cgroup v2 path of this process from `/proc/self/cgroup`: the `0::`

/// line, which is `/` inside a container with its own cgroup namespace and

/// the full `/kubepods.slice/.../cri-containerd-<id>.scope` path in one that

/// shares the host's, as a privileged container does.

pub fn parse_self_cgroup(self_cgroup: &str) -> Option<&str> {

    self_cgroup

        .lines()

        .find_map(|line| line.strip_prefix("0::"))

        .map(str::trim)

        .filter(|path| 
path71.0k
.
starts_with71.0k
('/'))

}

/// Where this process's own `cpu.max`, `memory.max` and `memory.current`

/// live: the mount root joined with the path from `/proc/self/cgroup` when

/// that directory exists under the mount, the mount root otherwise. A

/// container with a private cgroup namespace sees itself at the root, so its

/// path is `/` and both agree. A privileged container sees the host's tree,

/// where the root's files are the host's (or absent, as on a systemd host),

/// and its own numbers sit further down.

pub fn resolve_cgroup_dir(

    mount_root: &Path,

    self_cgroup: Option<&str>,

    is_dir: impl Fn(&Path) -> bool,

) -> PathBuf {

    if let Some(
path71.0k
) = self_cgroup.and_then(parse_self_cgroup) {

        let own = mount_root.join(path.trim_start_matches('/'));

        if is_dir(&own) {

            return own;

        }

    }

    mount_root.to_path_buf()

}

/// This process's cgroup directory on the live system.

fn cgroup_dir() -> PathBuf {

    let self_cgroup = std::fs::read_to_string("/proc/self/cgroup").ok();

    resolve_cgroup_dir(Path::new("/sys/fs/cgroup"), self_cgroup.as_deref(), |dir| {

        dir.join("memory.max").is_file()

    })

}

/// The `avg10` of the `full` line of a cgroup v2 `memory.pressure` file:

/// the share of the last ten seconds (0 to 100) during which every task in the

/// cgroup was stalled on memory. `None` when the text has no such line.

pub fn parse_pressure_full_avg10(text: &str) -> Option<f64> {

    text.lines().find_map(|line| {

        let 
rest2
 = line.strip_prefix("full")
?2
;

        rest.split_whitespace()

            .find_map(|field| field.strip_prefix("avg10="))

            .and_then(|value| value.parse().ok())

    })

}

/// The `memory.pressure` file that describes this worker: the one in the

/// nearest limited cgroup (the limit the kernel enforces, see

/// [`find_limit`]), else the one in the worker's own cgroup directory.

/// `None` where neither exists: cgroup v1, a kernel without PSI, or a

/// cgroup root, which has no `memory.pressure`.

#[cfg(target_os = "linux")]

pub fn memory_pressure_file() -> Option<PathBuf> {

    let mount_root = Path::new("/sys/fs/cgroup");

    let own = cgroup_dir();

    let limited = find_limit(&own, mount_root, "memory.max", read_file).map(|(dir, _)| dir);

    [limited, Some(own)]

        .into_iter()

        .flatten()

        .map(|dir| dir.join("memory.pressure"))

        .find(|path| path.is_file())

}

/// Whether a `cpu.max` or `memory.max` value is "no limit": `max`, or for

/// `cpu.max` `max <period>`.

fn is_unlimited(value: &str) -> bool {

    value.split_whitespace().next() == Some("max")

}

/// The first limit for `file` from `dir` up to and including `mount_root`,

/// with the directory it was found in. The process's own cgroup is often

/// an unlimited child of a limited parent (a systemd scope in a pod, a

/// split cgroup, a privileged container under `kubepods.slice`), and the

/// parent's limit is the one the kernel enforces on it. `None` when every

/// level says `max` or nothing is readable: no limit applies.

pub fn find_limit(

    dir: &Path,

    mount_root: &Path,

    file: &str,

    read: impl Fn(&Path) -> Option<String>,

) -> Option<(PathBuf, String)> {

    let mut at = dir;

    loop {

        if let Some(
value8
) = read(&at.join(file))

            && !is_unlimited(value.trim())

        {

            return Some((at.to_path_buf(), value));

        }

        if at == mount_root || 
!9
at9
.starts_with(mount_root) {

            return None;

        }

        at = at.parent()
?0
;

    }

}

fn read_file(path: &Path) -> Option<String> {

    std::fs::read_to_string(path).ok()

}

/// Reads the limits that apply to the worker: for CPU and for memory, the

/// nearest limit from its own cgroup up to the mount root. `None` where no

/// limit applies at any level (not Linux, cgroup v1, no permission, or a

/// host where nothing is limited), so the configured properties stand.

pub fn observe_cgroup() -> Option<ObservedCapacity> {

    let host_cpus = std::thread::available_parallelism().map_or(1, |n| n.get() as u64);

    let host_memory_kb = std::fs::read_to_string("/proc/meminfo")

        .ok()

        .and_then(|meminfo| parse_meminfo_total_kb(&meminfo))

        .unwrap_or(0);

    let mount_root = Path::new("/sys/fs/cgroup");

    let dir = cgroup_dir();

    let cpu = find_limit(&dir, mount_root, "cpu.max", read_file);

    let memory = find_limit(&dir, mount_root, "memory.max", read_file);

    if cpu.is_none() && memory.is_none() {

        warn!(

            cgroup = %dir.display(),

            "no cgroup v2 cpu.max or memory.max limit from this cgroup up to the mount root; keeping the configured platform_properties"

        );

        return None;

    }

    let cpu_max = cpu.map_or_else(|| "max".to_string(), |(_, value)| value);

    let memory_max = memory.map_or_else(|| "max".to_string(), |(_, value)| value);

    Some(ObservedCapacity {

        cpu_millicores: parse_cpu_max(&cpu_max, host_cpus)?,

        memory_kb: parse_memory_max(&memory_max, host_memory_kb)?,

    })

}

/// Sets the CPU and memory properties to what the cgroup allows. When

/// nothing can be read, the configured properties stand if they carry both,

/// and the worker fails to start if they do not. Returns what was

/// advertised.

pub fn apply<S: BuildHasher>(

    config: &CapacityConfig,

    memory_headroom_percent: u64,

    properties: &mut HashMap<String, WorkerProperty, S>,

) -> Result<Option<(u64, u64)>, Error> {

    let Some(observed) = observe_cgroup() else {

        without_cgroup(config, properties)?;

        return Ok(None);

    };

    let (cpu, memory_kb) = advertised(observed, config, memory_headroom_percent);

    properties.insert(

        config.cpu_property_name.clone(),

        WorkerProperty::Values(vec![cpu.to_string()]),

    );

    properties.insert(

        config.memory_property_name.clone(),

        WorkerProperty::Values(vec![memory_kb.to_string()]),

    );

    info!(

        observed_cpu_millicores = observed.cpu_millicores,

        observed_memory_kb = observed.memory_kb,

        cpu,

        cpu_unit = ?config.cpu_unit,

        memory_kb,

        memory_headroom_percent,

        "Advertising capacity from the cgroup"

    );

    Ok(Some((cpu, memory_kb)))

}

pub fn parse_memory_current_kb(memory_current: &str) -> Option<u64> {

    memory_current

        .trim()

        .parse::<u64>()

        .ok()

        .map(|bytes| 
bytes16
 / 1024)

}

/// `MemAvailable` from `/proc/meminfo`, in KiB.

pub fn parse_meminfo_available_kb(meminfo: &str) -> Option<u64> {

    keyed_value(meminfo, "MemAvailable:")

}

/// The page cache in a cgroup's `memory.stat` the kernel can drop at once,

/// in KiB: `inactive_file`, the cache it reclaims first and what kubelet

/// leaves out of the working set it evicts on, less `file_dirty` and

/// `file_writeback`, which cannot be dropped until they reach the disk

/// (kubelet does not take those off, so this reads up to that much lower).

/// Dirty pages may sit on the active list instead; taking all of them off

/// the inactive one errs towards less free, never more. `None` without an

/// `inactive_file` line.

pub fn parse_memory_stat_reclaimable_kb(memory_stat: &str) -> Option<u64> {

    let 
inactive_file9
 = keyed_value(memory_stat, "inactive_file")
?2
;

    let dirty = keyed_value(memory_stat, "file_dirty").unwrap_or(0);

    let writeback = keyed_value(memory_stat, "file_writeback").unwrap_or(0);

    Some(

        inactive_file

            .saturating_sub(dirty)

            .saturating_sub(writeback)

            / 1024,

    )

}

/// Memory the worker could still give an action: the cgroup limit less its

/// working set (current usage less the page cache the kernel can drop)

/// when there is a limit, the host's `MemAvailable` when there is not,

/// nothing when neither can be read. `memory.current` charges every page

/// of file cache the worker's I/O has touched, and the kernel only reclaims

/// it near the limit, so a busy worker's usage sits just under its limit

/// with little of it in use.

pub const fn free_memory_kb_from(

    limit_kb: Option<u64>,

    current_kb: Option<u64>,

    reclaimable_kb: Option<u64>,

    available_kb: Option<u64>,

) -> Option<u64> {

    match (limit_kb, current_kb) {

        (Some(limit), Some(current)) => {

            let reclaimable = match reclaimable_kb {

                Some(kb) => kb,

                None => 0,

            };

            Some(limit.saturating_sub(current.saturating_sub(reclaimable)))

        }

        _ => available_kb,

    }

}

/// Free memory measured against `limited`, the nearest limited cgroup and

/// its `memory.max` as `find_limit` returns them, with `read` for the

/// files; the host's `MemAvailable` when there is no limit.

pub fn free_memory_kb_with(

    limited: Option<&(PathBuf, String)>,

    read: impl Fn(&Path) -> Option<String>,

) -> Option<u64> {

    let Some((
dir4
, 
max4
)) = limited else {

        return read(Path::new("/proc/meminfo"))

            .and_then(|meminfo| parse_meminfo_available_kb(&meminfo));

    };

    let read_current =

        || read(&dir.join("memory.current")).and_then(|current| 
parse_memory_current_kb6
(
&current6
));

    // The usage is read on both sides of the cache, and the larger taken.

    // That mitigates the common races, cache that grows during the read or

    // is reclaimed during it, but the three reads are not one snapshot:

    // cache that grows at the `memory.stat` read and is reclaimed before

    // the second usage read is missed by both usage reads yet subtracted,

    // so a transient cache peak can report more free than there is. The

    // cache figure also comes from per-CPU counters the kernel flushes in

    // batches, so it can run ahead of what is there by a few MiB per CPU.

    let current_before = read_current();

    let reclaimable_kb =

        read(&dir.join("memory.stat")).and_then(|stat| parse_memory_stat_reclaimable_kb(&stat));

    let current_kb = current_before.max(read_current());

    match (parse_memory_current_kb(max.trim()), current_kb) {

        (Some(limit_kb), Some(current_kb)) => {

            free_memory_kb_from(Some(limit_kb), Some(current_kb), reclaimable_kb, None)

        }

        _ => read(Path::new("/proc/meminfo"))

            .and_then(|meminfo| parse_meminfo_available_kb(&meminfo)),

    }

}

/// Whether `dir`, the worker's own cgroup, carries a memory limit of its

/// own with readable usage: then nothing outside the worker gives memory

/// back, so what it reads free while idle is the most it will read. A limit

/// only on an ancestor (a shared `kubepods.slice`, say) is shared with

/// other workloads, which can release memory, and without readable usage

/// free memory falls back to the host's `MemAvailable`; neither counts.

pub fn own_memory_limit_in(dir: &Path, read: impl Fn(&Path) -> Option<String>) -> bool {

    let limited = read(&dir.join("memory.max"))

        .is_some_and(|max| 
parse_memory_current_kb3
(max.trim()).
is_some3
());

    let usage = read(&dir.join("memory.current"))

        .is_some_and(|current| 
parse_memory_current_kb(&current)3
.
is_some3
());

    limited && 
usage2

}

/// `own_memory_limit_in` for this process's cgroup.

#[cfg(target_os = "linux")]

pub fn memory_is_limited() -> bool {

    own_memory_limit_in(&cgroup_dir(), read_file)

}

#[cfg(not(target_os = "linux"))]

pub const fn memory_is_limited() -> bool {

    false

}

/// What the worker reports on each keepalive. Linux only; elsewhere the

/// worker reports nothing and is never vetoed.

#[cfg(target_os = "linux")]

pub fn free_memory_kb() -> Option<u64> {

    // The limit and the usage counted against it come from the same level:

    // the nearest limited ancestor, which charges everything below it.

    let limited = find_limit(

        &cgroup_dir(),

        Path::new("/sys/fs/cgroup"),

        "memory.max",

        read_file,

    );

    free_memory_kb_with(limited.as_ref(), read_file)

}

#[cfg(not(target_os = "linux"))]

pub const fn free_memory_kb() -> Option<u64> {

    None

}