Multi-Write in YugabyteDB: Testing Concurrent Writes, Conflicts, and Consistency

				
					#!/usr/bin/env bash
# YugabyteDB Tips: concurrent applications, shared keys, checked outcomes.
set -euo pipefail

NODES=${YB_DEMO_NODES:-"127.0.0.1 127.0.0.2 127.0.0.3"}
DB=${YB_DEMO_DB:-yugabyte}
DB_USER=${YB_DEMO_USER:-yugabyte}
PORT=${YB_DEMO_PORT:-5433}
NUM_TXNS=${YB_DEMO_TXNS:-10}
MAX_RETRIES=${YB_DEMO_RETRIES:-5}
CERTS=${YB_DEMO_CERTS:-}
SSL_MODE=${YB_DEMO_SSLMODE:-}
SSL_CERT=${YB_DEMO_SSLCERT:-}
SSL_KEY=${YB_DEMO_SSLKEY:-}
HOLD_MS=${YB_DEMO_HOLD_MS:-100}
LOCK_MS=${YB_DEMO_LOCK_TIMEOUT_MS:-5000}
CLIENT=${YB_DEMO_YSQLSH:-ysqlsh}

die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; }
usage() {
    cat <<'HELP'
Usage: yb_multi_app_simulator.sh [options]
  --nodes "host1 host2 ..."  One worker per endpoint; all in ONE universe
  --port N                  YSQL listener port (5433)
  --db NAME                 Database (yugabyte)
  --user NAME               Database role (yugabyte)
  --txns N                  Logical transactions per worker (10)
  --retries N               Extra attempts after a known aborted transaction (5)
  --certs FILE              Trusted CA file; defaults sslmode to verify-full
  --sslmode MODE            disable, require, verify-ca, or verify-full
  --sslcert FILE            Optional client certificate
  --sslkey FILE             Optional client private key
  --hold-ms N               Delay before commit to encourage overlap (100)
  --lock-timeout-ms N       Per-statement lock wait limit (5000)
  --ysqlsh PATH             Client executable (ysqlsh)
  -h, --help                Show help

Options override matching YB_DEMO_* environment variables; see the Tip.
Use .pgpass / PGPASSFILE / PGPASSWORD for authentication; no password prompts.
Without TLS options the local-demo default is sslmode=disable.
All runs use yb_multi_app_demo; each run has a separate UUID. No data is dropped.
HELP
}

while (($#)); do
    case $1 in
        -h|--help) usage; exit 0 ;;
        --nodes|--port|--db|--user|--txns|--retries|--certs|--sslmode|--sslcert|--sslkey|--hold-ms|--lock-timeout-ms|--ysqlsh)
            (($# >= 2)) && [[ -n $2 && $2 != --* ]] || die "Missing value for $1"
            case $1 in
                --nodes) NODES=$2 ;; --port) PORT=$2 ;; --db) DB=$2 ;;
                --user) DB_USER=$2 ;; --txns) NUM_TXNS=$2 ;;
                --retries) MAX_RETRIES=$2 ;; --certs) CERTS=$2 ;;
                --sslmode) SSL_MODE=$2 ;; --sslcert) SSL_CERT=$2 ;;
                --sslkey) SSL_KEY=$2 ;; --hold-ms) HOLD_MS=$2 ;;
                --lock-timeout-ms) LOCK_MS=$2 ;; --ysqlsh) CLIENT=$2 ;;
            esac
            shift 2 ;;
        *) die "Unknown option: $1" ;;
    esac
done
for value in "$PORT" "$NUM_TXNS" "$MAX_RETRIES" "$HOLD_MS" "$LOCK_MS"; do
    [[ $value =~ ^(0|[1-9][0-9]{0,8})$ ]] || die "Expected a nonnegative decimal integer: $value"
done
((PORT > 0 && PORT <= 65535 && NUM_TXNS > 0 && LOCK_MS > 0)) || die "Invalid port, transaction count, or lock timeout"
((MAX_RETRIES <= 20 && HOLD_MS <= 10000 && LOCK_MS <= 20000)) || die "Use retries <= 20, hold-ms <= 10000, lock-timeout-ms <= 20000"
read -r -a NODE_ARRAY <<< "$NODES"
APP_COUNT=${#NODE_ARRAY[@]}
((APP_COUNT > 0)) || die "At least one endpoint is required"
FIRST_NODE=${NODE_ARRAY[0]}
for node in "${NODE_ARRAY[@]}"; do
    [[ $node != -* && $node != *','* ]] || die "Supply one hostname or IP per endpoint"
done
if [[ -z $SSL_MODE ]]; then
    SSL_MODE=disable
    [[ -z $CERTS ]] || SSL_MODE=verify-full
fi
case $SSL_MODE in disable|require|verify-ca|verify-full) ;; *) die "Invalid sslmode: $SSL_MODE" ;; esac
for cert_file in "$CERTS" "$SSL_CERT" "$SSL_KEY"; do
    [[ -z $cert_file || -r $cert_file ]] || die "Cannot read certificate/key: $cert_file"
done
[[ -z $SSL_CERT && -z $SSL_KEY || -n $SSL_CERT && -n $SSL_KEY ]] || die "Provide both --sslcert and --sslkey"
command -v "$CLIENT" >/dev/null || die "Client not found: $CLIENT"

# Explicit connection settings override unrelated shell libpq defaults.
# Password / passfile settings are intentionally inherited.
sql() {
    local node=$1; shift
    env -u PGHOSTADDR -u PGSERVICE -u PGSERVICEFILE -u PGOPTIONS \
        PGSSLMODE="$SSL_MODE" PGSSLROOTCERT="$CERTS" \
        PGSSLCERT="$SSL_CERT" PGSSLKEY="$SSL_KEY" PGCONNECT_TIMEOUT=10 \
        PGAPPNAME=yb_multi_app_simulator PGTZ=UTC \
        "$CLIENT" -X -w -q -h "$node" -p "$PORT" -U "$DB_USER" -d "$DB" \
        -v ON_ERROR_STOP=1 -v VERBOSITY=verbose -P pager=off "$@"
}

printf 'YugabyteDB Multi-App Concurrency Simulator\n'
printf 'Workers: %s | Transactions per worker: %s | Hold: %s ms\n' "$APP_COUNT" "$NUM_TXNS" "$HOLD_MS"
sql "$FIRST_NODE" -c 'SELECT host, cloud, region, zone FROM yb_servers() ORDER BY host;'
for node in "${NODE_ARRAY[@]}"; do
    sql "$node" -c "SELECT version(), yb_server_region() AS server_region;" >/dev/null
done

# The narrow exception handler catches only the intended target-row collision.
sql "$FIRST_NODE" <<'SQL'
CREATE SCHEMA IF NOT EXISTS yb_multi_app_demo;
CREATE TABLE IF NOT EXISTS yb_multi_app_demo.target_rows (
    run_id uuid NOT NULL, id int NOT NULL,
    visit_count int NOT NULL CHECK (visit_count > 0),
    PRIMARY KEY ((run_id, id) HASH)
);
CREATE TABLE IF NOT EXISTS yb_multi_app_demo.events (
    run_id uuid NOT NULL, target_id int NOT NULL,
    step int NOT NULL, event_no int NOT NULL,
    app_id int NOT NULL, node_endpoint text NOT NULL,
    server_region text NOT NULL, operation text NOT NULL,
    event_time timestamptz NOT NULL DEFAULT clock_timestamp(),
    isolation_level text NOT NULL DEFAULT current_setting('transaction_isolation'),
    PRIMARY KEY ((run_id, target_id) HASH, step ASC, event_no ASC)
);
CREATE TABLE IF NOT EXISTS yb_multi_app_demo.attempt_errors (
    run_id uuid NOT NULL, app_id int NOT NULL, target_id int NOT NULL,
    attempt int NOT NULL, node_endpoint text NOT NULL,
    server_region text NOT NULL DEFAULT yb_server_region(),
    event_time timestamptz NOT NULL DEFAULT clock_timestamp(),
    sqlstate text NOT NULL, status text NOT NULL, detail text NOT NULL,
    PRIMARY KEY ((run_id, app_id) HASH, target_id ASC, attempt ASC)
);
CREATE OR REPLACE FUNCTION yb_multi_app_demo.apply_one(
    p_run uuid, p_app int, p_node text, p_id int, p_apps int, p_hold int
) RETURNS void LANGUAGE plpgsql AS $$
DECLARE
    v_count int := 1;
    v_collision boolean := false;
    v_region text := yb_server_region();
BEGIN
    BEGIN
        INSERT INTO yb_multi_app_demo.target_rows VALUES (p_run, p_id, 1);
    EXCEPTION WHEN unique_violation THEN
        v_collision := true;
    END;

    IF v_collision THEN
        UPDATE yb_multi_app_demo.target_rows
        SET visit_count = visit_count + 1
        WHERE run_id = p_run AND id = p_id
        RETURNING visit_count INTO v_count;
        IF NOT FOUND THEN
            RAISE EXCEPTION 'Expected target row is missing' USING ERRCODE = 'P0001';
        END IF;
        INSERT INTO yb_multi_app_demo.events
            (run_id, target_id, step, event_no, app_id, node_endpoint, server_region, operation)
        VALUES (p_run, p_id, v_count, 1, p_app, p_node, v_region, 'PK_COLLISION');
    END IF;

    INSERT INTO yb_multi_app_demo.events
        (run_id, target_id, step, event_no, app_id, node_endpoint, server_region, operation)
    VALUES (p_run, p_id, v_count, 2, p_app, p_node, v_region,
            CASE WHEN v_collision THEN 'UPDATE' ELSE 'INSERT' END);

    IF v_count = p_apps THEN
        DELETE FROM yb_multi_app_demo.target_rows WHERE run_id = p_run AND id = p_id;
        INSERT INTO yb_multi_app_demo.events
            (run_id, target_id, step, event_no, app_id, node_endpoint, server_region, operation)
        VALUES (p_run, p_id, v_count, 3, p_app, p_node, v_region, 'DELETE');
    ELSIF v_count > p_apps THEN
        RAISE EXCEPTION 'Visit count exceeded worker count';
    END IF;
    PERFORM pg_sleep(p_hold / 1000.0);
END;
$$;
SQL

RUN_ID=$(sql "$FIRST_NODE" -Atc 'SELECT gen_random_uuid();')
printf 'Run ID: %s\n' "$RUN_ID"
WORK_DIR=$(mktemp -d)
PIDS=()
cleanup() { rm -rf -- "$WORK_DIR"; }
stop_workers() {
    for pid in "${PIDS[@]}"; do kill "$pid" 2>/dev/null || true; done
    wait || true
    exit 130
}
trap cleanup EXIT
trap stop_workers INT TERM

worker() {
    local node=$1 app=$2 row attempt state status delay cap output
    while [[ ! -f $WORK_DIR/start ]]; do sleep 0.02; done
    for ((row=1; row<=NUM_TXNS; row++)); do
        for ((attempt=1; attempt<=MAX_RETRIES+1; attempt++)); do
            if output=$(sql "$node" -v run="$RUN_ID" -v app="$app" -v node="$node" \
                -v row="$row" -v apps="$APP_COUNT" -v hold="$HOLD_MS" -v lock_ms="$LOCK_MS" 2>&1 <<'SQL'
BEGIN ISOLATION LEVEL READ COMMITTED;
SET LOCAL statement_timeout = '30s';
SET LOCAL lock_timeout = :'lock_ms';
SELECT yb_multi_app_demo.apply_one(:'run'::uuid, :app, :'node', :row, :apps, :hold);
COMMIT;
SQL
            ); then
                break
            fi
            state=$(printf '%s\n' "$output" | sed -nE 's/^.*(ERROR|FATAL):[[:space:]]+([A-Z0-9]{5}):.*/\2/p' | head -n 1)
            state=${state:-UNKNOWN}
            status=FAILED
            case $state in
                40001|40P01|55P03) ((attempt > MAX_RETRIES)) || status=RETRY ;;
            esac
            # A fresh connection/transaction keeps this log after the failed attempt rolls back.
            sql "$node" -v run="$RUN_ID" -v app="$app" -v node="$node" -v row="$row" \
                -v attempt="$attempt" -v state="$state" -v status="$status" -v detail="$output" <<'SQL' || {
INSERT INTO yb_multi_app_demo.attempt_errors
    (run_id, app_id, target_id, attempt, node_endpoint, sqlstate, status, detail)
VALUES (:'run'::uuid, :app, :row, :attempt, :'node', :'state', :'status', :'detail');
SQL
                printf 'App %s: could not persist attempt error\n' "$app" >&2
            }
            printf 'App %s key %s attempt %s: %s [%s]\n' "$app" "$row" "$attempt" "$state" "$status" >&2
            if [[ $status != RETRY ]]; then
                printf '%s\n' "$output" >&2
                return 1
            fi
            cap=$((50 * (1 << (attempt < 6 ? attempt - 1 : 5))))
            printf -v delay '%d.%03d' "$((cap / 1000))" "$((cap % 1000 + RANDOM % 100))"
            sleep "$delay"
        done
    done
    printf 'App %s completed %s transactions through %s\n' "$app" "$NUM_TXNS" "$node"
}

START_SECONDS=$SECONDS
for ((app=1; app<=APP_COUNT; app++)); do
    worker "${NODE_ARRAY[app-1]}" "$app" &
    PIDS+=("$!")
done
touch "$WORK_DIR/start"
FAILED_WORKERS=0
for pid in "${PIDS[@]}"; do
    if ! wait "$pid"; then FAILED_WORKERS=$((FAILED_WORKERS + 1)); fi
done
printf 'Elapsed workload time: %s seconds | Failed workers: %s\n' "$((SECONDS - START_SECONDS))" "$FAILED_WORKERS"

sql "$FIRST_NODE" -v run="$RUN_ID" <<'SQL'
SELECT operation, count(*) AS committed_events
FROM yb_multi_app_demo.events WHERE run_id = :'run'::uuid
GROUP BY operation ORDER BY operation;

SELECT app_id, node_endpoint, server_region,
       count(*) FILTER (WHERE operation = 'INSERT') AS inserts,
       count(*) FILTER (WHERE operation = 'UPDATE') AS updates,
       count(*) FILTER (WHERE operation = 'DELETE') AS deletes,
       count(*) FILTER (WHERE operation = 'PK_COLLISION') AS pk_collisions
FROM yb_multi_app_demo.events WHERE run_id = :'run'::uuid
GROUP BY app_id, node_endpoint, server_region ORDER BY app_id, server_region;

SELECT target_id, step, event_no, app_id, server_region, operation, event_time
FROM yb_multi_app_demo.events WHERE run_id = :'run'::uuid AND target_id <= 3
ORDER BY target_id, step, event_no;

SELECT app_id, target_id, attempt, sqlstate, status, event_time
FROM yb_multi_app_demo.attempt_errors WHERE run_id = :'run'::uuid
ORDER BY app_id, target_id, attempt;
SQL

CHECK=$(sql "$FIRST_NODE" -At -v run="$RUN_ID" -v apps="$APP_COUNT" -v n="$NUM_TXNS" <<'SQL'
WITH e AS (
    SELECT * FROM yb_multi_app_demo.events WHERE run_id = :'run'::uuid
), per_key AS (
    SELECT target_id,
        count(*) FILTER (WHERE operation = 'INSERT') AS ins,
        count(*) FILTER (WHERE operation = 'UPDATE') AS upd,
        count(*) FILTER (WHERE operation = 'DELETE') AS del,
        count(*) FILTER (WHERE operation = 'PK_COLLISION') AS pk,
        count(DISTINCT app_id) FILTER (WHERE operation IN ('INSERT','UPDATE')) AS apps,
        count(DISTINCT step) FILTER (WHERE operation IN ('INSERT','UPDATE')) AS steps,
        min(step) AS first_step, max(step) AS last_step
    FROM e GROUP BY target_id
), per_app AS (
    SELECT app_id, count(*) AS completed
    FROM e WHERE operation IN ('INSERT','UPDATE') GROUP BY app_id
)
SELECT CASE WHEN
    (SELECT count(*) FROM per_key) = :n
    AND NOT EXISTS (SELECT 1 FROM per_key WHERE target_id NOT BETWEEN 1 AND :n
        OR ins <> 1 OR upd <> :apps - 1 OR del <> 1 OR pk <> :apps - 1
        OR apps <> :apps OR steps <> :apps OR first_step <> 1 OR last_step <> :apps)
    AND (SELECT count(*) FROM per_app) = :apps
    AND NOT EXISTS (SELECT 1 FROM per_app WHERE app_id NOT BETWEEN 1 AND :apps OR completed <> :n)
    AND NOT EXISTS (SELECT 1 FROM yb_multi_app_demo.target_rows WHERE run_id = :'run'::uuid)
    THEN 'PASS' ELSE 'FAIL' END;
SQL
)
if [[ $CHECK == PASS && $FAILED_WORKERS == 0 ]]; then
    printf 'PASS: %s worker transactions committed; every key reached %s visits and was deleted.\n' \
        "$((APP_COUNT * NUM_TXNS))" "$APP_COUNT"
    printf 'Expected/verified: %s INSERTs, %s UPDATEs, %s DELETEs, %s PK collisions, 0 remaining rows.\n' \
        "$NUM_TXNS" "$(((APP_COUNT - 1) * NUM_TXNS))" "$NUM_TXNS" "$(((APP_COUNT - 1) * NUM_TXNS))"
else
    die "FAIL: worker errors or workload invariants failed; inspect run $RUN_ID"
fi