cpu-lease.sh

gistfile1.txt

#!/usr/bin/env bash
# cpu-lease — cooperative CPU reservation broker for concurrent agents.
#
# Model: leases come out of a fixed pool of whole SMT pairs. Every lease is a
# transient scope under leases.slice; the slices that hold unleased work are
# narrowed to the complement of what's leased. They are cgroup siblings of
# leases.slice on purpose — cpuset is hierarchical, so a lease nested inside a
# narrowed slice could never exceed it.
#
# Requires cpuset delegated to the user manager; see `cpu-lease install`.

set -euo pipefail

readonly UID_NUM="$(id -u)"
readonly STATE_DIR="${XDG_STATE_HOME:-$HOME/.local/state}/cpu-lease"
readonly LOCK_FILE="$STATE_DIR/lock"
readonly STATE_FILE="$STATE_DIR/leases"
readonly QUEUE_FILE="$STATE_DIR/queue"
readonly CONFIG_FILE="${XDG_CONFIG_HOME:-$HOME/.config}/cpu-lease/config"
readonly SLICE_UNIT="$HOME/.config/systemd/user/leases.slice"
readonly DROPIN="/etc/systemd/system/user@.service.d/50-delegate-cpuset.conf"

# Tunables — override in $CONFIG_FILE.
POOL_VCPUS=8            # ceiling on how much may be leased at once
DEFAULT_VCPUS=2
DEFAULT_TTL=1800        # advisory only: status flags overruns, nothing is killed
WAIT_TIMEOUT=1800
POLL_INTERVAL=0.5
CGROUP_ROOT="${CGROUP_ROOT:-}"   # override the user manager's cgroup path if it differs

# Slices narrowed to the complement of what's leased. app.slice holds terminals
# and GUI apps (so, every agent); session.slice holds gnome-shell and the desktop
# services and is by far the biggest consumer of the three; background.slice is
# where systemd puts deliberately low-priority work. init.scope is left out on
# purpose — it is the user manager itself, and squeezing it buys nothing.
# Everything outside user@UID.service (system.slice, session-N.scope) and the
# kernel's per-CPU threads stay unconfined; see `cpu-lease help`.
CONFINED_SLICES=(app.slice session.slice background.slice)

[[ -r $CONFIG_FILE ]] && . "$CONFIG_FILE"

# Back-compat: configs written against the single-slice version.
if [[ -n ${CONFINED_SLICE:-} ]]; then CONFINED_SLICES=("$CONFINED_SLICE"); fi

die() { printf 'cpu-lease: %s\n' "$*" >&2; exit 1; }
warn() { printf 'cpu-lease: %s\n' "$*" >&2; }

# Option values end up in arithmetic contexts, where a non-numeric string reads
# as 0 — silently turning --timeout into --no-wait and --ttl into "always over".
require_uint() { [[ ${2-} =~ ^[0-9]+$ ]] || die "$1 needs a non-negative integer, got '${2-}'"; }

# ---------------------------------------------------------------- topology

# "0-3" or "0,1" -> "0,1,2,3"
expand_list() {
  local part a b i joined
  local -a out=()
  local IFS=,
  for part in $1; do
    if [[ $part == *-* ]]; then
      a=${part%%-*}; b=${part##*-}
      for ((i = a; i <= b; i++)); do out+=("$i"); done
    else
      out+=("$part")
    fi
  done
  printf -v joined '%s,' "${out[@]}"
  echo "${joined%,}"
}

# All SMT sibling groups, ascending: one "8,9" per line.
all_pairs() {
  local f sl norm
  for f in /sys/devices/system/cpu/cpu[0-9]*/topology/thread_siblings_list; do
    [[ -r $f ]] || continue
    read -r sl < "$f"
    norm="$(expand_list "$sl")"
    printf '%s %s\n' "${norm%%,*}" "$norm"
  done | sort -k1,1n -u | cut -d' ' -f2
}

# Leasable pairs, taken from the high end so low cores stay with the desktop.
pool_pairs() {
  local -a pairs
  mapfile -t pairs < <(all_pairs)
  ((${#pairs[@]})) || die "could not read CPU topology"
  local per want
  per=$(awk -F, '{print NF}' <<<"${pairs[0]}")
  want=$((POOL_VCPUS / per))
  ((want < 1)) && want=1
  ((want > ${#pairs[@]})) && want=${#pairs[@]}
  printf '%s\n' "${pairs[@]: -want}"
}

threads_per_core() { awk -F, '{print NF; exit}' < <(all_pairs); }

# Pool sizing, need=vcpus/per and every vCPU count in status assume all sibling
# groups are the same size. Hybrid CPUs (SMT P-cores + single-thread E-cores)
# break that, and pool_pairs takes from the high end — exactly where the E-cores
# are. Refuse rather than hand out leases that are a fraction of what was asked.
# Not callable from $(...): die there would only kill the subshell.
assert_uniform_topology() {
  local g n per=""
  while read -r g; do
    n=$(awk -F, '{print NF}' <<<"$g")
    [[ -n $per ]] || { per=$n; continue; }
    ((n == per)) || die "heterogeneous SMT topology (sibling groups of $per and $n threads) is not supported"
  done < <(all_pairs)
}

# systemd reports cpusets as "0-7 12-13"; flatten to the "0,1,2,…" form used here.
norm_cpus() {
  local s="${1// /,}"
  [[ -n $s ]] || return 0
  expand_list "$s" | tr ',' '\n' | sort -n | paste -sd,
}

all_cpus() {
  local joined
  printf -v joined '%s,' $(all_pairs)
  echo "${joined%,}" | tr ',' '\n' | sort -n | paste -sd,
}

pool_cpus() {
  local joined
  printf -v joined '%s,' $(pool_pairs)
  echo "${joined%,}" | tr ',' '\n' | sort -n | paste -sd,
}

# ---------------------------------------------------------------- readiness

# systemd's *config* view — it says cpuset is delegated the moment the drop-in
# lands, well before the controller actually reaches the cgroup.
cpuset_delegated() {
  [[ $(systemctl show "user@${UID_NUM}.service" -p DelegateControllers --value 2>/dev/null) == *cpuset* ]]
}

# The only view that matters: can we really write a cpuset for the slices we
# confine? Trusting cpuset_delegated here would hand out leases that enforce
# nothing — worse than refusing, because the caller believes it's isolated.
# Probing the first slice is enough: delegation is granted to user@UID.service
# as a whole, so it either reaches all its children or none.
cpuset_live() {
  [[ -f "$(user_cgroup)/${CONFINED_SLICES[0]}/cpuset.cpus" ]]
}

# leases.slice bakes AllowedCPUs at install time. If the pool has since moved
# (POOL_VCPUS changed in the config, or the topology did), a lease landing
# outside it is not clamped to nothing — cgroup v2 falls back to the parent's
# effective set, silently granting the lease the *whole* slice and overlapping
# every other lease. Verified: requesting 0,1 under leases.slice(8-15) yields
# Cpus_allowed_list 8-15.
pool_matches_slice() {
  [[ "$(norm_cpus "$(systemctl --user show leases.slice -p AllowedCPUs --value 2>/dev/null)")" == "$(pool_cpus)" ]]
}

require_ready() {
  assert_uniform_topology
  cpuset_live || die "cpuset is not active for ${CONFINED_SLICES[0]} — leases would not be enforced. Run: cpu-lease install"
  systemctl --user cat leases.slice >/dev/null 2>&1 || die "leases.slice is missing. Run: cpu-lease install"
  pool_matches_slice || die "leases.slice covers $(systemctl --user show leases.slice -p AllowedCPUs --value 2>/dev/null), but the pool is $(pool_cpus) — leases outside the slice would silently get all of it. Run: cpu-lease install"
}

# ---------------------------------------------------------------- state
#
# leases: id|scope|cpus|vcpus|pid:starttime|start_epoch|ttl|label
# queue:  ticket|vcpus|pid:starttime|epoch
#
# Field 5 packs the broker's pid with its /proc start time so a recycled pid
# can't resurrect a dead holder. Lines written by an older version carry a bare
# pid; those are accepted without the start-time check.

# /proc/PID/stat field 22, minus the two fields eaten by stripping comm (which
# may itself contain spaces and parentheses).
pid_starttime() {
  local line
  read -r line < "/proc/$1/stat" 2>/dev/null || return 1
  local -a fields
  read -ra fields <<<"${line#*") "}"
  printf '%s' "${fields[19]-}"
}

pid_tag() { printf '%s:%s' "$1" "$(pid_starttime "$1")"; }

holder_alive() {
  local pid=${1%%:*} want=${1#*:} got
  [[ -n $pid ]] || return 1
  kill -0 "$pid" 2>/dev/null || return 1
  [[ $1 == *:* ]] || return 0
  got="$(pid_starttime "$pid")" || return 1
  [[ $got == "$want" ]]
}

# Fail closed: only an explicit dead state counts. A systemctl that errors out
# (user D-Bus hiccup, manager busy) must not be read as "scope is gone" —
# that would hand the lease's CPUs back while the workload still runs on them.
scope_gone() {
  local state
  state="$(systemctl --user show "$1" -p ActiveState --value 2>/dev/null)" || return 1
  # 'deactivating' is deliberately not here: the workload may still be draining
  # between SIGTERM and SIGKILL, and holding the cores a few seconds too long
  # beats sharing them with a process that hasn't died yet.
  case $state in
    inactive|failed) return 0 ;;
    *)               return 1 ;;
  esac
}

# Rewrite a state file without the rows whose first field equals $2. Exact match,
# not a regex — `release '.*'` must not sweep every lease.
drop_row() {
  local f=$1 v=$2
  awk -F'|' -v v="$v" '$1 != v' "$f" >"$f.tmp.$$" && mv "$f.tmp.$$" "$f"
}

init_state() {
  mkdir -p "$STATE_DIR"
  : >>"$STATE_FILE"
  : >>"$QUEUE_FILE"
  exec 9>"$LOCK_FILE"
}

lock()   { flock 9; }
unlock() { flock -u 9; }

# Drop leases whose scope is gone and tickets whose holder died. Callers hold the lock.
# Sets GC_REAPED=1 if any lease disappeared, so callers know confinement is stale.
GC_REAPED=0
gc_locked() {
  local tmp id scope cpus vcpus pid start ttl label ticket epoch before after
  GC_REAPED=0
  before=$(wc -l <"$STATE_FILE")
  tmp="$STATE_FILE.tmp.$$"
  : >"$tmp"
  while IFS='|' read -r id scope cpus vcpus pid start ttl label; do
    [[ -n ${id:-} ]] || continue
    # The broker check comes first, and not just because it is cheaper: between
    # publishing the lease and systemd-run creating the scope there is a window
    # (~50ms measured) where the scope does not exist yet. Reaping on the scope
    # alone would drop a live lease and hand its CPUs to the very waiter whose
    # poll loop ran the gc — two workloads, same cores, nobody the wiser.
    if holder_alive "$pid" || ! scope_gone "$scope"; then
      printf '%s|%s|%s|%s|%s|%s|%s|%s\n' "$id" "$scope" "$cpus" "$vcpus" "$pid" "$start" "$ttl" "$label" >>"$tmp"
    fi
  done <"$STATE_FILE"
  mv "$tmp" "$STATE_FILE"
  after=$(wc -l <"$STATE_FILE")
  ((before != after)) && GC_REAPED=1

  tmp="$QUEUE_FILE.tmp.$$"
  : >"$tmp"
  while IFS='|' read -r ticket vcpus pid epoch; do
    [[ -n ${ticket:-} ]] || continue
    # A stale ticket at the head blocks the whole FIFO, so pid reuse matters here.
    holder_alive "$pid" && printf '%s|%s|%s|%s\n' "$ticket" "$vcpus" "$pid" "$epoch" >>"$tmp"
  done <"$QUEUE_FILE"
  mv "$tmp" "$QUEUE_FILE"
}

held_cpus_locked() {
  cut -d'|' -f3 "$STATE_FILE" | tr ',' '\n' | grep -E '^[0-9]+$' | sort -un || true
}

free_pairs_locked() {
  local -a held
  mapfile -t held < <(held_cpus_locked)
  local pair cpu busy
  while read -r pair; do
    busy=0
    for cpu in ${pair//,/ }; do
      for h in ${held[@]+"${held[@]}"}; do [[ $cpu == "$h" ]] && busy=1 && break; done
      ((busy)) && break
    done
    ((busy)) || echo "$pair"
  done < <(pool_pairs)
}

# The confined slices get everything that isn't leased right now — no lease, no
# penalty. Recomputed from scratch every time, so brokers may come and go in any
# order. Applies to inactive slices too: systemd stores the property and honours
# it when the slice starts, which closes the window where a unit could launch
# unconfined while a lease is held.
apply_locked() {
  local -a held
  mapfile -t held < <(held_cpus_locked)
  local -a commons=()
  local cpu h keep
  for cpu in $(all_cpus | tr ',' ' '); do
    keep=1
    for h in ${held[@]+"${held[@]}"}; do [[ $cpu == "$h" ]] && keep=0 && break; done
    ((keep)) && commons+=("$cpu")
  done
  ((${#commons[@]})) || die "refusing to leave ${CONFINED_SLICES[*]} with zero CPUs"
  local joined s
  printf -v joined '%s,' "${commons[@]}"
  joined=${joined%,}
  for s in "${CONFINED_SLICES[@]}"; do
    systemctl --user set-property --runtime "$s" "AllowedCPUs=$joined" \
      || die "could not narrow $s — unleased work there would keep running on leased cores"
  done
}

# ---------------------------------------------------------------- acquire

TICKET=""
LEASE_ID=""
ALLOC_PAIRS=""

cleanup() {
  local rc=$?
  if [[ -n $TICKET || -n $LEASE_ID ]]; then
    # Stop the scope before the CPUs go back to app.slice. On the normal path
    # the command already exited and this is a no-op, but if the broker is dying
    # while its workload lives on (signal, error), widening first would put
    # unleased work back on cores that are still busy.
    # `stop` on an already-collected scope exits 5; never let that abort cleanup.
    [[ -n $LEASE_ID ]] && { systemctl --user stop "lease-$LEASE_ID.scope" >/dev/null 2>&1 || true; }
    lock
    [[ -n $TICKET ]] && { drop_row "$QUEUE_FILE" "$TICKET" || true; }
    if [[ -n $LEASE_ID ]]; then
      drop_row "$STATE_FILE" "$LEASE_ID"
      apply_locked || true
    fi
    unlock
  fi
  return $rc
}

cmd_run() {
  local vcpus=$DEFAULT_VCPUS label="" ttl=$DEFAULT_TTL timeout=$WAIT_TIMEOUT no_smt=0
  while (($#)); do
    case $1 in
      -n|--vcpus)   require_uint "$1" "${2-}"; vcpus=$2; shift 2 ;;
      --label)      label=${2-}; shift 2 ;;
      --ttl)        require_uint "$1" "${2-}"; ttl=$2; shift 2 ;;
      --timeout)    require_uint "$1" "${2-}"; timeout=$2; shift 2 ;;
      --no-wait)    timeout=0; shift ;;
      --no-smt)     no_smt=1; shift ;;
      --)           shift; break ;;
      *)            die "unknown option for run: $1" ;;
    esac
  done
  (($#)) || die "nothing to run — usage: cpu-lease run [opts] -- <command>"

  require_ready
  init_state

  local per; per=$(threads_per_core)
  ((vcpus > 0)) || die "--vcpus must be positive"
  if ((vcpus % per)); then
    vcpus=$(((vcpus / per + 1) * per))
    warn "rounded up to $vcpus vCPU (whole SMT pairs — a shared core is a leaky lease)"
  fi
  ((vcpus <= POOL_VCPUS)) || die "$vcpus vCPU exceeds the pool ceiling of $POOL_VCPUS"

  local need=$((vcpus / per))
  label=${label:-$(basename -- "$1")}
  label=${label//|/ }; label=${label//$'\n'/ }

  TICKET="t$$-$(date +%s%N)"
  local me; me="$(pid_tag $$)"
  trap cleanup EXIT
  trap 'exit 130' INT
  trap 'exit 143' TERM

  lock; gc_locked
  printf '%s|%s|%s|%s\n' "$TICKET" "$vcpus" "$me" "$(date +%s)" >>"$QUEUE_FILE"
  unlock

  local deadline=$(( $(date +%s) + timeout ))
  local cpus="" announced=0
  while :; do
    lock
    gc_locked
    local head; head=$(head -n1 "$QUEUE_FILE" | cut -d'|' -f1)
    if [[ $head == "$TICKET" ]]; then
      local -a got
      mapfile -t got < <(free_pairs_locked)
      if ((${#got[@]} >= need)); then
        local joined
        printf -v joined '%s,' "${got[@]:0:need}"
        cpus="${joined%,}"
        ALLOC_PAIRS="${got[*]:0:need}"
        LEASE_ID="${TICKET#t}"
        printf '%s|%s|%s|%s|%s|%s|%s|%s\n' \
          "$LEASE_ID" "lease-$LEASE_ID.scope" "$cpus" "$vcpus" "$me" "$(date +%s)" "$ttl" "$label" >>"$STATE_FILE"
        drop_row "$QUEUE_FILE" "$TICKET"
        TICKET=""
        apply_locked
        unlock
        break
      fi
    fi
    unlock

    (($(date +%s) < deadline)) || die "timed out waiting for $vcpus vCPU (pool is $POOL_VCPUS)"
    if ((!announced)); then
      announced=1
      warn "waiting for $vcpus vCPU…"
    fi
    sleep "$POLL_INTERVAL"
  done

  local -a affinity_prefix=()
  if ((no_smt)); then
    # Hold the whole pair but run on one thread of it; the idle sibling is what
    # buys the low variance a benchmark needs.
    local pair sj
    local -a solo=()
    for pair in $ALLOC_PAIRS; do solo+=("${pair%%,*}"); done
    printf -v sj '%s,' "${solo[@]}"
    affinity_prefix=(taskset -c "${sj%,}")
  fi

  set +e
  systemd-run --user --scope --quiet --collect \
    --slice=leases.slice --unit="lease-$LEASE_ID" \
    -p "AllowedCPUs=$cpus" \
    -- ${affinity_prefix[@]+"${affinity_prefix[@]}"} "$@"
  local rc=$?
  set -e
  return $rc
}

# ---------------------------------------------------------------- reporting

cmd_status() {
  assert_uniform_topology
  init_state
  lock
  gc_locked
  # A reaped lease leaves app.slice narrower than it should be — widen it here
  # rather than making the user notice and run gc.
  ((GC_REAPED)) && cpuset_live && apply_locked
  local free; free=$(free_pairs_locked | wc -l)
  unlock

  local per; per=$(threads_per_core)
  printf 'pool     %s  (%s vCPU, %s pairs)\n' "$(pool_cpus)" "$POOL_VCPUS" "$((POOL_VCPUS / per))"
  printf 'free     %s vCPU\n' "$((free * per))"
  local s
  for s in "${CONFINED_SLICES[@]}"; do
    printf 'confined %-17s -> %s\n' "$s" \
      "$(systemctl --user show "$s" -p AllowedCPUs --value 2>/dev/null || echo '(unset)')"
  done
  cpuset_live || printf 'WARNING  cpuset not active — leases are NOT enforced. Run: cpu-lease install\n'

  if [[ -s $STATE_FILE ]]; then
    printf '\n%-12s %-8s %-14s %-7s %s\n' ID VCPU CPUS AGE LABEL
    local id scope cpus vcpus pid start ttl label now age flag
    now=$(date +%s)
    while IFS='|' read -r id scope cpus vcpus pid start ttl label; do
      [[ -n ${id:-} ]] || continue
      age=$((now - start))
      flag=""; ((age > ttl)) && flag=" [over TTL]"
      printf '%-12s %-8s %-14s %-7s %s%s\n' "${id:0:12}" "$vcpus" "$cpus" "${age}s" "$label" "$flag"
    done <"$STATE_FILE"
  else
    printf '\nno active leases\n'
  fi
}

cmd_gc() {
  init_state
  lock
  gc_locked
  cpuset_live && apply_locked
  unlock
}

cmd_release() {
  [[ $# -eq 1 ]] || die "usage: cpu-lease release <id>"
  init_state
  lock
  local scope; scope="$(awk -F'|' -v id="$1" '$1 == id { print $2; exit }' "$STATE_FILE")"
  [[ -n $scope ]] || { unlock; die "no such lease: $1"; }
  local pid; pid="$(awk -F'|' -v id="$1" '$1 == id { print $5; exit }' "$STATE_FILE")"
  systemctl --user stop "$scope" 2>/dev/null || true
  if holder_alive "$pid"; then
    # Signal the owner and let its own cleanup drop the row. Removing it here
    # would free the CPUs while a lease still inside its pending window goes on
    # to create its scope regardless.
    kill -TERM "${pid%%:*}" 2>/dev/null || true
  else
    drop_row "$STATE_FILE" "$1"
  fi
  gc_locked; apply_locked
  unlock
}

# ---------------------------------------------------------------- install

user_cgroup() {
  printf '%s' "${CGROUP_ROOT:-/sys/fs/cgroup/user.slice/user-$UID_NUM.slice/user@$UID_NUM.service}"
}

write_user_units() {
  local unitdir; unitdir="$(dirname "$SLICE_UNIT")"
  mkdir -p "$unitdir" "$STATE_DIR" "$(dirname "$CONFIG_FILE")"
  cat >"$SLICE_UNIT" <<EOF
# Generated by cpu-lease install. Sibling of app.slice, not a child — cpuset is
# hierarchical, so a lease nested under the narrowed slice could never exceed it.
[Unit]
Description=CPU lease pool

[Slice]
AllowedCPUs=$(pool_cpus)
EOF
  # Safety net: if every cpu-lease process dies at once, nobody widens app.slice
  # back. This reaps the stale narrowing within a minute.
  cat >"$unitdir/cpu-lease-gc.service" <<EOF
[Unit]
Description=Reap dead CPU leases

[Service]
Type=oneshot
ExecStart=$(readlink -f "$0") gc
EOF
  cat >"$unitdir/cpu-lease-gc.timer" <<'EOF'
[Unit]
Description=Reap dead CPU leases periodically

[Timer]
OnBootSec=30s
OnUnitActiveSec=60s
AccuracySec=10s

[Install]
WantedBy=timers.target
EOF
}

# Everything needing root, in one sudo call — sudo may prompt (password or
# fingerprint) and asking once is the difference between a smooth install and
# a half-applied one. Never feeds the drop-in via /dev/stdin: sudo takes stdin
# for its own auth prompt and the heredoc gets eaten.
root_stage() {
  local all="$1" tmpdir dropin script rc=0
  tmpdir="$(mktemp -d)"
  dropin="$tmpdir/dropin.conf"
  script="$tmpdir/root.sh"

  cat >"$dropin" <<EOF
# Generated by cpu-lease install.
#
# Delegate= hands the user manager the cpuset controller. AllowedCPUs is a
# no-op range covering every CPU — it exists so systemd keeps cpuset enabled in
# the parent slice's subtree_control; without a unit visibly asking for a cpuset
# property, the next daemon-reload drops it again.
[Service]
Delegate=cpu cpuset memory pids
AllowedCPUs=$all
EOF

  cat >"$script" <<EOF
set -eu
install -Dm644 '$dropin' '$DROPIN'
systemctl daemon-reload
systemctl set-property --runtime 'user@${UID_NUM}.service' 'AllowedCPUs=$all' 2>/dev/null || true
sub=/sys/fs/cgroup/user.slice/user-${UID_NUM}.slice/cgroup.subtree_control
if ! grep -qw cpuset "\$sub" 2>/dev/null; then
  echo +cpuset >"\$sub" 2>/dev/null || true
fi
EOF

  echo "==> root: installing $DROPIN and enabling cpuset for user-${UID_NUM}.slice"
  sudo bash "$script" || rc=$?
  rm -rf "$tmpdir"
  return $rc
}

cmd_install() {
  # cgroup.controllers only exists on the unified hierarchy. Don't probe this
  # with `stat -fc %T`: coreutils here reports cgroup2 as "UNKNOWN (0x63677270)".
  [[ -r /sys/fs/cgroup/cgroup.controllers ]] \
    || die "this needs a unified cgroup v2 hierarchy at /sys/fs/cgroup"
  grep -qw cpuset /sys/fs/cgroup/cgroup.controllers \
    || die "the kernel has no cpuset controller on the unified hierarchy"
  command -v systemd-run >/dev/null || die "systemd-run not found"
  systemctl --user show-environment >/dev/null 2>&1 || die "no systemd user manager for uid $UID_NUM"
  assert_uniform_topology

  local all cg; all="$(all_cpus)"; cg="$(user_cgroup)"

  if ! cpuset_delegated || ! grep -qw cpuset "$cg/cgroup.controllers" 2>/dev/null; then
    root_stage "$all" || die "root stage failed"
  else
    echo "==> root: already done, skipping sudo"
  fi

  grep -qw cpuset "$cg/cgroup.controllers" 2>/dev/null \
    || die "cpuset still absent from $cg/cgroup.controllers — see 'cpu-lease doctor'"

  echo "==> user: re-executing the user manager in place"
  # Not 'restart': re-exec keeps every unit and this session alive. Needed
  # because the manager cached its available-controller mask at startup.
  systemctl --user daemon-reexec || true

  echo "==> user: installing units"
  write_user_units
  systemctl --user daemon-reload
  # Starting a slice that declares AllowedCPUs is what makes the user manager
  # enable cpuset in its own subtree_control.
  systemctl --user start leases.slice 2>/dev/null || true
  if ! grep -qw cpuset "$cg/cgroup.subtree_control" 2>/dev/null; then
    echo +cpuset >"$cg/cgroup.subtree_control" 2>/dev/null || true
  fi
  systemctl --user enable --now cpu-lease-gc.timer >/dev/null 2>&1 || true

  echo
  cmd_doctor
}

cmd_uninstall() {
  local cg unitdir; cg="$(user_cgroup)"; unitdir="$(dirname "$SLICE_UNIT")"
  init_state
  lock; gc_locked; unlock
  systemctl --user disable --now cpu-lease-gc.timer >/dev/null 2>&1 || true
  systemctl --user stop leases.slice >/dev/null 2>&1 || true
  # Must hand back the full CPU list explicitly: an empty AllowedCPUs= does NOT
  # clear the restriction — systemd forgets the property but the cgroup keeps
  # whatever cpuset.cpus it last had.
  if cpuset_live; then
    local s
    for s in "${CONFINED_SLICES[@]}"; do
      systemctl --user set-property --runtime "$s" "AllowedCPUs=$(all_cpus)" 2>/dev/null || true
    done
  fi
  rm -f "$SLICE_UNIT" "$unitdir/cpu-lease-gc.service" "$unitdir/cpu-lease-gc.timer"
  systemctl --user daemon-reload
  echo "removed user units and lifted confinement from ${CONFINED_SLICES[*]}"
  echo "the root drop-in is left in place; remove it with:"
  echo "  sudo rm -f $DROPIN && sudo systemctl daemon-reload"
}

# ---------------------------------------------------------------- doctor

cmd_doctor() {
  local cg; cg="$(user_cgroup)"
  local fails=0
  chk() {
    local label="$1"; shift
    if "$@" >/dev/null 2>&1; then printf '  ok    %s\n' "$label"
    else printf '  FAIL  %s\n' "$label"; fails=$((fails + 1)); fi
  }

  echo "checks:"
  chk "drop-in installed"                   test -f "$DROPIN"
  chk "systemd reports cpuset delegated"    cpuset_delegated
  chk "cpuset reaches user@${UID_NUM}.service"    grep -qw cpuset "$cg/cgroup.controllers"
  chk "cpuset enabled for its children"     grep -qw cpuset "$cg/cgroup.subtree_control"
  local s
  for s in "${CONFINED_SLICES[@]}"; do
    chk "$s is cpuset-capable"              test -f "$cg/$s/cpuset.cpus"
  done
  chk "leases.slice active"                 systemctl --user is-active --quiet leases.slice
  chk "gc timer enabled"                    systemctl --user is-enabled --quiet cpu-lease-gc.timer

  if ((fails)); then
    echo
    echo "$fails check(s) failed — leases would NOT be enforced."
    echo "If cpuset never reaches user@${UID_NUM}.service, the delegation only applies"
    echo "at unit start: log out and back in (or reboot) to finish."
    return 1
  fi

  # The only proof that matters: does a lease actually confine the task?
  # One whole sibling group is the smallest real lease — asking for a fixed 2
  # vCPU would allocate two groups on a machine without SMT. And the check is a
  # subset test, not equality: which pairs are free depends on who else holds a
  # lease right now, so pinning it to the first pool pair fails on a busy pool
  # even though confinement worked perfectly.
  local per pool got probe rc=0 inner=""
  per="$(threads_per_core)"
  pool=",$(pool_cpus),"
  # While the lease is held, each confined slice must have lost exactly those
  # CPUs. Reading it from inside the lease is the only way to observe both sides
  # at once, and it is the tool's whole claim in one measurement.
  for s in "${CONFINED_SLICES[@]}"; do
    inner+="printf 'SLICE %s %s\\n' '$s' \"\$(cat '$cg/$s/cpuset.cpus.effective' 2>/dev/null)\"; "
  done
  # --no-wait: the probe must not inherit the 30-minute wait and look hung.
  probe="$("$(readlink -f "$0")" run -n "$per" --label doctor --no-wait -- \
            bash -c "grep -i '^Cpus_allowed_list' /proc/self/status; $inner" 2>&1)" || rc=$?
  got="$(awk '/[Cc]pus_allowed_list/ {print $2}' <<<"$probe")"
  echo
  if ((rc)) && [[ $probe == *"timed out waiting"* ]]; then
    echo "probe skipped: the pool is fully leased right now — run 'cpu-lease doctor' again when it frees up"
    echo "units are in place."
    return 0
  fi
  local ok=1 c n=0
  got="$(expand_list "${got:-}")"
  for c in ${got//,/ }; do
    n=$((n + 1))
    [[ $pool == *",$c,"* ]] || ok=0
  done
  if ((n != per)) || ((!ok)); then
    echo "probe returned '${got:-<nothing>}', expected $per CPU(s) drawn from the pool $(pool_cpus)"
    echo "units are in place but confinement is not taking effect."
    return 1
  fi

  # The other half: no confined slice may have kept the leased CPUs.
  # Count the reports too — a probe that printed nothing must not read as a pass.
  local leaked=0 seen=0 sname seff
  while read -r _ sname seff; do
    [[ -n ${sname:-} ]] || continue
    seen=$((seen + 1))
    seff=",$(norm_cpus "$seff"),"
    for c in ${got//,/ }; do
      if [[ $seff == *",$c,"* ]]; then
        printf 'LEAK  %s still allows CPU %s, which the lease held\n' "$sname" "$c"
        leaked=1
      fi
    done
  done < <(grep '^SLICE ' <<<"$probe" || true)
  ((leaked)) && { echo "confinement is incomplete — unleased work can share leased cores."; return 1; }
  if ((seen != ${#CONFINED_SLICES[@]})); then
    echo "probe reported $seen of ${#CONFINED_SLICES[@]} confined slices — cannot confirm the others were narrowed"
    return 1
  fi

  echo "verified: a $per-vCPU lease confined the task to CPUs $got (pool $(pool_cpus))"
  printf 'verified: %s lost those CPUs for the duration\n' "${CONFINED_SLICES[*]}"
  echo "ready — leases are enforced."
}

cmd_help() {
  cat <<'EOF'
cpu-lease — reserve CPUs for a task so concurrent agents can't steal them.

  cpu-lease run [opts] -- <command>   lease CPUs, run, release on exit
  cpu-lease status                    pool, free capacity, active leases
  cpu-lease release <id>              stop a lease that outlived its task
  cpu-lease gc                        reap dead leases, re-apply confinement
  cpu-lease install                   full install (asks for sudo once)
  cpu-lease doctor                    check the install, prove a lease confines
  cpu-lease uninstall                 remove units, lift confinement

run options:
  -n, --vcpus N     vCPUs to lease (default 2, rounded up to whole SMT pairs)
      --label TEXT  shown in status
      --ttl SEC     advisory; status flags overruns, nothing is killed
      --timeout SEC give up waiting (default 1800)
      --no-wait     fail immediately if the pool is full
      --no-smt      hold whole pairs, run on one thread each (benchmark mode)

How it works: leases come out of a fixed pool of whole SMT pairs and run in
transient scopes under leases.slice. app.slice, session.slice and background.slice
are narrowed to whatever isn't leased — so unleased work there physically cannot
touch a leased core, and gets the full machine back the moment the lease ends.
Waiters are served strictly FIFO, so a large request can't be starved by a stream
of small ones (at the cost of some idle capacity while it waits).

Scope: that covers terminals and GUI apps (app.slice), the desktop shell and its
services (session.slice) and low-priority work (background.slice). Still NOT
confined, and free to run on a leased core:

  system.slice          system daemons — dockerd, journald, NetworkManager
  session-N.scope       SSH and display-manager login sessions
  init.scope            the systemd user manager itself
  kernel threads        ksoftirqd/N, kworker — pinned per-CPU, cpuset cannot move them

'cpu-lease doctor' proves the confined half end to end: it takes a real lease and
checks that every confined slice lost exactly those CPUs while it was held.
EOF
}

case ${1:-help} in
  run)       shift; cmd_run "$@" ;;
  install)   shift; cmd_install "$@" ;;
  uninstall) shift; cmd_uninstall "$@" ;;
  doctor)    shift; cmd_doctor "$@" ;;
  status)  shift; cmd_status "$@" ;;
  release) shift; cmd_release "$@" ;;
  gc)      shift; cmd_gc "$@" ;;
  setup)     shift; cmd_install "$@" ;;
  help|-h|--help) cmd_help ;;
  *)       die "unknown command: $1 (try: cpu-lease help)" ;;
esac