#!/usr/bin/env bash
#
# Measures what slicing work by yb_hash_code() buys you on YugabyteDB:
#
#   1. where the tablets sit in the hash space, and which node leads each
#   2. the plans for full scans, aligned slices and straddling slices
#   3. a full table read by 1 reader against N readers on hash slices
#   4. a job queue polled by 12 workers, four ways: a status-hashed index,
#      a partial index hashed on id, and that index with a random slice of
#      6 or of 64
#   5. a hash code slice against an index hashed on a different column,
#      and whether yb_enable_parallel_scan_hash_sharded changes the plan
#
# Starts a disposable 3 node cluster plus a client container in Docker and
# removes them all on exit. Set SECTIONS (e.g. SECTIONS="3") to re-run a
# subset.

set -euo pipefail

IMAGE="${IMAGE:-yugabytedb/yugabyte:2026.1.2.0-b137}"
CLIENT_IMAGE="${CLIENT_IMAGE:-postgres:16}"
NODE_CPUS="${NODE_CPUS:-2}"
ROWS="${ROWS:-2000000}"
QUEUE_ROWS="${QUEUE_ROWS:-1000000}"
DURATION="${DURATION:-30}"
SECTIONS="${SECTIONS:-1 2 3 4 5}"
NET=ybfan
NODES=(ybfan-1 ybfan-2 ybfan-3)
CLIENT=ybfan-client

want() { [[ " $SECTIONS " == *" $1 "* ]]; }

cleanup() {
    docker rm -f "${NODES[@]}" "$CLIENT" >/dev/null 2>&1 || true
    docker network rm "$NET" >/dev/null 2>&1 || true
}
cleanup
# KEEP=1 leaves the cluster running afterwards, for poking at by hand.
[[ "${KEEP:-0}" == 1 ]] || trap cleanup EXIT

# Runs SQL from stdin or -c against a node, from the client container so
# the node's own CPUs are only spent on being a database.
psql_on() {
    local host=$1; shift
    docker exec -i "$CLIENT" psql -h "$host" -p 5433 -U yugabyte -d yugabyte -X -q "$@"
}
psql_() { psql_on "${NODES[0]}" "$@"; }
value() { psql_ -t -A -c "$1"; }

server_count() { value "SELECT count(*) FROM yb_servers()" 2>/dev/null || echo 0; }

wait_for_servers() {
    local want=$1
    for _ in $(seq 1 90); do
        [[ "$(server_count)" == "$want" ]] && return 0
        sleep 2
    done
    printf 'Cluster never reached %s servers\n' "$want" >&2
    exit 1
}

start_cluster() {
    docker network create "$NET" >/dev/null
    docker run -d --name "$CLIENT" --net "$NET" -e POSTGRES_HOST_AUTH_METHOD=trust \
        "$CLIENT_IMAGE" >/dev/null

    # Nodes join one at a time, because two masters joining at once race
    # and the loser exits with "Unable to add master".
    local n=0
    for node in "${NODES[@]}"; do
        local join=()
        [[ $n -gt 0 ]] && join=(--join="${NODES[0]}")
        docker run -d --name "$node" --hostname "$node" --net "$NET" --cpus "$NODE_CPUS" \
            "$IMAGE" bin/yugabyted start --background=false \
            --advertise_address="$node" ${join[@]+"${join[@]}"} >/dev/null
        n=$((n + 1))
        wait_for_servers "$n"
    done
}

seed_jobs() {
    psql_ <<SQL
CREATE TABLE jobs (
    id uuid DEFAULT gen_random_uuid(),
    created_at timestamptz NOT NULL,
    payload text NOT NULL,
    PRIMARY KEY (id HASH)
) SPLIT INTO 6 TABLETS;

INSERT INTO jobs (created_at, payload)
SELECT now() - (g || ' seconds')::interval, repeat(md5(g::text), 6)
FROM generate_series(1, $ROWS) g;

ANALYZE jobs;
SQL
}

# Reads every row of jobs using $1 readers, each owning an equal slice of
# the 65,536 hash codes and connected round robin to the three nodes. With
# 6 tablets split evenly, 1, 2, 3 and 6 readers line up with tablet
# boundaries, and 12 puts two readers on each tablet.
#
# The readers are launched and timed inside the client container, with
# their rows discarded there, since streaming 400 MB back out through
# docker exec turned out to be the bottleneck rather than the database.
read_all() {
    local readers=$1 script="start=\$(date +%s%3N)"
    for i in $(seq 0 $((readers - 1))); do
        local lo=$((i * 65536 / readers)) hi=$(((i + 1) * 65536 / readers))
        local host=${NODES[$((i % 3))]} query
        if [[ $readers -eq 1 ]]; then
            query="SELECT id, payload FROM jobs"
        else
            query="SELECT id, payload FROM jobs WHERE yb_hash_code(id) >= $lo AND yb_hash_code(id) < $hi"
        fi
        script+=$'\n'"psql -h $host -p 5433 -U yugabyte -X -q -c 'COPY ($query) TO STDOUT' >/dev/null &"
    done
    script+=$'\n'"wait; echo \$((\$(date +%s%3N) - start))"
    docker exec "$CLIENT" sh -c "$script"
}

median() { sort -n | awk '{ a[NR] = $1 } END { print a[int((NR + 1) / 2)] }'; }

seed_queue() {
    local index_sql=$1
    psql_ <<SQL
DROP TABLE IF EXISTS queue;
CREATE TABLE queue (
    id uuid DEFAULT gen_random_uuid(),
    status text NOT NULL,
    created_at timestamptz NOT NULL,
    PRIMARY KEY (id HASH)
) SPLIT INTO 6 TABLETS;

-- 1 row in 10 is pending, like a queue with history behind it. With every
-- row pending the planner rightly ignores a partial index on pending.
INSERT INTO queue (status, created_at)
SELECT CASE WHEN g % 10 = 0 THEN 'pending' ELSE 'done' END,
       now() - (g || ' seconds')::interval
FROM generate_series(1, $QUEUE_ROWS) g;

$index_sql
ANALYZE queue;
SQL
}

# Seeds the queue with index $1, then runs 12 workers for $DURATION
# seconds, 4 per node, each claiming 10 jobs per transaction with the
# SKIP LOCKED select $3 and marking them done. $2 is a pgbench \set line
# for :slice, or empty, and $4 an optional pg_hint_plan hint. Prints the
# summed claims per second, the summed failures, a CPU sample per node
# taken halfway through, and what one more claim costs in storage once the
# workers have stopped, with :slice = 0.
poll_queue() {
    local index_sql=$1 set_slice=$2 claim=$3 hint=${4:-} out
    seed_queue "$index_sql"
    docker exec -i "$CLIENT" sh -c "cat > /tmp/poll.sql" <<SQL
$set_slice
BEGIN;
$hint UPDATE queue SET status = 'done' WHERE id IN ($claim);
COMMIT;
SQL
    out=$(mktemp)
    for node in "${NODES[@]}"; do
        docker exec "$CLIENT" pgbench -h "$node" -p 5433 -U yugabyte -n -c 4 -j 4 \
            -T "$DURATION" --max-tries=10 -f /tmp/poll.sql yugabyte >>"$out" 2>&1 &
    done
    sleep $((DURATION / 2))
    docker stats --no-stream --format '{{.Name}} {{.CPUPerc}}' "${NODES[@]}" | sed 's/^/    cpu /'
    wait
    awk '/^tps = / { tps += $3 } /number of failed transactions/ { f += $5 }
         END { printf "    claims/s %.0f  failed txns %d\n", tps * 10, f }' "$out"
    rm -f "$out"
    psql_ -c "EXPLAIN (ANALYZE, DIST, COSTS OFF) $hint ${claim//:slice/0}" |
        grep -E 'Scan|Cond|Filter|Rows Scanned|Execution Time' | sed 's/^/   /'
}

start_cluster
seed_jobs

printf 'YugabyteDB %s, 3 nodes x %s CPUs, %s rows in jobs\n\n' \
    "$(value "SELECT split_part(version(), ' ', 2)")" "$NODE_CPUS" "$ROWS"

if want 1; then
    echo '== 1. Tablets of jobs in the hash space'
    psql_ -c "SELECT start_hash_code, end_hash_code, leader
              FROM yb_tablet_metadata WHERE relname = 'jobs' ORDER BY start_hash_code"
fi

if want 2; then
    echo '== 2. Plans'
    for where in "" \
        "WHERE yb_hash_code(id) >= 10922 AND yb_hash_code(id) < 21845" \
        "WHERE yb_hash_code(id) >= 5000 AND yb_hash_code(id) < 16000" \
        "WHERE yb_hash_code(id) % 6 = 1"; do
        printf -- '-- SELECT id, payload FROM jobs %s\n' "$where"
        psql_ -c "EXPLAIN (ANALYZE, DIST, COSTS OFF) SELECT id, payload FROM jobs $where" |
            grep -E 'Scan|Cond|Filter|Table Read (Requests|Ops)|Rows Scanned'
    done
    echo
fi

if want 3; then
    echo '== 3. Reading every row, median of 3 runs after a warm-up'
    read_all 6 >/dev/null
    for readers in 1 2 3 6 12; do
        ms=$(for _ in 1 2 3; do read_all "$readers"; done | median)
        printf '    %2s readers  %6s ms\n' "$readers" "$ms"
    done
    echo
fi


if want 4; then
    echo '== 4. Job queue, 12 workers'
    hashed_status="CREATE INDEX queue_status ON queue (status HASH, created_at ASC) SPLIT INTO 6 TABLETS;"
    hashed_id="CREATE INDEX queue_pending ON queue (id HASH) SPLIT INTO 6 TABLETS WHERE status = 'pending';"
    # FOR UPDATE needs the table row, so the planner costs the partial index
    # as a lookup plus a fetch and prefers the primary key with a filter,
    # which then scans past every claimed row. Pin the index instead.
    use_pending="/*+ IndexScan(q queue_pending) */"

    echo "-- (a) index on (status HASH, created_at ASC), oldest first"
    psql_ -c "SELECT yb_hash_code('pending'::text) AS pending_hash_code"
    poll_queue "$hashed_status" "" \
        "SELECT id FROM queue WHERE status = 'pending' ORDER BY created_at LIMIT 10 FOR UPDATE SKIP LOCKED"

    echo "-- (b) partial index on (id HASH) WHERE pending, no slice"
    poll_queue "$hashed_id" "" \
        "SELECT id FROM queue q WHERE status = 'pending' LIMIT 10 FOR UPDATE SKIP LOCKED" \
        "$use_pending"

    echo "-- (c) as (b), random slice of 6, one per tablet"
    poll_queue "$hashed_id" "\\set slice random(0, 5)" \
        "SELECT id FROM queue q WHERE status = 'pending' AND yb_hash_code(id) >= :slice * 65536 / 6 AND yb_hash_code(id) < (:slice + 1) * 65536 / 6 LIMIT 10 FOR UPDATE SKIP LOCKED" \
        "$use_pending"

    echo "-- (d) as (b), random slice of 64"
    poll_queue "$hashed_id" "\\set slice random(0, 63)" \
        "SELECT id FROM queue q WHERE status = 'pending' AND yb_hash_code(id) >= :slice * 1024 AND yb_hash_code(id) < (:slice + 1) * 1024 LIMIT 10 FOR UPDATE SKIP LOCKED" \
        "$use_pending"
fi

if want 5; then
    echo '== 5. Slicing an index hashed on something else, and built-in parallel scan'
    seed_queue "CREATE INDEX queue_status ON queue (status HASH, created_at ASC) SPLIT INTO 6 TABLETS;"
    psql_ -c "EXPLAIN (COSTS OFF) /*+ IndexScan(queue queue_status) */
              SELECT id FROM queue WHERE status = 'pending'
              AND yb_hash_code(id) >= 10922 AND yb_hash_code(id) < 21845 LIMIT 100"
    psql_ -c "SET yb_enable_parallel_scan_hash_sharded = on" \
          -c "EXPLAIN (COSTS OFF) SELECT id, payload FROM jobs"
fi