#!/usr/bin/env bash
# The statements of an application behind a connection pool that gives each request a backend of its own: every statement is the first of its backend, so nothing kept by
# a backend (the verdicts of the planner, the rows of lookup tables, the dictionaries, the statistics of an index that a new backend reads) helps it, and only what the
# server shares does. Each of N statements runs in a new session and its backend is terminated after it; the same statements in one session are the arm that shows what a
# backend keeps. Per statement: the blocks of its planning and of its execution (EXPLAIN BUFFERS), the milliseconds of both, and the user instructions and cycles of the
# backend from the moment before the statement to its end (perf stat attached to the backend; empty where perf cannot attach). The blocks and instructions are the measure,
# the milliseconds are secondary.
#
#   PGHOST=.. PGPORT=.. bench/steps/fresh_backend_statements.sh [STATEMENTS [KEYS]] > records.tsv       LOAD=0 reuses the tables of an earlier run
#
# Needs a scratch cluster with pg_sift (tmpfs, no autovacuum). Two tables: fb_ev (120,000 events of 1,000 parents, a sift_col index) and fb_map (300,000 rows, a unique
# btree on (COALESCE(scope, '{}'), kind, part, sub): the lookup table of a correlated LIMIT 1 child, as the items of an event are normalized through it). Statement i asks for
# the events of KEYS consecutive parents from parent i: the same statement with other constants, as an application sends it. Arms:
#   fresh    sift.enable_agg = on, each statement in a new backend
#   native   sift.enable_agg = off, each statement in a new backend (PostgreSQL's plan)
#   session  sift.enable_agg = on, the N statements one after another in one backend
set -euo pipefail
N=${1:-20} KEYS=${2:-5}
P="psql -X -q -At -v ON_ERROR_STOP=1"
D=$(mktemp -d /dev/shm/fresh-backend.XXXXXX)
trap 'rm -rf "$D"' EXIT

[ "${LOAD:-1}" = 1 ] && $P <<'SQL' >/dev/null
set client_min_messages = warning;
create extension if not exists pg_sift;
drop table if exists fb_ev, fb_map cascade;
create table fb_map (id bigint generated by default as identity primary key, scope int[], kind text not null, part text not null, sub text not null, canon text);
insert into fb_map (scope, kind, part, sub, canon)
  select case when g % 40 = 0 then array[g % 7] end, (array['type_a', 'type_b', 'type_c', 'type_d'])[1 + g % 4], 'p' || g % 52, md5(g::text), 'm' || g % 18
  from generate_series(1, 300000) g order by random();
create unique index fb_map_key on fb_map (coalesce(scope, '{}'::int[]), kind, part, sub);
create table fb_ev (id bigint generated by default as identity primary key, dim int not null, kind text not null, part text not null, sub text not null, score float8 not null);
insert into fb_ev (dim, kind, part, sub, score)
  select 1 + g % 1000, (array['type_a', 'type_a', 'type_a', 'type_b', 'type_c', 'type_d'])[1 + (g % 1000) % 6], 'p' || (r % 52), md5((r)::text), round((3 + random() * 7)::numeric, 1)
  from (select g, 1 + floor(random() * 12000)::int * 97 as r from generate_series(1, 120000) g) x;
create index fb_ev_col on fb_ev using sift_col (id, dim, kind, part, sub, score);
vacuum (freeze, analyze) fb_ev, fb_map;
SQL

stmt() {  # I: the events of KEYS consecutive parents from parent I
  echo "select m.canon, count(*) as n, avg(e.score) as score from fb_ev e left join lateral (select canon from fb_map t where coalesce(t.scope, '{}') = '{}' and t.kind = e.kind and t.part = e.part and t.sub = e.sub limit 1) m on true where e.dim between $1 and $(( $1 + KEYS - 1 )) group by 1 order by 1"
}
settings() { case $1 in native) echo "set sift.enable_agg = off;" ;; *) echo "set sift.enable_agg = on;" ;; esac; }

# the backend of a script waits at "ready N" until the harness has attached perf to it and says "go N"; after the statement it waits for "fin N" so that the exit of the process is outside the count
script() {  # ARM FIRST LAST: the statements FIRST..LAST of one session
  local arm=$1 first=$2 last=$3 i
  echo '\set ON_ERROR_STOP 1'
  echo 'select pg_backend_pid() \gset'
  echo "\\o $W/pid.$first"
  echo 'select :pg_backend_pid;'
  echo '\o'
  echo "set max_parallel_workers_per_gather = 0; $(settings "$arm")"
  for (( i = first; i <= last; i++ )); do
    echo "\\! touch $W/ready.$i"
    echo "\\! while [ ! -f $W/go.$i ]; do sleep 0.01; done"
    echo "\\o $W/plan.$i"
    echo "explain (analyze, buffers, costs off, format json) $(stmt "$i");"
    echo '\o'
    echo "\\! touch $W/done.$i"
    echo "\\! while [ ! -f $W/fin.$i ]; do sleep 0.01; done"
  done
}

waitfor() { while [ ! -f "$1" ]; do sleep 0.01; done; }
perf_ok() { command -v perf >/dev/null && perf stat -e instructions:u true >/dev/null 2>&1; }

measure() {  # ARM I PID: perf attached to PID around statement I; one record
  local arm=$1 i=$2 pid=$3
  waitfor "$W/ready.$i"
  if perf_ok; then
    perf stat -x, -e instructions:u,cycles:u -p "$pid" -o "$W/perf.$i" 2>/dev/null &
    local pp=$! k
    for (( k = 0; k < 200; k++ )); do grep -qs started "$W/perf.$i" && break; sleep 0.01; done   # perf writes its header once attached
  fi
  touch "$W/go.$i"
  waitfor "$W/done.$i"
  if [ -n "${pp:-}" ]; then kill -INT "$pp" 2>/dev/null || true; wait "$pp" 2>/dev/null || true; fi
  touch "$W/fin.$i"   # the backend stays until perf is done, so the exit of the process is not counted
  python3 -I - "$arm" "$i" "$W/plan.$i" "$W/perf.$i" <<'PY'
import json, os, sys
arm, i, plan, perf = sys.argv[1:]
j = json.load(open(plan))[0]
p, e = j['Planning'], j['Plan']
st = {}
if os.path.exists(perf):
    for line in open(perf):
        f = line.strip().split(',')
        if len(f) > 2 and f[0].replace('.', '').isdigit():
            st[f[2].split(':')[0]] = float(f[0])
m = lambda k: '%.1f' % (st[k] / 1e6) if k in st else ''
print('\t'.join(str(x) for x in [arm, i, p['Shared Read Blocks'], p['Shared Hit Blocks'], round(j['Planning Time'], 2), e['Shared Read Blocks'], e['Shared Hit Blocks'], round(j['Execution Time'], 2), m('instructions'), m('cycles')]))
PY
}

printf 'arm\tstatement\tplan_read_blocks\tplan_hit_blocks\tplan_ms\texec_read_blocks\texec_hit_blocks\texec_ms\tuser_minstr\tuser_mcycles\n'
for arm in fresh native; do
  W=$D/$arm; mkdir "$W"
  for (( i = 1; i <= N; i++ )); do
    script "$arm" "$i" "$i" > "$W/s.sql"
    $P -f "$W/s.sql" > /dev/null 2>&1 &
    sp=$!
    waitfor "$W/pid.$i"; sleep 0.05
    pid=$(tr -d ' \n' < "$W/pid.$i")
    measure "$arm" "$i" "$pid"
    $P -c "select pg_terminate_backend($pid)" > /dev/null 2>&1
    wait "$sp" 2>/dev/null || true
  done
done
W=$D/session; mkdir "$W"
script session 1 "$N" > "$W/s.sql"
$P -f "$W/s.sql" > /dev/null 2>&1 &
sp=$!
waitfor "$W/pid.1"; sleep 0.05
pid=$(tr -d ' \n' < "$W/pid.1")
for (( i = 1; i <= N; i++ )); do measure session "$i" "$pid"; done
wait "$sp" 2>/dev/null || true