leanrun

#!/bin/sh
set -eu

script_dir="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)"
repo_root="$(CDPATH= cd -- "$script_dir/.." && pwd)"
toolchain_name="$(sed -n '1p' "$repo_root/lean-toolchain")"
toolchain_key="$(printf '%s' "$toolchain_name" | sed 's,/,--,g; s,:,---,g')"
elan_root="${ELAN_HOME:-$HOME/.elan}"
toolchain="${LEANRUN_TOOLCHAIN:-$elan_root/toolchains/$toolchain_key}"
run_timeout="${LEANRUN_TIMEOUT:-900}"
lock_timeout="${LEANRUN_LOCK_TIMEOUT:-900}"
lock_dir="${LEANRUN_LOCKDIR:-${VQ_LOCKDIR:-${TMPDIR:-/tmp}/vq-leanrun.$(id -u)}}"
cache_dir="${LEANRUN_CACHE_HOME:-$repo_root/.lake/cache}"
local_mode="${LEANRUN_LOCAL:-0}"
memory_high="${LEANRUN_MEMORY_HIGH:-4G}"
cpu_count="${LEANRUN_CPUS:-}"
thread_count="${LEANRUN_THREADS:-}"
log_file=""

usage() {
  echo "usage: tools/leanrun [--timeout DURATION] [--lock-timeout SECONDS] [--memory-high SIZE] [--cpus N|auto] [--threads N] [--log FILE] [--toolchain DIR] <command> [args]" >&2
  echo "SIZE is the scope's MemoryHigh in whole gibibytes with the suffix G (default 4G), and MemoryMax is 2G more." >&2
  echo "--threads N sets LEAN_NUM_THREADS to a positive count (default 1, or the --cpus count)." >&2
  echo "--cpus N sets the scope's CPUQuota to N CPUs, or to all online CPUs for auto (default: the thread count)." >&2
  exit 2
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --timeout)
      [ "$#" -ge 2 ] || usage
      run_timeout="$2"
      shift 2
      ;;
    --lock-timeout)
      [ "$#" -ge 2 ] || usage
      lock_timeout="$2"
      shift 2
      ;;
    --memory-high)
      [ "$#" -ge 2 ] || usage
      memory_high="$2"
      shift 2
      ;;
    --cpus)
      [ "$#" -ge 2 ] && [ -n "$2" ] || usage
      cpu_count="$2"
      shift 2
      ;;
    --threads)
      [ "$#" -ge 2 ] && [ -n "$2" ] || usage
      thread_count="$2"
      shift 2
      ;;
    --log)
      [ "$#" -ge 2 ] && [ -n "$2" ] || usage
      log_file="$2"
      shift 2
      ;;
    --toolchain)
      [ "$#" -ge 2 ] || usage
      toolchain="$2"
      shift 2
      ;;
    --)
      shift
      break
      ;;
    -*)
      usage
      ;;
    *)
      break
      ;;
  esac
done

if [ "$#" -eq 0 ]; then
  usage
fi

if [ -n "$log_file" ]; then
  mkdir -p "$(dirname -- "$log_file")"
  printf 'leanrun: log %s\n' "$log_file" >&2
  exec >"$log_file" 2>&1
fi

case "$lock_timeout" in
  ''|*[!0-9]*)
    echo "leanrun: --lock-timeout must be a nonnegative integer" >&2
    exit 2
    ;;
esac

memory_high_gib="${memory_high%G}"
case "$memory_high_gib" in
  ''|0*|*[!0-9]*)
    memory_high_gib=""
    ;;
esac
if [ -z "$memory_high_gib" ] || [ "$memory_high" != "${memory_high_gib}G" ]; then
  echo "leanrun: --memory-high must be whole gibibytes with the suffix G, such as 10G" >&2
  exit 2
fi
memory_max_gib=$((memory_high_gib + 2))

if [ "$cpu_count" = "auto" ]; then
  if ! cpu_count="$(getconf _NPROCESSORS_ONLN)"; then
    echo "leanrun: cannot determine online CPU count with getconf" >&2
    exit 2
  fi
fi
thread_count="${thread_count:-${cpu_count:-1}}"
cpu_count="${cpu_count:-$thread_count}"
case "$cpu_count" in
  ''|0*|*[!0-9]*)
    echo "leanrun: --cpus/LEANRUN_CPUS must be a positive integer or auto" >&2
    exit 2
    ;;
esac
cpu_quota_percent="${cpu_count}00"
cpu_quota_usec="${cpu_count}00000"

case "$thread_count" in
  ''|0*|*[!0-9]*)
    echo "leanrun: --threads/LEANRUN_THREADS must be a positive integer" >&2
    exit 2
    ;;
esac

if [ -z "$run_timeout" ]; then
  echo "leanrun: --timeout must not be empty" >&2
  exit 2
fi

case "$local_mode" in
  0|1)
    ;;
  *)
    echo "leanrun: LEANRUN_LOCAL must be 0 or 1" >&2
    exit 2
    ;;
esac

if [ "${LEANRUN_LOCAL_IN_SCOPE:-0}" = "1" ]; then
  echo "leanrun: nested tools/leanrun is not supported in LEANRUN_LOCAL mode (LEANRUN_LOCAL=1); invoke the repository driver directly" >&2
  exit 1
fi

if [ ! -x "$toolchain/bin/lean" ]; then
  echo "leanrun: no Lean toolchain at $toolchain" >&2
  exit 2
fi

PATH="$toolchain/bin:$PATH"
LEAN_NUM_THREADS="$thread_count"
LEANRUN_CPUS="$cpu_count"
LEANRUN_THREADS="$thread_count"
NO_COLOR=1
XDG_CACHE_HOME="$cache_dir"
mkdir -p "$XDG_CACHE_HOME"
export PATH LEAN_NUM_THREADS LEANRUN_CPUS LEANRUN_THREADS NO_COLOR XDG_CACHE_HOME

if [ "${LEANRUN_IN_SCOPE:-}" = "1" ]; then
  cgroup_path="$(sed -n 's/^0:://p' /proc/self/cgroup)"
  cgroup_dir="/sys/fs/cgroup$cgroup_path"
  if [ ! -d "$cgroup_dir" ] ||
     [ "$(cat "$cgroup_dir/memory.high")" != "$((memory_high_gib * 1073741824))" ] ||
     [ "$(cat "$cgroup_dir/memory.max")" != "$((memory_max_gib * 1073741824))" ] ||
     [ "$(cat "$cgroup_dir/memory.swap.max")" != "1073741824" ] ||
     [ "$(cat "$cgroup_dir/cpu.max")" != "$cpu_quota_usec 100000" ]; then
    echo "leanrun: inherited scope does not have the required resource limits" >&2
    exit 1
  fi
  exec ionice -c 3 timeout "$run_timeout" "$@"
fi

if [ "$(uname -s)" = "Darwin" ]; then
  if [ "$local_mode" != "1" ]; then
    echo "leanrun: macOS requires explicit LEANRUN_LOCAL=1; systemd cgroups are unavailable" >&2
    exit 2
  fi
  helper_hash="$(shasum -a 256 "$script_dir/leanrun-macos.c")"
  helper="$cache_dir/leanrun-macos-${helper_hash%% *}"
  if [ ! -x "$helper" ]; then
    staging="$(mktemp -d "$cache_dir/leanrun-macos-build.XXXXXX")"
    cc -std=c11 -Wall -Wextra -Werror -O2 "$script_dir/leanrun-macos.c" -o "$staging/runner"
    # Publish without overwriting an existing helper; retain staging as local state.
    if ! ln "$staging/runner" "$helper"; then
      [ -x "$helper" ] || exit 2
    fi
  fi
  mkdir -p "$lock_dir"
  echo "leanrun: LEANRUN_LOCAL=1; macOS has no cgroup CPU/memory/swap limits or ionice" >&2
  exec "$helper" "$lock_dir/1" "$lock_timeout" "$run_timeout" "$@"
fi

mkdir -p "$lock_dir"
exec 9>"$lock_dir/1"
if ! flock -w "$lock_timeout" 9; then
  echo "leanrun: timed out waiting for the machine-wide Lean slot" >&2
  exit 75
fi

if [ "$local_mode" = "1" ]; then
  echo "leanrun: LEANRUN_LOCAL=1; cgroup CPU, memory, and swap limits are not enforced" >&2
  exec env LEANRUN_LOCAL_IN_SCOPE=1 nice -n 10 ionice -c 3 \
    timeout "$run_timeout" "$@"
fi

exec systemd-run --user --scope --quiet --collect \
  -p MemoryHigh="${memory_high_gib}G" \
  -p MemoryMax="${memory_max_gib}G" \
  -p MemorySwapMax=1G \
  -p CPUQuota="${cpu_quota_percent}%" \
  -p CPUQuotaPeriodSec=100ms \
  env LEANRUN_IN_SCOPE=1 LEANRUN_MEMORY_HIGH="$memory_high" LEANRUN_CPUS="$cpu_count" LEANRUN_THREADS="$thread_count" nice -n 10 ionice -c 3 \
  timeout "$run_timeout" "$@"