#!/usr/bin/env bash
# Memory bench for the forge's SQLite connections.
#
#   scripts/mem-bench.sh LABEL [path/to/quarry]
#
# Builds (once) three Git-derived stores of 2k, 8k and 30k check-ins (plus
# 24 copies of the small and 3 of the medium, as the live forge holds dozens) with
# scripts/mem-bench-stores.sh, starts the given quarry binary on them with
# a fresh data directory, drives a fixed read workload (timeline pages and
# a deep page, branch filter, file history, activity, branch and tag lists,
# search, a check-in and its diff, a file view; ROUNDS rounds, P requests
# in parallel), and writes OUT/LABEL.txt: RSS at start+60s, peak RSS
# before and during load, RSS after IDLE seconds, open index.db
# descriptors, and p50/p95 latency and non-200 count per route class.
#
# To judge one knob, build a binary with only that knob changed and run
# this with its own LABEL; compare the files with scripts/mem-bench.sh
# --compare LABEL...
#
# The forge, the sampler and the watchdog all die on every exit path, and
# the run carries its own deadline (DEADLINE seconds, default 1500).
set -u

if [ "${1:-}" = "--compare" ]; then
  shift
  out=${OUT:-/tmp/quarry-mem-bench/out}
  for l in "$@"; do
    echo "== $l"; grep -E '^(start60|peak_start|peak_load|idle|fds_|errors)' "$out/$l.txt"
    sort -t$'\t' -k4,4nr "$out/$l.txt" | awk -F'\t' '$1=="route"{print "  slow", $2, "p50="$3, "p95="$4; if(++n==3)exit}'
  done
  exit 0
fi

label=${1:?usage: mem-bench.sh LABEL [quarry-binary]}
here=$(cd "$(dirname "$0")" && pwd)
bin=${2:-quarry}
work=${WORK:-/tmp/quarry-mem-bench}
STORES=${STORES:-$work/stores}; export STORES
out=${OUT:-$work/out}
port=${PORT:-8961}
rounds=${ROUNDS:-20}
par=${P:-6}
idle=${IDLE:-180}
deadline=${DEADLINE:-1500}
mkdir -p "$out" "$work"
"$here/mem-bench-stores.sh" || exit 1

run=$(mktemp -d "$work/run.XXXXXX")
pid=""; samp=""; dog=""
cleanup() {
  trap - EXIT INT TERM
  [ -n "$samp" ] && kill "$samp" 2>/dev/null
  [ -n "$dog" ] && kill "$dog" 2>/dev/null
  [ -n "$pid" ] && kill "$pid" 2>/dev/null && wait "$pid" 2>/dev/null
  rm -rf "$run/data"
}
trap cleanup EXIT INT TERM

B=http://127.0.0.1:$port
names() { ls "$STORES" | grep -E '^(small|medium|large)(-[0-9]+)?$'; }
repos() { for n in $(names); do printf -- '--repo %s=%s ' "$n" "$STORES/$n"; done; }
nrepos=$(names | wc -l | tr -d ' ')
mkdir -p "$run/data"
"$bin" serve --addr 127.0.0.1:$port --data "$run/data" --repo-names=flat \
  $(repos) \
  >"$run/serve.log" 2>&1 &
pid=$!
t0=$(date +%s)
# Watchdog: a deadline the run enforces on itself, even if this shell dies.
( sleep "$deadline"; kill "$pid" 2>/dev/null ) >/dev/null 2>&1 &
dog=$!

rss() { ps -o rss= -p "$pid" 2>/dev/null | tr -d ' '; }
fds() { lsof -p "$pid" 2>/dev/null | grep -c 'index\.db'; }
# Sampler: epoch, rss KB, every second; the end is its own.
( end=$((t0 + deadline)); while [ "$(date +%s)" -lt "$end" ]; do
    echo "$(date +%s) $(rss)"; sleep 1; done ) >"$run/samples" 2>/dev/null &
samp=$!

for _ in $(seq 1 300); do
  [ "$(grep -c '^.* ingested [0-9]* artifacts from' "$run/serve.log")" -ge "$nrepos" ] && break
  kill -0 "$pid" 2>/dev/null || { echo "forge died at start"; tail "$run/serve.log"; exit 1; }
  sleep 1
done
ready=$(date +%s)
echo "ingested after $((ready - t0))s"
while [ "$(date +%s)" -lt $((t0 + 60)) ]; do sleep 1; done
start60=$(rss); fds_start=$(fds)

# Workload ---------------------------------------------------------------
urls="$run/urls"; : >"$urls"
for s in $(names); do
  k=${s%%-*}
  tl=$(curl -s "$B/r/$s/timeline")
  H=$(printf %s "$tl" | grep -o '/checkin/[0-9a-f]*' | head -1 | cut -d/ -f3)
  fpath=$(curl -s "$B/r/$s/checkin/$H" | grep -o "/file/$H/[^\"]*" | head -1 | cut -d/ -f5- | sed "s/%/%25/g")
  before=$(printf %s "$tl" | grep -o 'timeline?before=[^"]*' | head -1 | sed 's/&/\&/g')
  {
    echo "$k:home /r/$s"
    echo "$k:timeline /r/$s/timeline"
    echo "$k:branches /r/$s/branch"
    echo "$k:tags /r/$s/tag"
    echo "$k:search /r/$s/search?q=refactor"
    echo "$k:checkin /r/$s/checkin/$H"
    if [ "$k" != small ]; then
      echo "$k:timeline200 /r/$s/timeline?n=200"
      echo "$k:timelinedeep /r/$s/$before"
      echo "$k:branchfilter /r/$s/timeline?branch=feature-003"
      for f in src%2Ffile0000.go lib%2Ffile0004.go docs%2Ffile0005.md cmd%2Ffile0007.md tests%2Ffile0008.go; do
        echo "$k:filehistory /r/$s/timeline?p=$f"
      done
      echo "$k:activity /r/$s/activity"
      for q in alpha+beta index+query cache sort+merge diff fix+add remove; do echo "$k:search /r/$s/search?q=$q"; done
      echo "$k:diff /r/$s/checkin/$H/diff"
      echo "$k:file /r/$s/file/$H/$fpath"
    fi
  } >>"$urls"
done
: >"$run/all"; for _ in $(seq 1 "$rounds"); do sort -R "$urls" >>"$run/all"; done

loadstart=$(date +%s)
export B
xargs -P "$par" -L1 bash -c 'r=$(curl -s -o /dev/null -w "%{http_code} %{time_total}" --max-time 120 "$B$1"); echo "$0 $r"' <"$run/all" >"$run/lat" 2>/dev/null
loadend=$(date +%s)
fds_load=$(fds)
rss_after_load=$(rss)

while [ "$(date +%s)" -lt $((loadend + idle)) ]; do sleep 1; done
idle_rss=$(rss); fds_idle=$(fds)

peak_start=$(awk -v a="$t0" -v b="$loadstart" '$1>=a&&$1<b&&$2>m{m=$2}END{print m+0}' "$run/samples")
peak_load=$(awk -v a="$loadstart" -v b="$loadend" '$1>=a&&$1<=b+1&&$2>m{m=$2}END{print m+0}' "$run/samples")
{
  echo "label	$label"
  echo "start60	${start60} KB"
  echo "peak_start	${peak_start} KB"
  echo "peak_load	${peak_load} KB"
  echo "after_load	${rss_after_load} KB"
  echo "idle	${idle_rss} KB after ${idle}s"
  echo "fds_start60	$fds_start"
  echo "fds_load	$fds_load"
  echo "fds_idle	$fds_idle"
  echo "load_seconds	$((loadend - loadstart))"
  echo "errors	$(awk '$2!="200"' "$run/lat" | wc -l | tr -d ' ') of $(wc -l <"$run/lat" | tr -d ' ')"
  awk '$2!="200"{print "nonok", $1, $2}' "$run/lat" | sort | uniq -c
  sort -k1,1 -k3,3n "$run/lat" | awk '{c[$1]=c[$1]" "$3} END{for(k in c){n=split(c[k],a," "); i50=int((n+1)/2); i95=int(n*0.95+0.999); if(i95>n)i95=n; printf "route\t%s\t%.3f\t%.3f\t%d\n", k, a[i50], a[i95], n}}' | sort -t$'\t' -k4,4nr
} >"$out/$label.txt"
grep -v '^nonok' "$out/$label.txt" | head -22
grep -E 'out of memory|panic|SQLITE' "$run/serve.log" | head -5