#!/usr/bin/env bash
# cf6958-matrix.sh <bindir> <workdir>
#
# What pg_stat_progress_analyze and pg_stat_progress_copy show for a backend
# while it runs commands that nest one progress report inside another,
# mostly through file_fdw (ANALYZE of a foreign table, COPY TO over a foreign
# scan, a COPY FROM whose trigger reads a foreign table, two foreign scans in
# one query).
#
# It creates a throw-away cluster in <workdir> with the binaries in <bindir>
# (file_fdw must be installed there), writes the data files there (about
# 130 MB), and runs each command in one session while a second session
# samples both views for that backend every ~0.5 ms, calling
# pg_stat_clear_snapshot() between samples, until the backend goes idle.
# Takes about two minutes.  The cluster is stopped at the end.
#
# Output, per command: how many samples showed each row, where "(none)" means
# the backend was active but neither view had a row for it; the largest
# bytes_processed seen in pg_stat_progress_copy; the bytes the command sent to
# the client; stderr; and "SERVER CRASH" plus the first TRAP line if the
# backend died (the script waits for crash recovery and goes on).
set -euo pipefail
trap 'kill $(jobs -p) 2>/dev/null; exit 143' TERM INT
BIN=$1 W=$2; rm -rf "$W"; mkdir -p "$W"
PGDATA=$W/data; PORT=$((61000 + RANDOM % 4000)); export PGHOST=$W PGPORT=$PORT PGDATABASE=postgres
"$BIN/initdb" -A trust -N -D "$PGDATA" >/dev/null
"$BIN/pg_ctl" -D "$PGDATA" -o "-p $PORT -k $W -c autovacuum=off" -l "$W/log" -w start >/dev/null
trap '"$BIN/pg_ctl" -D "$PGDATA" -w -m immediate stop >/dev/null 2>&1 || true' EXIT
# big.txt: 1,500,000 rows x 8 columns; small.txt: 200,000 rows; bad.txt: one bad
# line at the end; three.txt: 3 rows.
awk 'BEGIN{for(i=1;i<=1500000;i++) printf "%d,%d,%d,%d,%d,x%d,y%d,z%d\n",i,i%997,i%101,i*7%100003,i%13,i%5003,i%211,i}' >"$W/big.txt"
head -200000 "$W/big.txt" >"$W/small.txt"
{ head -100000 "$W/big.txt"; echo 'oops,1,1,1,1,a,b,c'; } >"$W/bad.txt"
printf '1\n2\n3\n' >"$W/three.txt"
COLS='a int, b int, c int, d int, e int, f text, g text, h text'
SQ="$BIN/psql -X -q -v ON_ERROR_STOP=1"
$SQ <<SQL
CREATE EXTENSION file_fdw;
CREATE SERVER f FOREIGN DATA WRAPPER file_fdw;
CREATE FOREIGN TABLE ft ($COLS) SERVER f OPTIONS (filename '$W/big.txt', format 'csv');
CREATE FOREIGN TABLE ft_small ($COLS) SERVER f OPTIONS (filename '$W/small.txt', format 'csv');
CREATE FOREIGN TABLE ft_prog ($COLS) SERVER f OPTIONS (program 'cat $W/big.txt', format 'csv');
CREATE FOREIGN TABLE ft_bad ($COLS) SERVER f OPTIONS (filename '$W/bad.txt', format 'csv');
CREATE TABLE t ($COLS);
CREATE TABLE parted ($COLS) PARTITION BY RANGE (a);
CREATE TABLE p_heap PARTITION OF parted FOR VALUES FROM (2000000) TO (3000000);
INSERT INTO p_heap SELECT g, 1, 1, 1, 1, 'x', 'y', 'z' FROM generate_series(2000000, 2300000) g;
CREATE FOREIGN TABLE p_ft PARTITION OF parted FOR VALUES FROM (1) TO (2000000)
  SERVER f OPTIONS (filename '$W/big.txt', format 'csv');
CREATE TABLE t3 (a int);
CREATE FUNCTION read_ft() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN PERFORM count(*) FROM ft_small; RETURN NEW; END \$\$;
CREATE TRIGGER t3_read_ft BEFORE INSERT ON t3 FOR EACH ROW EXECUTE FUNCTION read_ft();
CREATE TABLE t4 (a int);
CREATE FUNCTION copy_ft() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN COPY (SELECT a FROM ft_small) TO '/dev/null'; RETURN NEW; END \$\$;
CREATE TRIGGER t4_copy_ft BEFORE INSERT ON t4 FOR EACH ROW EXECUTE FUNCTION copy_ft();
CREATE FUNCTION observe(app text) RETURNS TABLE (seen text, n bigint)
LANGUAGE plpgsql AS \$\$
DECLARE
  p int;
  t0 timestamptz := clock_timestamp();
  st text;
  o text;
  b bigint;
  maxb bigint := 0;
  arr text[] := '{}';
BEGIN
  LOOP
    PERFORM pg_stat_clear_snapshot();
    SELECT pid INTO p FROM pg_stat_activity
     WHERE application_name = app AND state = 'active';
    EXIT WHEN p IS NOT NULL;
    IF clock_timestamp() - t0 > interval '30 s' THEN
      RAISE EXCEPTION 'session % never became active', app;
    END IF;
    PERFORM pg_sleep(0.0005);
  END LOOP;
  LOOP
    PERFORM pg_stat_clear_snapshot();
    SELECT coalesce(
      (SELECT 'analyze: ' || phase || ' relid=' || relid::regclass ||
              CASE WHEN child_tables_total > 0
                   THEN ' children ' || child_tables_done || '/' || child_tables_total
                   ELSE '' END
         FROM pg_stat_progress_analyze WHERE pid = p),
      (SELECT 'copy: ' || command || ' ' || type || ' relid=' || relid::regclass ||
              CASE WHEN bytes_total > 0 AND bytes_processed > bytes_total
                   THEN ' bytes_processed>bytes_total' ELSE '' END
         FROM pg_stat_progress_copy WHERE pid = p),
      '(none)'),
      (SELECT state FROM pg_stat_activity WHERE pid = p),
      (SELECT bytes_processed FROM pg_stat_progress_copy WHERE pid = p)
      INTO o, st, b;
    EXIT WHEN st IS DISTINCT FROM 'active';
    arr := arr || o;
    maxb := greatest(maxb, b);
    PERFORM pg_sleep(0.0005);
  END LOOP;
  RETURN QUERY SELECT x, count(*) FROM unnest(arr) x GROUP BY x ORDER BY 2 DESC, 1;
  RETURN QUERY SELECT 'max bytes_processed seen in pg_stat_progress_copy: ' || maxb, NULL::bigint;
END \$\$;
SQL
cell() {			# cell <name> <sql>
  local name=$1 sql=$2
  $SQ -At -c "SELECT coalesce(n || '  ', '') || seen FROM observe('$name')" >"$W/obs.$name" & local o=$!
  PGAPPNAME=$name PGOPTIONS='-c default_statistics_target=1000' \
    $BIN/psql -X -q -c "$sql" 2>"$W/err.$name" | wc -c >"$W/out.$name" || true
  wait $o || true
  echo "== $name: $sql"
  sed 's/^/   /' "$W/obs.$name"
  echo "   bytes written to the client: $(cat "$W/out.$name")"
  grep -v '^$' "$W/err.$name" | sed 's/^/   stderr: /' || true
  if grep -q 'terminated by signal' "$W/log"; then
    echo "   SERVER CRASH: $(grep -m1 -E '^TRAP|terminated by signal' "$W/log")"
    : >"$W/log.seen"; cat "$W/log" >>"$W/log.seen"; : >"$W/log"
    for _ in $(seq 600); do $SQ -c 'SELECT 1' >/dev/null 2>&1 && break; sleep 0.1; done
  fi
}
cell analyze_ft          'ANALYZE ft'
cell analyze_partitioned 'ANALYZE parted'
cell vacuum_analyze_parted 'VACUUM ANALYZE parted'
cell analyze_all         'ANALYZE'
cell copy_to_select_ft   'COPY (SELECT a FROM ft) TO STDOUT'
cell copy_from_trigger   "COPY t3 FROM '$W/three.txt'"
cell copy_from_trigger_copy_to "COPY t4 FROM '$W/three.txt'"
cell copy_to_two_scans   'COPY (SELECT a FROM ft UNION ALL SELECT a FROM ft_small) TO STDOUT'
cell insert_select_ft    'INSERT INTO t SELECT * FROM ft'
cell ctas_ft             'CREATE TABLE ctas AS SELECT * FROM ft'
cell select_ft           'SELECT count(*) FROM ft'
cell explain_then_select 'EXPLAIN SELECT * FROM ft_small; SELECT count(*) FROM ft'
cell select_program      'SELECT count(*) FROM ft_prog'
cell copy_to_program     'COPY (SELECT a FROM ft_prog) TO STDOUT'
cell copy_to_plain       'COPY (SELECT a FROM t) TO STDOUT'
cell select_two_scans    'SELECT (SELECT count(*) FROM ft) + (SELECT count(*) FROM ft_small)'
cell select_rescan       'SELECT count(*) FROM generate_series(1, 3) g(i), LATERAL (SELECT a FROM ft WHERE a > i OFFSET 0) s'
cell select_error        'SELECT count(*) FROM ft_bad'
echo "== slot after an error, same session:"
$SQ -At <<'SQL' | sed 's/^/   /'
\set ON_ERROR_STOP 0
SELECT count(*) FROM ft_bad;
SELECT 'copy rows for own pid after error: ' || count(*) FROM pg_stat_progress_copy WHERE pid = pg_backend_pid();
SQL
echo "== server: $("$BIN/postgres" --version)"