recovery

Author: Shubham8287

crates/client-api/src/routes/database.rs

@@ -349,6 +349,31 @@ pub async fn abort_2pc<S: ControlStateDelegate + NodeDelegate>(
     Ok(StatusCode::OK)
 }
 
+/// 2PC coordinator status endpoint.
+///
+/// Returns `"commit"` if the coordinator has durably decided COMMIT for `prepare_id`,
+/// or `"abort"` otherwise.  Participant B polls this to recover from a timeout or crash.
+///
+/// `GET /v1/database/:name_or_identity/2pc/status/:prepare_id`
+pub async fn status_2pc<S: ControlStateDelegate + NodeDelegate>(
+    State(worker_ctx): State<S>,
+    Extension(_auth): Extension<SpacetimeAuth>,
+    Path(TwoPcParams {
+        name_or_identity,
+        prepare_id,
+    }): Path<TwoPcParams>,
+) -> axum::response::Result<impl IntoResponse> {
+    let (module, _database) = find_module_and_database(&worker_ctx, name_or_identity).await?;
+
+    let decision = if module.has_2pc_coordinator_commit(&prepare_id) {
+        "commit"
+    } else {
+        "abort"
+    };
+
+    Ok((StatusCode::OK, decision))
+}
+
 fn reducer_outcome_response(
     module: &ModuleHost,
     owner_identity: &Identity,
@@ -1361,6 +1386,8 @@ pub struct DatabaseRoutes<S> {
     pub commit_2pc_post: MethodRouter<S>,
     /// POST: /database/:name_or_identity/2pc/abort/:prepare_id
     pub abort_2pc_post: MethodRouter<S>,
+    /// GET: /database/:name_or_identity/2pc/status/:prepare_id
+    pub status_2pc_get: MethodRouter<S>,
 }
 
 impl<S> Default for DatabaseRoutes<S>
@@ -1389,6 +1416,7 @@ where
             prepare_post: post(prepare::<S>),
             commit_2pc_post: post(commit_2pc::<S>),
             abort_2pc_post: post(abort_2pc::<S>),
+            status_2pc_get: get(status_2pc::<S>),
         }
     }
 }
@@ -1416,7 +1444,8 @@ where
             .route("/reset", self.db_reset)
             .route("/prepare/:reducer", self.prepare_post)
             .route("/2pc/commit/:prepare_id", self.commit_2pc_post)
-            .route("/2pc/abort/:prepare_id", self.abort_2pc_post);
+            .route("/2pc/abort/:prepare_id", self.abort_2pc_post)
+            .route("/2pc/status/:prepare_id", self.status_2pc_get);
 
         axum::Router::new()
             .route("/", self.root_post)

crates/core/src/db/relational_db.rs

@@ -454,14 +454,25 @@ impl RelationalDB {
     }
 
     /// Read any 2PC participant transactions that were in PREPARE state when the database
-    /// last shut down (or crashed). Each returned string is a `prepare_id`.
+    /// last shut down (or crashed).
     ///
-    /// If non-empty, the caller must resume these transactions: retransmit PREPARED to
-    /// the coordinator and await a COMMIT or ABORT decision before allowing normal operation.
-    pub fn pending_2pc_prepares(&self) -> Result<Vec<String>, DBError> {
+    /// Each returned row contains all the information needed to resume the transaction:
+    /// the prepare_id, coordinator identity, reducer name/args, and caller context.
+    /// B never aborts on its own — it polls the coordinator for a decision.
+    pub fn pending_2pc_prepares(&self) -> Result<Vec<spacetimedb_datastore::system_tables::St2pcStateRow>, DBError> {
         self.with_auto_commit(Workload::Internal, |tx| tx.scan_st_2pc_state().map_err(DBError::from))
     }
 
+    /// Read any 2PC coordinator log entries that have not yet been acknowledged by their
+    /// participants. Used on coordinator crash-recovery to retransmit COMMIT decisions.
+    pub fn pending_2pc_coordinator_commits(
+        &self,
+    ) -> Result<Vec<spacetimedb_datastore::system_tables::St2pcCoordinatorLogRow>, DBError> {
+        self.with_auto_commit(Workload::Internal, |tx| {
+            tx.scan_st_2pc_coordinator_log().map_err(DBError::from)
+        })
+    }
+
     /// Read the set of clients currently connected to the database.
     pub fn connected_clients(&self) -> Result<ConnectedClients, DBError> {
         self.with_read_only(Workload::Internal, |tx| {

crates/core/src/host/host_controller.rs

@@ -1141,6 +1141,11 @@ impl Host {
         module_host.clear_all_clients().await?;
 
         scheduler_starter.start(&module_host)?;
+
+        // Crash recovery: retransmit any pending 2PC decisions from before the restart.
+        module_host.recover_2pc_coordinator();
+        module_host.recover_2pc_participant();
+
         let disk_metrics_recorder_task = tokio::spawn(metric_reporter(replica_ctx.clone())).abort_handle();
         let view_cleanup_task = spawn_view_cleanup_loop(replica_ctx.relational_db().clone());
 

crates/core/src/host/module_host.rs

@@ -1808,17 +1808,18 @@ impl ModuleHost {
         let this = self.clone();
         let reducer_name_owned = reducer_def.name.clone();
         let prepare_id_clone = prepare_id.clone();
+        let coordinator_identity = caller_identity;
         tokio::spawn(async move {
             let _ = this
                 .call(
                     &reducer_name_owned,
-                    (params, prepare_id_clone, prepared_tx, decision_rx),
-                    async |(p, pid, ptx, drx), inst| {
-                        inst.call_reducer_prepare_and_hold(p, pid, ptx, drx);
+                    (params, prepare_id_clone, coordinator_identity, prepared_tx, decision_rx),
+                    async |(p, pid, cid, ptx, drx), inst| {
+                        inst.call_reducer_prepare_and_hold(p, pid, cid, ptx, drx);
                         Ok::<(), ReducerCallError>(())
                     },
                     // JS modules: no 2PC support yet.
-                    async |(p, _pid, ptx, _drx), inst| {
+                    async |(p, _pid, _cid, ptx, _drx), inst| {
                         let (res, rv) = inst.call_reducer(p).await.map(|r| (r, None)).unwrap_or_else(|e| {
                             log::error!("prepare_reducer JS fallback: {e}");
                             (
@@ -1874,6 +1875,230 @@ impl ModuleHost {
         Ok(())
     }
 
+    /// Check whether `prepare_id` is present in the coordinator log of this database.
+    /// Used by participant B to ask coordinator A: "did you commit?"
+    pub fn has_2pc_coordinator_commit(&self, prepare_id: &str) -> bool {
+        let db = self.relational_db();
+        db.pending_2pc_coordinator_commits()
+            .map(|rows| rows.iter().any(|r| r.participant_prepare_id == prepare_id))
+            .unwrap_or(false)
+    }
+
+    /// Crash recovery for the **coordinator** role.
+    ///
+    /// Scans `st_2pc_coordinator_log` for participants that have not yet acked
+    /// COMMIT and retransmits the HTTP commit call.  Deletes the log entry on success.
+    pub fn recover_2pc_coordinator(&self) {
+        let db = self.relational_db().clone();
+        let rows = match db.pending_2pc_coordinator_commits() {
+            Ok(r) => r,
+            Err(e) => {
+                log::error!("recover_2pc_coordinator: scan failed: {e}");
+                return;
+            }
+        };
+        if rows.is_empty() {
+            return;
+        }
+        let replica_ctx = self.replica_ctx().clone();
+        let db2 = db.clone();
+        tokio::spawn(async move {
+            let client = replica_ctx.call_reducer_client.clone();
+            let router = replica_ctx.call_reducer_router.clone();
+            let auth_token = replica_ctx.call_reducer_auth_token.clone();
+            for row in rows {
+                let prepare_id = row.participant_prepare_id.clone();
+                let participant_identity = match Identity::from_hex(&row.participant_identity_hex) {
+                    Ok(id) => id,
+                    Err(e) => {
+                        log::error!("recover_2pc_coordinator: invalid participant identity hex {}: {e}", row.participant_identity_hex);
+                        continue;
+                    }
+                };
+                let base_url = match router.resolve_base_url(participant_identity).await {
+                    Ok(url) => url,
+                    Err(e) => {
+                        log::warn!("recover_2pc_coordinator: cannot resolve URL for {participant_identity}: {e}");
+                        continue;
+                    }
+                };
+                let url = format!(
+                    "{}/v1/database/{}/2pc/commit/{}",
+                    base_url,
+                    participant_identity.to_hex(),
+                    prepare_id,
+                );
+                let mut req = client.post(&url);
+                if let Some(ref token) = auth_token {
+                    req = req.header(http::header::AUTHORIZATION, format!("Bearer {token}"));
+                }
+                match req.send().await {
+                    Ok(resp) if resp.status().is_success() => {
+                        log::info!("recover_2pc_coordinator: re-committed {prepare_id} on {participant_identity}");
+                        if let Err(e) = db2.with_auto_commit::<_, _, anyhow::Error>(Workload::Internal, |tx| {
+                            Ok(tx.delete_st_2pc_coordinator_log(&prepare_id)?)
+                        }) {
+                            log::warn!("recover_2pc_coordinator: delete coordinator log failed for {prepare_id}: {e}");
+                        }
+                    }
+                    Ok(resp) => {
+                        log::warn!("recover_2pc_coordinator: commit for {prepare_id} returned {}", resp.status());
+                    }
+                    Err(e) => {
+                        log::warn!("recover_2pc_coordinator: transport error for {prepare_id}: {e}");
+                    }
+                }
+            }
+        });
+    }
+
+    /// Crash recovery for the **participant** role.
+    ///
+    /// Scans `st_2pc_state` for any prepared-but-not-decided transactions, re-runs
+    /// each reducer to reacquire the write lock, then polls the coordinator for a decision.
+    ///
+    /// **B never aborts on its own** — only the coordinator's response yields ABORT.
+    pub fn recover_2pc_participant(&self) {
+        let db = self.relational_db().clone();
+        let rows = match db.pending_2pc_prepares() {
+            Ok(r) => r,
+            Err(e) => {
+                log::error!("recover_2pc_participant: scan failed: {e}");
+                return;
+            }
+        };
+        if rows.is_empty() {
+            return;
+        }
+        let this = self.clone();
+        tokio::spawn(async move {
+            for row in rows {
+                let original_prepare_id = row.prepare_id.clone();
+                let coordinator_identity = match Identity::from_hex(&row.coordinator_identity_hex) {
+                    Ok(id) => id,
+                    Err(e) => {
+                        log::error!("recover_2pc_participant: invalid coordinator identity hex for {original_prepare_id}: {e}");
+                        continue;
+                    }
+                };
+                let caller_identity = match Identity::from_hex(&row.caller_identity_hex) {
+                    Ok(id) => id,
+                    Err(e) => {
+                        log::error!("recover_2pc_participant: invalid caller identity hex for {original_prepare_id}: {e}");
+                        continue;
+                    }
+                };
+                let caller_connection_id = u128::from_str_radix(&row.caller_connection_id_hex, 16)
+                    .map(ConnectionId::from_u128)
+                    .unwrap_or(ConnectionId::ZERO);
+                let args = FunctionArgs::Bsatn(row.args_bsatn.clone().into());
+
+                // Step 1: Re-run the reducer to reacquire the write lock.
+                let new_prepare_id = match this
+                    .prepare_reducer(caller_identity, Some(caller_connection_id), &row.reducer_name, args)
+                    .await
+                {
+                    Ok((pid, result, _rv)) if !pid.is_empty() => {
+                        log::info!(
+                            "recover_2pc_participant: re-prepared {original_prepare_id} as {pid}: {:?}",
+                            result.outcome
+                        );
+                        pid
+                    }
+                    Ok(_) => {
+                        // Reducer failed — treat as abort, clean up old marker.
+                        log::warn!("recover_2pc_participant: reducer failed on re-run for {original_prepare_id}");
+                        let _ = db.with_auto_commit::<_, _, anyhow::Error>(Workload::Internal, |tx| {
+                            Ok(tx.delete_st_2pc_state(&original_prepare_id)?)
+                        });
+                        continue;
+                    }
+                    Err(e) => {
+                        log::error!("recover_2pc_participant: prepare_reducer error for {original_prepare_id}: {e:?}");
+                        continue;
+                    }
+                };
+
+                // Step 2: Poll coordinator with the ORIGINAL prepare_id until we get a decision.
+                // We do this in a separate task so the loop can proceed to the next row.
+                let this2 = this.clone();
+                let db2 = db.clone();
+                let client = this.replica_ctx().call_reducer_client.clone();
+                let router = this.replica_ctx().call_reducer_router.clone();
+                let auth_token = this.replica_ctx().call_reducer_auth_token.clone();
+                tokio::spawn(async move {
+                    loop {
+                        let decision = Self::query_coordinator_status_with_client(
+                            &client,
+                            &router,
+                            auth_token.clone(),
+                            coordinator_identity,
+                            &original_prepare_id,
+                        ).await;
+                        match decision {
+                            Some(commit) => {
+                                if commit {
+                                    let _ = this2.commit_prepared(&new_prepare_id);
+                                } else {
+                                    let _ = this2.abort_prepared(&new_prepare_id);
+                                }
+                                // Clean up the old st_2pc_state entry.
+                                let _ = db2.with_auto_commit::<_, _, anyhow::Error>(Workload::Internal, |tx| {
+                                    Ok(tx.delete_st_2pc_state(&original_prepare_id)?)
+                                });
+                                break;
+                            }
+                            None => tokio::time::sleep(std::time::Duration::from_secs(5)).await,
+                        }
+                    }
+                });
+            }
+        });
+    }
+
+    /// Query `GET /v1/database/{coordinator}/2pc/status/{prepare_id}`.
+    ///
+    /// Returns `Some(true)` = COMMIT, `Some(false)` = ABORT, `None` = transient error (retry).
+    async fn query_coordinator_status_with_client(
+        client: &reqwest::Client,
+        router: &std::sync::Arc<dyn crate::host::reducer_router::ReducerCallRouter>,
+        auth_token: Option<String>,
+        coordinator_identity: Identity,
+        prepare_id: &str,
+    ) -> Option<bool> {
+        let base_url = match router.resolve_base_url(coordinator_identity).await {
+            Ok(url) => url,
+            Err(e) => {
+                log::warn!("2PC recovery status poll: cannot resolve coordinator URL: {e}");
+                return None;
+            }
+        };
+        let url = format!(
+            "{}/v1/database/{}/2pc/status/{}",
+            base_url,
+            coordinator_identity.to_hex(),
+            prepare_id,
+        );
+        let mut req = client.get(&url);
+        if let Some(token) = &auth_token {
+            req = req.header(http::header::AUTHORIZATION, format!("Bearer {token}"));
+        }
+        match req.send().await {
+            Ok(resp) if resp.status().is_success() => {
+                let body = resp.text().await.unwrap_or_default();
+                Some(body.trim() == "commit")
+            }
+            Ok(resp) => {
+                log::warn!("2PC recovery status poll: coordinator returned {}", resp.status());
+                None
+            }
+            Err(e) => {
+                log::warn!("2PC recovery status poll: transport error: {e}");
+                None
+            }
+        }
+    }
+
     pub async fn call_view_add_single_subscription(
         &self,
         sender: Arc<ClientConnectionSender>,