fix(control): track every vShard a transaction stages writes to

Sessions previously assumed a transaction stages writes to a single
"home" vShard, recording only the first one. Multi-core transactions
(e.g. two INSERTs to collections homed on different vShards) leaked
the staging overlay on every vShard after the first at ROLLBACK, and
savepoint mark/rewind only ever touched that one core.

Replace the single `tx_vshard: Option<VShardId>` with a `tx_vshards:
BTreeSet<VShardId>` recorded on every staged write. Overlay teardown
and per-vShard savepoint markers now fan out over the full set:
`MetaOp::DropTxnOverlay` targets every staged vShard, and each
savepoint stores a per-vShard `(value_marker, graph_marker)` map so
ROLLBACK TO rewinds all of them, defaulting to `(0, 0)` for a vShard
first staged after the savepoint was taken.

Also derive `PartialOrd`/`Ord` on `VShardId` so it can key a
`BTreeMap`/`BTreeSet`.

Author: farhan-syah

nodedb-types/src/id/vshard.rs

@@ -13,6 +13,8 @@ use serde::{Deserialize, Serialize};
     Copy,
     PartialEq,
     Eq,
+    PartialOrd,
+    Ord,
     Hash,
     Serialize,
     Deserialize,

nodedb/src/control/server/shared/session/lifecycle.rs

@@ -58,8 +58,9 @@ pub async fn run_rollback(
 ) {
     ddl_buffer::discard();
     // Snapshot the overlay identity BEFORE `rollback()` clears session state,
-    // so the staging overlay can be released on its home vShard.
-    let (overlay_txn_id, overlay_vshard) = sessions.txn_identity(addr);
+    // so the staging overlay can be released on EVERY vShard the transaction
+    // staged writes to (a transaction may span multiple cores).
+    let (overlay_txn_id, overlay_vshards) = sessions.txn_identity(addr);
     let reservations = sessions.rollback(addr).unwrap_or_default();
     for handle in &reservations {
         let key = &handle.sequence_key;
@@ -87,8 +88,11 @@ pub async fn run_rollback(
     // Discard NOTIFY messages buffered during this transaction.
     sessions.discard_pending_notifies(addr);
 
-    // Release any staging overlay populated by statement-time point writes.
-    if let (Some(txn_id), Some(vshard_id)) = (overlay_txn_id, overlay_vshard) {
-        drop_txn_overlay(dp, identity.tenant_id, vshard_id, txn_id).await;
+    // Release any staging overlay populated by statement-time point writes on
+    // every vShard the transaction staged to, so no core's overlay leaks.
+    if let Some(txn_id) = overlay_txn_id {
+        for vshard_id in overlay_vshards {
+            drop_txn_overlay(dp, identity.tenant_id, vshard_id, txn_id).await;
+        }
     }
 }

nodedb/src/control/server/shared/session/savepoint_ops.rs

@@ -8,10 +8,11 @@
 //! neutral `SessionStore` savepoint stack. Transports translate the returned
 //! [`SavepointError`] into their own SQLSTATE (`25P01` / `3B001`).
 
+use std::collections::BTreeMap;
 use std::net::SocketAddr;
 
 use crate::bridge::envelope::PhysicalPlan;
-use crate::types::TenantId;
+use crate::types::{TenantId, VShardId};
 use nodedb_physical::physical_plan::MetaOp;
 use nodedb_physical::physical_task::{PhysicalTask, PostSetOp};
 
@@ -20,6 +21,7 @@ use super::state::TransactionState;
 use super::store::SessionStore;
 
 /// Typed savepoint failure. Adapters map each variant to its SQLSTATE.
+#[derive(Debug)]
 pub enum SavepointError {
     /// A savepoint command was issued outside a transaction block. → `25P01`.
     NoActiveTransaction,
@@ -36,19 +38,14 @@ fn require_active_txn(sessions: &SessionStore, addr: &SocketAddr) -> Result<(),
     Ok(())
 }
 
-/// Dispatch a savepoint overlay meta-op to the transaction's home vShard and
-/// return the raw response payload bytes, or `None` when no staged write has
-/// homed a vShard yet (the overlay — and its journal — is empty, so there is
-/// nothing on the Data Plane to mark or rewind).
+/// Dispatch a savepoint overlay meta-op to a specific vShard's core and return
+/// the raw response payload bytes, or `None` on dispatch failure.
 async fn dispatch_overlay_savepoint(
-    sessions: &SessionStore,
-    addr: &SocketAddr,
     tenant_id: TenantId,
+    vshard_id: VShardId,
     dp: &impl TxnDataPlane,
     op: MetaOp,
 ) -> Option<Vec<u8>> {
-    let (txn_id, vshard) = sessions.txn_identity(addr);
-    let (_txn_id, vshard_id) = (txn_id?, vshard?);
     let task = PhysicalTask {
         tenant_id,
         vshard_id,
@@ -88,10 +85,11 @@ fn decode_markers(payload: Option<Vec<u8>>) -> (usize, usize) {
 
 /// Handle SAVEPOINT `<name>`.
 ///
-/// Captures the composite overlay undo-journal marker on the txn's home vShard
-/// so a later ROLLBACK TO reverts staged value/TTL AND graph state to exactly
-/// here. A missing/short payload (or no vShard yet) means empty journals →
-/// `(0, 0)`.
+/// Captures the composite overlay undo-journal marker on EVERY vShard the
+/// transaction has staged writes to, so a later ROLLBACK TO reverts staged
+/// value/TTL AND graph state on all of them to exactly here. A missing/short
+/// payload means empty journals → `(0, 0)`. With no vShard staged yet the
+/// marker map is empty.
 pub async fn run_savepoint(
     sessions: &SessionStore,
     addr: &SocketAddr,
@@ -100,21 +98,21 @@ pub async fn run_savepoint(
     name: &str,
 ) -> Result<(), SavepointError> {
     require_active_txn(sessions, addr)?;
-    let (value_marker, graph_marker) = match sessions.tx_id(addr) {
-        Some(txn_id) => {
+    let (txn_id, vshards) = sessions.txn_identity(addr);
+    let mut markers: BTreeMap<VShardId, (usize, usize)> = BTreeMap::new();
+    if let Some(txn_id) = txn_id {
+        for vshard_id in vshards {
             let payload = dispatch_overlay_savepoint(
-                sessions,
-                addr,
                 tenant_id,
+                vshard_id,
                 dp,
                 MetaOp::MarkSavepoint { txn_id },
             )
             .await;
-            decode_markers(payload)
+            markers.insert(vshard_id, decode_markers(payload));
         }
-        None => (0, 0),
-    };
-    sessions.create_savepoint(addr, name.to_string(), value_marker, graph_marker);
+    }
+    sessions.create_savepoint(addr, name.to_string(), markers);
     Ok(())
 }
 
@@ -139,7 +137,12 @@ pub fn run_release_savepoint(
 /// Handle ROLLBACK TO SAVEPOINT `<name>`.
 ///
 /// Truncates the write buffer to the saved position and rewinds BOTH the
-/// value/TTL overlay and the graph overlay to the marked journal points.
+/// value/TTL overlay and the graph overlay on every vShard the transaction has
+/// staged to. Iterates the CURRENT staged set (a superset of the savepoint's,
+/// since writes may have staged to NEW vShards after the savepoint): a vShard
+/// with a saved marker rewinds to it; a vShard first staged AFTER the savepoint
+/// has no saved marker and rewinds to `(0, 0)`, dropping ALL of its staged
+/// writes.
 pub async fn run_rollback_to_savepoint(
     sessions: &SessionStore,
     addr: &SocketAddr,
@@ -148,25 +151,28 @@ pub async fn run_rollback_to_savepoint(
     name: &str,
 ) -> Result<(), SavepointError> {
     require_active_txn(sessions, addr)?;
-    let (value_marker, graph_marker) =
+    let markers =
         sessions
             .rollback_to_savepoint(addr, name)
             .map_err(|e| SavepointError::NotFound {
                 message: e.to_string(),
             })?;
-    if let Some(txn_id) = sessions.tx_id(addr) {
-        dispatch_overlay_savepoint(
-            sessions,
-            addr,
-            tenant_id,
-            dp,
-            MetaOp::RollbackToSavepoint {
-                txn_id,
-                value_marker: value_marker as u64,
-                graph_marker: graph_marker as u64,
-            },
-        )
-        .await;
+    let (txn_id, vshards) = sessions.txn_identity(addr);
+    if let Some(txn_id) = txn_id {
+        for vshard_id in vshards {
+            let (value_marker, graph_marker) = markers.get(&vshard_id).copied().unwrap_or((0, 0));
+            dispatch_overlay_savepoint(
+                tenant_id,
+                vshard_id,
+                dp,
+                MetaOp::RollbackToSavepoint {
+                    txn_id,
+                    value_marker: value_marker as u64,
+                    graph_marker: graph_marker as u64,
+                },
+            )
+            .await;
+        }
     }
     Ok(())
 }

nodedb/src/control/server/shared/session/state.rs

@@ -2,11 +2,28 @@
 
 //! Per-connection session state types.
 
-use std::collections::HashMap;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 
 use crate::types::{DatabaseId, Lsn, TenantId, TxnId, VShardId};
 use nodedb_physical::physical_task::PhysicalTask;
 
+/// One entry on the transaction's savepoint stack.
+///
+/// A savepoint captures the write buffer length AND, for each vShard that had
+/// staged writes when the savepoint was established, that vShard's value/TTL and
+/// graph overlay undo-journal markers. On ROLLBACK TO, the buffer is truncated
+/// to `buffer_len` and every currently-staged vShard's overlays are rewound —
+/// to its saved marker if present, else to `(0, 0)` (a vShard first staged
+/// AFTER the savepoint must have ALL of its staged writes rewound).
+pub struct SavepointEntry {
+    /// User-visible savepoint name.
+    pub name: String,
+    /// `tx_buffer` length captured when the savepoint was established.
+    pub buffer_len: usize,
+    /// Per-vShard `(value_marker, graph_marker)` overlay journal markers.
+    pub markers: BTreeMap<VShardId, (usize, usize)>,
+}
+
 /// PostgreSQL transaction state for ReadyForQuery status byte.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum TransactionState {
@@ -73,23 +90,23 @@ pub struct ConnSession {
     /// and cleared on `COMMIT`/`ROLLBACK`. Keys the per-transaction staging
     /// overlay. `None` outside a transaction block.
     pub tx_id: Option<TxnId>,
-    /// vShard this transaction's staged writes were homed on, captured on the
-    /// first staged/buffered write. Under the single-vshard-per-transaction
-    /// assumption, this is where the staging overlay lives — used to target
-    /// `MetaOp::DropTxnOverlay` at COMMIT/ROLLBACK. `None` until the first write.
-    pub tx_vshard: Option<VShardId>,
+    /// Set of vShards this transaction has staged writes to, recorded on every
+    /// staged/buffered write. A transaction can stage to multiple vShards/cores
+    /// (e.g. two INSERTs to collections homed on different cores), so overlay
+    /// teardown (`MetaOp::DropTxnOverlay` at ROLLBACK) and per-vShard savepoint
+    /// mark/rewind must fan over ALL of them. Ordered (BTree) for deterministic
+    /// teardown. Empty until the first staged write.
+    pub tx_vshards: BTreeSet<VShardId>,
     /// Read-set: LSN-versioned, predicate-aware entries for write conflict
     /// detection, captured on the shared read seam by every transport. At
     /// COMMIT, each entry is checked — if the entry's collection has a current
     /// write-LSN past `read_lsn`, a concurrent write occurred and the
     /// transaction is rejected with SERIALIZATION_FAILURE.
     pub tx_read_set: Vec<super::read_set::ReadSetEntry>,
-    /// Savepoint stack: each entry is `(name, tx_buffer_len_at_savepoint,
-    /// value_overlay_marker, graph_overlay_marker)`.
-    /// On ROLLBACK TO, truncate tx_buffer to the saved length AND rewind both
-    /// Data-Plane staging overlays (value/TTL and GRAPH) to their saved
-    /// journal markers.
-    pub savepoints: Vec<(String, usize, usize, usize)>,
+    /// Savepoint stack. On ROLLBACK TO, truncate tx_buffer to the saved length
+    /// AND rewind each staged vShard's two Data-Plane staging overlays (value/TTL
+    /// and GRAPH) to their saved journal markers. See [`SavepointEntry`].
+    pub savepoints: Vec<SavepointEntry>,
     /// Pending consumer offset commits deferred until COMMIT.
     /// Each entry: (tenant_id, stream_name, group_name, partition_id, lsn).
     /// Flushed atomically on COMMIT, discarded on ROLLBACK.
@@ -159,7 +176,7 @@ impl ConnSession {
             tx_snapshot_lsn: None,
             tx_snapshot_epoch: None,
             tx_id: None,
-            tx_vshard: None,
+            tx_vshards: BTreeSet::new(),
             tx_read_set: Vec::new(),
             savepoints: Vec::new(),
             pending_offset_commits: Vec::new(),

nodedb/src/control/server/shared/session/store.rs

@@ -164,7 +164,7 @@ impl SessionStore {
             session.tx_snapshot_lsn = None;
             session.tx_snapshot_epoch = None;
             session.tx_id = None;
-            session.tx_vshard = None;
+            session.tx_vshards.clear();
             session.tx_read_set.clear();
             session.savepoints.clear();
             session.pending_offset_commits.clear();