Skip to content

fix(server): observe persisted sandbox watch updates #2257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
elezar wants to merge 1 commit into
base: main
Choose a base branch
from
+113 −1
Draft

fix(server): observe persisted sandbox watch updates #2257

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 113 additions & 1 deletion crates/openshell-server/src/grpc/sandbox.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use openshell_core::telemetry::{
LifecycleOperation, LifecycleResource, SandboxTemplateSource, TelemetryComputeDriver,
TelemetryOutcome,
};
use openshell_core::{ObjectId, ObjectName};
use openshell_core::{GetResourceVersion, ObjectId, ObjectName};
use prost::Message;
use std::net::IpAddr;
use std::pin::Pin;
Expand All @@ -53,6 +53,7 @@ use super::{MAX_PAGE_SIZE, MAX_PROVIDERS, clamp_limit};
use crate::persistence::current_time_ms;

const TCP_FORWARD_CHUNK_SIZE: usize = 64 * 1024;
const SANDBOX_STORE_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);

// ---------------------------------------------------------------------------
// Sandbox lifecycle handlers
Expand Down Expand Up @@ -638,9 +639,12 @@ pub(super) async fn handle_watch_sandbox(
None
};

let mut last_sandbox_resource_version;

// Re-read the snapshot now that we have subscriptions active.
match state.store.get_message::<Sandbox>(&sandbox_id).await {
Ok(Some(sandbox)) => {
last_sandbox_resource_version = Some(sandbox.get_resource_version());
state.sandbox_index.update_from_sandbox(&sandbox);
let _ = tx
.send(Ok(SandboxStreamEvent {
Expand Down Expand Up @@ -707,6 +711,17 @@ pub(super) async fn handle_watch_sandbox(
}
}

// The in-memory status bus only observes writes made by this gateway process.
// Reconcile status followers with the shared store at a bounded rate so writes
// from another process become visible too. Delay the first tick because the
// initial snapshot was just read, and skip missed ticks to avoid query bursts.
let mut store_poll = tokio::time::interval_at(
tokio::time::Instant::now() + SANDBOX_STORE_POLL_INTERVAL,
SANDBOX_STORE_POLL_INTERVAL,
);
store_poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut store_poll_error_reported = false;

loop {
tokio::select! {
res = async {
Expand All @@ -719,6 +734,8 @@ pub(super) async fn handle_watch_sandbox(
Ok(()) => {
match state.store.get_message::<Sandbox>(&sandbox_id).await {
Ok(Some(sandbox)) => {
last_sandbox_resource_version =
Some(sandbox.get_resource_version());
state.sandbox_index.update_from_sandbox(&sandbox);
if tx.send(Ok(SandboxStreamEvent { payload: Some(openshell_core::proto::sandbox_stream_event::Payload::Sandbox(sandbox.clone()))})).await.is_err() {
return;
Expand All @@ -745,6 +762,42 @@ pub(super) async fn handle_watch_sandbox(
}
}
}
_ = store_poll.tick(), if follow_status => {
match state.store.get_message::<Sandbox>(&sandbox_id).await {
Ok(Some(sandbox)) => {
store_poll_error_reported = false;
let resource_version = sandbox.get_resource_version();
if last_sandbox_resource_version == Some(resource_version) {
continue;
}
last_sandbox_resource_version = Some(resource_version);
state.sandbox_index.update_from_sandbox(&sandbox);
if tx.send(Ok(SandboxStreamEvent {
payload: Some(openshell_core::proto::sandbox_stream_event::Payload::Sandbox(sandbox.clone())),
})).await.is_err() {
return;
}
if stop_on_terminal {
let phase = SandboxPhase::try_from(sandbox.phase()).unwrap_or(SandboxPhase::Unknown);
if phase == SandboxPhase::Ready {
return;
}
}
}
Ok(None) => {
return;
}
Err(error) => {
// Polling supplements the local notification bus. A transient
// store error should not terminate a stream that can still
// receive local status, log, or platform events.
if !store_poll_error_reported {
warn!(sandbox_id = %sandbox_id, %error, "WatchSandbox store poll failed; retrying");
store_poll_error_reported = true;
}
}
}
}
res = async {
match log_rx.as_mut() {
Some(rx) => rx.recv().await,
Expand Down Expand Up @@ -2008,6 +2061,7 @@ mod tests {
use crate::grpc::test_support::test_server_state;
use openshell_core::proto::datamodel::v1::ObjectMeta;
use std::collections::HashMap;
use tokio_stream::StreamExt;

// ---- shell_escape ----

Expand Down Expand Up @@ -2337,6 +2391,64 @@ mod tests {
sandbox
}

#[tokio::test]
async fn watch_sandbox_observes_store_update_without_local_notification() {
let state = test_server_state().await;
let mut sandbox = test_sandbox("watch-store", Vec::new());
sandbox.set_phase(SandboxPhase::Provisioning as i32);
let sandbox_id = sandbox.object_id().to_string();
state.store.put_message(&sandbox).await.unwrap();

let response = handle_watch_sandbox(
&state,
Request::new(WatchSandboxRequest {
id: sandbox_id.clone(),
follow_status: true,
..Default::default()
}),
)
.await
.unwrap();
let mut stream = response.into_inner();

let first = tokio::time::timeout(std::time::Duration::from_secs(1), stream.next())
.await
.expect("initial watch snapshot should arrive")
.expect("stream should remain open")
.expect("initial watch snapshot should be valid");
let Some(openshell_core::proto::sandbox_stream_event::Payload::Sandbox(first)) =
first.payload
else {
panic!("expected initial sandbox snapshot");
};
assert_eq!(first.phase(), SandboxPhase::Provisioning as i32);
assert_eq!(first.get_resource_version(), 1);

// Persist the change without notifying this process's sandbox watch bus,
// matching a write performed by another gateway process.
state
.store
.update_message_cas::<Sandbox, _>(&sandbox_id, 0, |sandbox| {
sandbox.set_phase(SandboxPhase::Ready as i32);
})
.await
.unwrap();

let event = tokio::time::timeout(std::time::Duration::from_secs(3), stream.next())
.await
.expect("store reconciliation should observe the update")
.expect("stream should remain open")
.expect("reconciled watch event should be valid");
let Some(openshell_core::proto::sandbox_stream_event::Payload::Sandbox(updated)) =
event.payload
else {
panic!("expected updated sandbox snapshot");
};
assert_eq!(updated.object_id(), sandbox_id);
assert_eq!(updated.phase(), SandboxPhase::Ready as i32);
assert_eq!(updated.get_resource_version(), 2);
}

#[tokio::test]
async fn attach_sandbox_provider_persists_current_provider_list() {
let state = test_server_state().await;
Expand Down
Loading