Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: acquire SQLite write ownership before restore and deletion authority reads

Browse Source
This commit is contained in:
Hare
2026-10-12 07:57:46 +09:00
parent df9d056d42
commit f597f1b57a
3 changed files with 114 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/server_db_execution_tests.rs
+102 -24
View File
@@ -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<SqliteWorkspaceStore>),
}
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<BusyGate>,
resume: Option<mpsc::Sender<()>>,
}
impl BusyObservation {
fn install(database: crate::FeatureDatabase) -> (Self, oneshot::Receiver<()>) {
Self::install_on(BusyDatabase::Feature(database))
}
fn install_store(store: Arc<SqliteWorkspaceStore>) -> (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"
);
}
crates/workspace-server/src/store.rs
+7 -3
View File
@@ -2341,12 +2341,16 @@ impl SqliteWorkspaceStore {
f(&mut conn)
}
pub(crate) fn with_transaction<T>(
pub(crate) fn with_immediate_transaction<T>(
&self,
f: impl FnOnce(&rusqlite::Transaction<'_>) -> Result<T>,
) -> Result<T> {
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<server_api::WorkerRestoreRequest> {
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()?
crates/workspace-server/src/workspace_deletion.rs
+5 -5
View File
@@ -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<WorkspaceDeletionOperationResponse> {
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<WorkspaceDeletionOperationResponse> {
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<WorkspaceDeletionOperationResponse> {
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,