#!/bin/bash
#
# Tests for parallel heap vacuum.
#
#   SIZE_GB=10 PGBIN=/path/to/bin PGDATA=/path/to/data ./vacuum-parallel-heap.sh
#
set -euo pipefail

SIZE_GB=${SIZE_GB:-10}            # heap size, 226 rows per page
KIND=${KIND:-unlogged}            # empty for a logged table
INDEXES=${INDEXES:-"0 1 2 4"}     # btree indexes on the same int column, ascending
WORKERS=${WORKERS:-"0 2 4 8"}     # VACUUM (PARALLEL n)
PASSES=${PASSES:-3}               # number of runs, the summary takes the average
RING_BUFFER=${RING_BUFFER:-2MB}   # vacuum ring buffer, 0 = whole buffer pool
PGBIN=${PGBIN:-$HOME/pg/bin}
PGDATA=${PGDATA:-/pgdata/data}
OUT=${OUT:-$HOME/phv-results/x$SIZE_GB}
export PATH="$PGBIN:$PATH" PGPORT=${PGPORT:-5500} PGHOST=/tmp

ROWS=$((SIZE_GB * 131072 * 226))
NOPARALLEL=$((SIZE_GB * 2))GB
DB=phv$SIZE_GB$KIND
CSV=$OUT/results.csv
psql="psql -Xq -v ON_ERROR_STOP=1"
mkdir -p "$OUT"

if [ ! -s "$PGDATA/PG_VERSION" ]; then
    initdb -D "$PGDATA" --no-sync >/dev/null
    cat >> "$PGDATA/postgresql.conf" <<EOF
port = $PGPORT
unix_socket_directories = '/tmp'
listen_addresses = ''
shared_buffers = 32GB
maintenance_work_mem = 1GB
max_wal_size = 64GB
max_worker_processes = 32
max_parallel_workers = 32
max_parallel_maintenance_workers = 16
autovacuum = off
track_io_timing = on
log_min_messages = warning
EOF
fi
pg_ctl -D "$PGDATA" -l "$OUT/pg.log" start -w >/dev/null 2>&1 || true

# One base table, then one template per index count, copied before the delete.
$psql -d postgres -c "create database ${DB}_base"
$psql -d ${DB}_base -c "create $KIND table test (a int) with (autovacuum_enabled = off)"
for i in $(seq 0 7); do
    lo=$((i * ROWS / 8 + 1)); hi=$(((i + 1) * ROWS / 8))
    $psql -d ${DB}_base -c "insert into test select i from generate_series($lo, $hi) i" &
done
wait
$psql -d ${DB}_base -c "vacuum analyze test" -c "checkpoint"

built=0
for nidx in $INDEXES; do
    while [ "$built" -lt "$nidx" ]; do
        built=$((built + 1))
        $psql -d ${DB}_base -c "set maintenance_work_mem = '8GB'" \
                             -c "set max_parallel_maintenance_workers = 8" \
                             -c "create index idx$built on test (a)"
    done
    $psql -d postgres -c "create database ${DB}_d$nidx template ${DB}_base strategy = file_copy"
    $psql -d ${DB}_d$nidx -c "delete from test where a % 5 = 0" -c "checkpoint"
done

echo "nidx,workers,parallel_heap,pass,scan_s,index_s,heapvac_s,total_s" > "$CSV"
for nidx in $INDEXES; do
    for pass in $(seq 1 "$PASSES"); do
        for w in $(shuf -e $WORKERS); do
            for ph in $(shuf -e on off); do
                off=()
                [ "$ph" = off ] && off=(-c "set min_parallel_table_scan_size = '$NOPARALLEL'")
                $psql -d postgres -c "drop database if exists ${DB}_run" >/dev/null
                $psql -d postgres -c "create database ${DB}_run template ${DB}_d$nidx strategy = file_copy" >/dev/null
                pg_ctl -D "$PGDATA" -l "$OUT/pg.log" restart -m fast -w >/dev/null
                $psql -d ${DB}_run -c "set vacuum_buffer_usage_limit = '$RING_BUFFER'" "${off[@]}" \
                      -c '\timing on' -c "vacuum (parallel $w, verbose) test" \
                      > "$OUT/n$nidx-w$w-$ph-p$pass.txt" 2>&1
                cat "$OUT/n$nidx-w$w-$ph-p$pass.txt" >> "$OUT/vacuum-verbose.txt"
                awk -v n="$nidx" -v w="$w" -v h="$ph" -v p="$pass" '
                    /^phase times:/ { sc=$5/1000; iv=$9/1000; hv=$13/1000 }
                    /^Time:/ { printf "%s,%s,%s,%s,%.2f,%.2f,%.2f,%.2f\n", n,w,h,p,sc,iv,hv,$2/1000 }
                ' "$OUT/n$nidx-w$w-$ph-p$pass.txt" | tee -a "$CSV"
            done
        done
    done
done
$psql -d postgres -c "drop database if exists ${DB}_run" >/dev/null

avg () { awk '{s+=$1; n++} END {if (n) printf "%.2f\n", s/n}'; }
col () { grep "^$1,$2,$3," "$CSV" | cut -d, -f"$4" | avg; }
echo
echo "averages, ${SIZE_GB}GB ${KIND:-logged}, ring buffer $RING_BUFFER, seconds"
printf "%-6s %-8s %9s %9s %10s %10s\n" nidx workers scan_off scan_on total_off total_on
for nidx in $INDEXES; do
    for w in $WORKERS; do
        printf "%-6s %-8s %9s %9s %10s %10s\n" "$nidx" "$w" \
            "$(col "$nidx" "$w" off 5)" "$(col "$nidx" "$w" on 5)" \
            "$(col "$nidx" "$w" off 8)" "$(col "$nidx" "$w" on 8)"
    done
done
echo "output saved to $OUT"