Skip to content

Commit e57e68e

Browse files
committed
fix(control): reclaim native transaction overlays on session teardown
An abandoned NATIVE-protocol transaction's Data-Plane staging overlays leaked forever if the connection ended (abrupt disconnect, idle/absolute timeout, or clean EOF) without a COMMIT/ROLLBACK, since nothing dispatched a rollback for it. NativeSession::run now wraps the frame loop and reclaims any still-open transaction's overlay on every exit path. Also exposes shared Control-Plane state on the native test harness so tests can observe the active_txn_overlays gauge and verify reclamation rather than just row invisibility.
1 parent cc7bef1 commit e57e68e

4 files changed

Lines changed: 194 additions & 2 deletions

File tree

nodedb-test-support/src/native_harness/server.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ use nodedb::wal::WalManager;
1717
/// A running native-protocol test server.
1818
pub struct NativeTestServer {
1919
pub addr: std::net::SocketAddr,
20+
/// Shared Control-Plane state, exposed so tests can read metrics (e.g.
21+
/// `shared.system_metrics.active_txn_overlays`) that the Data-Plane
22+
/// core also updates via the same `Arc<SystemMetrics>`.
23+
pub shared: Arc<SharedState>,
2024
pub(super) shutdown_bus: nodedb::control::shutdown::ShutdownBus,
2125
pub(super) poller_shutdown_tx: tokio::sync::watch::Sender<bool>,
2226
pub(super) core_stop_tx: std::sync::mpsc::Sender<()>,
@@ -65,6 +69,7 @@ impl NativeTestServer {
6569
let core_dir = dir.path().to_path_buf();
6670
let event_producer = event_producers.into_iter().next().expect("event producer");
6771
let core_array_catalog = shared.array_catalog.clone();
72+
let core_metrics = shared.system_metrics.clone();
6873
let (core_stop_tx, core_stop_rx) = std::sync::mpsc::channel::<()>();
6974
let _core_handle = tokio::task::spawn_blocking(move || {
7075
let mut core = CoreLoop::open_with_array_catalog(
@@ -77,6 +82,9 @@ impl NativeTestServer {
7782
)
7883
.expect("open core");
7984
core.set_event_producer(event_producer);
85+
if let Some(m) = core_metrics {
86+
core.set_metrics(m);
87+
}
8088
while matches!(
8189
core_stop_rx.try_recv(),
8290
Err(std::sync::mpsc::TryRecvError::Empty)
@@ -145,6 +153,7 @@ impl NativeTestServer {
145153

146154
Self {
147155
addr,
156+
shared: Arc::clone(&shared),
148157
shutdown_bus,
149158
poller_shutdown_tx,
150159
core_stop_tx,

nodedb/src/control/server/native/dispatch/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,4 @@ pub(crate) use direct_ops::{handle_direct_op, handle_graph_match};
2727
pub(crate) use session_ops::{handle_reset, handle_set, handle_show};
2828
pub(crate) use sql::{handle_sql, handle_sql_streaming};
2929
pub(crate) use streaming::{SqlOutcome, SqlStream};
30-
pub(crate) use transaction::{handle_begin, handle_commit, handle_rollback};
30+
pub(crate) use transaction::{NativeTxnDp, handle_begin, handle_commit, handle_rollback};

nodedb/src/control/server/native/session/run.rs

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,45 @@ use super::dispatch;
1515
use super::{NativeSession, chunk_large_response};
1616

1717
impl NativeSession {
18+
/// Run the session: drives the frame loop, then reclaims any
19+
/// still-open transaction's Data-Plane overlays on every exit path
20+
/// (clean EOF, idle/absolute timeout, or abrupt disconnect/error).
21+
pub async fn run(mut self) -> crate::Result<()> {
22+
let result = self.run_loop().await;
23+
self.reclaim_open_txn().await;
24+
result
25+
}
26+
27+
/// Reclaim a still-open transaction's Data-Plane overlays if the
28+
/// connection ended mid-transaction (no COMMIT/ROLLBACK). Idempotent:
29+
/// a no-op when the session is idle, so a graceful end does not
30+
/// double-drop.
31+
async fn reclaim_open_txn(&self) {
32+
use crate::control::server::native::dispatch::NativeTxnDp;
33+
use crate::control::server::shared::session::{TransactionState, lifecycle};
34+
35+
let Some(identity) = self.identity.as_ref() else {
36+
return;
37+
};
38+
if self.sessions.transaction_state(&self.peer_addr) == TransactionState::Idle {
39+
return;
40+
}
41+
let dp = NativeTxnDp {
42+
state: self.state.as_ref(),
43+
};
44+
lifecycle::run_rollback(
45+
&self.sessions,
46+
&self.peer_addr,
47+
identity,
48+
self.state.as_ref(),
49+
&dp,
50+
)
51+
.await;
52+
}
53+
1854
/// Run the session loop: read frames, route by opcode, write responses.
1955
#[instrument(skip(self), fields(peer = %self.peer_addr))]
20-
pub async fn run(mut self) -> crate::Result<()> {
56+
async fn run_loop(&mut self) -> crate::Result<()> {
2157
// Perform the version-negotiation handshake before any frame exchange.
2258
let limits = self.state.limits.clone();
2359
self.proto_ver =
Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
// SPDX-License-Identifier: BUSL-1.1
2+
3+
//! Regression test: an abandoned NATIVE-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. `NativeSession::run`
8+
//! now wraps the frame loop and reclaims any still-open transaction's
9+
//! overlay via `lifecycle::run_rollback` on every exit path.
10+
//!
11+
//! This test proves reclamation (rather than just "row not visible", which
12+
//! cannot distinguish reclaimed from leaked-but-invisible) via the
13+
//! `active_txn_overlays` gauge on `SystemMetrics`: it rises when the
14+
//! transaction's first staged write materializes the overlay, and must
15+
//! fall back to baseline once the abandoned connection is torn down.
16+
17+
mod common;
18+
19+
use std::sync::atomic::Ordering;
20+
use std::time::{Duration, Instant};
21+
22+
use common::native_harness::{NativeTestServer, do_handshake, send_sql};
23+
24+
use nodedb_types::protocol::HelloFrame;
25+
use nodedb_types::protocol::opcodes::ResponseStatus;
26+
27+
/// Poll `active_txn_overlays` until `pred` is satisfied or `deadline` elapses.
28+
/// Returns the last observed value; callers assert against it so a timeout
29+
/// panics with the actual value instead of a bare "timed out".
30+
async fn poll_gauge(
31+
server: &NativeTestServer,
32+
deadline: Duration,
33+
pred: impl Fn(u64) -> bool,
34+
) -> u64 {
35+
let metrics = server
36+
.shared
37+
.system_metrics
38+
.as_ref()
39+
.expect("system_metrics must be wired into the native harness");
40+
let start = Instant::now();
41+
loop {
42+
let value = metrics.active_txn_overlays.load(Ordering::Relaxed);
43+
if pred(value) || start.elapsed() >= deadline {
44+
return value;
45+
}
46+
tokio::time::sleep(Duration::from_millis(20)).await;
47+
}
48+
}
49+
50+
#[tokio::test]
51+
async fn native_abandoned_txn_overlay_reclaimed_on_teardown() {
52+
let server = NativeTestServer::start().await;
53+
54+
let baseline = server
55+
.shared
56+
.system_metrics
57+
.as_ref()
58+
.expect("system_metrics must be wired into the native harness")
59+
.active_txn_overlays
60+
.load(Ordering::Relaxed);
61+
assert_eq!(baseline, 0, "gauge must start at zero: {baseline}");
62+
63+
{
64+
let (mut stream, _ack) = do_handshake(server.addr, &HelloFrame::current())
65+
.await
66+
.expect("handshake");
67+
68+
let create_resp = send_sql(
69+
&mut stream,
70+
1,
71+
"CREATE COLLECTION native_txn_overlay_teardown (id STRING PRIMARY KEY, n INT) \
72+
WITH (engine='document_schemaless')",
73+
)
74+
.await;
75+
assert_ne!(
76+
create_resp.status,
77+
ResponseStatus::Error,
78+
"CREATE must succeed: {create_resp:?}"
79+
);
80+
81+
let begin_resp = send_sql(&mut stream, 2, "BEGIN").await;
82+
assert_ne!(
83+
begin_resp.status,
84+
ResponseStatus::Error,
85+
"BEGIN must succeed: {begin_resp:?}"
86+
);
87+
88+
let staged_insert = send_sql(
89+
&mut stream,
90+
3,
91+
"INSERT INTO native_txn_overlay_teardown (id, n) VALUES ('a', 1)",
92+
)
93+
.await;
94+
assert_ne!(
95+
staged_insert.status,
96+
ResponseStatus::Error,
97+
"in-tx INSERT must succeed: {staged_insert:?}"
98+
);
99+
100+
// The staged write must have materialized the overlay, bumping the
101+
// gauge above baseline -- confirms there is something to reclaim.
102+
let after_stage = poll_gauge(&server, Duration::from_secs(5), |v| v > baseline).await;
103+
assert!(
104+
after_stage > baseline,
105+
"staged write must raise active_txn_overlays above baseline {baseline}, \
106+
got {after_stage}"
107+
);
108+
109+
// Abruptly abandon the connection -- no COMMIT/ROLLBACK. Dropping the
110+
// raw TcpStream (returned by `do_handshake`) closes the socket,
111+
// triggering the server-side EOF/error teardown path.
112+
drop(stream);
113+
}
114+
115+
// The abandoned transaction's overlay must be reclaimed on teardown,
116+
// bringing the gauge back down to baseline.
117+
let after_teardown = poll_gauge(&server, Duration::from_secs(5), |v| v == baseline).await;
118+
assert_eq!(
119+
after_teardown, baseline,
120+
"abandoned txn overlay must be reclaimed on connection teardown, \
121+
active_txn_overlays still {after_teardown} (baseline {baseline})"
122+
);
123+
124+
// Belt-and-suspenders: a fresh autocommit connection must not see the
125+
// staged row -- it was never committed.
126+
let (mut stream2, _ack2) = do_handshake(server.addr, &HelloFrame::current())
127+
.await
128+
.expect("handshake");
129+
let check = send_sql(
130+
&mut stream2,
131+
1,
132+
"SELECT * FROM native_txn_overlay_teardown WHERE id = 'a'",
133+
)
134+
.await;
135+
assert_ne!(
136+
check.status,
137+
ResponseStatus::Error,
138+
"post-teardown SELECT must succeed: {check:?}"
139+
);
140+
assert_eq!(
141+
check.rows.as_ref().map(Vec::len).unwrap_or(0),
142+
0,
143+
"the never-committed staged row must not be visible: {check:?}"
144+
);
145+
146+
server.shutdown().await;
147+
}

0 commit comments

Comments
 (0)