fix(pegboard): delete branch-backed sqlite on actor destroy

engine/packages/pegboard/src/actor_sqlite.rs

@@ -48,6 +48,67 @@ pub fn clear_v2_storage_for_destroy(tx: &universaldb::Transaction, actor_id: Id)
 	}
 }
 
+pub async fn clear_branch_storage_for_destroy(
+	db: &universaldb::Database,
+	namespace_id: Id,
+	actor_id: Id,
+) -> Result<()> {
+	let bucket = BucketId::from_gas_id(namespace_id);
+	let database_name = actor_id.to_string();
+
+	let branch_id = match resolve_branch_for_destroy(db, bucket, &database_name).await? {
+		Some(id) => id,
+		None => return Ok(()),
+	};
+
+	match depot_branch::delete_database(db, bucket, branch_id).await {
+		Ok(()) => Ok(()),
+		Err(err) if is_database_not_found(&err) => {
+			tracing::debug!(?branch_id, %database_name, "branch already deleted on destroy");
+			Ok(())
+		}
+		Err(err) => Err(err),
+	}
+}
+
+async fn resolve_branch_for_destroy(
+	db: &universaldb::Database,
+	bucket: BucketId,
+	database_name: &str,
+) -> Result<Option<depot::conveyer::types::DatabaseBranchId>> {
+	let database_name_owned = database_name.to_string();
+	let database_name_for_txn = database_name_owned.clone();
+	let res = db
+		.txn("pegboard_actor_sqlite_resolve_branch", move |tx| {
+			let database_name = database_name_for_txn.clone();
+			async move {
+				depot_branch::resolve_database_branch(
+					&tx,
+					bucket,
+					&database_name,
+					universaldb::utils::IsolationLevel::Snapshot,
+				)
+				.await
+			}
+		})
+		.await;
+
+	match res {
+		Ok(branch) => Ok(branch),
+		Err(err) if is_database_not_found(&err) => {
+			tracing::debug!(%database_name_owned, "database not found while resolving branch for destroy");
+			Ok(None)
+		}
+		Err(err) => Err(err),
+	}
+}
+
+fn is_database_not_found(err: &anyhow::Error) -> bool {
+	err.chain()
+		.find_map(|c| c.downcast_ref::<depot::conveyer::error::SqliteStorageError>())
+		.is_some_and(|e| matches!(e, depot::conveyer::error::SqliteStorageError::DatabaseNotFound))
+}
+
 fn prefix_range(prefix: &[u8]) -> (Vec<u8>, Vec<u8>) {
 	universaldb::tuple::Subspace::from_bytes(prefix.to_vec()).range()
 }

engine/packages/pegboard/src/workflows/actor/destroy.rs

@@ -180,8 +180,10 @@ struct ClearKvOutput {
 
 #[activity(ClearKv)]
 async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result<ClearKvOutput> {
-	let final_size = ctx
-		.udb()?
+	let state = ctx.state::<State>()?;
+	let namespace_id = state.namespace_id;
+	let udb = ctx.udb()?;
+	let final_size = udb
 		.txn("pegboard_actor_clear_kv", |tx| async move {
 			let subspace = keys::actor_kv::subspace(input.actor_id);
 
@@ -197,6 +199,10 @@ async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result<ClearKvOutp
 		.custom_instrument(tracing::info_span!("actor_clear_kv_tx"))
 		.await?;
 
+	// Keep outside txn - branch delete has its own txn.
+	crate::actor_sqlite::clear_branch_storage_for_destroy(&udb, namespace_id, input.actor_id)
+		.await?;
+
 	Ok(ClearKvOutput { final_size })
 }
 

engine/packages/pegboard/src/workflows/actor2/mod.rs

@@ -1240,8 +1240,9 @@ async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result<ClearKvOutp
 	let state = ctx.state::<State>()?;
 
 	let actor_id = state.actor_id;
-	let final_size = ctx
-		.udb()?
+	let namespace_id = state.namespace_id;
+	let udb = ctx.udb()?;
+	let final_size = udb
 		.txn("pegboard_actor2_clear_kv", |tx| async move {
 			let subspace = crate::keys::actor_kv::subspace(actor_id);
 
@@ -1257,6 +1258,9 @@ async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result<ClearKvOutp
 		.custom_instrument(tracing::info_span!("actor_clear_kv_tx"))
 		.await?;
 
+	// Keep outside txn - branch delete has its own txn.
+	crate::actor_sqlite::clear_branch_storage_for_destroy(&udb, namespace_id, actor_id).await?;
+
 	Ok(ClearKvOutput { final_size })
 }
 

engine/packages/pegboard/tests/actor_sqlite_destroy.rs

@@ -1,11 +1,16 @@
 use std::sync::Arc;
 
 use anyhow::{Result, anyhow};
-use depot::keys::{
-	delta_chunk_key, meta_compact_key, meta_compactor_lease_key, meta_head_key, meta_quota_key,
-	pidx_delta_key, shard_key,
+use depot::{
+	conveyer::{Db, branch},
+	keys::{
+		branches_refcount_key, delta_chunk_key, meta_compact_key, meta_compactor_lease_key,
+		meta_head_key, meta_quota_key, pidx_delta_key, shard_key,
+	},
+	types::{BucketId, DatabaseBranchId, DirtyPage},
 };
 use gas::prelude::Id;
+use rivet_pools::NodeId;
 use tempfile::Builder;
 use universaldb::utils::IsolationLevel::Snapshot;
 
@@ -98,3 +103,135 @@ async fn actor_destroy_in_one_tx() -> Result<()> {
 
 	Ok(())
 }
+
+async fn read_refcount(db: &universaldb::Database, branch: DatabaseBranchId) -> Result<i64> {
+	let bytes = db
+		.txn("test_pegboardactor_sqlite_destroy", move |tx| async move {
+			Ok(tx.informal().get(&branches_refcount_key(branch), Snapshot).await?)
+		})
+		.await?
+		.expect("branch refcount should exist");
+	let arr: [u8; 8] = bytes.as_slice().try_into().expect("refcount i64 LE");
+	Ok(i64::from_le_bytes(arr))
+}
+
+fn branch_page(pgno: u32) -> DirtyPage {
+	DirtyPage {
+		pgno,
+		bytes: vec![0xAA; depot::keys::PAGE_SIZE as usize],
+	}
+}
+
+#[tokio::test]
+async fn actor_destroy_clears_branch_backed_database() -> Result<()> {
+	let db = test_db().await?;
+	let namespace_id = Id::new_v1(0x1111);
+	let actor_id = Id::new_v1(0x2222);
+	let bucket = BucketId::from_gas_id(namespace_id);
+	let db_arc = Arc::new(db.clone());
+	let depot_db = Db::new(db_arc.clone(), namespace_id, actor_id.to_string(), NodeId::new());
+	depot_db.commit(vec![branch_page(1)], 1, 1_000).await?;
+
+	let branch = db
+		.txn("test_resolve", |tx| {
+			let actor_id = actor_id.to_string();
+			async move {
+				branch::resolve_database_branch(&tx, bucket, &actor_id, Snapshot).await
+			}
+		})
+		.await?
+		.expect("branch should exist before destroy");
+	assert_eq!(read_refcount(&db, branch).await?, 1);
+	assert_eq!(branch::list_databases(&db, bucket).await?, vec![branch]);
+
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+
+	assert_eq!(read_refcount(&db, branch).await?, 0);
+	assert_eq!(
+		branch::list_databases(&db, bucket).await?,
+		Vec::<DatabaseBranchId>::new()
+	);
+
+	Ok(())
+}
+
+#[tokio::test]
+async fn actor_destroy_branch_is_idempotent() -> Result<()> {
+	let db = test_db().await?;
+	let namespace_id = Id::new_v1(0x3333);
+	let actor_id = Id::new_v1(0x4444);
+	let bucket = BucketId::from_gas_id(namespace_id);
+	let db_arc = Arc::new(db.clone());
+	let depot_db = Db::new(db_arc, namespace_id, actor_id.to_string(), NodeId::new());
+	depot_db.commit(vec![branch_page(1)], 1, 1_000).await?;
+
+	let branch = db
+		.txn("test_resolve", |tx| {
+			let actor_id = actor_id.to_string();
+			async move {
+				branch::resolve_database_branch(&tx, bucket, &actor_id, Snapshot).await
+			}
+		})
+		.await?
+		.expect("branch should exist");
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+	// Second destroy is a no-op.
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+
+	assert_eq!(read_refcount(&db, branch).await?, 0);
+	Ok(())
+}
+
+#[tokio::test]
+async fn actor_destroy_branch_noop_when_no_database() -> Result<()> {
+	let db = test_db().await?;
+	let namespace_id = Id::new_v1(0x5555);
+	let actor_id = Id::new_v1(0x6666);
+
+	// No depot database was ever created.
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+
+	// Also works when bucket itself was never created.
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+	Ok(())
+}
+
+#[tokio::test]
+async fn actor_destroy_clears_both_layouts_like_workflow() -> Result<()> {
+	// Simulates ClearKv wiring.
+	let db = test_db().await?;
+	let namespace_id = Id::new_v1(0x7777);
+	let actor_id = Id::new_v1(0x8888);
+	let bucket = BucketId::from_gas_id(namespace_id);
+
+	// Legacy keys.
+	let legacy_keys = sqlite_keys(actor_id);
+	seed(&db, &legacy_keys).await?;
+
+	// Branch-backed DB.
+	let depot_db = Db::new(Arc::new(db.clone()), namespace_id, actor_id.to_string(), NodeId::new());
+	depot_db.commit(vec![branch_page(1)], 1, 1_000).await?;
+	let branch = db
+		.txn("test_resolve", |tx| {
+			let actor_id = actor_id.to_string();
+			async move { branch::resolve_database_branch(&tx, bucket, &actor_id, Snapshot).await }
+		})
+		.await?
+		.expect("branch should exist");
+
+	db.txn("test_workflow_clear_kv", |tx| async move {
+		pegboard::actor_sqlite::clear_v2_storage_for_destroy(&tx, actor_id);
+		Ok(())
+	})
+	.await?;
+	pegboard::actor_sqlite::clear_branch_storage_for_destroy(&db, namespace_id, actor_id).await?;
+
+	// Both layouts gone.
+	for key in legacy_keys {
+		assert!(!value_exists(&db, key).await?);
+	}
+	assert_eq!(read_refcount(&db, branch).await?, 0);
+	assert_eq!(branch::list_databases(&db, bucket).await?, Vec::<DatabaseBranchId>::new());
+
+	Ok(())
+}