#!/usr/bin/env bash
# buildslot.sh — machine-global admission gate for heavy builds/tests.
# At most N jobs run concurrently; excess invocations BLOCK until a slot frees.
# Admission is strict FIFO via an on-disk ticket queue: only the queue head may
# probe slots, so a waiter can never lose a freed slot to a later arrival
# (non-FIFO polling starved long waiters indefinitely under sustained load).
# Dead ticket PIDs are filtered on every head computation, so a kill -9'd
# waiter never wedges the queue.
# N adapts to real-time interactive demand measured from /proc/stat: non-nice
# CPU time (user+system+irq+softirq+steal) is "demand"; nice time is NOT —
# queued builds run at nice 19, so a box saturated only by its own builds still
# reads as idle and admits up to BUILD_SLOTS. A box busy with interactive work
# shrinks to BUILD_SLOT_FLOOR. Holders are never preempted; adaptation only
# gates NEW admissions.
# Slot lock = flock on an fd held open across exec: the kernel releases it when
# the command exits or crashes — no PID files, no daemon.
#
# Tunables:
#   BUILD_SLOTS=4          max concurrent slots (adaptive ceiling)
#   BUILD_SLOT_FLOOR=1     slots still admitted under full interactive load
#   BUILD_SLOT_FIXED=1     disable adaptation; always exactly BUILD_SLOTS
#   BUILD_SLOT_TIMEOUT=    seconds to wait before giving up (default: forever)
#   BUILD_SLOT_DIR=        lock directory override (tests)
#   BUILD_SLOT_STAT=       /proc/stat override (tests)
set -euo pipefail

if (( $# == 0 )); then
  printf 'usage: %s COMMAND [ARG...]\n' "$0" >&2
  exit 64
fi

# Reentrancy: a nested invocation inside a slot-holding tree inherits the slot.
if [[ -n "${BUILD_SLOT_HELD:-}" ]]; then
  exec "$@"
fi

slots_max=${BUILD_SLOTS:-1}
if ! [[ "$slots_max" =~ ^[1-9][0-9]*$ ]]; then
  printf '[buildslot] BUILD_SLOTS must be a positive integer: %s\n' "$slots_max" >&2
  exit 64
fi

floor=${BUILD_SLOT_FLOOR:-1}
if ! [[ "$floor" =~ ^[1-9][0-9]*$ ]] || (( floor > slots_max )); then
  printf '[buildslot] BUILD_SLOT_FLOOR must be an integer in [1, BUILD_SLOTS]: %s\n' "$floor" >&2
  exit 64
fi

timeout=${BUILD_SLOT_TIMEOUT:-}
if [[ -n "$timeout" ]] && ! [[ "$timeout" =~ ^[0-9]+$ ]]; then
  printf '[buildslot] BUILD_SLOT_TIMEOUT must be a non-negative integer: %s\n' "$timeout" >&2
  exit 64
fi

lock_dir=${BUILD_SLOT_DIR:-/run/dev-build-slot}
if ! mkdir -p "$lock_dir" 2>/dev/null || [[ ! -d "$lock_dir" || ! -w "$lock_dir" || ! -x "$lock_dir" ]]; then
  printf '[buildslot] unable to use lock directory: %s\n' "$lock_dir" >&2
  exit 73
fi

# Reentrancy fallback: env-filtering runners (turbo passthrough, env -i) strip
# BUILD_SLOT_HELD from descendants of a slot holder. A slot-lock fd is inherited
# across fork/exec, so every process in a holder's tree appears in lsof on the
# slot file — if any of OUR ancestors holds a slot, this tree already owns one.
# Without this, holder -> env-scrubbed child -> child queues behind its own
# ancestor = machine-wide deadlock. lsof missing -> fall through to queuing.
ancestor_holds_slot() {
  command -v lsof >/dev/null 2>&1 || return 1
  local holders
  holders=$(lsof -t "$lock_dir"/slot-*.lock 2>/dev/null) || true
  [[ -n "$holders" ]] || return 1
  local pid=$$ stat rest ppid
  while [[ "$pid" =~ ^[0-9]+$ ]] && (( pid > 1 )); do
    if grep -qx -- "$pid" <<< "$holders"; then
      return 0
    fi
    [[ -r "/proc/$pid/stat" ]] || return 1
    stat=$(<"/proc/$pid/stat")
    rest=${stat##*) }
    read -r _ ppid _ <<< "$rest"
    pid=$ppid
  done
  return 1
}

if ancestor_holds_slot; then
  export BUILD_SLOT_HELD=1
  exec "$@"
fi

stat_file=${BUILD_SLOT_STAT:-/proc/stat}
adaptive=1
if [[ "${BUILD_SLOT_FIXED:-}" == 1 ]] || [[ ! -r "$stat_file" ]]; then
  adaptive=0
fi

# Prints "busy total" from the aggregate cpu line. busy excludes the nice field
# so CPU consumed by nice>0 tasks (the queued builds themselves) never counts
# as demand — otherwise a build-saturated idle box would choke its own queue.
read_stat() {
  local _c user nice system idle iowait irq softirq steal _rest
  read -r _c user nice system idle iowait irq softirq steal _rest < "$stat_file" || return 1
  printf '%s %s\n' "$((user + system + irq + softirq + steal))" \
    "$((user + nice + system + idle + iowait + irq + softirq + steal))"
}

# Map interactive-demand percent over the last sample window to a slot count.
effective_slots() {
  local prev_busy=$1 prev_total=$2 cur_busy=$3 cur_total=$4
  local db=$((cur_busy - prev_busy)) dt=$((cur_total - prev_total))
  if (( dt <= 0 )); then
    printf '%s\n' "$floor"
    return
  fi
  local pct=$((100 * db / dt))
  if (( pct < 30 )); then
    printf '%s\n' "$slots_max"
  elif (( pct < 60 )); then
    local half=$(((slots_max + 1) / 2))
    (( half < floor )) && half=$floor
    printf '%s\n' "$half"
  else
    printf '%s\n' "$floor"
  fi
}

queue_file="$lock_dir/queue"
queue_lock="$lock_dir/queue.lock"
exec {queue_fd}>"$queue_lock"

enqueue() {
  flock "$queue_fd"
  printf '%s\n' "$$" >>"$queue_file"
  flock -u "$queue_fd"
}

dequeue_self() {
  flock "$queue_fd"
  if [[ -f "$queue_file" ]]; then
    grep -Fxv -- "$$" "$queue_file" >"$queue_file.new" || true
    mv "$queue_file.new" "$queue_file"
  fi
  flock -u "$queue_fd"
}

# Head = first still-alive PID; dead tickets are pruned in place so a crashed
# waiter never blocks the line. Returns 0 when we are head.
at_head() {
  local head="" pid alive=()
  flock "$queue_fd"
  if [[ -f "$queue_file" ]]; then
    while read -r pid; do
      [[ "$pid" =~ ^[0-9]+$ ]] || continue
      if kill -0 "$pid" 2>/dev/null; then
        alive+=("$pid")
      fi
    done <"$queue_file"
    if (( ${#alive[@]} )); then
      printf '%s\n' "${alive[@]}" >"$queue_file.new"
      mv "$queue_file.new" "$queue_file"
      head=${alive[0]}
    else
      : >"$queue_file"
    fi
  fi
  flock -u "$queue_fd"
  [[ "$head" == "$$" ]]
}

trap 'dequeue_self' EXIT
enqueue

prev_busy="" prev_total=""
if (( adaptive )); then
  if sample="$(read_stat)"; then
    read -r prev_busy prev_total <<< "$sample"
    sleep 0.1
  else
    printf '[buildslot] cannot read %s; falling back to fixed %s slots\n' "$stat_file" "$slots_max" >&2
    adaptive=0
  fi
fi

started=$SECONDS
announced=0

while :; do
  if at_head; then
    eff=$slots_max
    if (( adaptive )); then
      if sample="$(read_stat)"; then
        read -r cur_busy cur_total <<< "$sample"
        eff="$(effective_slots "$prev_busy" "$prev_total" "$cur_busy" "$cur_total")"
        prev_busy=$cur_busy prev_total=$cur_total
      else
        printf '[buildslot] cannot read %s; falling back to fixed %s slots\n' "$stat_file" "$slots_max" >&2
        adaptive=0
      fi
    fi
    for (( index = 0; index < eff; index++ )); do
      lock_file="$lock_dir/slot-${index}.lock"
      if ! exec {lock_fd}>"$lock_file"; then
        printf '[buildslot] unable to use lock directory: %s\n' "$lock_dir" >&2
        exit 73
      fi
      if flock -n "$lock_fd"; then
        dequeue_self
        trap - EXIT
        export BUILD_SLOT_HELD=1
        if (( announced )); then
          printf '[buildslot] acquired build slot after %ss\n' "$((SECONDS - started))" >&2
        fi
        exec "$@"
      fi
      exec {lock_fd}>&-
    done
  fi

  elapsed=$((SECONDS - started))
  if [[ -n "$timeout" ]] && (( elapsed >= timeout )); then
    printf '[buildslot] timed out waiting for build slot after %ss\n' "$elapsed" >&2
    exit 75
  fi
  if (( elapsed > 2 && ! announced )); then
    printf '[buildslot] waiting for build slot\n' >&2
    announced=1
  fi
  sleep 0.1
done