fix(supervisor): report runtime readiness after admission

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
Piotr Mlocek
2026-09-18 16:30:58 -07:00
parent f855310360
commit c10f3724cf
8 changed files with 1090 additions and 796 deletions
+7 -2
View File
@@ -664,12 +664,16 @@ sequenceDiagram
GW->>SUP: ConfigurationAdmission(durable state)
alt Configuration accepted
SUP->>SUP: Launch workload
SUP->>GW: SupervisorRuntimeReady
GW->>DB: Promote sandbox to Ready
else Configuration rejected
GW->>SUP: ConfigUpdate(repaired snapshot)
SUP->>GW: ConfigUpdateResult(component outcome, admission)
GW->>DB: Persist accepted repair
GW->>SUP: ConfigurationAdmission(accepted)
SUP->>SUP: Launch workload on the same supervisor
SUP->>GW: SupervisorRuntimeReady
GW->>DB: Promote sandbox to Ready
end
```
@@ -699,8 +703,9 @@ The gateway validates generation identity, persists admission, and returns that
durable state on the stream. The supervisor holds the workload boundary until
it receives an accepted acknowledgement. Rejection leaves the stream and
supervisor alive so a later complete replacement can repair the generation and
release the same workload. Compatibility protocol revisions continue using
polling.
release the same workload. After launch, the supervisor separately reports
runtime readiness once its relay plane is usable; admission alone never promotes
the sandbox to `Ready`. Compatibility protocol revisions continue using polling.
The gateway serializes construction per sandbox and component, and coalesces
repeated mutations into the latest full snapshot. An enqueue result means only
+124 -37
View File
@@ -654,6 +654,25 @@ impl fmt::Debug for ComputeRuntime {
}
}
#[derive(Clone, Copy)]
enum SupervisorSessionStateUpdate<'a> {
Connected {
instance_id: &'a str,
admission: Option<&'a openshell_core::proto::SandboxConfigurationAdmission>,
readiness: SupervisorRuntimeReadiness,
},
Disconnected {
terminal_delivery_finalized: bool,
},
}
#[derive(Clone, Copy)]
enum SupervisorRuntimeReadiness {
Initializing,
Ready,
ReadyAfterAdmission,
}
impl ComputeRuntime {
#[allow(clippy::too_many_arguments)]
#[tracing::instrument(
@@ -1610,7 +1629,7 @@ impl ComputeRuntime {
) -> Option<Sandbox> {
let sandbox_id = transition.object_id().to_string();
let expected_resource_version = sandbox_resource_version(transition);
let session_connected = self.supervisor_sessions.has_session(&sandbox_id);
let session_connected = self.supervisor_sessions.is_runtime_ready(&sandbox_id);
match self
.store
.update_message_cas::<Sandbox, _>(&sandbox_id, expected_resource_version, |sandbox| {
@@ -2100,7 +2119,7 @@ impl ComputeRuntime {
match observed {
Ok(Some(snapshot)) if snapshot.id == sandbox_id && snapshot.status.is_some() => {
let session_connected = self.supervisor_sessions.has_session(sandbox_id);
let session_connected = self.supervisor_sessions.is_runtime_ready(sandbox_id);
self.write_delete_recovery_with_retry(
sandbox_id,
deleting_resource_version,
@@ -3312,7 +3331,7 @@ impl ComputeRuntime {
expected_resource_version: u64,
existing_phase: SandboxPhase,
) -> Result<(), String> {
let session_connected = self.supervisor_sessions.has_session(&incoming.id);
let session_connected = self.supervisor_sessions.is_runtime_ready(&incoming.id);
let sandbox = self
.store
.update_message_cas::<Sandbox, _>(
@@ -3354,8 +3373,31 @@ impl ComputeRuntime {
sandbox_id: &str,
instance_id: &str,
) -> Result<(), String> {
self.set_supervisor_session_state(sandbox_id, true, Some(instance_id), None, false)
.await
self.set_supervisor_session_state(
sandbox_id,
SupervisorSessionStateUpdate::Connected {
instance_id,
admission: None,
readiness: SupervisorRuntimeReadiness::Ready,
},
)
.await
}
pub async fn supervisor_runtime_ready(
&self,
sandbox_id: &str,
instance_id: &str,
) -> Result<(), String> {
self.set_supervisor_session_state(
sandbox_id,
SupervisorSessionStateUpdate::Connected {
instance_id,
admission: None,
readiness: SupervisorRuntimeReadiness::ReadyAfterAdmission,
},
)
.await
}
pub async fn supervisor_session_admission(
@@ -3366,10 +3408,15 @@ impl ComputeRuntime {
) -> Result<(), String> {
self.set_supervisor_session_state(
sandbox_id,
true,
Some(instance_id),
Some(admission),
false,
SupervisorSessionStateUpdate::Connected {
instance_id,
admission: Some(admission),
readiness: if self.supervisor_sessions.is_runtime_ready(sandbox_id) {
SupervisorRuntimeReadiness::Ready
} else {
SupervisorRuntimeReadiness::Initializing
},
},
)
.await
}
@@ -3381,10 +3428,9 @@ impl ComputeRuntime {
) -> Result<(), String> {
self.set_supervisor_session_state(
sandbox_id,
false,
None,
None,
terminal_delivery_finalized,
SupervisorSessionStateUpdate::Disconnected {
terminal_delivery_finalized,
},
)
.await
}
@@ -3392,10 +3438,7 @@ impl ComputeRuntime {
async fn set_supervisor_session_state(
&self,
sandbox_id: &str,
connected: bool,
instance_id: Option<&str>,
admission: Option<&openshell_core::proto::SandboxConfigurationAdmission>,
terminal_delivery_finalized: bool,
update: SupervisorSessionStateUpdate<'_>,
) -> Result<(), String> {
let _guard = self.sync_lock.lock().await;
let existing = self
@@ -3403,32 +3446,54 @@ impl ComputeRuntime {
.get_message::<Sandbox>(sandbox_id)
.await
.map_err(|err| err.to_string())?;
self.set_supervisor_session_state_from_snapshot(
sandbox_id,
connected,
instance_id,
admission,
terminal_delivery_finalized,
existing,
)
.await
self.set_supervisor_session_state_from_snapshot(sandbox_id, update, existing)
.await
}
async fn set_supervisor_session_state_from_snapshot(
&self,
sandbox_id: &str,
connected: bool,
instance_id: Option<&str>,
admission: Option<&openshell_core::proto::SandboxConfigurationAdmission>,
terminal_delivery_finalized: bool,
update: SupervisorSessionStateUpdate<'_>,
mut existing: Option<Sandbox>,
) -> Result<(), String> {
let (connected, instance_id, admission, terminal_delivery_finalized, readiness) =
match update {
SupervisorSessionStateUpdate::Connected {
instance_id,
admission,
readiness,
} => (true, Some(instance_id), admission, false, Some(readiness)),
SupervisorSessionStateUpdate::Disconnected {
terminal_delivery_finalized,
} => (false, None, None, terminal_delivery_finalized, None),
};
let runtime_ready = matches!(
readiness,
Some(
SupervisorRuntimeReadiness::Ready | SupervisorRuntimeReadiness::ReadyAfterAdmission
)
);
let require_activated_configuration = matches!(
readiness,
Some(SupervisorRuntimeReadiness::ReadyAfterAdmission)
);
for attempt in 1..=SUPERVISOR_SESSION_CAS_RETRY_LIMIT {
let Some(current) = existing else {
return Ok(());
};
let current_phase =
SandboxPhase::try_from(current.phase()).unwrap_or(SandboxPhase::Unknown);
if connected
&& runtime_ready
&& require_activated_configuration
&& current
.status
.as_ref()
.and_then(|status| status.configuration_activated)
!= Some(true)
{
return Err("sandbox configuration is not durably admitted".to_string());
}
if connected
&& (provisioning_deadline::timed_out(&current)
|| matches!(
@@ -3468,7 +3533,6 @@ impl ComputeRuntime {
expected_resource_version,
|sandbox| {
if connected {
ensure_supervisor_ready_status(&mut sandbox.status);
let status = sandbox.status.get_or_insert_with(Default::default);
status.main_process_instance_id =
instance_id.unwrap_or_default().to_string();
@@ -3480,7 +3544,10 @@ impl ComputeRuntime {
== i32::from(openshell_core::proto::ConfigurationAdmissionState::Accepted),
);
}
sandbox.set_phase(SandboxPhase::Ready as i32);
if runtime_ready {
ensure_supervisor_ready_status(&mut sandbox.status);
sandbox.set_phase(SandboxPhase::Ready as i32);
}
} else {
ensure_supervisor_not_ready_status(&mut sandbox.status);
sandbox.set_phase(SandboxPhase::Provisioning as i32);
@@ -10541,6 +10608,13 @@ mod tests {
rejected.status.as_ref().unwrap().configuration_activated,
Some(false)
);
assert!(
runtime
.supervisor_runtime_ready("sb-1", "supervisor-1")
.await
.unwrap_err()
.contains("not durably admitted")
);
admission.state = ConfigurationAdmissionState::Accepted.into();
admission.error.clear();
@@ -10554,7 +10628,7 @@ mod tests {
.await
.unwrap()
.unwrap();
assert_eq!(repaired.phase(), SandboxPhase::Ready as i32);
assert_eq!(repaired.phase(), SandboxPhase::Provisioning as i32);
assert_eq!(
repaired.status.as_ref().unwrap().main_process_instance_id,
"supervisor-1"
@@ -10563,6 +10637,18 @@ mod tests {
repaired.status.as_ref().unwrap().configuration_activated,
Some(true)
);
runtime
.supervisor_runtime_ready("sb-1", "supervisor-1")
.await
.unwrap();
let ready = runtime
.store
.get_message::<Sandbox>("sb-1")
.await
.unwrap()
.unwrap();
assert_eq!(ready.phase(), SandboxPhase::Ready as i32);
}
#[tokio::test]
@@ -10604,10 +10690,11 @@ mod tests {
runtime
.set_supervisor_session_state_from_snapshot(
"sb-1",
true,
Some("test-generation"),
None,
false,
SupervisorSessionStateUpdate::Connected {
instance_id: "test-generation",
admission: None,
readiness: SupervisorRuntimeReadiness::Ready,
},
stale,
)
.await
+2 -2
View File
@@ -162,7 +162,7 @@ mod tests {
const STORAGE_V1_SCHEMA_SHA256: &str =
"d68401809d8cea445c35233ef32412bbd041cb2ac5acaf368a0d0bf74d2ddf17";
const PUBLIC_RPC_SCHEMA_SHA256: &str =
"ab978d653fc8bff9f7dac633dd651b015fc73aecf3704d1a43f3b32e0143167f";
"aec69234398a97e094eaeb7214803dadd8beb38eeffce25b34f08a20a8108652";
const DURABLE_SCHEMA_SHA256: &str =
"965a8a09fa80168906a4f9cc6169d3623b98f5281c367695f2f2898230915dd0";
const PUBLIC_DURABLE_OVERLAP_SHA256: &str =
@@ -619,7 +619,7 @@ mod tests {
overlap_hash.as_str(),
),
(
(314, 25),
(315, 25),
(93, 19),
(80, 19),
PUBLIC_RPC_SCHEMA_SHA256,
+104 -15
View File
@@ -102,6 +102,8 @@ struct LiveSession {
/// Installation evidence belongs to this connection and is never restored
/// from persistence or inherited by a replacement supervisor session.
provider_readiness: Option<ProviderReadinessEvidence>,
/// True only after the admitted workload and relay plane are usable.
runtime_ready: bool,
#[allow(dead_code)]
connected_at: Instant,
}
@@ -391,6 +393,27 @@ impl SupervisorSessionRegistry {
session_id: String,
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
) -> bool {
self.register_with_runtime_state(sandbox_id, session_id, tx, shutdown, true)
}
fn register_initializing(
&self,
sandbox_id: String,
session_id: String,
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
) -> bool {
self.register_with_runtime_state(sandbox_id, session_id, tx, shutdown, false)
}
fn register_with_runtime_state(
&self,
sandbox_id: String,
session_id: String,
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
runtime_ready: bool,
) -> bool {
let mut sessions = self.sessions.lock().unwrap();
let previous = sessions.remove(&sandbox_id);
@@ -406,6 +429,7 @@ impl SupervisorSessionRegistry {
endpoint_status_initialized: false,
endpoint_report_cursor: None,
provider_readiness: None,
runtime_ready,
connected_at: Instant::now(),
},
);
@@ -493,6 +517,26 @@ impl SupervisorSessionRegistry {
self.sessions.lock().unwrap().contains_key(sandbox_id)
}
pub fn is_runtime_ready(&self, sandbox_id: &str) -> bool {
self.sessions
.lock()
.unwrap()
.get(sandbox_id)
.is_some_and(|session| session.runtime_ready)
}
fn mark_runtime_ready(&self, sandbox_id: &str, session_id: &str) -> bool {
let mut sessions = self.sessions.lock().unwrap();
let Some(session) = sessions.get_mut(sandbox_id) else {
return false;
};
if session.session_id != session_id {
return false;
}
session.runtime_ready = true;
true
}
pub fn terminal_delivery_finalized(&self, sandbox_id: &str) -> bool {
self.sessions
.lock()
@@ -1418,12 +1462,21 @@ async fn accept_supervisor_session(
};
accepted_payload.bootstrap = None;
}
let superseded = state.supervisor_sessions.register(
sandbox_id.clone(),
session_id.clone(),
session_tx.clone(),
shutdown_tx,
);
let superseded = if stream_applies_config {
state.supervisor_sessions.register_initializing(
sandbox_id.clone(),
session_id.clone(),
session_tx.clone(),
shutdown_tx,
)
} else {
state.supervisor_sessions.register(
sandbox_id.clone(),
session_id.clone(),
session_tx.clone(),
shutdown_tx,
)
};
if superseded {
info!(
sandbox_id = %sandbox_id,
@@ -1467,7 +1520,7 @@ async fn accept_supervisor_session(
}
if !stream_applies_config
&& !mark_supervisor_initialized(&state, &sandbox_id, &session_id, &instance_id).await
&& !mark_supervisor_initialized(&state, &sandbox_id, &session_id, &instance_id, false).await
{
state
.supervisor_sessions
@@ -2231,6 +2284,22 @@ async fn handle_supervisor_message(
}
}
}
Some(supervisor_message::Payload::RuntimeReady(_)) => {
if !stream_applies_config {
debug!(
sandbox_id,
session_id, "ignored runtime-ready from compatibility supervisor"
);
return;
}
if !mark_supervisor_initialized(state, sandbox_id, session_id, instance_id, true).await
{
warn!(
sandbox_id,
session_id, "failed to persist supervisor runtime readiness"
);
}
}
_ => {
warn!(
sandbox_id = %sandbox_id,
@@ -2312,6 +2381,7 @@ async fn mark_supervisor_initialized(
sandbox_id: &str,
session_id: &str,
instance_id: &str,
require_activated_configuration: bool,
) -> bool {
if !state
.supervisor_sessions
@@ -2319,11 +2389,18 @@ async fn mark_supervisor_initialized(
{
return false;
}
if let Err(err) = state
.compute
.supervisor_session_connected(sandbox_id, instance_id)
.await
{
let persisted = if require_activated_configuration {
state
.compute
.supervisor_runtime_ready(sandbox_id, instance_id)
.await
} else {
state
.compute
.supervisor_session_connected(sandbox_id, instance_id)
.await
};
if let Err(err) = persisted {
warn!(
sandbox_id,
session_id,
@@ -2331,9 +2408,14 @@ async fn mark_supervisor_initialized(
"supervisor session: failed to mark sandbox initialized"
);
false
} else {
} else if state
.supervisor_sessions
.mark_runtime_ready(sandbox_id, session_id)
{
state.telemetry.sandbox_session_connected(sandbox_id);
true
} else {
false
}
}
@@ -2592,8 +2674,15 @@ mod tests {
},
)),
};
let decoded = SupervisorMessage::decode(result.encode_to_vec().as_slice()).unwrap();
assert_eq!(decoded, result);
let runtime_ready = SupervisorMessage {
payload: Some(supervisor_message::Payload::RuntimeReady(
openshell_core::proto::SupervisorRuntimeReady {},
)),
};
for original in [result, runtime_ready] {
let decoded = SupervisorMessage::decode(original.encode_to_vec().as_slice()).unwrap();
assert_eq!(decoded, original);
}
}
#[test]
@@ -62,6 +62,24 @@ pub struct PrestartedSupervisorSession {
task: Option<tokio::task::JoinHandle<()>>,
readiness: tokio::sync::watch::Receiver<bool>,
loopback: Arc<DeferredLoopbackConnector>,
outbound: tokio::sync::mpsc::Sender<openshell_core::proto::SupervisorMessage>,
runtime_ready: Arc<AtomicBool>,
}
impl PrestartedSupervisorSession {
async fn report_runtime_ready(&self) -> Result<()> {
self.runtime_ready.store(true, Ordering::Release);
self.outbound
.send(openshell_core::proto::SupervisorMessage {
payload: Some(
openshell_core::proto::supervisor_message::Payload::RuntimeReady(
openshell_core::proto::SupervisorRuntimeReady {},
),
),
})
.await
.map_err(|_| miette::miette!("supervisor session ended before runtime readiness"))
}
}
impl Drop for PrestartedSupervisorSession {
@@ -89,7 +107,7 @@ pub async fn start_prepared_supervisor_session(
|| std::path::PathBuf::from(openshell_core::container_paths::SSH_SOCKET_PATH),
std::path::PathBuf::from,
);
let (task, mut readiness) = crate::supervisor_session::spawn_prepared(
let (task, mut readiness, outbound, runtime_ready) = crate::supervisor_session::spawn_prepared(
prepared,
bootstrap_result,
target,
@@ -111,6 +129,8 @@ pub async fn start_prepared_supervisor_session(
task: Some(task),
readiness,
loopback,
outbound,
runtime_ready,
}),
Ok(Err(_)) => {
task.abort();
@@ -201,6 +221,9 @@ pub async fn start_boundary_access(
|prestarted| prestarted.terminating.clone(),
);
let Some(ssh_socket_path) = ssh_socket_path.map(std::path::PathBuf::from) else {
if let Some(prestarted) = prestarted_supervisor_session.as_ref() {
prestarted.report_runtime_ready().await?;
}
let (session_task, session_readiness) = match prestarted_supervisor_session.as_mut() {
Some(prestarted) => (prestarted.task.take(), Some(prestarted.readiness.clone())),
None => (None, None),
@@ -311,6 +334,10 @@ pub async fn start_boundary_access(
},
};
if let Some(prestarted) = prestarted_supervisor_session.as_ref() {
prestarted.report_runtime_ready().await?;
}
Ok(BoundaryAccess {
instance_id,
terminating,
@@ -434,6 +434,7 @@ pub fn spawn_with_readiness(
session_id_updates: runtime.session_id_updates,
config_apply_tx: runtime.config_apply_tx,
ready_tx,
runtime_ready: Arc::new(AtomicBool::new(true)),
};
(tokio::spawn(run_session_loop(config, None)), ready_rx)
}
@@ -477,8 +478,15 @@ pub fn spawn_prepared(
terminating: Arc<AtomicBool>,
config_apply_tx: mpsc::Sender<ConfigApplyRequest>,
session_id_updates: Option<watch::Sender<Option<String>>>,
) -> (tokio::task::JoinHandle<()>, watch::Receiver<bool>) {
) -> (
tokio::task::JoinHandle<()>,
watch::Receiver<bool>,
mpsc::Sender<SupervisorMessage>,
Arc<AtomicBool>,
) {
let (ready_tx, ready_rx) = watch::channel(false);
let outbound = prepared.tx.clone();
let runtime_ready = Arc::new(AtomicBool::new(false));
let config = SessionConfig {
endpoint: prepared.endpoint.clone(),
sandbox_id: prepared.sandbox_id.clone(),
@@ -490,11 +498,10 @@ pub fn spawn_prepared(
config_apply_tx: Some(config_apply_tx),
session_id_updates,
ready_tx,
runtime_ready: runtime_ready.clone(),
};
(
tokio::spawn(run_session_loop(config, Some((prepared, bootstrap_result)))),
ready_rx,
)
let task = tokio::spawn(run_session_loop(config, Some((prepared, bootstrap_result))));
(task, ready_rx, outbound, runtime_ready)
}
struct SessionConfig {
@@ -509,6 +516,7 @@ struct SessionConfig {
/// Publishes the currently accepted session to sibling control-plane reporters.
session_id_updates: Option<watch::Sender<Option<String>>>,
ready_tx: watch::Sender<bool>,
runtime_ready: Arc<AtomicBool>,
}
async fn run_session_loop(
@@ -765,6 +773,7 @@ async fn run_prepared_session(
config_apply_tx: config.config_apply_tx.as_ref(),
config_sequences: &config_sequences,
ready_tx: &config.ready_tx,
runtime_ready: &config.runtime_ready,
};
handle_gateway_message(
&msg,
@@ -937,6 +946,7 @@ struct GatewayMessageContext<'a> {
config_apply_tx: Option<&'a mpsc::Sender<ConfigApplyRequest>>,
config_sequences: &'a Arc<Mutex<ConfigSequenceWatermarks>>,
ready_tx: &'a watch::Sender<bool>,
runtime_ready: &'a Arc<AtomicBool>,
}
fn handle_gateway_message(msg: &GatewayMessage, context: &GatewayMessageContext<'_>) {
@@ -948,6 +958,18 @@ fn handle_gateway_message(msg: &GatewayMessage, context: &GatewayMessageContext<
let accepted = admission.state
== i32::from(openshell_core::proto::ConfigurationAdmissionState::Accepted);
context.ready_tx.send_replace(accepted);
if accepted && context.runtime_ready.load(Ordering::Acquire) {
let tx = context.tx.clone();
tokio::spawn(async move {
let _ = tx
.send(SupervisorMessage {
payload: Some(supervisor_message::Payload::RuntimeReady(
openshell_core::proto::SupervisorRuntimeReady {},
)),
})
.await;
});
}
}
Some(gateway_message::Payload::ConfigUpdate(update)) => {
let update = update.clone();
+5
View File
@@ -2862,9 +2862,14 @@ message SupervisorMessage {
ConfigUpdateResult config_update_result = 5;
ConfigBootstrapResult config_bootstrap_result = 6;
StartupConfigPrepared startup_config_prepared = 7;
// Sent after accepted admission once the workload relay is usable.
SupervisorRuntimeReady runtime_ready = 8;
}
}
// Confirms that the admitted workload and its gateway relay are usable.
message SupervisorRuntimeReady {}
// Envelope for gateway-to-supervisor messages on the ConnectSupervisor stream.
message GatewayMessage {
oneof payload {
File diff suppressed because it is too large Load Diff