#!/bin/bash
# Reproduces every factual claim in
# https://www.heatware.net/postgresql/fix-high-cpu-usage/
#
# Creates a throwaway PostgreSQL cluster in a temp directory on port 55731, loads
# pg_stat_statements into it, builds a 2,000,000-row table, runs a genuinely CPU-bound query,
# and asserts what pg_stat_activity, pg_stat_statements, EXPLAIN and the autovacuum log say
# about it. Then it stops the cluster and deletes everything.
# It does not touch any existing PostgreSQL installation, data directory, log or role.
#
# It deliberately saturates several CPUs for a few seconds during the connection test.
#
# Assertions are on RELATIONSHIPS -- a cached scan reads zero blocks; an indexed expression is
# faster than a regex over the same rows; idle connections burn less CPU than active queries --
# never on hard-coded millisecond values, which vary by machine.
#
# Requires: initdb, pg_ctl, psql on PATH, pg_stat_statements in the server's library directory
# (ships with PostgreSQL), and about 400 MB of temp space.
# Runtime: about 3 minutes.
# Usage: bash verify-fix-high-cpu-usage.sh

set -u

PORT=55731
ROOT="$(mktemp -d -t hwcpu)"
# the Unix socket path is capped at 103 bytes, so keep the socket dir short
SOCK="$(mktemp -d /tmp/hwcpu.XXXXXX)"
DATA="$ROOT/data"
PGLOG="$DATA/log/postgresql.log"

cleanup(){
  pg_ctl -D "$DATA" stop -m immediate >/dev/null 2>&1
  rm -rf "$ROOT" "$SOCK"
}
trap cleanup EXIT

PASSES=0; FAILS=0; SKIPS=0
pass(){ PASSES=$((PASSES+1)); printf 'PASS %s\n' "$1"; }
fail(){ FAILS=$((FAILS+1));   printf 'FAIL %s\n     expected: %s\n     actual:   %s\n' "$1" "$2" "$3"; }
skip(){ SKIPS=$((SKIPS+1));   printf 'SKIP %s (needs %s)\n' "$1" "$2"; }
check(){ [ "$2" = "$3" ] && pass "$1" || fail "$1" "$2" "$3"; }
contains(){ case "$3" in *"$2"*) pass "$1";; *) fail "$1" "text containing: $2" "$3";; esac; }
lt(){ if awk "BEGIN{exit !($2 < $3)}" 2>/dev/null; then pass "$1"; else fail "$1" "$2 < $3" "it is not"; fi; }

for bin in initdb pg_ctl psql; do
  command -v "$bin" >/dev/null 2>&1 || { echo "SKIP all checks (needs $bin on PATH)"; echo "PASS=0 FAIL=0 SKIP=1"; exit 0; }
done

initdb -D "$DATA" -U postgres --no-locale -E UTF8 >"$ROOT/initdb.log" 2>&1 \
  || { echo "SKIP all checks (initdb failed; see $ROOT/initdb.log)"; echo "PASS=0 FAIL=0 SKIP=1"; exit 0; }

cat >> "$DATA/postgresql.conf" <<'EOF'
shared_preload_libraries = 'pg_stat_statements'
pg_stat_statements.track = 'all'
shared_buffers = '1GB'
track_io_timing = on
logging_collector = on
log_directory = 'log'
log_filename = 'postgresql.log'
log_autovacuum_min_duration = 0
log_line_prefix = '%m [%p] %u@%d '
max_connections = 200
autovacuum_naptime = '5s'
autovacuum_vacuum_scale_factor = 0.02
EOF

pg_ctl -D "$DATA" -o "-p $PORT -k $SOCK -c listen_addresses=127.0.0.1" \
       -l "$ROOT/startup.log" start >/dev/null 2>&1
for _ in $(seq 20); do psql -h 127.0.0.1 -p $PORT -U postgres -X -At -c 'SELECT 1' >/dev/null 2>&1 && break; sleep 1; done
q(){ psql -h 127.0.0.1 -p $PORT -U postgres -d postgres -X -At -c "$1" 2>&1; }

q 'SELECT 1' >/dev/null 2>&1 || {
  echo "SKIP all checks (server did not start; see $ROOT/startup.log)"; echo "PASS=0 FAIL=0 SKIP=1"; exit 0; }

echo "server: $(q 'SELECT version()')"
echo "cpus:   $( (sysctl -n hw.logicalcpu 2>/dev/null || nproc 2>/dev/null || echo '?') )"
echo

if [ -z "$(q "CREATE EXTENSION IF NOT EXISTS pg_stat_statements; SELECT extversion FROM pg_extension WHERE extname='pg_stat_statements'" | tail -1)" ]; then
  echo "SKIP all checks (pg_stat_statements not available in this build)"; echo "PASS=0 FAIL=0 SKIP=1"; exit 0
fi
pass "pg_stat_statements loads (version $(q "SELECT extversion FROM pg_extension WHERE extname='pg_stat_statements'" | tail -1))"

q "CREATE TABLE users (id bigserial primary key, email text not null,
       created_at timestamptz not null default now(), payload text);
   INSERT INTO users (email, payload)
   SELECT 'user'||g||'@'||(ARRAY['example.com','gmail.com','yahoo.com','proton.me'])[1+(g%4)], md5(g::text)
   FROM generate_series(1,2000000) g;" >/dev/null
q "VACUUM ANALYZE users" >/dev/null
check "the table holds 2,000,000 rows" "2000000" "$(q 'SELECT count(*) FROM users')"
q "SELECT count(*) FROM users" >/dev/null; q "SELECT count(*) FROM users" >/dev/null   # warm the cache

# ---- 1. the query is CPU-bound: every block is a shared_buffers hit, none read from disk
PLAN="$(q "SET max_parallel_workers_per_gather=0;
           EXPLAIN (ANALYZE, BUFFERS, COSTS OFF) SELECT count(*) FROM users WHERE email ~ '@example\.com\$'")"
contains "the regex filter forces a Seq Scan" "Seq Scan on users" "$PLAN"
contains "it discards 1,500,000 rows in the filter" "Rows Removed by Filter: 1500000" "$PLAN"
case "$PLAN" in
  *"Buffers: shared hit="*read=*) fail "the whole table is served from shared_buffers (read=0)" "no 'read=' in the Buffers line" "$(echo "$PLAN" | grep -m1 Buffers)";;
  *"Buffers: shared hit="*)       pass "the whole table is served from shared_buffers (read=0)";;
  *)                              fail "the whole table is served from shared_buffers (read=0)" "a Buffers line" "$PLAN";;
esac

# ---- 2. an active CPU-bound backend has NO wait event, and is a live OS process
q "SELECT pg_stat_statements_reset()" >/dev/null
( for _ in 1 2 3; do q "SET max_parallel_workers_per_gather=0; SELECT count(*) FROM users WHERE email ~ '@example\.com\$'" >/dev/null; done ) &
BG=$!
sleep 2
ROW="$(q "SELECT pid||'|'||state||'|'||coalesce(wait_event_type,'NONE')
          FROM pg_stat_activity
          WHERE state='active' AND backend_type='client backend' AND pid<>pg_backend_pid() LIMIT 1")"
case "$ROW" in
  *"|active|NONE") pass "the running backend is active with an empty wait_event_type (on CPU)";;
  "")              fail "the running backend is active with an empty wait_event_type (on CPU)" "one active client backend" "none found";;
  *)               fail "the running backend is active with an empty wait_event_type (on CPU)" "<pid>|active|NONE" "$ROW";;
esac
BPID="${ROW%%|*}"
if [ -n "$BPID" ] && kill -0 "$BPID" 2>/dev/null; then
  pass "pg_stat_activity.pid ($BPID) is a live operating-system process"
else
  fail "pg_stat_activity.pid is a live operating-system process" "a pid that ps/kill -0 can see" "${BPID:-<none>}"
fi
wait $BG

# ---- 3. pg_stat_statements attributes the time, with zero disk reads
STAT="$(q "SELECT calls||'|'||shared_blks_read||'|'||round(mean_exec_time::numeric,1)
           FROM pg_stat_statements WHERE query LIKE '%users WHERE email ~%' ORDER BY total_exec_time DESC LIMIT 1")"
check "pg_stat_statements recorded 3 calls with 0 blocks read from disk" \
      "3|0" "$(echo "$STAT" | cut -d'|' -f1,2)"
REGEX_MEAN="$(echo "$STAT" | cut -d'|' -f3)"

# ---- 4. the fix: a functional index on what is actually filtered
q "CREATE INDEX users_domain_idx ON users ((split_part(email,'@',2))); ANALYZE users" >/dev/null
PLAN2="$(q "EXPLAIN (ANALYZE, BUFFERS, COSTS OFF) SELECT count(*) FROM users WHERE split_part(email,'@',2) = 'example.com'")"
contains "the functional index is used" "Bitmap Index Scan on users_domain_idx" "$PLAN2"
q "SELECT pg_stat_statements_reset()" >/dev/null
for _ in 1 2 3; do q "SET max_parallel_workers_per_gather=0; SELECT count(*) FROM users WHERE email ~ '@example\.com\$'" >/dev/null; done
for _ in 1 2 3; do q "SET max_parallel_workers_per_gather=0; SELECT count(*) FROM users WHERE split_part(email,'@',2) = 'example.com'" >/dev/null; done
A="$(q "SELECT round(mean_exec_time::numeric,1) FROM pg_stat_statements WHERE query LIKE '%email ~%' LIMIT 1")"
B="$(q "SELECT round(mean_exec_time::numeric,1) FROM pg_stat_statements WHERE query LIKE '%split_part%' LIMIT 1")"
echo "    (measured on this machine: regex ${A} ms mean, indexed expression ${B} ms mean)"
lt "the indexed expression is faster than the regex scan" "${B:-0}" "${A:-0}"
check "both forms return the same 500,000 rows" "500000|500000" \
      "$(q "SELECT (SELECT count(*) FROM users WHERE email ~ '@example\.com\$')||'|'||(SELECT count(*) FROM users WHERE split_part(email,'@',2)='example.com')")"

# ---- 5. autovacuum reports its own CPU cost
q "UPDATE users SET payload = md5(payload) WHERE id % 3 = 0" >/dev/null
FOUND=""
for _ in $(seq 120); do
  grep -q 'automatic vacuum of table "postgres.public.users"' "$PGLOG" 2>/dev/null && { FOUND=1; break; }
  sleep 1
done
if [ -n "$FOUND" ]; then
  ENTRY="$(awk '/automatic vacuum of table "postgres.public.users"/{f=1} f{print; if (/system usage/) exit}' "$PGLOG")"
  contains "the autovacuum log entry reports its own CPU usage" "system usage: CPU: user:" "$ENTRY"
  check "pg_stat_user_tables records the autovacuum" "1" \
        "$(q "SELECT CASE WHEN autovacuum_count > 0 THEN 1 ELSE 0 END FROM pg_stat_user_tables WHERE relname='users'")"
else
  skip "the autovacuum log entry reports its own CPU usage" "autovacuum to fire within 120s"
  skip "pg_stat_user_tables records the autovacuum" "autovacuum to fire within 120s"
fi

# ---- 6. idle connections vs active queries, measured in CPU seconds
PM="$(head -1 "$DATA/postmaster.pid")"
cpusecs(){
  local pids; pids="$PM $(pgrep -P "$PM" 2>/dev/null | tr '\n' ' ')"
  ps -o time= -p "$(echo $pids | tr ' ' ',')" 2>/dev/null | awk -F: '{s += $1*60 + $2} END {printf "%.2f", s}'
}
if [ -z "$(cpusecs)" ]; then
  skip "100 idle connections cost no measurable CPU" "ps -o time= to report cumulative CPU"
  skip "10 concurrent CPU-bound queries cost far more CPU than 100 idle connections" "ps -o time="
else
  sleep 15                      # let any autovacuum worker exit so the process set is stable
  A0=$(cpusecs); sleep 10; A1=$(cpusecs)
  IDLE_BASE=$(awk "BEGIN{printf \"%.2f\", $A1-$A0}")
  for _ in $(seq 100); do ( sleep 120 | psql -h 127.0.0.1 -p $PORT -U postgres -X -q >/dev/null 2>&1 & ); done
  sleep 5
  NCONN="$(q "SELECT count(*) FROM pg_stat_activity WHERE backend_type='client backend' AND state='idle'")"
  B0=$(cpusecs); sleep 10; B1=$(cpusecs)
  IDLE_100=$(awk "BEGIN{printf \"%.2f\", $B1-$B0}")
  { echo "SET max_parallel_workers_per_gather=0;"
    for _ in $(seq 20); do echo "SELECT count(*) FROM users WHERE email ~ '@example\\.com\$';"; done
  } > "$ROOT/burn.sql"
  for _ in $(seq 10); do ( psql -h 127.0.0.1 -p $PORT -U postgres -X -qAt -f "$ROOT/burn.sql" >/dev/null 2>&1 & ); done
  sleep 3
  C0=$(cpusecs); sleep 8; C1=$(cpusecs)
  ACTIVE=$(awk "BEGIN{printf \"%.2f\", $C1-$C0}")
  echo "    (idle cluster ${IDLE_BASE}s CPU/10s; ${NCONN} idle connections ${IDLE_100}s CPU/10s; 10 active queries ${ACTIVE}s CPU/8s)"
  check "the 100 extra connections really were idle" "100" "$NCONN"
  lt "100 idle connections cost less than 1 CPU-second over 10 seconds" "$IDLE_100" "1.0"
  lt "10 concurrent CPU-bound queries cost far more CPU than 100 idle connections" \
     "$(awk "BEGIN{printf \"%.2f\", $IDLE_100 * 10}")" "$ACTIVE"
fi

echo
echo "PASS=$PASSES FAIL=$FAILS SKIP=$SKIPS"
[ "$FAILS" -eq 0 ] || exit 1