fix(sqlite): stop actors on head fence mismatch

Author: NathanFlurry

Cargo.lock

@@ -16,4 +16,5 @@ async-trait.workspace = true
 depot.workspace = true
 depot-client.workspace = true
 rivet-envoy-protocol.workspace = true
+rivet-error.workspace = true
 tokio.workspace = true

engine/artifacts/errors/depot.head_fence_mismatch.json

@@ -8,6 +8,7 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use async_trait::async_trait;
+use depot::error::SqliteStorageError;
 use depot_client::{
 	database::{NativeDatabaseHandle, open_database_from_transport},
 	vfs::{SqliteTransport, SqliteVfsMetrics},
@@ -48,22 +49,31 @@ impl SqliteTransport for EmbeddedDepotSqliteTransport {
 		&self,
 		request: protocol::SqliteGetPagesRequest,
 	) -> Result<protocol::SqliteGetPagesResponse> {
-		match self.db.get_pages(request.pgnos).await {
-			Ok(pages) => Ok(protocol::SqliteGetPagesResponse::SqliteGetPagesOk(
+		match self
+			.db
+			.get_pages_with_options(
+				request.pgnos,
+				depot::types::GetPagesOptions {
+					expected_head_txid: request.expected_head_txid,
+				},
+			)
+			.await
+		{
+			Ok(result) => Ok(protocol::SqliteGetPagesResponse::SqliteGetPagesOk(
 				protocol::SqliteGetPagesOk {
-					pages: pages
+					pages: result
+						.pages
 						.into_iter()
 						.map(|page| protocol::SqliteFetchedPage {
 							pgno: page.pgno,
 							bytes: page.bytes,
 						})
 						.collect(),
+					head_txid: Some(result.head_txid),
 				},
 			)),
 			Err(err) => Ok(protocol::SqliteGetPagesResponse::SqliteErrorResponse(
-				protocol::SqliteErrorResponse {
-					message: sqlite_error_reason(&err),
-				},
+				sqlite_error_response(&err),
 			)),
 		}
 	}
@@ -74,7 +84,7 @@ impl SqliteTransport for EmbeddedDepotSqliteTransport {
 	) -> Result<protocol::SqliteCommitResponse> {
 		match self
 			.db
-			.commit(
+			.commit_with_options(
 				request
 					.dirty_pages
 					.into_iter()
@@ -85,15 +95,20 @@ impl SqliteTransport for EmbeddedDepotSqliteTransport {
 					.collect(),
 				request.db_size_pages,
 				request.now_ms,
+				depot::types::CommitOptions {
+					expected_head_txid: request.expected_head_txid,
+				},
 			)
 			.await
 		{
-			Ok(()) => Ok(protocol::SqliteCommitResponse::SqliteCommitOk),
-			Err(err) => Ok(protocol::SqliteCommitResponse::SqliteErrorResponse(
-				protocol::SqliteErrorResponse {
-					message: sqlite_error_reason(&err),
+			Ok(result) => Ok(protocol::SqliteCommitResponse::SqliteCommitOk(
+				protocol::SqliteCommitOk {
+					head_txid: Some(result.head_txid),
 				},
 			)),
+			Err(err) => Ok(protocol::SqliteCommitResponse::SqliteErrorResponse(
+				sqlite_error_response(&err),
+			)),
 		}
 	}
 }
@@ -104,3 +119,19 @@ fn sqlite_error_reason(err: &anyhow::Error) -> String {
 		.collect::<Vec<_>>()
 		.join(": ")
 }
+
+fn sqlite_error_response(err: &anyhow::Error) -> protocol::SqliteErrorResponse {
+	let structured = depot_error(err)
+		.map(|err| rivet_error::RivetError::extract(&err.clone().build()))
+		.unwrap_or_else(|| rivet_error::RivetError::extract(err));
+	protocol::SqliteErrorResponse {
+		group: structured.group().to_string(),
+		code: structured.code().to_string(),
+		message: sqlite_error_reason(err),
+	}
+}
+
+fn depot_error(err: &anyhow::Error) -> Option<&SqliteStorageError> {
+	err.chain()
+		.find_map(|source| source.downcast_ref::<SqliteStorageError>())
+}

engine/packages/depot-client-embedded/Cargo.toml

@@ -1,5 +1,12 @@
 //! Shared SQLite execution types for local and remote depot client backends.
 
+pub const HEAD_FENCE_MISMATCH_GROUP: &str = "depot";
+pub const HEAD_FENCE_MISMATCH_CODE: &str = "head_fence_mismatch";
+
+pub fn is_head_fence_mismatch(group: &str, code: &str) -> bool {
+	group == HEAD_FENCE_MISMATCH_GROUP && code == HEAD_FENCE_MISMATCH_CODE
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub enum BindParam {
 	Null,

engine/packages/depot-client-embedded/src/lib.rs

@@ -29,6 +29,7 @@ depot = { workspace = true, features = ["test-faults"] }
 futures-util.workspace = true
 gas.workspace = true
 rivet-config.workspace = true
+rivet-error.workspace = true
 rivet-pools.workspace = true
 rivet-test-deps.workspace = true
 sha2.workspace = true

engine/packages/depot-client-types/src/lib.rs

@@ -10,7 +10,7 @@ use crate::{
 		SqliteVfsMetricsSnapshot, VfsConfig, VfsPreloadHintSnapshot,
 		fetch_initial_pages_for_registration,
 	},
-	worker::SqliteWorkerHandle,
+	worker::{SqliteWorkerFatalError, SqliteWorkerHandle},
 };
 
 #[derive(Clone)]
@@ -70,7 +70,8 @@ impl NativeDatabaseHandle {
 	}
 
 	pub async fn exec(&self, sql: String) -> Result<QueryResult> {
-		self.worker.exec(sql).await
+		self.check_fatal_error()?;
+		self.map_worker_result(self.worker.exec(sql).await)
 	}
 
 	pub async fn query(&self, sql: String, params: Option<Vec<BindParam>>) -> Result<QueryResult> {
@@ -91,11 +92,15 @@ impl NativeDatabaseHandle {
 		sql: String,
 		params: Option<Vec<BindParam>>,
 	) -> Result<ExecuteResult> {
-		self.worker.execute(sql, params).await
+		self.check_fatal_error()?;
+		self.map_worker_result(self.worker.execute(sql, params).await)
 	}
 
 	pub async fn close(&self) -> Result<()> {
-		self.worker.close().await
+		match self.worker.close().await {
+			Ok(()) => Ok(()),
+			Err(error) => Err(self.fatal_error().unwrap_or(error)),
+		}
 	}
 
 	pub async fn wait_for_worker_failure(&self) -> bool {
@@ -106,6 +111,10 @@ impl NativeDatabaseHandle {
 		self.vfs.take_last_error()
 	}
 
+	pub fn clone_fatal_error(&self) -> Option<String> {
+		self.vfs.clone_fatal_error()
+	}
+
 	pub fn snapshot_preload_hints(&self) -> VfsPreloadHintSnapshot {
 		self.vfs.snapshot_preload_hints()
 	}
@@ -130,7 +139,30 @@ impl NativeDatabaseHandle {
 	}
 
 	async fn initialize(&self) -> Result<()> {
-		self.worker.wait_ready().await
+		self.map_worker_result(self.worker.wait_ready().await)
+	}
+
+	fn check_fatal_error(&self) -> Result<()> {
+		if let Some(error) = self.fatal_error() {
+			return Err(error);
+		}
+
+		Ok(())
+	}
+
+	fn map_worker_result<T>(&self, result: Result<T>) -> Result<T> {
+		match result {
+			Ok(value) => {
+				self.check_fatal_error()?;
+				Ok(value)
+			}
+			Err(error) => Err(self.fatal_error().unwrap_or(error)),
+		}
+	}
+
+	fn fatal_error(&self) -> Option<anyhow::Error> {
+		self.clone_fatal_error()
+			.map(|message| SqliteWorkerFatalError::new(message).into())
 	}
 }