#!/usr/bin/env bash
# spin_ab_urus.sh — A/B the RFC 004 spinning policy on the *realistic* urus
# workload (not the multi_scheduler microbench).
#
# It does exactly what the handoff asks: check out the pre-spin smarm commit,
# run the urus benches; check out the latest smarm commit, run them again;
# then print a before/after comparison.
#
#   BEFORE = b0c9685   (pre-spin: no spinning code exists at all)
#   AFTER  = current smarm HEAD   (spin defaults on + N=1 fix; e.g. 2f011f6)
#
# smarm is a *path dependency* of urus, which is a path dependency of
# urus-benches (all three are siblings). So we A/B by checking out smarm at
# each sha and rebuilding urus-benches; urus itself is held fixed — only the
# smarm under it changes between A and B. Each (commit, scenario, cores) run
# goes through the existing runner.sh (saturation probe -> warmup -> wrk2
# latency-honest measure) and lands a result.json; analyze_spin_ab_urus.py
# pairs them up at the end.
#
# Scenarios (per RECON §6.4): s2 (router+auth, no store) and s3 (memory store,
# the headline). Server is urus-mem for both. We do NOT touch s4/sqlite.
#
# Core sweep: the spin gate is active at 2-4 schedulers and inert at 1 and
# >=5. We sweep 1/2/4/6 server cores: 2 and 4 are gate-ACTIVE, 1 and 6 are
# inert controls -- 6 also being a *saturating* control, since the old 8-core
# point couldn't be saturated by an 8-core loadgen (it plateaued ~102k rps,
# loadgen-bound). Server is pinned to cores 0..N-1; the load generator gets
# its own disjoint cores (LOADGEN_CPUS), sized ~>=2x the server cores per the
# saturation goal while leaving cores free for the OS.
#
# Usage:
#   ./spin_ab_urus.sh                      # full A/B, defaults
#   ./spin_ab_urus.sh --quick              # short timings (15s phases)
#   ./spin_ab_urus.sh --dry-run            # print the plan, run nothing
#   CORES="2 4" SCENARIOS="s3" ./spin_ab_urus.sh
#   BEFORE=b0c9685 AFTER=2f011f6 ./spin_ab_urus.sh
#   ./spin_ab_urus.sh --with-axum-control  # also run axum-mem (smarm-free
#                                          # invariant; should be flat A vs B)
#
# Env knobs (all overridable):
#   BEFORE / AFTER     smarm commit-ishes              (default b0c9685 / HEAD)
#   SCENARIOS          space-separated                 (default "s2 s3")
#   CORES              server core counts to sweep     (default "1 2 4 6")
#   LOADGEN_CPUS       loadgen taskset range           (default "8-23" = 16 cores)
#   WRK_THREADS        loadgen wrk/wrk2 threads         (default 16)
#   WRK_CONNS          loadgen open connections         (default 512)
#   SERVER             memory server backend           (default urus-mem)
#   PROBE_SEC/WARMUP_SEC/MEASURE_SEC/SAT_RATIO/PREPOP/WRK_CONNS/WRK_THREADS
#                      passed straight through to runner.sh
#   SMARM_DIR          override smarm location (default: sibling ../smarm)
#   OUTDIR             results dir (default: results/spin_ab_urus_<ts>)

set -euo pipefail

# ---------------------------------------------------------------------------
# Locate the pieces
# ---------------------------------------------------------------------------

BENCH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"   # urus-benches
RUNNER="${BENCH_DIR}/runner.sh"
ANALYZER="${BENCH_DIR}/analyze_spin_ab_urus.py"
SMARM_DIR="${SMARM_DIR:-$(cd "${BENCH_DIR}/../smarm" 2>/dev/null && pwd || true)}"

if [[ ! -x "$RUNNER" ]]; then
  echo "ERROR: runner.sh not found/executable at $RUNNER" >&2; exit 1
fi
if [[ -z "$SMARM_DIR" || ! -d "$SMARM_DIR/.git" ]]; then
  echo "ERROR: smarm git repo not found (looked for sibling ../smarm). Set SMARM_DIR." >&2
  exit 1
fi

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

BEFORE="${BEFORE:-b0c9685}"
# AFTER defaults to whatever smarm is currently on ("the latest").
AFTER_DEFAULT="$(git -C "$SMARM_DIR" rev-parse --short HEAD)"
AFTER="${AFTER:-$AFTER_DEFAULT}"

SCENARIOS="${SCENARIOS:-s2 s3}"
CORES="${CORES:-1 2 4 6}"
SERVER="${SERVER:-urus-mem}"
LOADGEN_CPUS="${LOADGEN_CPUS:-8-23}"

# Loadgen sizing: ~2x loadgen cores per server core. 16 cores covers the
# 6-core server at ~2.7x and the 4-core at 4x. Threads track the core count;
# connections bumped so the probe can actually find the server's peak.
export WRK_THREADS="${WRK_THREADS:-16}"
export WRK_CONNS="${WRK_CONNS:-512}"

QUICK=0
DRY=0
WITH_AXUM=0
for arg in "$@"; do
  case "$arg" in
    --quick)             QUICK=1 ;;
    --dry-run|-n)        DRY=1 ;;
    --with-axum-control) WITH_AXUM=1 ;;
    -h|--help)           sed -n '2,/^set -euo/p' "$0" | sed 's/^# \{0,1\}//; $d'; exit 0 ;;
    *) echo "unknown arg: $arg" >&2; exit 2 ;;
  esac
done

if [[ "$QUICK" == "1" ]]; then
  export PROBE_SEC="${PROBE_SEC:-15}"
  export WARMUP_SEC="${WARMUP_SEC:-15}"
  export MEASURE_SEC="${MEASURE_SEC:-15}"
fi

TS="$(date +%Y%m%d-%H%M%S)"
OUTDIR="${OUTDIR:-${BENCH_DIR}/results/spin_ab_urus_${TS}}"

NPROC="$(nproc)"
# Loadgen low core, for an overlap sanity check against the server range.
LG_LO="${LOADGEN_CPUS%%-*}"

SERVERS_TO_RUN=("$SERVER")
[[ "$WITH_AXUM" == "1" ]] && SERVERS_TO_RUN+=("axum-mem")

echo ">>> smarm:     $SMARM_DIR"
echo ">>> BEFORE=$BEFORE  AFTER=$AFTER"
echo ">>> servers:   ${SERVERS_TO_RUN[*]}"
echo ">>> scenarios: $SCENARIOS"
echo ">>> cores:     $CORES"
echo ">>> loadgen:   cpus=$LOADGEN_CPUS threads=$WRK_THREADS conns=$WRK_CONNS"
echo ">>> machine:   $NPROC cores"
echo ">>> outdir:    $OUTDIR"
[[ "$QUICK" == "1" ]] && echo ">>> --quick: probe/warmup/measure = ${PROBE_SEC}/${WARMUP_SEC}/${MEASURE_SEC}s"
echo

# server core range for a given count: "0" for 1, else "0-(N-1)"
server_cpus() { (( $1 == 1 )) && echo "0" || echo "0-$(( $1 - 1 ))"; }

# ---------------------------------------------------------------------------
# Safety: refuse a dirty smarm tree; restore original ref on exit
# ---------------------------------------------------------------------------

if ! git -C "$SMARM_DIR" diff-index --quiet HEAD -- 2>/dev/null; then
  echo "ERROR: smarm working tree is dirty. Commit/stash before A/B so checkouts are safe." >&2
  git -C "$SMARM_DIR" status --short >&2
  exit 1
fi
ORIG_REF="$(git -C "$SMARM_DIR" symbolic-ref --quiet --short HEAD || git -C "$SMARM_DIR" rev-parse HEAD)"
restore_smarm() {
  echo
  echo ">>> restoring smarm to $ORIG_REF"
  git -C "$SMARM_DIR" checkout --quiet "$ORIG_REF" 2>/dev/null \
    || echo "WARN: could not restore $ORIG_REF — smarm is on $(git -C "$SMARM_DIR" rev-parse --short HEAD)" >&2
}
[[ "$DRY" == "0" ]] && trap restore_smarm EXIT

# Overlap sanity: server top core must be below the loadgen range.
MAX_SRV=0; for c in $CORES; do (( c > MAX_SRV )) && MAX_SRV=$c; done
if [[ "$LG_LO" =~ ^[0-9]+$ ]] && (( MAX_SRV - 1 >= LG_LO )); then
  echo "WARN: server cores reach $((MAX_SRV-1)) which overlaps LOADGEN_CPUS=$LOADGEN_CPUS." >&2
  echo "      set LOADGEN_CPUS to a disjoint range (two-process discipline)." >&2
fi
if (( MAX_SRV > NPROC )); then
  echo "WARN: sweep asks for up to ${MAX_SRV} server cores but machine has ${NPROC}." >&2
fi

mkdir -p "$OUTDIR"

# ---------------------------------------------------------------------------
# Run one commit across the (server x scenario x cores) sub-matrix
# ---------------------------------------------------------------------------

run_commit() {
  local sha="$1" label="$2"
  echo "============================================================"
  echo ">>> $label: checking out smarm $sha and rebuilding urus-benches"
  echo "============================================================"
  if [[ "$DRY" == "1" ]]; then
    echo "    [dry] git -C $SMARM_DIR checkout $sha"
    echo "    [dry] (cd $BENCH_DIR && cargo build --release)"
  else
    git -C "$SMARM_DIR" checkout --quiet "$sha"
    # urus is a path dep of urus-benches and smarm a path dep of urus, so a
    # workspace release build picks up the swapped smarm. cargo refingerprints
    # on the changed sources.
    ( cd "$BENCH_DIR" && cargo build --release ) 2>&1 | tail -3
  fi

  for srv in "${SERVERS_TO_RUN[@]}"; do
    for scn in $SCENARIOS; do
      for c in $CORES; do
        local scpus; scpus="$(server_cpus "$c")"
        local run_id="${label}__${srv}__${scn}__c${c}"
        local dest="${OUTDIR}/${label}/${srv}-${scn}-c${c}"
        echo ">>> $label  $srv  $scn  server-cores=$c (taskset $scpus)  loadgen=$LOADGEN_CPUS"
        if [[ "$DRY" == "1" ]]; then
          echo "    [dry] SERVER_CPUS=$scpus LOADGEN_CPUS=$LOADGEN_CPUS $RUNNER $srv $scn $run_id"
          continue
        fi
        mkdir -p "$dest"
        if SERVER_CPUS="$scpus" LOADGEN_CPUS="$LOADGEN_CPUS" \
             "$RUNNER" "$srv" "$scn" "$run_id" > "${dest}/runner.log" 2>&1; then
          if [[ -d "${BENCH_DIR}/results/${run_id}" ]]; then
            mv "${BENCH_DIR}/results/${run_id}/"* "$dest/" 2>/dev/null || true
            rmdir "${BENCH_DIR}/results/${run_id}" 2>/dev/null || true
          fi
          local rps p99 non2xx ms errpct
          rps=$(awk -F'[:,]' '/"sustained_rps"/ {gsub(/[ "]/,"",$2); print $2; exit}' "${dest}/result.json" 2>/dev/null || echo "?")
          p99=$(awk -F'"' '/"p99"/ {print $4; exit}' "${dest}/result.json" 2>/dev/null || echo "?")
          non2xx=$(awk -F'[:,]' '/"non_2xx"/ {gsub(/[ "]/,"",$2); print $2; exit}' "${dest}/result.json" 2>/dev/null || echo 0)
          ms="${MEASURE_SEC:-60}"
          # error rate vs approx total requests (sustained_rps * measure_sec).
          errpct=$(awk -v n="$non2xx" -v r="$rps" -v s="$ms" \
                     'BEGIN { t=r*s; if (t>0) printf "%.1f", 100*n/t; else print "?" }')
          if awk -v e="$errpct" 'BEGIN { exit !(e+0 > 2.0) }' 2>/dev/null; then
            echo "      -> VOID: ${errpct}% non-2xx (${non2xx}) — server is erroring, not serving. rps=${rps} is meaningless."
          else
            echo "      -> sustained=${rps} rps  p99=${p99}  (non-2xx ${errpct}%)"
          fi
        else
          echo "      -> FAILED (see ${dest}/runner.log)"
        fi
      done
    done
  done
}

run_commit "$BEFORE" "before"
run_commit "$AFTER"  "after"

[[ "$DRY" == "1" ]] && { echo; echo ">>> dry run complete (nothing executed, smarm untouched)."; exit 0; }

# ---------------------------------------------------------------------------
# Analyze
# ---------------------------------------------------------------------------

echo
echo "============================================================"
echo ">>> A/B comparison"
echo "============================================================"
if command -v python3 >/dev/null && [[ -f "$ANALYZER" ]]; then
  python3 "$ANALYZER" "$OUTDIR" --before before --after after | tee "${OUTDIR}/comparison.txt"
else
  echo "(python3 or analyzer missing; raw result.json files are under $OUTDIR)"
fi

echo
echo ">>> DONE. results: $OUTDIR"