#!/usr/bin/env bash
#
# Measures what DROP COLUMN actually costs in PostgreSQL:
#
#   1. the catalog flip itself, against table size
#   2. the cascading index drop, across a range of index sizes
#   3. the lock queue which forms behind AccessExclusiveLock
#   4. lock_timeout as a mitigation for (3)
#   5. attnum exhaustion, and whether VACUUM FULL reclaims the slots
#   6. tuple widths either side of the drop
#
# Creates a disposable cluster under /private/tmp and removes it on exit.
# Set SECTIONS (e.g. SECTIONS="3 4") to re-run a subset without disturbing
# the numbers from the others.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PG_BIN="${PG_BIN:-/Applications/Postgres.app/Contents/Versions/latest/bin}"
PORT="${PORT:-56434}"
DATA_DIR="${DATA_DIR:-/private/tmp/drop-column-bench}"
CHURN_CYCLES="${CHURN_CYCLES:-1599}"
SECTIONS="${SECTIONS:-1 2 3 4 5 6}"

for program in initdb pg_ctl psql; do
    if [[ ! -x "$PG_BIN/$program" ]]; then
        printf 'Missing required PostgreSQL program: %s/%s\n' "$PG_BIN" "$program" >&2
        exit 1
    fi
done

cleanup() {
    "$PG_BIN/pg_ctl" -D "$DATA_DIR" -m immediate stop >/dev/null 2>&1 || true
    rm -rf "$DATA_DIR"
}
trap cleanup EXIT

rm -rf "$DATA_DIR"
"$PG_BIN/initdb" -D "$DATA_DIR" -U postgres --no-sync >/dev/null
"$PG_BIN/pg_ctl" -D "$DATA_DIR" -o "-p $PORT -c fsync=off -c full_page_writes=off" \
    -l "$DATA_DIR/server.log" -w start >/dev/null

psql() { "$PG_BIN/psql" -h 127.0.0.1 -p "$PORT" -U postgres -X -q "$@"; }
quiet() { psql -t -c "$1" >/dev/null; }
value() { psql -t -A -c "$1"; }
want() { [[ " $SECTIONS " == *" $1 "* ]]; }

# Returns the milliseconds a single statement took, as measured by psql's
# \timing around the round trip, so client startup is excluded.
timed() {
    psql -t -A -c '\timing on' -c "$1" 2>&1 | grep -o '[0-9.]* ms' | head -1
}

seed_table() {
    local table=$1 rows=$2
    quiet "DROP TABLE IF EXISTS $table"
    quiet "CREATE TABLE $table (id bigint, notes text, status text)"
    quiet "INSERT INTO $table SELECT g, md5(g::text) || repeat('y', 100), 'open'
           FROM generate_series(1, $rows) g"
    quiet "CREATE INDEX ${table}_notes_idx ON $table (notes)"
    quiet "CREATE INDEX ${table}_notes_status_idx ON $table (notes, status)"
    quiet "VACUUM ANALYZE $table"
}

if want 1; then
    printf '=== 1. The catalog flip ===\n\n'

    quiet "DROP TABLE IF EXISTS flip"
    quiet "CREATE TABLE flip (id bigint, customer_id bigint, notes text, status text)"
    quiet "INSERT INTO flip SELECT g, g % 1000, repeat('x', 200), 'open'
           FROM generate_series(1, 500000) g"
    quiet "VACUUM ANALYZE flip"

    before=$(value "SELECT pg_size_pretty(pg_relation_size('flip'))")
    elapsed=$(timed "ALTER TABLE flip DROP COLUMN notes")
    after=$(value "SELECT pg_size_pretty(pg_relation_size('flip'))")

    printf 'ALTER TABLE flip DROP COLUMN notes: %s\n' "$elapsed"
    printf 'heap size before: %s, after: %s\n\n' "$before" "$after"
    printf 'pg_attribute now holds:\n'
    psql -c "SELECT attname, attnum, attisdropped, atttypid
             FROM pg_attribute WHERE attrelid = 'flip'::regclass AND attnum > 0
             ORDER BY attnum"

    vacuum_elapsed=$(timed "VACUUM FULL flip")
    printf 'VACUUM FULL flip: %s, heap size now %s\n\n' \
        "$vacuum_elapsed" "$(value "SELECT pg_size_pretty(pg_relation_size('flip'))")"
fi

if want 2; then
    printf '=== 2. The cascading index drop ===\n\n'
    printf '%-12s %-14s %s\n' 'rows' 'index size' 'ALTER TABLE ... DROP COLUMN'

    for rows in 250000 500000 1000000 2000000; do
        for run in 1 2 3; do
            seed_table cascade "$rows"
            index_size=$(value "SELECT pg_size_pretty(sum(pg_relation_size(indexrelid)))
                                FROM pg_index WHERE indrelid = 'cascade'::regclass")
            printf '%-12s %-14s %s\n' "$rows" "$index_size" \
                "$(timed 'ALTER TABLE cascade DROP COLUMN notes')"
        done
    done

    printf '\nSame table, but with the indexes removed first:\n'
    printf '%-12s %-14s %s\n' 'rows' 'index size' 'ALTER TABLE ... DROP COLUMN'
    for rows in 250000 2000000; do
        seed_table predrop "$rows"
        index_size=$(value "SELECT pg_size_pretty(sum(pg_relation_size(indexrelid)))
                            FROM pg_index WHERE indrelid = 'predrop'::regclass")
        quiet "DROP INDEX CONCURRENTLY predrop_notes_idx"
        quiet "DROP INDEX CONCURRENTLY predrop_notes_status_idx"
        printf '%-12s %-14s %s\n' "$rows" "$index_size" \
            "$(timed 'ALTER TABLE predrop DROP COLUMN notes')"
    done
    printf '\n'
fi

if want 3; then
    printf '=== 3. The lock queue ===\n\n'

    quiet "DROP TABLE IF EXISTS queued"
    quiet "CREATE TABLE queued (id bigint, notes text)"
    quiet "INSERT INTO queued SELECT g, md5(g::text) FROM generate_series(1, 100000) g"

    # A reporting query holds AccessShareLock for ten seconds. DROP COLUMN queues
    # behind it, and every query arriving afterwards queues behind DROP COLUMN.
    quiet "BEGIN; SELECT count(*) FROM queued; SELECT pg_sleep(10); COMMIT;" &
    sleep 1
    ( printf 'DROP COLUMN               waited %s\n' \
        "$( { time quiet 'ALTER TABLE queued DROP COLUMN notes'; } 2>&1 | awk '/real/ {print $2}')" ) &
    sleep 1
    ( printf 'SELECT ... LIMIT 1 behind it waited %s\n' \
        "$( { time quiet 'SELECT id FROM queued LIMIT 1'; } 2>&1 | awk '/real/ {print $2}')" ) &

    sleep 2
    printf '\nLock graph while all three are in flight:\n'
    psql -c "SELECT a.pid, left(a.query, 38) AS query, l.mode, l.granted
             FROM pg_locks l JOIN pg_stat_activity a USING (pid)
             WHERE l.relation = 'queued'::regclass
             ORDER BY l.granted DESC, a.pid"
    wait
    printf '\n'
fi

if want 4; then
    printf '=== 4. lock_timeout as the mitigation ===\n\n'

    # Same ten second reporting query as section 3, but the SELECT arrives
    # whilst the ALTER is still waiting, so this measures how long lock_timeout
    # holds the queue up rather than how fast a query runs once the ALTER has
    # already given up.
    for timeout in 2s 500ms; do
        quiet "DROP TABLE IF EXISTS guarded"
        quiet "CREATE TABLE guarded (id bigint, notes text)"
        quiet "INSERT INTO guarded SELECT g, md5(g::text) FROM generate_series(1, 100000) g"

        quiet "BEGIN; SELECT count(*) FROM guarded; SELECT pg_sleep(10); COMMIT;" &
        sleep 1
        # Expected to fail: the point of the section is that it gives up rather than queue.
        (
            output=$( { time psql -c "SET lock_timeout = '$timeout'" \
                                  -c "ALTER TABLE guarded DROP COLUMN notes" 2>&1; } 2>&1 || true )
            printf 'lock_timeout=%-5s ALTER gave up after    %s  %s\n' "$timeout" \
                "$(awk '/real/ {print $2}' <<<"$output")" "$(grep -o 'ERROR:.*' <<<"$output")"
        ) &
        sleep 0.2
        ( printf 'lock_timeout=%-5s SELECT behind it waited %s\n' "$timeout" \
            "$( { time quiet 'SELECT id FROM guarded LIMIT 1'; } 2>&1 | awk '/real/ {print $2}')" ) &
        wait
    done
    printf '\n'
fi

if want 5; then
    printf '=== 5. attnum exhaustion ===\n\n'

    quiet "DROP TABLE IF EXISTS churn"
    quiet "CREATE TABLE churn (id bigint)"
    quiet "INSERT INTO churn SELECT generate_series(1, 1000)"
    for _ in $(seq 1 "$CHURN_CYCLES"); do
        quiet "ALTER TABLE churn ADD COLUMN tmp int; ALTER TABLE churn DROP COLUMN tmp;"
    done

    slots() {
        psql -c "SELECT count(*) FILTER (WHERE attisdropped)     AS ghost_slots,
                        count(*) FILTER (WHERE NOT attisdropped) AS live_columns,
                        max(attnum)                              AS highest_attnum
                 FROM pg_attribute WHERE attrelid = '$1'::regclass AND attnum > 0"
    }

    printf 'After %s add/drop cycles on a table with one real column:\n' "$CHURN_CYCLES"
    slots churn

    printf '\nVACUUM FULL, which rewrites the heap:\n'
    quiet "VACUUM FULL churn"
    slots churn

    printf '\nAdding one more column:\n'
    add_result=$(psql -c "ALTER TABLE churn ADD COLUMN one_too_many int" 2>&1 | head -1 || true)
    printf '%s\n' "${add_result:-ALTER TABLE (succeeded)}"

    printf '\nCREATE TABLE ... AS SELECT, which builds fresh catalog rows:\n'
    quiet "CREATE TABLE churn_rebuilt AS SELECT * FROM churn"
    slots churn_rebuilt
    printf '\n'
fi

if want 6; then
    printf '=== 6. Tuple widths either side of the drop ===\n\n'

    quiet "CREATE EXTENSION IF NOT EXISTS pageinspect"
    quiet "DROP TABLE IF EXISTS widths"
    quiet "CREATE TABLE widths (id bigint, pad text, keep text)"
    quiet "INSERT INTO widths SELECT g, repeat('z', 500), 'keep' FROM generate_series(1, 20) g"
    quiet "ALTER TABLE widths DROP COLUMN pad"
    quiet "INSERT INTO widths (id, keep) SELECT g, 'keep' FROM generate_series(21, 40) g"

    psql -c "SELECT p AS page, count(*) AS tuples,
                    min(lp_len) AS min_tuple_bytes, max(lp_len) AS max_tuple_bytes
             FROM generate_series(0, (pg_relation_size('widths') / 8192)::int - 1) p,
                  LATERAL heap_page_items(get_raw_page('widths', p))
             WHERE lp_len > 0 GROUP BY p ORDER BY p"
    printf '\n'
fi

printf 'PostgreSQL version: %s\n' "$(value 'SHOW server_version')"