Skip to content

Commit b960125

Browse files
drewelezar
andauthored
fix(gateway): batch SSH session cleanup writes (#3000)
* fix(gateway): batch SSH session cleanup writes Part of #2999. Signed-off-by: Drew Newberry <anewberry@nvidia.com> * test(server): reproduce SSH reaper pagination race Signed-off-by: Evan Lezar <elezar@nvidia.com> * fix(server): use keyset SSH session cleanup Signed-off-by: Evan Lezar <elezar@nvidia.com> --------- Signed-off-by: Drew Newberry <anewberry@nvidia.com> Signed-off-by: Evan Lezar <elezar@nvidia.com> Co-authored-by: Evan Lezar <elezar@nvidia.com>
1 parent 7ceea96 commit b960125

7 files changed

Lines changed: 566 additions & 48 deletions

File tree

‎architecture/compute-runtimes.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,9 @@ backend is already absent (`deleted = false`), the request removes gateway state
208208
synchronously. Sandbox row removal remains bound to the stable ID and resource
209209
version. Settings retain their existing best-effort name-based cleanup; SSH
210210
sessions, indexes, and watch/log buses are cleaned after confirmed removal.
211+
Owned-record cleanup discovers records before mutating them and uses bounded
212+
set-based deletes so teardown cannot amplify one sandbox into an unbounded
213+
sequence of individual persistence writes.
211214

212215
The request acquires both locks before starting owned work, so cancellation
213216
while queued does not leave a delete armed. After that commitment point, the

‎crates/openshell-server/src/compute/mod.rs‎

Lines changed: 91 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ pub use vm::VmComputeConfig;
2020
use crate::grpc::policy::SANDBOX_SETTINGS_OBJECT_TYPE;
2121
use crate::otel_tracing::TraceContextInterceptor;
2222
use crate::persistence::{
23-
DRAFT_CHUNK_OBJECT_TYPE, ObjectId, ObjectName, ObjectRecord, ObjectType, POLICY_OBJECT_TYPE,
24-
Store, WriteCondition,
23+
DRAFT_CHUNK_OBJECT_TYPE, ObjectCursor, ObjectId, ObjectName, ObjectRecord, ObjectType,
24+
POLICY_OBJECT_TYPE, Store, WriteCondition,
2525
};
2626
use crate::sandbox_index::SandboxIndex;
2727
use crate::sandbox_watch::SandboxWatchBus;
@@ -70,7 +70,7 @@ use std::net::SocketAddr;
7070
use std::path::{Path, PathBuf};
7171
use std::pin::Pin;
7272
use std::sync::{Arc, Mutex as StdMutex, Weak};
73-
use std::time::Duration;
73+
use std::time::{Duration, Instant};
7474
#[cfg(unix)]
7575
use tokio::net::UnixStream;
7676
use tokio::sync::{Mutex, watch};
@@ -3305,21 +3305,60 @@ impl ComputeRuntime {
33053305
sandbox_id: &str,
33063306
workspace: &str,
33073307
) -> Result<(), String> {
3308-
let records = self
3308+
let started = Instant::now();
3309+
let mut cursor = None;
3310+
let mut scanned = 0_usize;
3311+
let mut decode_failures = 0_usize;
3312+
let mut session_ids = Vec::new();
3313+
3314+
loop {
3315+
let records = self
3316+
.store
3317+
.list_after(
3318+
SshSession::object_type(),
3319+
workspace,
3320+
cursor.as_ref(),
3321+
LIFECYCLE_SWEEP_PAGE_SIZE,
3322+
)
3323+
.await
3324+
.map_err(|e| format!("list SSH sessions: {e}"))?;
3325+
let page_len = records.len();
3326+
scanned += page_len;
3327+
3328+
cursor = records.last().map(ObjectCursor::from);
3329+
for record in records {
3330+
match SshSession::decode(record.payload.as_slice()) {
3331+
Ok(session) if session.sandbox_id == sandbox_id => {
3332+
session_ids.push(session.object_id().to_string());
3333+
}
3334+
Ok(_) => {}
3335+
Err(_) => decode_failures += 1,
3336+
}
3337+
}
3338+
3339+
if page_len < LIFECYCLE_SWEEP_PAGE_SIZE as usize {
3340+
break;
3341+
}
3342+
}
3343+
3344+
let matched = session_ids.len();
3345+
let deleted = self
33093346
.store
3310-
.list(SshSession::object_type(), workspace, 1000, 0)
3347+
.delete_many(SshSession::object_type(), &session_ids)
33113348
.await
3312-
.map_err(|e| format!("list SSH sessions: {e}"))?;
3349+
.map_err(|e| format!("delete sandbox SSH sessions: {e}"))?;
33133350

3314-
for record in records {
3315-
if let Ok(session) = SshSession::decode(record.payload.as_slice())
3316-
&& session.sandbox_id == sandbox_id
3317-
{
3318-
self.store
3319-
.delete(SshSession::object_type(), session.object_id())
3320-
.await
3321-
.map_err(|e| format!("delete SSH session {}: {e}", session.object_id()))?;
3322-
}
3351+
if matched > 0 || decode_failures > 0 {
3352+
debug!(
3353+
sandbox_id,
3354+
workspace,
3355+
scanned,
3356+
matched,
3357+
deleted,
3358+
decode_failures,
3359+
elapsed_ms = started.elapsed().as_millis(),
3360+
"Sandbox SSH session cleanup complete"
3361+
);
33233362
}
33243363

33253364
Ok(())
@@ -7803,6 +7842,43 @@ mod tests {
78037842
assert_eq!(driver.delete_calls(), 1);
78047843
}
78057844

7845+
#[tokio::test]
7846+
async fn sandbox_ssh_session_cleanup_batches_across_list_pages() {
7847+
let runtime = test_runtime(ControlledDriver::new()).await;
7848+
for idx in 0..(LIFECYCLE_SWEEP_PAGE_SIZE + 5) {
7849+
let session = ssh_session_record(&format!("owned-{idx:04}"), "sb-owned");
7850+
runtime.store.put_message(&session).await.unwrap();
7851+
}
7852+
for idx in 0..7 {
7853+
let session = ssh_session_record(&format!("unrelated-{idx:04}"), "sb-unrelated");
7854+
runtime.store.put_message(&session).await.unwrap();
7855+
}
7856+
7857+
runtime
7858+
.cleanup_sandbox_ssh_sessions("sb-owned", "default")
7859+
.await
7860+
.unwrap();
7861+
7862+
assert_eq!(
7863+
runtime
7864+
.store
7865+
.count_in_workspace(SshSession::object_type(), "default")
7866+
.await
7867+
.unwrap(),
7868+
7
7869+
);
7870+
for idx in 0..7 {
7871+
assert!(
7872+
runtime
7873+
.store
7874+
.get_message::<SshSession>(&format!("unrelated-{idx:04}"))
7875+
.await
7876+
.unwrap()
7877+
.is_some()
7878+
);
7879+
}
7880+
}
7881+
78067882
#[tokio::test]
78077883
async fn already_absent_driver_resource_is_removed_synchronously() {
78087884
let driver = ControlledDriver::new();

‎crates/openshell-server/src/persistence/mod.rs‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,12 @@ pub const DRAFT_CHUNK_OBJECT_TYPE: &str = "draft_policy_chunk";
2626

2727
pub type PersistenceResult<T> = Result<T, PersistenceError>;
2828

29+
/// Maximum number of object ids sent in one set-based delete statement.
30+
///
31+
/// Keep this well below `SQLite`'s bind-variable limit. Backends split larger
32+
/// requests into independently retryable, bounded write statements.
33+
pub const DELETE_MANY_BATCH_SIZE: usize = 128;
34+
2935
/// Persistence-layer error type.
3036
#[derive(Debug, Error, Clone)]
3137
pub enum PersistenceError {
@@ -101,6 +107,30 @@ pub struct ObjectRecord {
101107
pub resource_version: u64,
102108
}
103109

110+
/// Stable position in the global object-listing order.
111+
///
112+
/// Keyset consumers must use the matching store method for the order encoded
113+
/// here: workspace-scoped lists use `created_at_ms`, `name`, and `id`; global
114+
/// lists additionally include `workspace`.
115+
#[derive(Debug, Clone)]
116+
pub struct ObjectCursor {
117+
pub created_at_ms: i64,
118+
pub name: String,
119+
pub workspace: String,
120+
pub id: String,
121+
}
122+
123+
impl From<&ObjectRecord> for ObjectCursor {
124+
fn from(record: &ObjectRecord) -> Self {
125+
Self {
126+
created_at_ms: record.created_at_ms,
127+
name: record.name.clone(),
128+
workspace: record.workspace.clone(),
129+
id: record.id.clone(),
130+
}
131+
}
132+
}
133+
104134
/// Write condition for compare-and-swap operations.
105135
#[derive(Debug, Clone, Copy)]
106136
pub enum WriteCondition {
@@ -415,6 +445,22 @@ impl Store {
415445
store_dispatch_traced!(self.delete(object_type, id))
416446
}
417447

448+
/// Delete objects of one type by id in bounded, set-based statements.
449+
#[tracing::instrument(
450+
name = "store",
451+
skip_all,
452+
fields(
453+
otel.name = "store.delete_many",
454+
otel.status_code = tracing::field::Empty,
455+
object_type = %object_type,
456+
object_count = ids.len(),
457+
batch_count = ids.len().div_ceil(DELETE_MANY_BATCH_SIZE),
458+
)
459+
)]
460+
pub async fn delete_many(&self, object_type: &str, ids: &[String]) -> PersistenceResult<u64> {
461+
store_dispatch_traced!(self.delete_many(object_type, ids))
462+
}
463+
418464
/// Count objects of a given type within a workspace.
419465
#[tracing::instrument(
420466
name = "store",
@@ -499,6 +545,46 @@ impl Store {
499545
store_dispatch_traced!(self.list_by_type(object_type, limit, offset))
500546
}
501547

548+
/// List workspace objects after a stable cursor, without offset drift.
549+
#[tracing::instrument(
550+
name = "store",
551+
skip_all,
552+
fields(
553+
otel.name = "store.list_after",
554+
otel.status_code = tracing::field::Empty,
555+
object_type = %object_type,
556+
workspace = %workspace,
557+
)
558+
)]
559+
pub async fn list_after(
560+
&self,
561+
object_type: &str,
562+
workspace: &str,
563+
after: Option<&ObjectCursor>,
564+
limit: u32,
565+
) -> PersistenceResult<Vec<ObjectRecord>> {
566+
store_dispatch_traced!(self.list_after(object_type, workspace, after, limit))
567+
}
568+
569+
/// List objects across workspaces after a stable cursor, without offset drift.
570+
#[tracing::instrument(
571+
name = "store",
572+
skip_all,
573+
fields(
574+
otel.name = "store.list_by_type_after",
575+
otel.status_code = tracing::field::Empty,
576+
object_type = %object_type,
577+
)
578+
)]
579+
pub async fn list_by_type_after(
580+
&self,
581+
object_type: &str,
582+
after: Option<&ObjectCursor>,
583+
limit: u32,
584+
) -> PersistenceResult<Vec<ObjectRecord>> {
585+
store_dispatch_traced!(self.list_by_type_after(object_type, after, limit))
586+
}
587+
502588
/// List objects by type and application-owned scope.
503589
///
504590
/// Workspace filtering is intentionally omitted: scope values are sandbox

‎crates/openshell-server/src/persistence/postgres.rs‎

Lines changed: 53 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22
// SPDX-License-Identifier: Apache-2.0
33

44
use super::{
5-
DraftChunkRecord, ObjectRecord, PersistenceError, PersistenceResult, PolicyRecord,
6-
WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error,
5+
DraftChunkRecord, ObjectCursor, ObjectRecord, PersistenceError, PersistenceResult,
6+
PolicyRecord, WriteCondition, WriteResult, current_time_ms, map_db_error, map_migrate_error,
77
};
88
use crate::policy_store::{
99
AtomicPolicyRevisionWrite, draft_chunk_payload_from_record, draft_chunk_record_from_parts,
@@ -14,11 +14,11 @@ use openshell_core::SetResourceVersion;
1414
use openshell_core::proto::Sandbox;
1515
use prost::Message;
1616
use sqlx::postgres::PgPoolOptions;
17-
use sqlx::{Connection, PgPool, Row};
17+
use sqlx::{Connection, PgPool, Postgres, QueryBuilder, Row};
1818

1919
static POSTGRES_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/postgres");
2020

21-
use super::{DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE};
21+
use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE};
2222

2323
#[derive(Debug, Clone)]
2424
pub struct PostgresStore {
@@ -391,6 +391,28 @@ WHERE object_type = $1 AND workspace = $2 AND name = $3
391391
Ok(result.rows_affected() > 0)
392392
}
393393

394+
pub async fn delete_many(&self, object_type: &str, ids: &[String]) -> PersistenceResult<u64> {
395+
let mut deleted = 0_u64;
396+
for ids in ids.chunks(DELETE_MANY_BATCH_SIZE) {
397+
let mut query =
398+
QueryBuilder::<Postgres>::new("DELETE FROM objects WHERE object_type = ");
399+
query.push_bind(object_type).push(" AND id IN (");
400+
let mut separated = query.separated(", ");
401+
for id in ids {
402+
separated.push_bind(id);
403+
}
404+
separated.push_unseparated(")");
405+
406+
deleted += query
407+
.build()
408+
.execute(&self.pool)
409+
.await
410+
.map_err(|e| map_db_error(&e))?
411+
.rows_affected();
412+
}
413+
Ok(deleted)
414+
}
415+
394416
pub async fn count_in_workspace(
395417
&self,
396418
object_type: &str,
@@ -500,6 +522,33 @@ LIMIT $2 OFFSET $3
500522

501523
Ok(rows.into_iter().map(row_to_object_record).collect())
502524
}
525+
pub async fn list_after(
526+
&self,
527+
object_type: &str,
528+
workspace: &str,
529+
after: Option<&ObjectCursor>,
530+
limit: u32,
531+
) -> PersistenceResult<Vec<ObjectRecord>> {
532+
let rows = if let Some(cursor) = after {
533+
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
534+
} else {
535+
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
536+
}.map_err(|e| map_db_error(&e))?;
537+
Ok(rows.into_iter().map(row_to_object_record).collect())
538+
}
539+
pub async fn list_by_type_after(
540+
&self,
541+
object_type: &str,
542+
after: Option<&ObjectCursor>,
543+
limit: u32,
544+
) -> PersistenceResult<Vec<ObjectRecord>> {
545+
let rows = if let Some(cursor) = after {
546+
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
547+
} else {
548+
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
549+
}.map_err(|e| map_db_error(&e))?;
550+
Ok(rows.into_iter().map(row_to_object_record).collect())
551+
}
503552

504553
pub async fn list_with_membership(
505554
&self,

0 commit comments

Comments
 (0)