diff --git a/crates/workspace-server/src/server_db_execution_tests.rs b/crates/workspace-server/src/server_db_execution_tests.rs index 1db761f13..e9a6538f2 100644 --- a/crates/workspace-server/src/server_db_execution_tests.rs +++ b/crates/workspace-server/src/server_db_execution_tests.rs @@ -226,13 +226,41 @@ struct BusyGate { /// rusqlite's fn-pointer busy_handler cannot capture a fixture. SQLite's /// userdata API can, without process globals or changing production hooks. +enum BusyDatabase { + Feature(crate::FeatureDatabase), + Store(Arc), +} +impl BusyDatabase { + fn access(&self, access: impl FnOnce(&Connection)) { + match self { + Self::Feature(database) => database + .with_connection(|conn| { + access(conn); + Ok(()) + }) + .unwrap(), + Self::Store(store) => store + .with_conn(|conn| { + access(conn); + Ok(()) + }) + .unwrap(), + } + } +} struct BusyObservation { - database: crate::FeatureDatabase, + database: BusyDatabase, gate: Box, resume: Option>, } impl BusyObservation { fn install(database: crate::FeatureDatabase) -> (Self, oneshot::Receiver<()>) { + Self::install_on(BusyDatabase::Feature(database)) + } + fn install_store(store: Arc) -> (Self, oneshot::Receiver<()>) { + Self::install_on(BusyDatabase::Store(store)) + } + fn install_on(database: BusyDatabase) -> (Self, oneshot::Receiver<()>) { unsafe extern "C" fn busy(context: *mut c_void, _: i32) -> i32 { // SAFETY: owned Box remains registered until Drop unregisters it, // serialized by the managed connection mutex. @@ -249,20 +277,17 @@ impl BusyObservation { entered: Some(entered), resume: permit, }); - database - .with_connection(|conn| { - // SAFETY: handle belongs to this connection; the Box is stable. - let code = unsafe { - ffi::sqlite3_busy_handler( - conn.handle(), - Some(busy), - (&mut *gate as *mut BusyGate).cast(), - ) - }; - assert_eq!(code, ffi::SQLITE_OK); - Ok(()) - }) - .unwrap(); + database.access(|conn| { + // SAFETY: handle belongs to this connection; the Box is stable. + let code = unsafe { + ffi::sqlite3_busy_handler( + conn.handle(), + Some(busy), + (&mut *gate as *mut BusyGate).cast(), + ) + }; + assert_eq!(code, ffi::SQLITE_OK); + }); ( Self { database, @@ -281,15 +306,12 @@ impl BusyObservation { impl Drop for BusyObservation { fn drop(&mut self) { self.resume(); - self.database - .with_connection(|conn| { - // SAFETY: unregister before the boxed userdata is freed. - unsafe { - ffi::sqlite3_busy_handler(conn.handle(), None, std::ptr::null_mut()); - } - Ok(()) - }) - .unwrap(); + self.database.access(|conn| { + // SAFETY: unregister before the boxed userdata is freed. + unsafe { + ffi::sqlite3_busy_handler(conn.handle(), None, std::ptr::null_mut()); + } + }); let _ = &self.gate; } } @@ -689,3 +711,59 @@ async fn saturated_cancelled_host_jobs_keep_their_budget_but_not_other_fixture_o host.run(|| Ok::<_, Error>(())).await.unwrap(); host.drain().await.unwrap(); } + +#[tokio::test(flavor = "current_thread")] +async fn restore_intent_pins_once_after_write_ownership_without_a_deferred_read_upgrade() { + let fixture = Fixture::new().await; + let store = fixture.api.config_store.clone(); + let (mut busy, waiting) = BusyObservation::install_store(store.clone()); + let mut lock = WriteLock::acquire(fixture.api.config.database_path.clone()); + let worker = RuntimeWorkerRef::new("restore-runtime", "restore-worker"); + let pin_store = store.clone(); + let pin_worker = worker.clone(); + let mut pinned = Box::pin(fixture.executor.run(move || { + pin_store.pin_internal_worker_restore_intent( + WORKSPACE, + &pin_worker, + "restore-domain", + Some("observed-before-dispatch"), + Some("restore-request"), + ) + })); + tokio::select! { + result = &mut pinned => panic!("restore intent must wait for write ownership before reading, not fail a deferred read upgrade: {result:?}"), + arrived = observed(waiting) => arrived.unwrap(), + } + runtime_progress_while_locked(&lock).await; + assert!(futures::poll!(&mut pinned).is_pending()); + lock.release(); + busy.resume(); + let request = observed(pinned).await.unwrap(); + assert_eq!(request.request_id, "restore-request"); + assert_eq!( + request.expected_observation_token, + "observed-before-dispatch" + ); + fixture.executor.drain().await.unwrap(); + let saved = store + .get_internal_worker_restore_intent(WORKSPACE, &worker, &request.request_id) + .unwrap() + .unwrap(); + assert_eq!(saved.request_id, request.request_id); + assert_eq!( + saved.expected_observation_token, + request.expected_observation_token + ); + let count: i64 = store + .with_conn(|conn| { + Ok(conn.query_row( + "SELECT count(*) FROM worker_restore_intents WHERE workspace_id=?1 AND request_id=?2", + rusqlite::params![WORKSPACE, request.request_id], |row| row.get(0) + )?) + }) + .unwrap(); + assert_eq!( + count, 1, + "one admitted intent must persist exactly once, with no operation retry" + ); +} diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 2d954fa27..90a040b77 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -2341,12 +2341,16 @@ impl SqliteWorkspaceStore { f(&mut conn) } - pub(crate) fn with_transaction( + pub(crate) fn with_immediate_transaction( &self, f: impl FnOnce(&rusqlite::Transaction<'_>) -> Result, ) -> Result { self.with_conn_mut(|conn| { - let tx = conn.transaction()?; + // These callers read authority/journal state before mutating it. + // Acquire write ownership before the first read: a DEFERRED read + // upgrade against another connection's writer can return BUSY + // without invoking the busy handler. Never retry the domain mutation. + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let value = f(&tx)?; tx.commit()?; Ok(value) @@ -2407,7 +2411,7 @@ impl SqliteWorkspaceStore { observed_token: Option<&str>, explicit_request_id: Option<&str>, ) -> Result { - self.with_transaction(|tx| { + self.with_immediate_transaction(|tx| { let existing: Option<(String, String)> = if let Some(id) = explicit_request_id { tx.query_row("SELECT request_id, expected_token FROM worker_restore_intents WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND request_id=?4", params![workspace_id, worker.runtime_id, worker.worker_id, id], |row| Ok((row.get(0)?, row.get(1)?))).optional()? diff --git a/crates/workspace-server/src/workspace_deletion.rs b/crates/workspace-server/src/workspace_deletion.rs index b68c7e4af..5003d4153 100644 --- a/crates/workspace-server/src/workspace_deletion.rs +++ b/crates/workspace-server/src/workspace_deletion.rs @@ -236,7 +236,7 @@ impl WorkspaceDeletionStore for SqliteWorkspaceStore { "Workspace deletion request exceeds bounded field limits".to_string(), )); } - self.with_transaction(|tx| { + self.with_immediate_transaction(|tx| { if let Some(existing) = read_operation(tx, &request.operation_id)? { if existing.actor_account_id != actor_account_id { return Err(Error::WorkspacePermissionDenied( @@ -369,7 +369,7 @@ impl WorkspaceDeletionStore for SqliteWorkspaceStore { child_operation_id: &str, ) -> Result { validate_operation_id(child_operation_id)?; - self.with_transaction(|tx| { + self.with_immediate_transaction(|tx| { let operation = read_operation(tx, operation_id)? .ok_or_else(|| Error::InvalidInput("Workspace deletion operation".to_string()))? .response; @@ -411,7 +411,7 @@ impl WorkspaceDeletionStore for SqliteWorkspaceStore { failure_category: Option<&str>, ) -> Result { validate_operation_projection(child_operation_ids, blockers)?; - self.with_transaction(|tx| { + self.with_immediate_transaction(|tx| { let current = read_operation(tx, operation_id)? .ok_or_else(|| Error::InvalidInput("Workspace deletion operation".to_string()))? .response; @@ -462,7 +462,7 @@ impl WorkspaceDeletionStore for SqliteWorkspaceStore { &self, operation_id: &str, ) -> Result { - self.with_transaction(|tx| { + self.with_immediate_transaction(|tx| { tx.execute_batch("PRAGMA defer_foreign_keys = ON;")?; let operation = read_operation(tx, operation_id)? .ok_or_else(|| Error::InvalidInput("Workspace deletion operation".to_string()))?; @@ -1019,7 +1019,7 @@ mod tests { // Seed domain handling evidence under a transaction. This fixture // exercises owner purge, not singleton authentication or delivery. store - .with_transaction(|tx| { + .with_immediate_transaction(|tx| { ticket::sqlite_update_ticket_assignment_tracking_in_transaction( tx, workspace,