fix(pegboard-envoy): allow sqlite requests across lifecycle states

Author: NathanFlurry

engine/packages/pegboard-envoy/src/actor_lifecycle.rs

@@ -305,30 +305,3 @@ async fn close_actor_on_shutdown(
 		}
 	}
 }
-
-pub async fn assert_sqlite_actor_active(
-	conn: &Conn,
-	actor_id: &str,
-	sqlite_generation: u64,
-) -> Result<ActiveActor> {
-	// Stopping is accepted in addition to Running: the actor still owns its sqlite
-	// generation until actor_stopped runs, and may flush a final commit while draining.
-	let active = conn
-		.active_actors
-		.get_async(actor_id)
-		.await
-		.map(|entry| entry.get().clone())
-		.context("sqlite actor is not active on envoy connection")?;
-
-	let active_sqlite_generation = active
-		.sqlite_generation
-		.context("sqlite actor is still starting")?;
-	ensure!(
-		active_sqlite_generation == sqlite_generation,
-		"sqlite request generation {} did not match active generation {}",
-		sqlite_generation,
-		active_sqlite_generation
-	);
-
-	Ok(active)
-}

engine/packages/pegboard-envoy/src/sqlite_runtime.rs

@@ -3,35 +3,10 @@ use std::sync::Arc;
 use anyhow::Result;
 use gas::prelude::StandaloneCtx;
 use rivet_envoy_protocol as protocol;
-use sqlite_storage::{compaction::CompactionCoordinator, engine::SqliteEngine, open::OpenResult};
-use tokio::sync::OnceCell;
-use universaldb::Subspace;
-
-static SQLITE_ENGINE: OnceCell<Arc<SqliteEngine>> = OnceCell::const_new();
+use sqlite_storage::{engine::SqliteEngine, open::OpenResult};
 
 pub async fn shared_engine(ctx: &StandaloneCtx) -> Result<Arc<SqliteEngine>> {
-	let db = (*ctx.udb()?).clone();
-	let subspace = sqlite_subspace();
-
-	SQLITE_ENGINE
-		.get_or_try_init(|| async move {
-			tracing::info!("initializing shared sqlite dispatch runtime");
-
-			let (engine, compaction_rx) = SqliteEngine::new(db, subspace.clone());
-			let engine = Arc::new(engine);
-			tokio::spawn(CompactionCoordinator::run(
-				compaction_rx,
-				Arc::clone(&engine),
-			));
-
-			Ok(engine)
-		})
-		.await
-		.cloned()
-}
-
-fn sqlite_subspace() -> Subspace {
-	pegboard::keys::subspace().subspace(&("sqlite-storage",))
+	pegboard::actor_sqlite::shared_engine(ctx).await
 }
 
 pub fn protocol_sqlite_startup_data(startup: OpenResult) -> protocol::SqliteStartupData {

engine/packages/pegboard-envoy/src/ws_to_tunnel_task.rs

@@ -703,8 +703,6 @@ async fn handle_sqlite_get_pages(
 ) -> Result<protocol::SqliteGetPagesResponse> {
 	validate_sqlite_get_pages_request(&request)?;
 	validate_sqlite_actor(ctx, conn, &request.actor_id).await?;
-	actor_lifecycle::assert_sqlite_actor_active(conn, &request.actor_id, request.generation)
-		.await?;
 
 	match conn
 		.sqlite_engine
@@ -749,8 +747,6 @@ async fn handle_sqlite_commit(
 	let decode_request_start = Instant::now();
 	validate_sqlite_dirty_pages("sqlite commit", &request.dirty_pages)?;
 	validate_sqlite_actor(ctx, conn, &request.actor_id).await?;
-	actor_lifecycle::assert_sqlite_actor_active(conn, &request.actor_id, request.generation)
-		.await?;
 	let decode_request_duration = decode_request_start.elapsed();
 	conn.sqlite_engine.metrics().observe_commit_phase(
 		"fast",
@@ -817,8 +813,6 @@ async fn handle_sqlite_commit_stage(
 	request: protocol::SqliteCommitStageRequest,
 ) -> Result<protocol::SqliteCommitStageResponse> {
 	validate_sqlite_actor(ctx, conn, &request.actor_id).await?;
-	actor_lifecycle::assert_sqlite_actor_active(conn, &request.actor_id, request.generation)
-		.await?;
 
 	match conn
 		.sqlite_engine
@@ -856,8 +850,6 @@ async fn handle_sqlite_commit_stage_begin(
 	request: protocol::SqliteCommitStageBeginRequest,
 ) -> Result<protocol::SqliteCommitStageBeginResponse> {
 	validate_sqlite_actor(ctx, conn, &request.actor_id).await?;
-	actor_lifecycle::assert_sqlite_actor_active(conn, &request.actor_id, request.generation)
-		.await?;
 
 	match conn
 		.sqlite_engine
@@ -892,8 +884,6 @@ async fn handle_sqlite_commit_finalize(
 ) -> Result<protocol::SqliteCommitFinalizeResponse> {
 	let decode_request_start = Instant::now();
 	validate_sqlite_actor(ctx, conn, &request.actor_id).await?;
-	actor_lifecycle::assert_sqlite_actor_active(conn, &request.actor_id, request.generation)
-		.await?;
 	conn.sqlite_engine.metrics().observe_commit_phase(
 		"slow",
 		"decode_request",

engine/packages/pegboard-outbound/src/lib.rs

@@ -10,15 +10,14 @@ use rivet_runtime::TermSignal;
 use rivet_types::actor::RunnerPoolError;
 use rivet_types::runner_configs::RunnerConfigKind;
 use sqlite_storage::{
-	compaction::CompactionCoordinator,
 	engine::SqliteEngine,
 	open::{OpenConfig, OpenResult},
 	types::{FetchedPage, SqliteMeta},
 };
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
-use tokio::{sync::OnceCell, task::JoinHandle};
+use tokio::task::JoinHandle;
 use universalpubsub::NextOutput;
 use vbare::OwnedVersionedData;
 
@@ -29,25 +28,9 @@ const X_RIVET_POOL_NAME: HeaderName = HeaderName::from_static("x-rivet-pool-name
 const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-token");
 const X_RIVET_NAMESPACE_NAME: HeaderName = HeaderName::from_static("x-rivet-namespace-name");
 const SHUTDOWN_PROGRESS_INTERVAL: Duration = Duration::from_secs(7);
-static SQLITE_ENGINE: OnceCell<Arc<SqliteEngine>> = OnceCell::const_new();
 
 async fn shared_sqlite_engine(ctx: &StandaloneCtx) -> Result<Arc<SqliteEngine>> {
-	let db = (*ctx.udb()?).clone();
-	let subspace = pegboard::keys::subspace().subspace(&("sqlite-storage",));
-
-	SQLITE_ENGINE
-		.get_or_try_init(|| async move {
-			let (engine, compaction_rx) = SqliteEngine::new(db, subspace.clone());
-			let engine = Arc::new(engine);
-			tokio::spawn(CompactionCoordinator::run(
-				compaction_rx,
-				Arc::clone(&engine),
-			));
-
-			Ok(engine)
-		})
-		.await
-		.cloned()
+	pegboard::actor_sqlite::shared_engine(ctx).await
 }
 
 fn protocol_sqlite_startup_data(startup: OpenResult) -> protocol::SqliteStartupData {

engine/packages/pegboard/src/actor_sqlite.rs

@@ -1,15 +1,17 @@
-use std::time::Instant;
+use std::{sync::Arc, time::Instant};
 
 use anyhow::{Context, Result, ensure};
-use gas::prelude::{Id, util::timestamp};
+use gas::prelude::{Id, StandaloneCtx, util::timestamp};
 use rivet_envoy_protocol as protocol;
 use sqlite_storage::{
+	compaction::CompactionCoordinator,
 	commit::{CommitFinalizeRequest, CommitStageBeginRequest, CommitStageRequest},
 	engine::SqliteEngine,
 	ltx::{LtxHeader, encode_ltx_v3},
 	open::OpenConfig,
 	types::{DirtyPage, SQLITE_PAGE_SIZE, SqliteOrigin},
 };
+use tokio::sync::OnceCell;
 use universaldb::Subspace;
 
 use crate::{actor_kv::Recipient, metrics};
@@ -28,6 +30,7 @@ const FILE_TAG_JOURNAL: u8 = 0x01;
 const FILE_TAG_WAL: u8 = 0x02;
 const FILE_TAG_SHM: u8 = 0x03;
 const SQLITE_MAGIC: &[u8; 16] = b"SQLite format 3\0";
+static SQLITE_ENGINE: OnceCell<Arc<SqliteEngine>> = OnceCell::const_new();
 
 pub fn sqlite_subspace() -> Subspace {
 	crate::keys::subspace().subspace(&("sqlite-storage",))
@@ -39,6 +42,26 @@ pub fn new_engine(
 	SqliteEngine::new(db, sqlite_subspace())
 }
 
+pub async fn shared_engine(ctx: &StandaloneCtx) -> Result<Arc<SqliteEngine>> {
+	let db = (*ctx.udb()?).clone();
+
+	SQLITE_ENGINE
+		.get_or_try_init(|| async move {
+			tracing::info!("initializing shared sqlite dispatch runtime");
+
+			let (engine, compaction_rx) = new_engine(db);
+			let engine = Arc::new(engine);
+			tokio::spawn(CompactionCoordinator::run(
+				compaction_rx,
+				Arc::clone(&engine),
+			));
+
+			Ok(engine)
+		})
+		.await
+		.cloned()
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Hash)]
 pub struct MigrateV1ToV2Input {
 	pub actor_id: Id,

engine/sdks/rust/envoy-client/src/actor.rs

@@ -199,7 +199,8 @@ async fn actor_inner(
 		.await;
 
 	if let Err(error) = start_result {
-		tracing::error!(?error, "actor start failed");
+		let error_chain = error.chain().map(ToString::to_string).collect::<Vec<_>>();
+		tracing::error!(?error, error_chain = ?error_chain, "actor start failed");
 		send_event(
 			&mut ctx,
 			protocol::Event::EventActorStateUpdate(protocol::EventActorStateUpdate {

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -1398,7 +1398,7 @@ impl ActorTask {
 		let clean_exit = match outcome {
 			Ok(Ok(())) => true,
 			Ok(Err(error)) => {
-				tracing::error!(?error, "actor run handler failed");
+				log_actor_error(&error, "actor run handler failed");
 				false
 			}
 			Err(error) => {
@@ -1529,7 +1529,7 @@ impl ActorTask {
 		match (&mut run_handle).await {
 			Ok(Ok(())) => {}
 			Ok(Err(error)) => {
-				tracing::error!(?error, "actor run handler failed during shutdown");
+				log_actor_error(&error, "actor run handler failed during shutdown");
 			}
 			Err(error) => {
 				if !error.is_cancelled() {
@@ -2129,6 +2129,18 @@ fn clone_shutdown_result(result: &Result<()>) -> Result<()> {
 	}
 }
 
+fn log_actor_error(error: &anyhow::Error, log_message: &'static str) {
+	let structured = rivet_error::RivetError::extract(error);
+	tracing::error!(
+		?error,
+		group = structured.group(),
+		code = structured.code(),
+		message = %structured.message(),
+		metadata = ?structured.metadata(),
+		"{log_message}"
+	);
+}
+
 fn result_outcome<T>(result: &Result<T>) -> &'static str {
 	match result {
 		Ok(_) => "ok",

rivetkit-rust/packages/rivetkit-sqlite/src/vfs.rs

@@ -1308,6 +1308,15 @@ impl VfsContext {
 		{
 			Ok(outcome) => outcome,
 			Err(err) => {
+				tracing::error!(
+					actor_id = %self.actor_id,
+					generation = request.generation,
+					expected_head_txid = request.expected_head_txid,
+					new_db_size_pages = request.new_db_size_pages,
+					dirty_pages = request.dirty_pages.len(),
+					?err,
+					"sqlite flush commit failed"
+				);
 				mark_dead_for_non_fence_commit_error(self, &err);
 				return Err(err);
 			}
@@ -1387,6 +1396,15 @@ impl VfsContext {
 		{
 			Ok(outcome) => outcome,
 			Err(err) => {
+				tracing::error!(
+					actor_id = %self.actor_id,
+					generation = request.generation,
+					expected_head_txid = request.expected_head_txid,
+					new_db_size_pages = request.new_db_size_pages,
+					dirty_pages = request.dirty_pages.len(),
+					?err,
+					"sqlite atomic commit failed"
+				);
 				mark_dead_for_non_fence_commit_error(self, &err);
 				return Err(err);
 			}
@@ -1460,13 +1478,28 @@ fn assert_batch_atomic_probe(db: *mut sqlite3, vfs: &SqliteVfs) -> std::result::
 	";
 
 	if let Err(err) = sqlite_exec(db, probe_sql) {
+		let last_error = vfs.clone_last_error();
+		tracing::error!(
+			%err,
+			last_error = ?last_error,
+			commit_atomic_before,
+			"sqlite batch atomic probe failed"
+		);
 		cleanup_batch_atomic_probe(db);
+		if let Some(last_error) = last_error {
+			return Err(format!(
+				"batch atomic probe failed: {err}; vfs last_error: {last_error}"
+			));
+		}
 		return Err(format!("batch atomic probe failed: {err}"));
 	}
 
 	let commit_atomic_after = vfs.commit_atomic_count();
 	if commit_atomic_after == commit_atomic_before {
 		tracing::error!(
+			commit_atomic_before,
+			commit_atomic_after,
+			last_error = ?vfs.clone_last_error(),
 			"batch atomic writes not active for sqlite, SQLITE_ENABLE_BATCH_ATOMIC_WRITE may be missing"
 		);
 		cleanup_batch_atomic_probe(db);
@@ -2215,6 +2248,12 @@ unsafe extern "C" fn io_sync(p_file: *mut sqlite3_file, _flags: c_int) -> c_int
 		match ctx.flush_dirty_pages() {
 			Ok(_) => SQLITE_OK,
 			Err(err) => {
+				tracing::error!(
+					actor_id = %ctx.actor_id,
+					last_error = ?ctx.clone_last_error(),
+					?err,
+					"sqlite sync failed"
+				);
 				mark_dead_from_fence_commit_error(ctx, &err);
 				SQLITE_IOERR_FSYNC
 			}
@@ -2280,6 +2319,12 @@ unsafe extern "C" fn io_file_control(
 					SQLITE_OK
 				}
 				Err(err) => {
+					tracing::error!(
+						actor_id = %ctx.actor_id,
+						last_error = ?ctx.clone_last_error(),
+						?err,
+						"sqlite atomic write file control failed"
+					);
 					mark_dead_from_fence_commit_error(ctx, &err);
 					SQLITE_IOERR
 				}
@@ -2532,6 +2577,10 @@ impl SqliteVfs {
 		unsafe { (*self.ctx_ptr).take_last_error() }
 	}
 
+	fn clone_last_error(&self) -> Option<String> {
+		unsafe { (*self.ctx_ptr).clone_last_error() }
+	}
+
 	fn register_with_transport(
 		name: &str,
 		transport: SqliteTransport,
@@ -2659,6 +2708,13 @@ pub fn open_database(
 	};
 	if rc != SQLITE_OK {
 		let message = sqlite_error_message(db);
+		tracing::error!(
+			file_name,
+			rc,
+			%message,
+			last_error = ?vfs.clone_last_error(),
+			"failed to open sqlite database with custom VFS"
+		);
 		if !db.is_null() {
 			unsafe {
 				sqlite3_close(db);
@@ -2676,6 +2732,13 @@ pub fn open_database(
 		"PRAGMA locking_mode = EXCLUSIVE;",
 	] {
 		if let Err(err) = sqlite_exec(db, pragma) {
+			tracing::error!(
+				file_name,
+				pragma,
+				%err,
+				last_error = ?vfs.clone_last_error(),
+				"failed to configure sqlite database"
+			);
 			unsafe {
 				sqlite3_close(db);
 			}
@@ -2684,6 +2747,12 @@ pub fn open_database(
 	}
 
 	if let Err(err) = assert_batch_atomic_probe(db, &vfs) {
+		tracing::error!(
+			file_name,
+			%err,
+			last_error = ?vfs.clone_last_error(),
+			"failed to verify sqlite batch atomic writes"
+		);
 		unsafe {
 			sqlite3_close(db);
 		}