diff --git a/architecture/compute-runtimes.md b/architecture/compute-runtimes.md index 2a36073486..cd697c5d5d 100644 --- a/architecture/compute-runtimes.md +++ b/architecture/compute-runtimes.md @@ -191,6 +191,9 @@ backend is already absent (`deleted = false`), the request removes gateway state synchronously. Sandbox row removal remains bound to the stable ID and resource version. Settings retain their existing best-effort name-based cleanup; SSH sessions, indexes, and watch/log buses are cleaned after confirmed removal. +Owned-record cleanup discovers records before mutating them and uses bounded +set-based deletes so teardown cannot amplify one sandbox into an unbounded +sequence of individual persistence writes. The request acquires both locks before starting owned work, so cancellation while queued does not leave a delete armed. After that commitment point, the diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index b85641c986..41f5352812 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -69,7 +69,7 @@ use std::net::SocketAddr; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::{Arc, Mutex as StdMutex, Weak}; -use std::time::Duration; +use std::time::{Duration, Instant}; #[cfg(unix)] use tokio::net::UnixStream; use tokio::sync::{Mutex, watch}; @@ -3084,21 +3084,64 @@ impl ComputeRuntime { sandbox_id: &str, workspace: &str, ) -> Result<(), String> { - let records = self + let started = Instant::now(); + let mut offset = 0_u32; + let mut scanned = 0_usize; + let mut decode_failures = 0_usize; + let mut session_ids = Vec::new(); + + loop { + let records = self + .store + .list( + SshSession::object_type(), + workspace, + LIFECYCLE_SWEEP_PAGE_SIZE, + offset, + ) + .await + .map_err(|e| format!("list SSH sessions: {e}"))?; + let page_len = records.len(); + scanned += page_len; + + for record in records { + match SshSession::decode(record.payload.as_slice()) { + Ok(session) if session.sandbox_id == sandbox_id => { + session_ids.push(session.object_id().to_string()); + } + Ok(_) => {} + Err(_) => decode_failures += 1, + } + } + + if page_len < LIFECYCLE_SWEEP_PAGE_SIZE as usize { + break; + } + let page_len = u32::try_from(page_len) + .map_err(|_| "SSH session cleanup page length overflow".to_string())?; + offset = offset + .checked_add(page_len) + .ok_or_else(|| "SSH session cleanup pagination overflow".to_string())?; + } + + let matched = session_ids.len(); + let deleted = self .store - .list(SshSession::object_type(), workspace, 1000, 0) + .delete_many(SshSession::object_type(), &session_ids) .await - .map_err(|e| format!("list SSH sessions: {e}"))?; + .map_err(|e| format!("delete sandbox SSH sessions: {e}"))?; - for record in records { - if let Ok(session) = SshSession::decode(record.payload.as_slice()) - && session.sandbox_id == sandbox_id - { - self.store - .delete(SshSession::object_type(), session.object_id()) - .await - .map_err(|e| format!("delete SSH session {}: {e}", session.object_id()))?; - } + if matched > 0 || decode_failures > 0 { + debug!( + sandbox_id, + workspace, + scanned, + matched, + deleted, + decode_failures, + elapsed_ms = started.elapsed().as_millis(), + "Sandbox SSH session cleanup complete" + ); } Ok(()) @@ -7151,6 +7194,43 @@ mod tests { assert_eq!(driver.delete_calls(), 1); } + #[tokio::test] + async fn sandbox_ssh_session_cleanup_batches_across_list_pages() { + let runtime = test_runtime(ControlledDriver::new()).await; + for idx in 0..(LIFECYCLE_SWEEP_PAGE_SIZE + 5) { + let session = ssh_session_record(&format!("owned-{idx:04}"), "sb-owned"); + runtime.store.put_message(&session).await.unwrap(); + } + for idx in 0..7 { + let session = ssh_session_record(&format!("unrelated-{idx:04}"), "sb-unrelated"); + runtime.store.put_message(&session).await.unwrap(); + } + + runtime + .cleanup_sandbox_ssh_sessions("sb-owned", "default") + .await + .unwrap(); + + assert_eq!( + runtime + .store + .count_in_workspace(SshSession::object_type(), "default") + .await + .unwrap(), + 7 + ); + for idx in 0..7 { + assert!( + runtime + .store + .get_message::(&format!("unrelated-{idx:04}")) + .await + .unwrap() + .is_some() + ); + } + } + #[tokio::test] async fn already_absent_driver_resource_is_removed_synchronously() { let driver = ControlledDriver::new(); diff --git a/crates/openshell-server/src/persistence/mod.rs b/crates/openshell-server/src/persistence/mod.rs index 516faf4fe4..9c63e7d351 100644 --- a/crates/openshell-server/src/persistence/mod.rs +++ b/crates/openshell-server/src/persistence/mod.rs @@ -26,6 +26,12 @@ pub const DRAFT_CHUNK_OBJECT_TYPE: &str = "draft_policy_chunk"; pub type PersistenceResult = Result; +/// Maximum number of object ids sent in one set-based delete statement. +/// +/// Keep this well below `SQLite`'s bind-variable limit. Backends split larger +/// requests into independently retryable, bounded write statements. +pub const DELETE_MANY_BATCH_SIZE: usize = 128; + /// Persistence-layer error type. #[derive(Debug, Error, Clone)] pub enum PersistenceError { @@ -415,6 +421,22 @@ impl Store { store_dispatch_traced!(self.delete(object_type, id)) } + /// Delete objects of one type by id in bounded, set-based statements. + #[tracing::instrument( + name = "store", + skip_all, + fields( + otel.name = "store.delete_many", + otel.status_code = tracing::field::Empty, + object_type = %object_type, + object_count = ids.len(), + batch_count = ids.len().div_ceil(DELETE_MANY_BATCH_SIZE), + ) + )] + pub async fn delete_many(&self, object_type: &str, ids: &[String]) -> PersistenceResult { + store_dispatch_traced!(self.delete_many(object_type, ids)) + } + /// Count objects of a given type within a workspace. #[tracing::instrument( name = "store", diff --git a/crates/openshell-server/src/persistence/postgres.rs b/crates/openshell-server/src/persistence/postgres.rs index 8bab0ada96..f88ce81b99 100644 --- a/crates/openshell-server/src/persistence/postgres.rs +++ b/crates/openshell-server/src/persistence/postgres.rs @@ -14,11 +14,11 @@ use openshell_core::SetResourceVersion; use openshell_core::proto::Sandbox; use prost::Message; use sqlx::postgres::PgPoolOptions; -use sqlx::{Connection, PgPool, Row}; +use sqlx::{Connection, PgPool, Postgres, QueryBuilder, Row}; static POSTGRES_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/postgres"); -use super::{DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; +use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; #[derive(Debug, Clone)] pub struct PostgresStore { @@ -391,6 +391,28 @@ WHERE object_type = $1 AND workspace = $2 AND name = $3 Ok(result.rows_affected() > 0) } + pub async fn delete_many(&self, object_type: &str, ids: &[String]) -> PersistenceResult { + let mut deleted = 0_u64; + for ids in ids.chunks(DELETE_MANY_BATCH_SIZE) { + let mut query = + QueryBuilder::::new("DELETE FROM objects WHERE object_type = "); + query.push_bind(object_type).push(" AND id IN ("); + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + + deleted += query + .build() + .execute(&self.pool) + .await + .map_err(|e| map_db_error(&e))? + .rows_affected(); + } + Ok(deleted) + } + pub async fn count_in_workspace( &self, object_type: &str, diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index 658530f753..a1688b4c49 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -15,13 +15,13 @@ use openshell_core::paths::set_file_owner_only; use openshell_core::proto::Sandbox; use prost::Message; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; -use sqlx::{Connection, Row, SqlitePool}; +use sqlx::{Connection, QueryBuilder, Row, Sqlite, SqlitePool}; use std::path::{Path, PathBuf}; use std::str::FromStr; static SQLITE_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/sqlite"); -use super::{DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; +use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; #[derive(Debug, Clone)] pub struct SqliteStore { @@ -416,6 +416,27 @@ WHERE "object_type" = ?1 AND "id" = ?2 Ok(result.rows_affected() > 0) } + pub async fn delete_many(&self, object_type: &str, ids: &[String]) -> PersistenceResult { + let mut deleted = 0_u64; + for ids in ids.chunks(DELETE_MANY_BATCH_SIZE) { + let mut query = QueryBuilder::::new("DELETE FROM objects WHERE object_type = "); + query.push_bind(object_type).push(" AND id IN ("); + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + + deleted += query + .build() + .execute(&self.pool) + .await + .map_err(|e| map_db_error(&e))? + .rows_affected(); + } + Ok(deleted) + } + pub async fn count_in_workspace( &self, object_type: &str, diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 8802ac8d10..479f46e493 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -366,6 +366,116 @@ async fn sqlite_delete_behavior() { assert!(!deleted_again); } +#[tokio::test] +async fn delete_many_is_bounded_idempotent_and_type_scoped() { + let store = test_store().await; + let mut ids = Vec::new(); + for idx in 0..(super::DELETE_MANY_BATCH_SIZE + 12) { + let id = format!("sandbox-{idx}"); + store + .put( + "sandbox", + &id, + &format!("name-{idx}"), + "default", + b"payload", + None, + ) + .await + .unwrap(); + ids.push(id); + } + store + .put( + "provider", + "other-type", + "other-type", + "default", + b"payload", + None, + ) + .await + .unwrap(); + + ids.extend([ + "missing".to_string(), + "other-type".to_string(), + "sandbox-0".to_string(), + ]); + let expected = u64::try_from(super::DELETE_MANY_BATCH_SIZE + 12).unwrap(); + assert_eq!(store.delete_many("sandbox", &ids).await.unwrap(), expected); + assert_eq!(store.delete_many("sandbox", &ids).await.unwrap(), 0); + assert_eq!(store.delete_many("sandbox", &[]).await.unwrap(), 0); + assert!(store.get("provider", "other-type").await.unwrap().is_some()); +} + +#[tokio::test] +async fn file_backed_sqlite_bulk_delete_allows_concurrent_control_reads() { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + let tmp = tempfile::tempdir().expect("tempdir"); + let url = format!("sqlite:{}?mode=rwc", tmp.path().join("bulk.db").display()); + let store = Store::connect(&url) + .await + .expect("connect file-backed store"); + store + .put( + "provider", + "control-row", + "control-row", + "default", + b"control", + None, + ) + .await + .unwrap(); + + let mut ids = Vec::new(); + for idx in 0..500 { + let id = format!("session-{idx}"); + store + .put("ssh_session", &id, &id, "default", b"payload", None) + .await + .unwrap(); + ids.push(id); + } + + let stop = Arc::new(AtomicBool::new(false)); + let read_store = store.clone(); + let read_stop = stop.clone(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let reader = tokio::spawn(async move { + let mut reads = 0_usize; + let mut started_tx = Some(started_tx); + while !read_stop.load(Ordering::Relaxed) { + read_store + .get("provider", "control-row") + .await + .map_err(|error| error.to_string())? + .ok_or_else(|| "control row disappeared".to_string())?; + reads += 1; + if let Some(started_tx) = started_tx.take() { + let _ = started_tx.send(()); + } + tokio::task::yield_now().await; + } + Ok::(reads) + }); + started_rx.await.expect("reader started"); + + assert_eq!(store.delete_many("ssh_session", &ids).await.unwrap(), 500); + stop.store(true, Ordering::Relaxed); + assert!(reader.await.unwrap().unwrap() > 0); + assert!( + store + .get("provider", "control-row") + .await + .unwrap() + .is_some() + ); +} + #[tokio::test] async fn sqlite_protobuf_round_trip() { let store = test_store().await; diff --git a/crates/openshell-server/src/ssh_sessions.rs b/crates/openshell-server/src/ssh_sessions.rs index 490f0dbb25..85f7c2446f 100644 --- a/crates/openshell-server/src/ssh_sessions.rs +++ b/crates/openshell-server/src/ssh_sessions.rs @@ -13,6 +13,8 @@ use tracing::{info, warn}; use crate::persistence::{ObjectType, Store}; +const SESSION_REAPER_PAGE_SIZE: u32 = 1000; + impl ObjectType for SshSession { fn object_type() -> &'static str { "ssh_session" @@ -35,36 +37,54 @@ pub fn spawn_session_reaper(store: Arc, interval: Duration) { async fn reap_expired_sessions(store: &Store) -> Result<(), String> { let now_ms = now_ms(); - - let records = store - .list_by_type(SshSession::object_type(), 1000, 0) - .await - .map_err(|e| e.to_string())?; - - let mut reaped = 0u32; - for record in records { - let session: SshSession = match Message::decode(record.payload.as_slice()) { - Ok(s) => s, - Err(_) => continue, - }; - - let should_delete = - (session.expires_at_ms > 0 && now_ms > session.expires_at_ms) || session.revoked; - - if should_delete { - if let Err(e) = store - .delete(SshSession::object_type(), session.object_id()) - .await - { - warn!(session_id = %session.object_id(), error = %e, "Failed to reap SSH session"); - } else { - reaped += 1; + let started = std::time::Instant::now(); + let mut offset = 0_u32; + let mut scanned = 0_usize; + let mut decode_failures = 0_usize; + let mut session_ids = Vec::new(); + + loop { + let records = store + .list_by_type(SshSession::object_type(), SESSION_REAPER_PAGE_SIZE, offset) + .await + .map_err(|e| e.to_string())?; + let page_len = records.len(); + scanned += page_len; + + for record in records { + let Ok(session) = SshSession::decode(record.payload.as_slice()) else { + decode_failures += 1; + continue; + }; + if (session.expires_at_ms > 0 && now_ms > session.expires_at_ms) || session.revoked { + session_ids.push(session.object_id().to_string()); } } + + if page_len < SESSION_REAPER_PAGE_SIZE as usize { + break; + } + let page_len = u32::try_from(page_len) + .map_err(|_| "SSH session reaper page length overflow".to_string())?; + offset = offset + .checked_add(page_len) + .ok_or_else(|| "SSH session reaper pagination overflow".to_string())?; } - if reaped > 0 { - info!(count = reaped, "SSH session reaper: cleaned up sessions"); + let matched = session_ids.len(); + let deleted = store + .delete_many(SshSession::object_type(), &session_ids) + .await + .map_err(|e| e.to_string())?; + if matched > 0 || decode_failures > 0 { + info!( + scanned, + matched, + deleted, + decode_failures, + elapsed_ms = started.elapsed().as_millis(), + "SSH session reaper sweep complete" + ); } Ok(()) } @@ -175,4 +195,38 @@ mod tests { "session with no expiry should be preserved" ); } + + #[tokio::test] + async fn reaper_batches_expired_and_revoked_sessions() { + let store = test_store().await; + let session_count = crate::persistence::DELETE_MANY_BATCH_SIZE + 9; + for idx in 0..session_count { + let session = make_session( + &format!("reap-{idx}"), + "sbx1", + if idx % 2 == 0 { now_ms() - 1 } else { 0 }, + idx % 2 != 0, + ); + store.put_message(&session).await.unwrap(); + } + let active = make_session("keep", "sbx1", now_ms() + 60_000, false); + store.put_message(&active).await.unwrap(); + + reap_expired_sessions(&store).await.unwrap(); + + assert_eq!( + store + .count_in_workspace(SshSession::object_type(), "default") + .await + .unwrap(), + 1 + ); + assert!( + store + .get_message::("keep") + .await + .unwrap() + .is_some() + ); + } }