diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9be05e8..543d8b7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -205,6 +205,10 @@ jobs: id: shutdown_test run: ./scripts/test-shutdown.sh --pg-version ${{ matrix.pg_version }} --cycles 1 + - name: Run epoch-race regression test + id: epoch_race_test + run: ./scripts/test-epoch-race.sh --pg-version ${{ matrix.pg_version }} + - name: Upload PostgreSQL logs on E2E failure if: failure() uses: actions/upload-artifact@v4 diff --git a/scripts/test-epoch-race.sh b/scripts/test-epoch-race.sh new file mode 100755 index 0000000..57ca49b --- /dev/null +++ b/scripts/test-epoch-race.sh @@ -0,0 +1,278 @@ +#!/bin/bash +# Copyright (c) Microsoft Corporation. +# Licensed under the PostgreSQL License. + +# test-epoch-race.sh - Regression test for the extension epoch race that could +# certify a stale duroxide runtime across DROP/CREATE EXTENSION (GitHub issue #333). +# +# The background worker used to write its readiness record and epoch sentinel only +# after runtime initialization. If a DROP EXTENSION / CREATE EXTENSION replaced the +# provider objects mid-initialization, the sentinel landed in the NEW epoch's `df` +# schema, so the stale runtime — bound to provider objects that no longer existed — +# was incorrectly certified as current and never restarted. Readiness then timed out. +# +# This test makes the race deterministic using a build-in test hook +# (PG_DURABLE_TEST_PAUSE_BEFORE_READY_MS) that pauses the worker between epoch +# capture / runtime initialization and readiness publication. During that pause the +# test drops and recreates the extension, then verifies: +# 1. The worker tears down the stale runtime and reinitializes (log evidence). +# 2. The worker eventually becomes ready. +# 3. `_worker_ready` describes the CURRENT provider epoch (a durable function runs). +# +# Usage: ./scripts/test-epoch-race.sh [options] +# +# Options: +# --pg-version VER PostgreSQL major version (default: 17) +# --pause-ms MS Worker init->ready pause window (default: 4000) +# --verbose, -v Show detailed output + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" + +PG_VERSION="${PG_VERSION:-17}" +PAUSE_MS=4000 +VERBOSE=false + +while [[ $# -gt 0 ]]; do + case "$1" in + --pg-version) + PG_VERSION="$2"; shift 2 ;; + --pause-ms) + PAUSE_MS="$2"; shift 2 ;; + --verbose|-v) + VERBOSE=true; shift ;; + --help|-h) + sed -n '/^# Usage:/,/^[^#]/{ /^[^#]/d; s/^# \{0,1\}//; p }' "$0" + exit 0 ;; + *) + echo "Unknown option: $1"; exit 1 ;; + esac +done + +PG_PORT="$((28800 + PG_VERSION))" +PGRX_HOME="$HOME/.pgrx" +DATA_DIR="$PGRX_HOME/data-$PG_VERSION" +LOG_FILE="$PGRX_HOME/$PG_VERSION.log" + +shopt -s nullglob +PGRX_CANDIDATES=("$PGRX_HOME"/"$PG_VERSION".*/pgrx-install/bin) +shopt -u nullglob +if [ "${#PGRX_CANDIDATES[@]}" -eq 0 ]; then + echo "Error: pgrx PostgreSQL $PG_VERSION not installed (run: cargo pgrx init)" + exit 1 +fi + +PGRX_BIN="${PGRX_CANDIDATES[0]}" +PSQL="$PGRX_BIN/psql" +PG_CTL="$PGRX_BIN/pg_ctl" +PG_ISREADY="$PGRX_BIN/pg_isready" +PG_CONFIG="$PGRX_BIN/pg_config" + +RED='\033[0;31m' +GREEN='\033[0;32m' +CYAN='\033[0;36m' +NC='\033[0m' + +PASS=0 +FAIL=0 + +log() { echo -e "$*"; } +info() { log "${CYAN}$*${NC}"; } +ok() { log "${GREEN}[PASS]${NC} $*"; PASS=$((PASS + 1)); } +fail() { log "${RED}[FAIL]${NC} $*"; FAIL=$((FAIL + 1)); } + +# --------------------------------------------------------------------------- +# Server helpers +# --------------------------------------------------------------------------- + +configure_standard() { + local conf="$DATA_DIR/postgresql.conf" + sed -i.bak '/^[#[:space:]]*shared_preload_libraries/d' "$conf" + sed -i.bak '/^[#[:space:]]*pg_durable\./d' "$conf" + rm -f "$conf.bak" + : > "$DATA_DIR/postgresql.auto.conf" + cat >> "$conf" </dev/null 2>&1 + local attempts=0 + until "$PG_ISREADY" -h localhost -p "$PG_PORT" -U postgres -q >/dev/null 2>&1; do + attempts=$((attempts + 1)) + if [ "$attempts" -ge 60 ]; then + echo "PostgreSQL did not become ready on port $PG_PORT" + return 1 + fi + sleep 0.5 + done +} + +stop_server() { + "$PG_CTL" -D "$DATA_DIR" stop -m fast -t 20 >/dev/null 2>&1 || \ + "$PG_CTL" -D "$DATA_DIR" stop -m immediate >/dev/null 2>&1 || true +} + +wait_for_worker() { + local attempts=0 + local dx_schema ready + while [ "$attempts" -lt 120 ]; do + dx_schema=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres \ + -Atqc "SELECT df.duroxide_schema();" 2>/dev/null | tr -d ' \n' || echo "_duroxide") + ready=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres \ + -Atqc "SELECT COALESCE((SELECT TRUE FROM ${dx_schema}._worker_ready WHERE schema_version >= 1), FALSE);" \ + 2>/dev/null | tr -d ' \n' || echo "f") + [ "$ready" = "t" ] && return 0 + sleep 0.5 + attempts=$((attempts + 1)) + done + echo "Worker did not become ready within 60 s" + return 1 +} + +ensure_extension() { + "$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres >/dev/null 2>&1 <<'SQL' +BEGIN; +DROP EXTENSION IF EXISTS pg_durable CASCADE; +CREATE EXTENSION pg_durable; +COMMIT; +SQL +} + +# Run a durable function end-to-end. Success proves the worker's runtime is bound to +# the CURRENT provider epoch: the fetch_work_item / fetch_orchestration_item provider +# functions and their backing tables must exist and be the ones this runtime polls. +run_durable_function() { + local inst status attempts=0 + inst=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT df.start(df.sql('SELECT 42'), 'epoch-race-check');" 2>/dev/null | tr -d ' \n') + [ -n "$inst" ] || return 1 + while [ "$attempts" -lt 300 ]; do + status=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT lower(s) FROM df.status('$inst') s;" 2>/dev/null | tr -d ' \n') + case "$status" in + completed) return 0 ;; + failed|cancelled) return 1 ;; + esac + sleep 0.1 + attempts=$((attempts + 1)) + done + return 1 +} + +cleanup() { + stop_server +} +trap cleanup EXIT + +# --------------------------------------------------------------------------- +# Build +# --------------------------------------------------------------------------- + +info "Building pg_durable extension..." +cd "$PROJECT_DIR" +if ! cargo pgrx install --pg-config="$PG_CONFIG" --features http-allow-test-domains \ + > /tmp/pg_durable-epoch-race-build.log 2>&1; then + echo -e "${RED}Build failed:${NC}" + cat /tmp/pg_durable-epoch-race-build.log + exit 1 +fi +info "Build complete" + +if [ ! -d "$DATA_DIR" ]; then + "$PGRX_BIN/initdb" -D "$DATA_DIR" -U postgres --no-locale -E UTF8 >/dev/null 2>&1 +fi +configure_standard + +# --------------------------------------------------------------------------- +# Scenario: DROP/CREATE EXTENSION during the worker's init->ready pause +# --------------------------------------------------------------------------- + +info "=== Epoch race: drop/recreate extension during worker initialization ===" + +# 1. Start without the pause and install the extension so the worker is healthy. +info "Starting server and installing extension..." +start_server "" +ensure_extension +if ! wait_for_worker; then + fail "Baseline worker readiness failed before the race scenario" + stop_server + exit 1 +fi +ok "Baseline worker is ready" + +# 2. Restart with the pause enabled. On restart the worker captures the current +# epoch, initializes a runtime, then pauses before publishing readiness. +info "Restarting server with a ${PAUSE_MS}ms init->ready pause..." +stop_server +# Record where the current log ends so we only inspect messages from this restart. +LOG_MARK=$(wc -l < "$LOG_FILE" 2>/dev/null || echo 0) +start_server "$PAUSE_MS" + +# 3. During the pause window, replace the extension epoch. The just-initialized +# runtime is now stale: its provider objects have been dropped and recreated. +info "Dropping and recreating the extension mid-initialization..." +# Small delay so the worker has reached the pause; well within the pause window. +sleep 1 +ensure_extension + +# 4. The worker must tear down the stale runtime and reinitialize against the new +# epoch, then become ready. +info "Waiting for worker to recover and become ready..." +if wait_for_worker; then + ok "Worker became ready after drop/recreate during initialization" +else + fail "Worker did NOT become ready after drop/recreate (stale-runtime race)" +fi + +# 5. Log evidence that the stale runtime was detected and torn down. +if tail -n +"$((LOG_MARK + 1))" "$LOG_FILE" 2>/dev/null \ + | grep -q "extension epoch changed"; then + ok "Worker logged stale-runtime teardown ('extension epoch changed')" +else + fail "Expected 'extension epoch changed' teardown log not found" +fi + +# 6. Prove readiness describes the CURRENT epoch: a durable function must complete, +# which requires the runtime to poll live provider functions/tables. +info "Running a durable function against the current epoch..." +if run_durable_function; then + ok "Durable function completed — readiness reflects the current provider epoch" +else + fail "Durable function did not complete — runtime not bound to current epoch" +fi + +if [ "$VERBOSE" = true ]; then + info "--- Worker log (this restart) ---" + tail -n +"$((LOG_MARK + 1))" "$LOG_FILE" 2>/dev/null | grep "pg_durable:" || true +fi + +stop_server + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- + +echo +info "=== Summary ===" +log " Passed: $PASS" +log " Failed: $FAIL" + +if [ "$FAIL" -gt 0 ]; then + echo -e "${RED}EPOCH RACE TEST FAILED${NC}" + exit 1 +fi +echo -e "${GREEN}EPOCH RACE TEST PASSED${NC}" diff --git a/src/worker.rs b/src/worker.rs index 934f4e4..df992f8 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -252,6 +252,13 @@ async fn run_duroxide_runtime() { duroxide_schema ); + // Capture the extension epoch identity BEFORE initializing the runtime. + // A DROP/CREATE during initialization installs a new epoch; readiness must + // only be published for a runtime initialized against this same epoch, + // otherwise a stale runtime would poll provider objects that no longer + // exist while being certified as current. + let epoch_oid = capture_extension_epoch(&poll_pool).await; + let Some((duroxide_runtime, duroxide_store)) = initialize_duroxide_runtime( &pg_conn_str, INIT_RETRY_INTERVAL, @@ -264,6 +271,28 @@ async fn run_duroxide_runtime() { continue; }; + // Test-only pause that widens the window between runtime initialization + // and readiness publication so a regression test can drop/recreate the + // extension deterministically. No-op in production. + test_pause_before_ready().await; + + // Revalidate the epoch AFTER initialization and BEFORE publishing + // readiness. If the extension was dropped/recreated during init (or the + // pre-init capture failed), tear down this now-stale runtime and retry + // initialization against the current epoch. This prevents certifying a + // runtime bound to provider objects that no longer exist. + let current_epoch = capture_extension_epoch(&poll_pool).await; + if epoch_oid.is_none() || current_epoch != epoch_oid { + log!( + "pg_durable: extension epoch changed during initialization \ + (before={:?}, after={:?}) — tearing down stale runtime and retrying", + epoch_oid, + current_epoch + ); + teardown_runtime(duroxide_runtime, duroxide_store).await; + continue; + } + // Write the worker readiness record so backend sessions know the // duroxide schema is fully initialized for this schema version. // Skipped if the row already has the current WORKER_SCHEMA_VERSION. @@ -284,6 +313,20 @@ async fn run_duroxide_runtime() { } }; + // Final epoch check after readiness/sentinel writes and before entering the + // processing loop. A drop/recreate in this narrow window would have written + // the readiness record and sentinel into the new epoch's `df` schema, so the + // processing loop's sentinel check could not detect the replacement. Catch it + // here: tear down and retry against the current epoch. + if capture_extension_epoch(&poll_pool).await != epoch_oid { + log!( + "pg_durable: extension epoch changed after readiness publication \ + — tearing down stale runtime and retrying" + ); + teardown_runtime(duroxide_runtime, duroxide_store).await; + continue; + } + run_until_extension_dropped_or_shutdown( &poll_pool, &mgmt_pool, @@ -341,6 +384,38 @@ async fn check_extension_exists(pool: &sqlx::PgPool) -> bool { result.map(|(exists,)| exists).unwrap_or(false) } +/// Identifies a single extension install ("epoch"). `pg_extension.oid` is a fresh +/// value for every CREATE EXTENSION, so it changes across a DROP/CREATE cycle even +/// when the extension appears continuously present between polls. Returns `None` +/// when the extension is absent or the query fails. +async fn capture_extension_epoch(pool: &sqlx::PgPool) -> Option { + let result: Result<(i64,), sqlx::Error> = + sqlx::query_as("SELECT oid::bigint FROM pg_extension WHERE extname = 'pg_durable'") + .fetch_one(pool) + .await; + + result.map(|(oid,)| oid).ok() +} + +/// Test-only hook: when the `PG_DURABLE_TEST_PAUSE_BEFORE_READY_MS` environment +/// variable is set to a positive integer, sleep that many milliseconds between +/// runtime initialization and readiness publication. This creates a deterministic +/// window for a regression test to DROP/CREATE the extension mid-initialization +/// and exercise the stale-runtime detection path. No effect in production. +async fn test_pause_before_ready() { + if let Ok(val) = std::env::var("PG_DURABLE_TEST_PAUSE_BEFORE_READY_MS") { + if let Ok(ms) = val.parse::() { + if ms > 0 { + log!( + "pg_durable: TEST hook — pausing {}ms before readiness publication", + ms + ); + tokio::time::sleep(Duration::from_millis(ms)).await; + } + } + } +} + /// Returns true if the `duroxide` schema exists AND is owned by the `pg_durable` /// extension (dependency type 'e' in pg_depend). /// @@ -947,8 +1022,20 @@ async fn run_until_extension_dropped_or_shutdown( } } - log!("pg_durable: initiating duroxide runtime shutdown..."); + teardown_runtime(duroxide_runtime, duroxide_store).await; + log!("pg_durable: duroxide runtime shutdown complete"); +} +/// Shut down a duroxide runtime and close its store pool. +/// +/// Used both when exiting the processing loop and when tearing down a runtime that +/// was initialized against a stale extension epoch (drop/recreate mid-init). The +/// live path (extension restart) drains and reclaims connections; the shutdown +/// path closes first because PostgreSQL is terminating this process's backends. +async fn teardown_runtime( + duroxide_runtime: Arc, + duroxide_store: Arc, +) { if is_shutdown_requested() { // Close before aborting — the reverse of the branch below. PostgreSQL is // killing this process's backends, so duroxide's dispatcher tasks stay @@ -982,7 +1069,6 @@ async fn run_until_extension_dropped_or_shutdown( log!("pg_durable: duroxide store pool close timed out — forcing shutdown"); } } - log!("pg_durable: duroxide runtime shutdown complete"); } /// Maximum orphans one reconciliation pass reclaims, bounding a single tick's work.