#!/usr/bin/env bash
# Run a command with its disk-write rate held under a ceiling.
#
# Cargo on this box can saturate the virtual disk badly enough to take the
# whole machine out. The usual levers are unavailable here and it is worth
# recording why, so nobody burns an afternoon rediscovering it:
#
#   cgroup v2 io.max  -- /sys/fs/cgroup is mounted ro and cannot be remounted
#                        or re-mounted elsewhere; the container has no
#                        CAP_SYS_ADMIN, so this fails even under sudo.
#   vm.dirty_bytes    -- /proc/sys is read-only.
#   ionice            -- the scheduler is mq-deadline, where IO priority is a
#                        no-op, and /sys/block/*/queue/scheduler is read-only
#                        so BFQ cannot be selected.
#   nice              -- CPU only. Writeback happens in kernel flusher threads
#                        that never see the nice value, so a niced build
#                        saturates the disk exactly as fast as an un-niced one.
#
# What is left is duty-cycling from userspace: sample the block device's
# written-sectors counter, and when the observed rate exceeds the ceiling,
# SIGSTOP the whole process group until it drops. Crude, but it acts on
# measured throughput rather than on a scheduler hint, which is the only
# property that actually matters here.
#
#   scripts/paced.sh -- make build
#   MBPS=30 scripts/paced.sh -- env CARGO_CMD=test PKG=xai-grok-pager ./scripts/build.sh --lib
set -euo pipefail

# Measured on this box, idle, 1 GB transfers:
#
#   read   210 MB/s  (O_DIRECT and buffered alike)
#   write   50 MB/s  (identical for buffered+fdatasync, O_DIRECT 1M, O_DIRECT 64k)
#
# Write is pinned at ~50 MB/s regardless of block size or IO mode, so it is a
# volume-level cap rather than disk physics, and it is 4x slower than read.
# A short burst can reach ~150 MB/s on what look like burst credits -- do not
# calibrate against that, it is the number that made an earlier 40 MB/s ceiling
# useless. 15 MB/s is under a third of sustained capacity, which leaves the rest
# of the machine usable while a build runs.
MBPS="${MBPS:-15}"        # write ceiling in MB/s, averaged over one window
DEV="${DEV:-vdb}"         # block device to watch; the disk, not the partition
WINDOW="${WINDOW:-1}"     # sampling window in seconds

[[ "${1:-}" == "--" ]] && shift
if [[ $# -eq 0 ]]; then
  echo "usage: [MBPS=60] [DEV=vdb] scripts/paced.sh -- <command...>" >&2
  exit 2
fi

stats_field() {
  # /proc/diskstats field 10 is sectors written, in 512-byte units.
  awk -v dev="$DEV" '$3 == dev { print $10; exit }' /proc/diskstats
}

if [[ -z "$(stats_field)" ]]; then
  echo "paced: no such device in /proc/diskstats: $DEV" >&2
  exit 1
fi

# Its own process group, so one signal reaches cargo and every rustc it spawned.
#
# `set -m` rather than setsid: setsid only forks when it is already a group
# leader, so under `cmd &` it usually execs in place -- but when it does fork,
# $! is setsid's pid, which exits immediately. The loop below then sees the
# child "finish" instantly and reports success while the real build is still
# running. Job control gives the background job its own pgid deterministically.
set -m
"$@" &
child=$!
pgid=$child
set +m

cleanup() {
  # Never leave the build stopped: a SIGSTOPped process group that outlives
  # this script looks exactly like a hang.
  kill -CONT -- "-$pgid" 2>/dev/null || true
  kill -TERM -- "-$pgid" 2>/dev/null || true
}
trap cleanup INT TERM

ceiling_sectors=$(( MBPS * 1024 * 1024 / 512 * WINDOW ))
prev=$(stats_field)
stalled=0

while kill -0 "$child" 2>/dev/null; do
  sleep "$WINDOW"
  now=$(stats_field)
  delta=$(( now - prev ))
  prev=$now

  if (( delta > ceiling_sectors )); then
    # Stop for as long as the overshoot implies, capped so the build still
    # makes progress even when something else on the box is writing hard.
    pause=$(( delta / ceiling_sectors ))
    (( pause > 4 )) && pause=4
    stalled=$(( stalled + pause ))
    kill -STOP -- "-$pgid" 2>/dev/null || true
    sleep "$pause"
    kill -CONT -- "-$pgid" 2>/dev/null || true
    prev=$(stats_field)
  fi
done

rc=0
wait "$child" || rc=$?
trap - INT TERM
echo "paced: done rc=$rc (throttled ${stalled}s at ${MBPS}MB/s ceiling on /dev/$DEV)" >&2
# Propagate the wrapped command's status. Swallowing it once already produced a
# confident "build finished, exit 0" for a build that had in fact been killed.
exit "$rc"