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..c6bc153bbf 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -20,8 +20,8 @@ pub use vm::VmComputeConfig; use crate::grpc::policy::SANDBOX_SETTINGS_OBJECT_TYPE; use crate::otel_tracing::TraceContextInterceptor; use crate::persistence::{ - DRAFT_CHUNK_OBJECT_TYPE, ObjectId, ObjectName, ObjectRecord, ObjectType, POLICY_OBJECT_TYPE, - Store, WriteCondition, + DRAFT_CHUNK_OBJECT_TYPE, ObjectCursor, ObjectId, ObjectName, ObjectRecord, ObjectType, + POLICY_OBJECT_TYPE, Store, WriteCondition, }; use crate::sandbox_index::SandboxIndex; use crate::sandbox_watch::SandboxWatchBus; @@ -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,60 @@ impl ComputeRuntime { sandbox_id: &str, workspace: &str, ) -> Result<(), String> { - let records = self + let started = Instant::now(); + let mut cursor = None; + let mut scanned = 0_usize; + let mut decode_failures = 0_usize; + let mut session_ids = Vec::new(); + + loop { + let records = self + .store + .list_after( + SshSession::object_type(), + workspace, + cursor.as_ref(), + LIFECYCLE_SWEEP_PAGE_SIZE, + ) + .await + .map_err(|e| format!("list SSH sessions: {e}"))?; + let page_len = records.len(); + scanned += page_len; + + cursor = records.last().map(ObjectCursor::from); + 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 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 +7190,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..6934630110 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 { @@ -101,6 +107,30 @@ pub struct ObjectRecord { pub resource_version: u64, } +/// Stable position in the global object-listing order. +/// +/// Keyset consumers must use the matching store method for the order encoded +/// here: workspace-scoped lists use `created_at_ms`, `name`, and `id`; global +/// lists additionally include `workspace`. +#[derive(Debug, Clone)] +pub struct ObjectCursor { + pub created_at_ms: i64, + pub name: String, + pub workspace: String, + pub id: String, +} + +impl From<&ObjectRecord> for ObjectCursor { + fn from(record: &ObjectRecord) -> Self { + Self { + created_at_ms: record.created_at_ms, + name: record.name.clone(), + workspace: record.workspace.clone(), + id: record.id.clone(), + } + } +} + /// Write condition for compare-and-swap operations. #[derive(Debug, Clone, Copy)] pub enum WriteCondition { @@ -415,6 +445,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", @@ -499,6 +545,46 @@ impl Store { store_dispatch_traced!(self.list_by_type(object_type, limit, offset)) } + /// List workspace objects after a stable cursor, without offset drift. + #[tracing::instrument( + name = "store", + skip_all, + fields( + otel.name = "store.list_after", + otel.status_code = tracing::field::Empty, + object_type = %object_type, + workspace = %workspace, + ) + )] + pub async fn list_after( + &self, + object_type: &str, + workspace: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + store_dispatch_traced!(self.list_after(object_type, workspace, after, limit)) + } + + /// List objects across workspaces after a stable cursor, without offset drift. + #[tracing::instrument( + name = "store", + skip_all, + fields( + otel.name = "store.list_by_type_after", + otel.status_code = tracing::field::Empty, + object_type = %object_type, + ) + )] + pub async fn list_by_type_after( + &self, + object_type: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + store_dispatch_traced!(self.list_by_type_after(object_type, after, limit)) + } + /// List objects by type and application-owned scope. /// /// Workspace filtering is intentionally omitted: scope values are sandbox diff --git a/crates/openshell-server/src/persistence/postgres.rs b/crates/openshell-server/src/persistence/postgres.rs index 8bab0ada96..89fe442695 100644 --- a/crates/openshell-server/src/persistence/postgres.rs +++ b/crates/openshell-server/src/persistence/postgres.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - DraftChunkRecord, ObjectRecord, PersistenceError, PersistenceResult, PolicyRecord, - WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error, + DraftChunkRecord, ObjectCursor, ObjectRecord, PersistenceError, PersistenceResult, + PolicyRecord, WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error, }; use crate::policy_store::{ AtomicPolicyRevisionWrite, draft_chunk_payload_from_record, draft_chunk_record_from_parts, @@ -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, @@ -500,6 +522,33 @@ LIMIT $2 OFFSET $3 Ok(rows.into_iter().map(row_to_object_record).collect()) } + pub async fn list_after( + &self, + object_type: &str, + workspace: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + let rows = if let Some(cursor) = after { + sqlx::query("SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version FROM objects WHERE object_type = $1 AND workspace = $2 AND (created_at_ms, name, id) > ($3, $4, $5) ORDER BY created_at_ms, name, id LIMIT $6").bind(object_type).bind(workspace).bind(cursor.created_at_ms).bind(&cursor.name).bind(&cursor.id).bind(i64::from(limit)).fetch_all(&self.pool).await + } else { + sqlx::query("SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version FROM objects WHERE object_type = $1 AND workspace = $2 ORDER BY created_at_ms, name, id LIMIT $3").bind(object_type).bind(workspace).bind(i64::from(limit)).fetch_all(&self.pool).await + }.map_err(|e| map_db_error(&e))?; + Ok(rows.into_iter().map(row_to_object_record).collect()) + } + pub async fn list_by_type_after( + &self, + object_type: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + let rows = if let Some(cursor) = after { + sqlx::query("SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version FROM objects WHERE object_type = $1 AND (created_at_ms, name, workspace, id) > ($2, $3, $4, $5) ORDER BY created_at_ms, name, workspace, id LIMIT $6").bind(object_type).bind(cursor.created_at_ms).bind(&cursor.name).bind(&cursor.workspace).bind(&cursor.id).bind(i64::from(limit)).fetch_all(&self.pool).await + } else { + sqlx::query("SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version FROM objects WHERE object_type = $1 ORDER BY created_at_ms, name, workspace, id LIMIT $2").bind(object_type).bind(i64::from(limit)).fetch_all(&self.pool).await + }.map_err(|e| map_db_error(&e))?; + Ok(rows.into_iter().map(row_to_object_record).collect()) + } pub async fn list_with_membership( &self, diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index 658530f753..c28318d261 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - DraftChunkRecord, ObjectRecord, PersistenceError, PersistenceResult, PolicyRecord, - WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error, + DraftChunkRecord, ObjectCursor, ObjectRecord, PersistenceError, PersistenceResult, + PolicyRecord, WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error, }; use crate::policy_store::{ AtomicPolicyRevisionWrite, draft_chunk_payload_from_record, draft_chunk_record_from_parts, @@ -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, @@ -541,6 +562,77 @@ LIMIT ?2 OFFSET ?3 Ok(rows.into_iter().map(row_to_object_record).collect()) } + pub async fn list_after( + &self, + object_type: &str, + workspace: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + let rows = if let Some(cursor) = after { + sqlx::query( + r#" +SELECT "object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version" +FROM "objects" +WHERE "object_type" = ?1 AND "workspace" = ?2 + AND ("created_at_ms", "name", "id") > (?3, ?4, ?5) +ORDER BY "created_at_ms" ASC, "name" ASC, "id" ASC +LIMIT ?6 +"#, + ) + .bind(object_type) + .bind(workspace) + .bind(cursor.created_at_ms) + .bind(&cursor.name) + .bind(&cursor.id) + .bind(i64::from(limit)) + .fetch_all(&self.pool) + .await + } else { + sqlx::query( + r#" +SELECT "object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version" +FROM "objects" +WHERE "object_type" = ?1 AND "workspace" = ?2 +ORDER BY "created_at_ms" ASC, "name" ASC, "id" ASC +LIMIT ?3 +"#, + ) + .bind(object_type) + .bind(workspace) + .bind(i64::from(limit)) + .fetch_all(&self.pool) + .await + } + .map_err(|e| map_db_error(&e))?; + Ok(rows.into_iter().map(row_to_object_record).collect()) + } + pub async fn list_by_type_after( + &self, + object_type: &str, + after: Option<&ObjectCursor>, + limit: u32, + ) -> PersistenceResult> { + let rows = if let Some(cursor) = after { + sqlx::query(r#" +SELECT "object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version" +FROM "objects" +WHERE "object_type" = ?1 + AND ("created_at_ms", "name", "workspace", "id") > (?2, ?3, ?4, ?5) +ORDER BY "created_at_ms" ASC, "name" ASC, "workspace" ASC, "id" ASC +LIMIT ?6 +"#).bind(object_type).bind(cursor.created_at_ms).bind(&cursor.name).bind(&cursor.workspace).bind(&cursor.id).bind(i64::from(limit)).fetch_all(&self.pool).await + } else { + sqlx::query(r#" +SELECT "object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version" +FROM "objects" +WHERE "object_type" = ?1 +ORDER BY "created_at_ms" ASC, "name" ASC, "workspace" ASC, "id" ASC +LIMIT ?2 +"#).bind(object_type).bind(i64::from(limit)).fetch_all(&self.pool).await + }.map_err(|e| map_db_error(&e))?; + Ok(rows.into_iter().map(row_to_object_record).collect()) + } pub async fn list_with_membership( &self, 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..12081c217d 100644 --- a/crates/openshell-server/src/ssh_sessions.rs +++ b/crates/openshell-server/src/ssh_sessions.rs @@ -7,11 +7,14 @@ use openshell_core::ObjectId; use openshell_core::proto::SshSession; use openshell_core::time::now_ms; use prost::Message; +use std::future::Future; use std::sync::Arc; use std::time::Duration; use tracing::{info, warn}; -use crate::persistence::{ObjectType, Store}; +use crate::persistence::{ObjectCursor, ObjectType, Store}; + +const SESSION_REAPER_PAGE_SIZE: u32 = 1000; impl ObjectType for SshSession { fn object_type() -> &'static str { @@ -34,37 +37,69 @@ 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())?; + reap_expired_sessions_after_page(store, |_| std::future::ready(())).await +} - let mut reaped = 0u32; - for record in records { - let session: SshSession = match Message::decode(record.payload.as_slice()) { - Ok(s) => s, - Err(_) => continue, - }; +async fn reap_expired_sessions_after_page( + store: &Store, + mut after_page: F, +) -> Result<(), String> +where + F: FnMut(usize) -> Fut, + Fut: Future, +{ + let now_ms = now_ms(); + let started = std::time::Instant::now(); + let mut cursor = None; + let mut page_number = 0_usize; + let mut scanned = 0_usize; + let mut decode_failures = 0_usize; + let mut session_ids = Vec::new(); - let should_delete = - (session.expires_at_ms > 0 && now_ms > session.expires_at_ms) || session.revoked; + loop { + let records = store + .list_by_type_after( + SshSession::object_type(), + cursor.as_ref(), + SESSION_REAPER_PAGE_SIZE, + ) + .await + .map_err(|e| e.to_string())?; + let page_len = records.len(); + scanned += page_len; - 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; + cursor = records.last().map(ObjectCursor::from); + 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; + } + page_number += 1; + after_page(page_number).await; } - 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 +210,71 @@ 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() + ); + } + + #[tokio::test] + async fn reaper_removes_every_expired_session_when_an_earlier_page_row_is_deleted() { + let store = test_store().await; + for idx in 0..=SESSION_REAPER_PAGE_SIZE { + let session = make_session(&format!("reap-{idx:04}"), "sbx1", now_ms() - 1, false); + store.put_message(&session).await.unwrap(); + } + + let delete_store = store.clone(); + reap_expired_sessions_after_page(&store, move |page_number| { + let delete_store = delete_store.clone(); + async move { + if page_number == 1 { + delete_store + .delete(SshSession::object_type(), "reap-0000") + .await + .unwrap(); + } + } + }) + .await + .unwrap(); + + assert_eq!( + store + .count_in_workspace(SshSession::object_type(), "default") + .await + .unwrap(), + 0, + "the reaper must not leave an expired session behind" + ); + } }