#!/usr/bin/env bash
# Buffer header state contention: base vs. half-state CAS patches.
#
# Every backend hammers the same one or two buffers.  Pin (CAS, low half) and
# content-lock acquire (CAS, high half) race with unpin / unlock (atomic sub)
# of the other backends; on base any change to either half fails the CAS.
#
#   index   : nested-loop index probes on a one-page table, per probe the
#             index leaf is pinned, share-locked, unlocked, unpinned and the
#             heap page is share-locked and unlocked.
#   seqscan : nested-loop rescans of a one-page heap, per rescan the page is
#             pinned, share-locked, unlocked, unpinned.
#   point   : pgbench -S style point select (realistic, protocol-bound).
#
# Builds land in $WORK/{base,patched}; set REBUILD=1 to rebuild.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PG_SRC="${PG_SRC:-$(git -C "$HERE" rev-parse --show-toplevel)}"
BASE_REF="${BASE_REF:-73e69062c12}"
WORK="${WORK:-$PG_SRC/builds/half-state}"
NCPU="$(sysctl -n hw.ncpu 2>/dev/null || nproc)"
CLIENTS="${CLIENTS:-1 $NCPU}"
WORKLOADS="${WORKLOADS:-index seqscan point}"
DURATION="${DURATION:-15}"
ROUNDS="${ROUNDS:-3}"
PROBES="${PROBES:-1000}"
PORT="${PORT:-55440}"
REBUILD="${REBUILD:-0}"
DATA="$WORK/data"
RESULTS="$WORK/results-$(date +%Y%m%d-%H%M%S).csv"

build() { # name [patch...]
  local name=$1 src="$WORK/src-$1" prefix="$WORK/$1"
  shift
  [[ $REBUILD != 1 && -x $prefix/bin/postgres ]] && return
  if [[ ! -d $src ]]; then
    git -C "$PG_SRC" worktree add -f --detach "$src" "$BASE_REF"
    for p in "$@"; do git -C "$src" am -q "$p"; done
  fi
  [[ -d $src/build ]] || meson setup "$src/build" "$src" --buildtype=release \
    -Dcassert=false -Dauto_features=disabled --prefix="$prefix"
  ninja -C "$src/build" install >/dev/null
}

start() { # name
  "$WORK/$1/bin/pg_ctl" -D "$DATA" -l "$WORK/$1.log" -w -o \
    "-p $PORT -k /tmp -c listen_addresses= -c shared_buffers=128MB -c fsync=off \
     -c autovacuum=off -c jit=off -c max_parallel_workers_per_gather=0" start >/dev/null
}

stop() { # name
  "$WORK/$1/bin/pg_ctl" -D "$DATA" -m fast stop >/dev/null 2>&1 || true
}

trap 'stop base; stop patched' EXIT

mkdir -p "$WORK"
build base
build patched "$HERE"/v1-0001-*.patch "$HERE"/v1-0002-*.patch

# Force nested loops with a parameterised inner scan; no memoize/material caching.
export PGOPTIONS="-c enable_seqscan=off -c enable_bitmapscan=off -c enable_memoize=off -c enable_material=off"

cat >"$WORK/index.sql" <<EOF
SELECT sum(s.v) FROM generate_series(1, $PROBES) g(i),
  LATERAL (SELECT v FROM hot WHERE id = g.i % 10 + 1 OFFSET 0) s;
EOF
cat >"$WORK/seqscan.sql" <<EOF
SELECT sum(s.v) FROM generate_series(1, $PROBES) g(i),
  LATERAL (SELECT v FROM hot WHERE v = g.i % 10 + 1 OFFSET 0) s;
EOF
cat >"$WORK/point.sql" <<EOF
\set id random(1, 10)
SELECT v FROM hot WHERE id = :id;
EOF

if [[ ! -d $DATA ]]; then
  "$WORK/base/bin/initdb" -D "$DATA" -A trust >/dev/null
  start base
  "$WORK/base/bin/createdb" -h /tmp -p "$PORT" bench
  "$WORK/base/bin/psql" -qX -h /tmp -p "$PORT" -d bench -v ON_ERROR_STOP=1 <<EOF
CREATE TABLE hot(id int PRIMARY KEY, v int);
INSERT INTO hot SELECT g, g FROM generate_series(1, 10) g;
VACUUM ANALYZE hot;
EOF
  for w in index seqscan; do
    echo "== plan: $w"
    "$WORK/base/bin/psql" -qXAt -h /tmp -p "$PORT" -d bench -c "EXPLAIN (COSTS OFF) $(cat "$WORK/$w.sql")"
  done
  stop base
fi

echo "round,build,workload,clients,tps" >"$RESULTS"
for ((r = 1; r <= ROUNDS; r++)); do
  # alternate order so drift/thermal throttling does not favour one build
  builds="base patched"
  ((r % 2 == 1)) && builds="patched base"
  for b in $builds; do
    start "$b"
    for w in $WORKLOADS; do
      for c in $CLIENTS; do
        tps=$("$WORK/$b/bin/pgbench" -n -M prepared -c "$c" -j "$c" -T "$DURATION" \
          -h /tmp -p "$PORT" -f "$WORK/$w.sql" bench 2>&1 |
          sed -n 's/^tps = \([0-9.]*\).*/\1/p')
        echo "$r,$b,$w,$c,$tps" | tee -a "$RESULTS"
      done
    done
    stop "$b"
  done
done

echo
echo "== median tps over $ROUNDS rounds ($RESULTS)"
printf "%-8s %7s %12s %12s %8s\n" workload clients base patched change
awk -F, 'NR > 1 { v[$3 "," $4 "," $2] = v[$3 "," $4 "," $2] " " $5 }
  function median(s,  a, n, i, j, t) {
    n = split(s, a, " ")
    for (i = 2; i <= n; i++)
      for (j = i; j > 1 && a[j - 1] + 0 > a[j] + 0; j--) 
      {
        t = a[j];
        a[j] = a[j - 1];
        a[j - 1] = t
      }
    return n % 2 ? a[(n + 1) / 2] : (a[n / 2] + a[n / 2 + 1]) / 2
  }
  END {
    for (k in v) if (k ~ /,base$/) {
      split(k, f, ","); kb = f[1] "," f[2]
      b = median(v[kb ",base"]); p = median(v[kb ",patched"])
      printf "%-8s %7s %12.1f %12.1f %+7.1f%%\n", f[1], f[2], b, p, 100 * (p - b) / b
    }
  }' "$RESULTS" | sort -k1,1 -k2,2n