#!/bin/bash
# Stress the InnoDB parallel-scan chunking logic: concurrent writers reshape the
# B-tree (page splits, merges, level changes) while a reader runs long parallel
# grouped queries.  Fails on a crash, a server error, or a wrong result.
#
#   ./pq_stress.sh [--rows=200000] [--secs=300]
#                  [--writers=4] [--errlog=/path/mysqld.err]
#
# --errlog defaults to whatever the server reports in @@log_error.  Give it
# explicitly only when the server logs to stderr (@@log_error empty) and you
# captured that stream to a file yourself.
#
# TWO CORRECTNESS CHECKS, because under churn there is no single fixed "right
# answer":
#
#   A. INVARIANT.  Base rows have id % 10 == 0 and are never written; writers
#      only ever touch ids with id % 10 != 0.  A query restricted to
#      id % 10 = 0 therefore has ONE correct answer for the whole run, taken
#      once while the table is quiet.  The filter is not an index range, so the
#      scan still crosses every page the writers are splitting and merging.
#
#   B. SNAPSHOT.  Inside one REPEATABLE READ transaction, run the same query
#      serially and then in parallel.  Both see the same read view, so they
#      must agree no matter what the writers do.  This one also covers the
#      churned rows, which check A deliberately excludes.
#
# NOTE: the parallel-query gate refuses several plan shapes, so a workload can
# silently run serially and "pass".  check_parallel() below verifies with
# ANALYZE FORMAT=JSON that every variant really is chunked, and aborts if not.

set -u
SOCKET=; ROWS=200000; SECS=300; WRITERS=4; ERRLOG=; DB=pqstress
for arg in "$@"; do case $arg in
  --socket=*)  SOCKET=${arg#*=} ;;
  --rows=*)    ROWS=${arg#*=} ;;
  --secs=*)    SECS=${arg#*=} ;;
  --writers=*) WRITERS=${arg#*=} ;;
  --errlog=*)  ERRLOG=${arg#*=} ;;
  *) echo "unknown option: $arg" >&2; exit 2 ;;
esac; done

WORK=$(mktemp -d); trap 'rm -rf "$WORK"' EXIT
FAILURES=$WORK/failures; : > "$FAILURES"
CLIENT=(./client/mariadb --batch --raw --skip-column-names)
sql()  { "${CLIENT[@]}" "$DB" -e "$1"; }
fail() { echo "[FAIL] $*" | tee -a "$FAILURES"; }

# Where to look for assertions and crashes.  Ask the server rather than guess:
# a wrong guess is worse than none, because an unreadable file would just skip
# the check and the run would still look clean.
if [ -z "$ERRLOG" ]; then
  ERRLOG=$("${CLIENT[@]}" -e "SELECT @@log_error" 2>/dev/null)
  case "$ERRLOG" in
    ""|stderr) ERRLOG= ;;                       # server logs to stderr
    /*) ;;                                      # absolute, use as is
    *) ERRLOG="$("${CLIENT[@]}" -e "SELECT @@datadir" 2>/dev/null)$ERRLOG" ;;
  esac
fi
if [ -n "$ERRLOG" ] && [ -r "$ERRLOG" ]; then
  echo "== watching error log: $ERRLOG =="
  ERRLOG_BASE=$(wc -l < "$ERRLOG")              # ignore anything already there
else
  echo "== WARNING: no readable error log (@@log_error='${ERRLOG:-}')." >&2
  echo "==          Assertions and crashes will NOT be detected." >&2
  echo "==          Pass --errlog=PATH, or start the server with --log-error." >&2
  ERRLOG=; ERRLOG_BASE=0
fi

# ---------------------------------------------------------------- setup ----
echo "== loading $ROWS base rows =="
"${CLIENT[@]}" -e "DROP DATABASE IF EXISTS $DB; CREATE DATABASE $DB;" || exit 1
sql "
CREATE TABLE t (
  id  BIGINT NOT NULL,
  grp INT    NOT NULL,
  val BIGINT NOT NULL,
  pad CHAR(200) NOT NULL,
  PRIMARY KEY (id),
  KEY k_val (val)
) ENGINE=InnoDB;
INSERT INTO t SELECT seq*10, seq % 97, seq*2,
                     CONCAT('base-', seq, REPEAT('x', 180))
              FROM seq_1_to_$ROWS;" || exit 1

# One variant per access path the partitioner has to handle.  GROUP BY must be
# a plain column, and must NOT be the order an index already supplies, or the
# gate declines the plan and the whole run degenerates to a serial test.
declare -A Q
Q[pk_full]="SELECT grp, COUNT(*), SUM(val), SUM(id), MIN(id), MAX(id) FROM t FORCE INDEX (PRIMARY)
            WHERE id % 10 = 0 GROUP BY grp ORDER BY grp"
Q[pk_range]="SELECT grp, COUNT(*), SUM(val), SUM(id), MIN(id), MAX(id) FROM t FORCE INDEX (PRIMARY)
            WHERE id BETWEEN $((ROWS*2)) AND $((ROWS*8)) AND id % 10 = 0
            GROUP BY grp ORDER BY grp"
Q[sec_range]="SELECT grp, COUNT(*), SUM(val), SUM(id), MIN(id), MAX(id) FROM t FORCE INDEX (k_val)
            WHERE val > 0 AND id % 10 = 0 GROUP BY grp ORDER BY grp"

check_parallel() {                  # abort unless the plan is really chunked
  local k=$1 json
  json=$(sql "SET SESSION parallel_worker_threads=8;
              ANALYZE FORMAT=JSON ${Q[$k]}" | tr ',' '\n')
  if ! grep -q '_parallel' <<< "$json"; then
    echo "ABORT: '$k' does not run in parallel - the gate refused it:" >&2
    sql "SET optimizer_trace=1; SET SESSION parallel_worker_threads=8; ${Q[$k]};
         SELECT JSON_EXTRACT(trace,'\$**.parallel_scan_declined_because')
         FROM information_schema.optimizer_trace" >&2
    exit 1
  fi
  echo "  $k: $(grep -oE '"(access_type|r_chunks|r_chunks_resplit)": "?[a-zA-Z_0-9]*"?' <<< "$json" | tr -d '"' | tr '\n' ' ')"
}

echo "== recording expected results, and checking each plan is chunked =="
for k in "${!Q[@]}"; do
  sql "SET SESSION parallel_worker_threads=0; ${Q[$k]}" > "$WORK/expect.$k" || exit 1
  [ -s "$WORK/expect.$k" ] || { echo "empty baseline for $k" >&2; exit 1; }
  check_parallel "$k"
done

# --------------------------------------------------------------- writers ----
STOP=$WORK/stop
writer_split() {            # fill the gaps between base rows -> page splits
  local seed=$1
  while [ ! -e "$STOP" ]; do
    "${CLIENT[@]}" "$DB" -e "INSERT IGNORE INTO t
      SELECT seq*10 + $seed, seq % 97, seq*2, CONCAT('fill', REPEAT('y',180))
      FROM seq_1_to_$ROWS WHERE seq % 7 = $((seed % 7))" 2>>"$WORK/writer.err"
  done
}
writer_merge() {            # remove them again -> page merges
  local seed=$1
  while [ ! -e "$STOP" ]; do
    "${CLIENT[@]}" "$DB" -e "DELETE FROM t WHERE id % 10 = $seed" \
      2>>"$WORK/writer.err"
  done
}
writer_levels() {           # bulk inflate / deflate -> tree levels come and go
  while [ ! -e "$STOP" ]; do
    "${CLIENT[@]}" "$DB" -e "INSERT IGNORE INTO t
      SELECT seq*10 + 9, seq % 97, seq*2, CONCAT('bulk', REPEAT('z',180))
      FROM seq_1_to_$ROWS" 2>>"$WORK/writer.err"
    "${CLIENT[@]}" "$DB" -e "DELETE FROM t WHERE id % 10 = 9" 2>>"$WORK/writer.err"
  done
}

echo "== starting $((WRITERS*2+1)) writers =="
for ((i=1; i<=WRITERS; i++)); do writer_split  $i & done
for ((i=1; i<=WRITERS; i++)); do writer_merge  $i & done
writer_levels &

# ---------------------------------------------------------------- reader ----
echo "== reader running for ${SECS}s =="
END=$((SECONDS + SECS)); ROUND=0; THREADS_SET=(2 4 8 16)
while [ $SECONDS -lt $END ]; do
  ROUND=$((ROUND+1))
  T=${THREADS_SET[$((ROUND % 4))]}          # vary the worker count per round
  for k in "${!Q[@]}"; do

    # -- check A: invariant rows, answer fixed for the whole run
    if ! sql "SET SESSION parallel_worker_threads=$T; ${Q[$k]}" \
             > "$WORK/got.$k" 2>"$WORK/err.$k"; then
      fail "round $ROUND $k (threads=$T): query error: $(tr '\n' ' ' < "$WORK/err.$k")"
    elif ! diff -q "$WORK/expect.$k" "$WORK/got.$k" >/dev/null; then
      fail "round $ROUND $k (threads=$T): WRONG RESULT under churn"
      diff "$WORK/expect.$k" "$WORK/got.$k" | head -20 >> "$FAILURES"
    fi

    # -- check B: serial vs parallel in one read view, churned rows included
    rm -f "$WORK/snap.1" "$WORK/snap.2"
    if ! sql "START TRANSACTION WITH CONSISTENT SNAPSHOT;
              SET SESSION parallel_worker_threads=0;
              SELECT 'SERIAL';   ${Q[$k]//id % 10 = 0/1=1};
              SET SESSION parallel_worker_threads=$T;
              SELECT 'PARALLEL'; ${Q[$k]//id % 10 = 0/1=1};
              COMMIT;" > "$WORK/snap.$k" 2>"$WORK/snaperr.$k"; then
      fail "round $ROUND $k: snapshot query error: $(tr '\n' ' ' < "$WORK/snaperr.$k")"
    else
      awk -v d="$WORK" '/^SERIAL$/{s=1;next} /^PARALLEL$/{s=2;next}
                        s{print > (d "/snap." s)}' "$WORK/snap.$k"
      if [ ! -s "$WORK/snap.1" ] || [ ! -s "$WORK/snap.2" ]; then
        fail "round $ROUND $k: snapshot comparison produced no rows"
      elif ! diff -q "$WORK/snap.1" "$WORK/snap.2" >/dev/null; then
        fail "round $ROUND $k (threads=$T): serial and parallel DISAGREE in one snapshot"
        diff "$WORK/snap.1" "$WORK/snap.2" | head -20 >> "$FAILURES"
      fi
    fi
  done

  sql "SELECT 1" >/dev/null 2>&1 || { fail "round $ROUND: server is gone (crash?)"; break; }
done

# ----------------------------------------------------------------- report ----
touch "$STOP"; wait 2>/dev/null
echo; echo "== $ROUND rounds, $(wc -l < "$WORK/writer.err" 2>/dev/null || echo 0) writer errors =="
sql "SHOW GLOBAL STATUS WHERE Variable_name IN
     ('Handler_write','Handler_delete','Handler_read_next','Handler_read_rnd_next')" |
  sed 's/^/  /'
sql "SELECT COUNT(*) AS rows_now, SUM(id%10=0) AS base_intact FROM t" | sed 's/^/  rows: /'

if [ -n "$ERRLOG" ]; then
  if tail -n +$((ERRLOG_BASE + 1)) "$ERRLOG" \
       | grep -nE "Assertion|signal 1[12]|InnoDB: Failing|corrupt|\[ERROR\]" \
       | grep -v "Aborted connection" > "$WORK/errlog.hits"; then
    fail "server error log has new entries:"; head -20 "$WORK/errlog.hits" >> "$FAILURES"
  fi
else
  echo "  (error log was not scanned - crashes and assertions unchecked)"
fi

if [ -s "$FAILURES" ]; then
  echo; echo "FAILURES:"; cat "$FAILURES"; cp "$FAILURES" ./pq_stress_failures.txt
  echo "(saved to ./pq_stress_failures.txt)"; exit 1
fi
echo "OK - no crashes, no wrong results"