|
| 1 | +// SPDX-License-Identifier: BUSL-1.1 |
| 2 | + |
| 3 | +//! Regression test: an abandoned PGWIRE-protocol transaction's Data-Plane |
| 4 | +//! staging overlays (`txn_overlays` / `graph_txn_overlays`, keyed by |
| 5 | +//! `txn_id`) leaked forever if the connection ended (abrupt disconnect, |
| 6 | +//! idle/absolute timeout, or clean EOF) without a `COMMIT`/`ROLLBACK` -- |
| 7 | +//! nothing dispatched `MetaOp::DropTxnOverlay` for it. The pgwire listener |
| 8 | +//! now calls `NodeDbPgHandlerFactory::on_connection_end` on every connection |
| 9 | +//! exit, which reclaims any still-open transaction's overlay via |
| 10 | +//! `lifecycle::run_rollback` and drops the shared session entry. |
| 11 | +//! |
| 12 | +//! This test proves reclamation (rather than just "row not visible", which |
| 13 | +//! cannot distinguish reclaimed from leaked-but-invisible) via the |
| 14 | +//! `active_txn_overlays` gauge on `SystemMetrics`: it rises when the |
| 15 | +//! transaction's first staged write materializes the overlay, and must fall |
| 16 | +//! back to baseline once the abandoned connection is torn down. |
| 17 | +
|
| 18 | +mod common; |
| 19 | + |
| 20 | +use std::sync::atomic::Ordering; |
| 21 | +use std::time::{Duration, Instant}; |
| 22 | + |
| 23 | +use common::pgwire_harness::TestServer; |
| 24 | + |
| 25 | +/// Poll `active_txn_overlays` until `pred` is satisfied or `deadline` elapses. |
| 26 | +/// Returns the last observed value; callers assert against it so a timeout |
| 27 | +/// panics with the actual value instead of a bare "timed out". |
| 28 | +async fn poll_gauge(server: &TestServer, deadline: Duration, pred: impl Fn(u64) -> bool) -> u64 { |
| 29 | + let metrics = server |
| 30 | + .shared |
| 31 | + .system_metrics |
| 32 | + .as_ref() |
| 33 | + .expect("system_metrics must be wired into the pgwire harness"); |
| 34 | + let start = Instant::now(); |
| 35 | + loop { |
| 36 | + let value = metrics.active_txn_overlays.load(Ordering::Relaxed); |
| 37 | + if pred(value) || start.elapsed() >= deadline { |
| 38 | + return value; |
| 39 | + } |
| 40 | + tokio::time::sleep(Duration::from_millis(20)).await; |
| 41 | + } |
| 42 | +} |
| 43 | + |
| 44 | +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] |
| 45 | +async fn pgwire_abandoned_txn_overlay_reclaimed_on_teardown() { |
| 46 | + let server = TestServer::start().await; |
| 47 | + |
| 48 | + let baseline = server |
| 49 | + .shared |
| 50 | + .system_metrics |
| 51 | + .as_ref() |
| 52 | + .expect("system_metrics must be wired into the pgwire harness") |
| 53 | + .active_txn_overlays |
| 54 | + .load(Ordering::Relaxed); |
| 55 | + assert_eq!(baseline, 0, "gauge must start at zero: {baseline}"); |
| 56 | + |
| 57 | + // Open a SEPARATE connection the test OWNS (not `server.client`, which the |
| 58 | + // harness owns) so we can close it mid-transaction. tokio-postgres runs the |
| 59 | + // socket in a spawned task; we keep its JoinHandle to abort it later. |
| 60 | + let conn_str = format!( |
| 61 | + "host=127.0.0.1 port={} user=nodedb dbname=nodedb", |
| 62 | + server.pg_port |
| 63 | + ); |
| 64 | + let (client, connection) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls) |
| 65 | + .await |
| 66 | + .expect("owned pgwire connection must connect"); |
| 67 | + let conn_handle = tokio::spawn(async move { |
| 68 | + let _ = connection.await; |
| 69 | + }); |
| 70 | + |
| 71 | + client |
| 72 | + .batch_execute( |
| 73 | + "CREATE COLLECTION pg_txn_overlay_teardown (id STRING PRIMARY KEY, n INT) \ |
| 74 | + WITH (engine='document_schemaless')", |
| 75 | + ) |
| 76 | + .await |
| 77 | + .expect("CREATE must succeed"); |
| 78 | + |
| 79 | + client |
| 80 | + .batch_execute("BEGIN") |
| 81 | + .await |
| 82 | + .expect("BEGIN must succeed"); |
| 83 | + client |
| 84 | + .batch_execute("INSERT INTO pg_txn_overlay_teardown (id, n) VALUES ('a', 1)") |
| 85 | + .await |
| 86 | + .expect("in-tx INSERT must succeed"); |
| 87 | + |
| 88 | + // The staged write must have materialized the overlay, bumping the gauge |
| 89 | + // above baseline -- confirms there is something to reclaim. |
| 90 | + let after_stage = poll_gauge(&server, Duration::from_secs(5), |v| v > baseline).await; |
| 91 | + assert!( |
| 92 | + after_stage > baseline, |
| 93 | + "staged write must raise active_txn_overlays above baseline {baseline}, got {after_stage}" |
| 94 | + ); |
| 95 | + |
| 96 | + // Abruptly abandon the connection -- no COMMIT/ROLLBACK. Dropping the |
| 97 | + // Client resolves the Connection future; aborting its task guarantees the |
| 98 | + // socket is dropped so the server sees EOF and runs `on_connection_end`. |
| 99 | + drop(client); |
| 100 | + conn_handle.abort(); |
| 101 | + let _ = conn_handle.await; |
| 102 | + |
| 103 | + // The abandoned transaction's overlay must be reclaimed on teardown, |
| 104 | + // bringing the gauge back down to baseline. |
| 105 | + let after_teardown = poll_gauge(&server, Duration::from_secs(5), |v| v == baseline).await; |
| 106 | + assert_eq!( |
| 107 | + after_teardown, baseline, |
| 108 | + "abandoned txn overlay must be reclaimed on connection teardown, \ |
| 109 | + active_txn_overlays still {after_teardown} (baseline {baseline})" |
| 110 | + ); |
| 111 | + |
| 112 | + // Belt-and-suspenders: the harness-owned autocommit connection must not see |
| 113 | + // the staged row -- it was never committed. |
| 114 | + let rows = server |
| 115 | + .query_text("SELECT n FROM pg_txn_overlay_teardown WHERE id = 'a'") |
| 116 | + .await |
| 117 | + .expect("post-teardown SELECT must succeed"); |
| 118 | + assert!( |
| 119 | + rows.is_empty(), |
| 120 | + "the never-committed staged row must not be visible: {rows:?}" |
| 121 | + ); |
| 122 | +} |
0 commit comments