mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-29 07:23:25 +08:00
fix(server): prevent unbounded bus entry growth for sandbox IDs (#138)
Closes #27 Add remove() methods to TracingLogBus, PlatformEventBus, and SandboxWatchBus to clean up entries when sandboxes are deleted. Wire cleanup into both handle_deleted (K8s reconciler) and delete_sandbox (gRPC handler). Reorder watch_sandbox to validate sandbox existence before subscribing to buses, preventing entries for non-existent IDs. Add one-time sandbox validation at stream open in push_sandbox_logs. Co-authored-by: John Myers <johntmyers@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
co-authored by
John Myers
parent
024150e5c6
commit
ae9e766567
@@ -340,6 +340,10 @@ All buses use `tokio::sync::broadcast` channels keyed by sandbox ID. Buffer size
|
||||
|
||||
Broadcast lag is translated to `Status::resource_exhausted` via `broadcast_to_status()`.
|
||||
|
||||
**Cleanup:** Each bus exposes a `remove(sandbox_id)` method that drops the broadcast sender (closing active receivers with `RecvError::Closed`) and frees internal map entries. Cleanup is wired into both the `handle_deleted` reconciler (Kubernetes watcher) and the `delete_sandbox` gRPC handler to prevent unbounded memory growth from accumulated entries for deleted sandboxes.
|
||||
|
||||
**Validation:** `WatchSandbox` validates that the sandbox exists before subscribing to any bus, preventing entries from being created for non-existent IDs. `PushSandboxLogs` validates sandbox existence once on the first batch of the stream.
|
||||
|
||||
## Remote Exec via SSH
|
||||
|
||||
The `ExecSandbox` RPC (`crates/navigator-server/src/grpc.rs`) executes a command inside a sandbox pod over SSH and streams stdout/stderr/exit back to the client.
|
||||
|
||||
@@ -205,7 +205,23 @@ impl Navigator for NavigatorService {
|
||||
|
||||
// Spawn producer task.
|
||||
tokio::spawn(async move {
|
||||
// Subscribe to all buses BEFORE reading the initial snapshot to avoid
|
||||
// Validate that the sandbox exists BEFORE subscribing to any buses.
|
||||
// This prevents creating bus entries for non-existent sandbox IDs.
|
||||
match state.store.get_message::<Sandbox>(&sandbox_id).await {
|
||||
Ok(Some(_)) => {} // sandbox exists, proceed
|
||||
Ok(None) => {
|
||||
let _ = tx.send(Err(Status::not_found("sandbox not found"))).await;
|
||||
return;
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx
|
||||
.send(Err(Status::internal(format!("fetch sandbox failed: {e}"))))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Subscribe to all buses BEFORE reading the snapshot to avoid
|
||||
// missing notifications that fire between the snapshot read and subscribe.
|
||||
let mut status_rx = if follow_status {
|
||||
Some(state.sandbox_watch_bus.subscribe(&sandbox_id))
|
||||
@@ -228,7 +244,8 @@ impl Navigator for NavigatorService {
|
||||
None
|
||||
};
|
||||
|
||||
// Always start with a snapshot if present.
|
||||
// Re-read the snapshot now that we have subscriptions active
|
||||
// (avoids missing notifications between validate and subscribe).
|
||||
match state.store.get_message::<Sandbox>(&sandbox_id).await {
|
||||
Ok(Some(sandbox)) => {
|
||||
state.sandbox_index.update_from_sandbox(&sandbox);
|
||||
@@ -253,6 +270,7 @@ impl Navigator for NavigatorService {
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
// Sandbox was deleted between validate and subscribe — end stream.
|
||||
let _ = tx.send(Err(Status::not_found("sandbox not found"))).await;
|
||||
return;
|
||||
}
|
||||
@@ -481,6 +499,11 @@ impl Navigator for NavigatorService {
|
||||
warn!(sandbox_id = %id, error = %e, "Failed to clean up store after delete");
|
||||
}
|
||||
|
||||
// Clean up bus entries to prevent unbounded memory growth.
|
||||
self.state.tracing_log_bus.remove(&id);
|
||||
self.state.tracing_log_bus.platform_event_bus.remove(&id);
|
||||
self.state.sandbox_watch_bus.remove(&id);
|
||||
|
||||
info!(
|
||||
sandbox_id = %id,
|
||||
sandbox_name = %sandbox.name,
|
||||
@@ -1147,6 +1170,7 @@ impl Navigator for NavigatorService {
|
||||
request: Request<tonic::Streaming<PushSandboxLogsRequest>>,
|
||||
) -> Result<Response<PushSandboxLogsResponse>, Status> {
|
||||
let mut stream = request.into_inner();
|
||||
let mut validated = false;
|
||||
|
||||
while let Some(batch) = stream
|
||||
.message()
|
||||
@@ -1157,6 +1181,20 @@ impl Navigator for NavigatorService {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Validate sandbox existence once at stream open (first batch).
|
||||
// Subsequent batches trust the validated sandbox_id. If the sandbox
|
||||
// is deleted mid-stream, bus remove() drops the sender and publish
|
||||
// silently discards via `let _ = tx.send(...)`.
|
||||
if !validated {
|
||||
self.state
|
||||
.store
|
||||
.get_message::<Sandbox>(&batch.sandbox_id)
|
||||
.await
|
||||
.map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))?
|
||||
.ok_or_else(|| Status::not_found("sandbox not found"))?;
|
||||
validated = true;
|
||||
}
|
||||
|
||||
// Cap lines per batch to prevent abuse.
|
||||
for log in batch.logs.into_iter().take(100) {
|
||||
let mut log = log;
|
||||
|
||||
@@ -128,6 +128,7 @@ pub async fn run_server(config: Config, tracing_log_bus: TracingLogBus) -> Resul
|
||||
state.sandbox_client.clone(),
|
||||
state.sandbox_index.clone(),
|
||||
state.sandbox_watch_bus.clone(),
|
||||
state.tracing_log_bus.clone(),
|
||||
);
|
||||
spawn_kube_event_tailer(state.clone());
|
||||
|
||||
|
||||
@@ -217,6 +217,7 @@ pub fn spawn_sandbox_watcher(
|
||||
client: SandboxClient,
|
||||
index: crate::sandbox_index::SandboxIndex,
|
||||
watch_bus: crate::sandbox_watch::SandboxWatchBus,
|
||||
tracing_log_bus: crate::tracing_bus::TracingLogBus,
|
||||
) {
|
||||
let namespace = client.namespace().to_string();
|
||||
info!(namespace = %namespace, "Starting sandbox watcher");
|
||||
@@ -240,7 +241,9 @@ pub fn spawn_sandbox_watcher(
|
||||
Event::Deleted(obj) => {
|
||||
let obj_name = obj.metadata.name.clone().unwrap_or_default();
|
||||
debug!(sandbox_name = %obj_name, "Received Deleted event from Kubernetes");
|
||||
if let Err(err) = handle_deleted(&store, &index, &watch_bus, obj).await {
|
||||
if let Err(err) =
|
||||
handle_deleted(&store, &index, &watch_bus, &tracing_log_bus, obj).await
|
||||
{
|
||||
warn!(sandbox_name = %obj_name, error = %err, "Failed to delete sandbox record");
|
||||
}
|
||||
}
|
||||
@@ -363,6 +366,7 @@ async fn handle_deleted(
|
||||
store: &Store,
|
||||
index: &crate::sandbox_index::SandboxIndex,
|
||||
watch_bus: &crate::sandbox_watch::SandboxWatchBus,
|
||||
tracing_log_bus: &crate::tracing_bus::TracingLogBus,
|
||||
obj: DynamicObject,
|
||||
) -> Result<(), String> {
|
||||
let id = sandbox_id_from_object(&obj)?;
|
||||
@@ -373,6 +377,12 @@ async fn handle_deleted(
|
||||
debug!(sandbox_id = %id, deleted, "Deleted sandbox record");
|
||||
index.remove_sandbox(&id);
|
||||
watch_bus.notify(&id);
|
||||
|
||||
// Clean up bus entries to prevent unbounded memory growth.
|
||||
tracing_log_bus.remove(&id);
|
||||
tracing_log_bus.platform_event_bus.remove(&id);
|
||||
watch_bus.remove(&id);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -57,6 +57,15 @@ impl SandboxWatchBus {
|
||||
pub fn subscribe(&self, sandbox_id: &str) -> broadcast::Receiver<()> {
|
||||
self.sender_for(sandbox_id).subscribe()
|
||||
}
|
||||
|
||||
/// Remove the bus entry for the given sandbox id.
|
||||
///
|
||||
/// This drops the broadcast sender, closing any active receivers with
|
||||
/// `RecvError::Closed`.
|
||||
pub fn remove(&self, sandbox_id: &str) {
|
||||
let mut inner = self.inner.lock().expect("sandbox watch bus lock poisoned");
|
||||
inner.remove(sandbox_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a background Kubernetes Event tailer.
|
||||
@@ -160,6 +169,53 @@ fn map_kube_event_to_platform(
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn sandbox_watch_bus_remove_cleans_up() {
|
||||
let bus = SandboxWatchBus::new();
|
||||
let sandbox_id = "sb-1";
|
||||
|
||||
let mut rx = bus.subscribe(sandbox_id);
|
||||
|
||||
// Notify and receive
|
||||
bus.notify(sandbox_id);
|
||||
assert!(rx.try_recv().is_ok());
|
||||
|
||||
// Remove
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// Receiver should be closed
|
||||
match rx.try_recv() {
|
||||
Err(broadcast::error::TryRecvError::Closed) => {} // expected
|
||||
other => panic!("expected Closed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sandbox_watch_bus_subscribe_after_remove_creates_fresh_channel() {
|
||||
let bus = SandboxWatchBus::new();
|
||||
let sandbox_id = "sb-2";
|
||||
|
||||
let _old_rx = bus.subscribe(sandbox_id);
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// New subscription should work
|
||||
let mut new_rx = bus.subscribe(sandbox_id);
|
||||
bus.notify(sandbox_id);
|
||||
assert!(new_rx.try_recv().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sandbox_watch_bus_remove_nonexistent_is_noop() {
|
||||
let bus = SandboxWatchBus::new();
|
||||
// Should not panic
|
||||
bus.remove("nonexistent");
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to translate broadcast lag into a gRPC status.
|
||||
pub fn broadcast_to_status(err: broadcast::error::RecvError) -> Status {
|
||||
match err {
|
||||
|
||||
@@ -75,6 +75,16 @@ impl TracingLogBus {
|
||||
self.sender_for(sandbox_id).subscribe()
|
||||
}
|
||||
|
||||
/// Remove all bus entries for the given sandbox id.
|
||||
///
|
||||
/// This drops the broadcast sender (closing any active receivers with
|
||||
/// `RecvError::Closed`) and frees the tail buffer.
|
||||
pub fn remove(&self, sandbox_id: &str) {
|
||||
let mut inner = self.inner.lock().expect("tracing bus lock poisoned");
|
||||
inner.per_id.remove(sandbox_id);
|
||||
inner.tails.remove(sandbox_id);
|
||||
}
|
||||
|
||||
pub fn tail(&self, sandbox_id: &str, max: usize) -> Vec<SandboxStreamEvent> {
|
||||
let inner = self.inner.lock().expect("tracing bus lock poisoned");
|
||||
inner
|
||||
@@ -186,6 +196,129 @@ fn current_time_ms() -> Option<i64> {
|
||||
i64::try_from(now.as_millis()).ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_log_event(sandbox_id: &str, message: &str) -> SandboxLogLine {
|
||||
SandboxLogLine {
|
||||
sandbox_id: sandbox_id.to_string(),
|
||||
timestamp_ms: 1000,
|
||||
level: "INFO".to_string(),
|
||||
target: "test".to_string(),
|
||||
message: message.to_string(),
|
||||
source: "gateway".to_string(),
|
||||
fields: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_log_bus_remove_cleans_up_all_maps() {
|
||||
let bus = TracingLogBus::new();
|
||||
let sandbox_id = "sb-1";
|
||||
|
||||
// Create entries via subscribe and publish
|
||||
let _rx = bus.subscribe(sandbox_id);
|
||||
bus.publish_external(make_log_event(sandbox_id, "hello"));
|
||||
|
||||
// Verify entries exist
|
||||
assert_eq!(bus.tail(sandbox_id, 10).len(), 1);
|
||||
|
||||
// Remove
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// Verify entries are gone
|
||||
assert!(bus.tail(sandbox_id, 10).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_log_bus_subscribe_after_remove_creates_fresh_channel() {
|
||||
let bus = TracingLogBus::new();
|
||||
let sandbox_id = "sb-2";
|
||||
|
||||
// Create and remove
|
||||
bus.publish_external(make_log_event(sandbox_id, "old message"));
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// Subscribe again — should get a fresh channel with no history
|
||||
let mut rx = bus.subscribe(sandbox_id);
|
||||
assert!(bus.tail(sandbox_id, 10).is_empty());
|
||||
|
||||
// New publish should reach the new subscriber
|
||||
bus.publish_external(make_log_event(sandbox_id, "new message"));
|
||||
let evt = rx.try_recv().expect("should receive new event");
|
||||
assert!(evt.payload.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_log_bus_remove_closes_active_receivers() {
|
||||
let bus = TracingLogBus::new();
|
||||
let sandbox_id = "sb-3";
|
||||
|
||||
let mut rx = bus.subscribe(sandbox_id);
|
||||
|
||||
// Remove drops the sender
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// Existing receiver should get Closed error
|
||||
match rx.try_recv() {
|
||||
Err(broadcast::error::TryRecvError::Closed) => {} // expected
|
||||
other => panic!("expected Closed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_log_bus_remove_nonexistent_is_noop() {
|
||||
let bus = TracingLogBus::new();
|
||||
// Should not panic
|
||||
bus.remove("nonexistent");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn platform_event_bus_remove_cleans_up() {
|
||||
let bus = PlatformEventBus::new();
|
||||
let sandbox_id = "sb-4";
|
||||
|
||||
let mut rx = bus.subscribe(sandbox_id);
|
||||
|
||||
// Publish an event
|
||||
let evt = SandboxStreamEvent { payload: None };
|
||||
bus.publish(sandbox_id, evt);
|
||||
assert!(rx.try_recv().is_ok());
|
||||
|
||||
// Remove
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// Receiver should be closed
|
||||
match rx.try_recv() {
|
||||
Err(broadcast::error::TryRecvError::Closed) => {} // expected
|
||||
other => panic!("expected Closed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn platform_event_bus_subscribe_after_remove_creates_fresh_channel() {
|
||||
let bus = PlatformEventBus::new();
|
||||
let sandbox_id = "sb-5";
|
||||
|
||||
let _old_rx = bus.subscribe(sandbox_id);
|
||||
bus.remove(sandbox_id);
|
||||
|
||||
// New subscription should work
|
||||
let mut new_rx = bus.subscribe(sandbox_id);
|
||||
let evt = SandboxStreamEvent { payload: None };
|
||||
bus.publish(sandbox_id, evt);
|
||||
assert!(new_rx.try_recv().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn platform_event_bus_remove_nonexistent_is_noop() {
|
||||
let bus = PlatformEventBus::new();
|
||||
// Should not panic
|
||||
bus.remove("nonexistent");
|
||||
}
|
||||
}
|
||||
|
||||
/// Separate bus for platform event stream events.
|
||||
///
|
||||
/// This keeps platform events isolated from tracing capture.
|
||||
@@ -220,4 +353,12 @@ impl PlatformEventBus {
|
||||
let tx = self.sender_for(sandbox_id);
|
||||
let _ = tx.send(event);
|
||||
}
|
||||
|
||||
/// Remove the bus entry for the given sandbox id.
|
||||
///
|
||||
/// This drops the broadcast sender, closing any active receivers.
|
||||
pub(crate) fn remove(&self, sandbox_id: &str) {
|
||||
let mut inner = self.inner.lock().expect("platform event bus lock poisoned");
|
||||
inner.remove(sandbox_id);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user