trending/OpenShell
1
0
Fork 0
mirror of https://github.com/NVIDIA/OpenShell.git synced 2026-11-02 08:04:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(supervisor): keep polling startup and readiness intact without apply

Browse Source 
- Close the startup session when the gateway does not enable apply and
  reconnect after the workload starts, as polling supervisors always have.
  Retry transient connection failures while preparing the startup session.
- Mark apply-capable supervisors ready only on SupervisorRuntimeReady, in
  any delivery mode, and resend it after a reconnect without apply.
- Resume polling in a stream-started runtime whenever its current session
  does not apply configuration, such as after a rollback to poll mode.
- Leave restart bookkeeping to runtime readiness instead of admission.
- Keep the owner reconciler local to each replica and suppress unchanged
  snapshots for shadow sessions.

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
Piotr Mlocek
2026-10-06 13:55:16 -07:00
parent fde409b7a7
commit addd0132be
6 changed files with 335 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/openshell-server/src/compute/mod.rs
+19 -8
View File
@@ -4848,11 +4848,6 @@ impl ComputeRuntime {
|sandbox| {
if connected {
let status = sandbox.status.get_or_insert_with(Default::default);
let same_instance =
Some(status.main_process_instance_id.as_str()) == instance_id;
status.main_process_instance_id =
instance_id.unwrap_or_default().to_string();
status.exit_code = None;
if let Some(admission) = admission {
status.configuration_admission = Some(admission.clone());
status.configuration_activated = Some(
@@ -4860,7 +4855,14 @@ impl ComputeRuntime {
== i32::from(openshell_core::proto::ConfigurationAdmissionState::Accepted),
);
}
// An admission-only update leaves restart state to
// the instance that last reported runtime readiness.
if runtime_ready {
let same_instance =
Some(status.main_process_instance_id.as_str()) == instance_id;
status.main_process_instance_id =
instance_id.unwrap_or_default().to_string();
status.exit_code = None;
set_next_restart_at_ms(status, 0);
if !same_instance || main_process_started_at_ms(status) == 0 {
set_main_process_started_at_ms(
@@ -13865,9 +13867,14 @@ mod tests {
.unwrap()
.unwrap();
assert_eq!(repaired.phase(), SandboxPhase::Provisioning as i32);
assert_eq!(
repaired.status.as_ref().unwrap().main_process_instance_id,
"supervisor-1"
// Admission alone leaves restart bookkeeping to runtime readiness.
assert!(
repaired
.status
.as_ref()
.unwrap()
.main_process_instance_id
.is_empty()
);
assert_eq!(
repaired.status.as_ref().unwrap().configuration_activated,
@@ -13885,6 +13892,10 @@ mod tests {
.unwrap()
.unwrap();
assert_eq!(ready.phase(), SandboxPhase::Ready as i32);
assert_eq!(
ready.status.as_ref().unwrap().main_process_instance_id,
"supervisor-1"
);
}
#[tokio::test]
crates/openshell-server/src/config_delivery/mod.rs
+21 -5
View File
@@ -942,7 +942,11 @@ fn register_test_session(
control.clone(),
tokio::sync::oneshot::channel().0,
Some(Arc::clone(&slots)),
config_apply,
if config_apply {
crate::supervisor_session::SessionMode::ConfigApply
} else {
crate::supervisor_session::SessionMode::Legacy
},
);
register_session(
state,
@@ -994,16 +998,28 @@ fn record_delivery(component: ConfigComponentKind, outcome: &'static str) {
.increment(1);
}
/// Periodically rebuild current snapshots for every locally routable session.
/// Unchanged snapshots are suppressed per session, so this only repairs missed
/// publications. Dormant in poll mode.
/// Periodically rebuild current snapshots for this replica's sessions.
/// Streamed-apply sessions suppress unchanged snapshots, so this only repairs
/// missed publications. Every replica runs its own pass, so it never hints
/// peers. Dormant in poll mode.
pub fn spawn_owner_reconciler(state: Arc<ServerState>, interval: Duration) {
tokio::spawn(async move {
let mut timer = tokio::time::interval(interval);
timer.tick().await;
loop {
timer.tick().await;
publish_all_connected(&state, ConfigComponents::ALL);
if !push_enabled(&state) {
continue;
}
state.config_delivery.with_scheduler(|scheduler| {
scheduler.publish_fanout(
&FanoutScope::AllConnected,
ConfigComponents::ALL,
false,
Instant::now(),
);
});
kick(&state);
}
});
}
crates/openshell-server/src/supervisor_session.rs
+108 -36
View File
@@ -58,6 +58,19 @@ const HEARTBEAT_INTERVAL_SECS: u32 = 15;
const OWNER_RENEW_TIMEOUT: Duration = Duration::from_secs(5);
const STARTUP_POLICY_REPAIR_TIMEOUT: Duration = Duration::from_mins(5);
const RELAY_PENDING_TIMEOUT: Duration = Duration::from_secs(10);
/// How a supervisor session takes part in configuration delivery and readiness.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum SessionMode {
/// Predates streamed apply. Ready as soon as the session is accepted.
Legacy,
/// Can apply streamed configuration but polls on this gateway. Ready after
/// `SupervisorRuntimeReady`, since it connects before its workload starts.
ReportsReadiness,
/// Applies and acknowledges streamed configuration. Ready after
/// `SupervisorRuntimeReady` with an accepted admission.
ConfigApply,
}
/// How long a streamed-apply session holds newer updates for a component
/// while waiting for the supervisor to acknowledge the one it already has.
const CONFIG_ACK_TIMEOUT: Duration = Duration::from_mins(1);
@@ -725,14 +738,20 @@ impl SupervisorSessionRegistry {
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
) -> bool {
self.register_with_config_slots(sandbox_id, session_id, tx, shutdown, None, false)
self.register_with_config_slots(
sandbox_id,
session_id,
tx,
shutdown,
None,
SessionMode::Legacy,
)
}
/// Register a session that may receive pushed configuration.
///
/// A `config_apply` session acknowledges each update before the next one
/// for that component is released, and is not runtime-ready until the
/// supervisor sends `SupervisorRuntimeReady`.
/// A `ConfigApply` session acknowledges each update before the next one
/// for that component is released.
pub(crate) fn register_with_config_slots(
&self,
sandbox_id: String,
@@ -740,7 +759,7 @@ impl SupervisorSessionRegistry {
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
config_slots: Option<Arc<ConfigSlots>>,
config_apply: bool,
mode: SessionMode,
) -> bool {
let mut sessions = self.sessions.lock().unwrap();
let previous = sessions.remove(&sandbox_id);
@@ -752,7 +771,7 @@ impl SupervisorSessionRegistry {
config_slots,
tx,
config_sequences: ConfigSequences {
acknowledged: config_apply,
acknowledged: mode == SessionMode::ConfigApply,
..ConfigSequences::default()
},
shutdown,
@@ -760,7 +779,7 @@ impl SupervisorSessionRegistry {
endpoint_status_initialized: false,
endpoint_report_cursor: None,
provider_readiness: None,
runtime_ready: !config_apply,
runtime_ready: mode == SessionMode::Legacy,
connected_at: Instant::now(),
_gauge_slot: GaugeSlot::supervisor_session(),
},
@@ -918,13 +937,15 @@ impl SupervisorSessionRegistry {
&mut session.config_sequences.provider_environment
}
};
let fingerprint = config_message_fingerprint(&message);
// A shadow session never acknowledges, so its last delivered snapshot
// stands in for the acknowledged one.
if delivery_state.in_flight.is_none()
&& delivery_state.last_acknowledged_fingerprint.as_ref() == Some(&fingerprint)
{
return DeliveryDisposition::SuppressedUnchanged;
}
if config_apply {
if delivery_state.in_flight.is_none()
&& delivery_state.last_acknowledged_fingerprint.as_ref()
== Some(&config_message_fingerprint(&message))
{
return DeliveryDisposition::SuppressedUnchanged;
}
// Hold the component until the supervisor acknowledges the update
// it already has, bounded so a lost result cannot stall it.
if delivery_state
@@ -946,6 +967,8 @@ impl SupervisorSessionRegistry {
}
if config_apply {
delivery_state.in_flight = Some(in_flight);
} else {
delivery_state.last_acknowledged_fingerprint = Some(fingerprint);
}
drop(sessions);
@@ -2568,7 +2591,7 @@ struct SessionSetup {
sandbox_id: String,
instance_id: String,
connection_epoch: u64,
stream_applies_config: bool,
mode: SessionMode,
bootstrap: Option<ConfigBootstrap>,
provider_readiness: ProviderReadinessEvidence,
config_push: Option<ConfigPushSetup>,
@@ -2584,7 +2607,7 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> {
sandbox_id,
instance_id,
connection_epoch,
stream_applies_config,
mode,
bootstrap,
provider_readiness,
config_push,
@@ -2593,6 +2616,7 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> {
mut inbound,
session_lifetime,
} = setup;
let stream_applies_config = mode == SessionMode::ConfigApply;
let expected_bootstrap_admission = bootstrap
.as_ref()
.and_then(|bootstrap| bootstrap.sandbox_config.as_ref())
@@ -2666,7 +2690,7 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> {
session_tx.clone(),
shutdown_tx,
config_slots,
stream_applies_config,
mode,
);
if superseded {
info!(
@@ -2724,9 +2748,9 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> {
// Do not expose SessionAccepted to the supervisor when the gateway could
// not durably record the connection. Dropping the buffered response forces
// a reconnect, which gives the state transition a fresh chance. A
// streamed-apply session becomes ready only after SupervisorRuntimeReady.
if !stream_applies_config
// a reconnect, which gives the state transition a fresh chance. Other
// sessions become ready only after SupervisorRuntimeReady.
if mode == SessionMode::Legacy
&& !mark_supervisor_initialized(&state, &sandbox_id, &session_id, &instance_id, false).await
{
state
@@ -3039,6 +3063,13 @@ pub async fn handle_connect_supervisor(
// The stock supervisor includes discovery only on its first connection.
// Reconnects from the same process omit it and proceed directly to the
// current authoritative bootstrap.
let mode = if stream_applies_config {
SessionMode::ConfigApply
} else if hello.supports_config_apply {
SessionMode::ReportsReadiness
} else {
SessionMode::Legacy
};
let prepares_startup_policy = stream_applies_config && hello.image_policy_discovery.is_some();
let image_policy_admission = if stream_applies_config {
image_policy_admission(&hello)?
@@ -3088,7 +3119,7 @@ pub async fn handle_connect_supervisor(
sandbox_id,
instance_id: hello.instance_id,
connection_epoch: hello.connection_epoch,
stream_applies_config,
mode,
bootstrap,
provider_readiness,
config_push,
@@ -3222,7 +3253,7 @@ pub async fn handle_connect_supervisor(
sandbox_id: sandbox_id.clone(),
instance_id: hello.instance_id,
connection_epoch: hello.connection_epoch,
stream_applies_config,
mode,
bootstrap: Some(bootstrap),
provider_readiness,
config_push,
@@ -3688,14 +3719,15 @@ 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
// A polling session has no streamed admission to require.
if !mark_supervisor_initialized(
state,
sandbox_id,
session_id,
instance_id,
stream_applies_config,
)
.await
{
warn!(
sandbox_id,
@@ -4153,7 +4185,7 @@ mod tests {
tx.clone(),
make_shutdown(),
Some(Arc::clone(&slots)),
false,
SessionMode::Legacy,
);
(tx, SessionOutbound::new(rx, slots))
}
@@ -4194,18 +4226,38 @@ mod tests {
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (_tx, mut outbound) = register_push_session(&registry, "sb-1", "session-1");
for expected in [DeliveryDisposition::Queued, DeliveryDisposition::Replaced] {
for (config_revision, expected) in [
(1, DeliveryDisposition::Queued),
(2, DeliveryDisposition::Replaced),
] {
assert_eq!(
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
SupervisorConfigMessage::SandboxConfig(Box::new(SandboxConfigSnapshot {
config_revision,
..Default::default()
})),
)
.await,
expected
);
}
// An identical snapshot is not sent again.
assert_eq!(
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::SandboxConfig(Box::new(SandboxConfigSnapshot {
config_revision: 2,
..Default::default()
})),
)
.await,
DeliveryDisposition::SuppressedUnchanged
);
assert_eq!(
router
.deliver(
@@ -4333,7 +4385,7 @@ mod tests {
tx,
make_shutdown(),
Some(Arc::clone(&slots)),
true,
SessionMode::ConfigApply,
);
slots
}
@@ -6738,7 +6790,7 @@ mod tests {
tx.clone(),
shutdown_tx,
Some(Arc::clone(&slots)),
true,
SessionMode::ConfigApply,
);
let snapshot = SandboxConfigSnapshot {
configuration_instance_id: "configuration-1".into(),
@@ -7072,7 +7124,7 @@ mod tests {
tx.clone(),
shutdown_tx,
Some(Arc::clone(&slots)),
true,
SessionMode::ConfigApply,
);
let snapshot = ProviderEnvironmentSnapshot {
provider_env_revision: 11,
@@ -7410,7 +7462,7 @@ mod tests {
tx,
make_shutdown(),
Some(Arc::clone(&slots)),
false,
SessionMode::Legacy,
);
let update = |revision| {
SupervisorConfigMessage::ProviderEnvironment(ProviderEnvironmentSnapshot {
@@ -7467,6 +7519,26 @@ mod tests {
};
assert!(!accepted.config_apply_enabled);
assert!(accepted.bootstrap.is_none());
// This supervisor connects before its workload starts, so acceptance
// alone must not report it ready.
assert!(!state.supervisor_sessions.is_runtime_ready("sb-poll-apply"));
harness
.outbound
.send(SupervisorMessage {
payload: Some(supervisor_message::Payload::RuntimeReady(
openshell_core::proto::SupervisorRuntimeReady {},
)),
})
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(5), async {
while !state.supervisor_sessions.is_runtime_ready("sb-poll-apply") {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("runtime readiness marks the polling session ready");
}
#[tokio::test]
crates/openshell-supervisor-process/src/delegated.rs
+2
View File
@@ -100,6 +100,7 @@ pub async fn start_prepared_supervisor_session(
ssh_socket_path: Option<&str>,
config_apply_tx: tokio::sync::mpsc::Sender<crate::supervisor_session::ConfigApplyRequest>,
supervisor_session_updates: Option<tokio::sync::watch::Sender<Option<String>>>,
config_apply_updates: tokio::sync::watch::Sender<bool>,
) -> Result<PrestartedSupervisorSession> {
let terminating = Arc::new(AtomicBool::new(false));
let loopback = Arc::new(DeferredLoopbackConnector::default());
@@ -116,6 +117,7 @@ pub async fn start_prepared_supervisor_session(
terminating.clone(),
config_apply_tx,
supervisor_session_updates,
config_apply_updates,
);
let ready = tokio::time::timeout(
crate::supervisor_session::SESSION_PREPARE_TIMEOUT,
crates/openshell-supervisor-process/src/supervisor_session.rs
+108 -21
View File
@@ -450,6 +450,7 @@ pub fn spawn_with_readiness(
config_apply_tx: runtime.config_apply_tx,
ready_tx,
runtime_ready: Arc::new(AtomicBool::new(true)),
config_apply_updates: None,
};
(tokio::spawn(run_session_loop(config, None)), ready_rx)
}
@@ -464,17 +465,35 @@ pub async fn prepare(
image_policy_discovery: ImagePolicyDiscovery,
prepare_policy: impl FnOnce(SandboxPolicy) -> Result<Option<SandboxPolicy>, String> + Send + 'static,
) -> Result<PreparedSupervisorSession, Box<dyn std::error::Error + Send + Sync>> {
let prepared = tokio::time::timeout(
SESSION_PREPARE_TIMEOUT,
open_session(
endpoint,
sandbox_id,
instance_id,
1,
Some(image_policy_discovery),
Some(Box::new(prepare_policy)),
),
)
let prepared = tokio::time::timeout(SESSION_PREPARE_TIMEOUT, async {
let mut prepare_policy: Option<StartupPolicyPreparer> = Some(Box::new(prepare_policy));
let mut backoff = INITIAL_BACKOFF;
let mut connection_epoch = 1;
loop {
match open_session(
endpoint.clone(),
sandbox_id.clone(),
instance_id.clone(),
connection_epoch,
Some(image_policy_discovery.clone()),
prepare_policy.take(),
)
.await
{
Ok(prepared) => return Ok(prepared),
// Retry only failures before the gateway saw this hello, so
// the startup preparer is still unused.
Err(OpenSessionError::Connect(error, preparer)) => {
warn!(error = %error, "supervisor session: startup connection failed; retrying");
prepare_policy = preparer;
connection_epoch += 1;
tokio::time::sleep(backoff).await;
backoff = (backoff * 2).min(MAX_BACKOFF);
}
Err(OpenSessionError::Session(error)) => return Err(error),
}
}
})
.await
.map_err(|_| "timed out waiting for supervisor session bootstrap")??;
if prepared.config_apply_enabled && prepared.bootstrap.is_none() {
@@ -495,6 +514,7 @@ pub fn spawn_prepared(
terminating: Arc<AtomicBool>,
config_apply_tx: mpsc::Sender<ConfigApplyRequest>,
session_id_updates: Option<watch::Sender<Option<String>>>,
config_apply_updates: watch::Sender<bool>,
) -> (
tokio::task::JoinHandle<()>,
watch::Receiver<bool>,
@@ -516,6 +536,7 @@ pub fn spawn_prepared(
session_id_updates,
ready_tx,
runtime_ready: runtime_ready.clone(),
config_apply_updates: Some(config_apply_updates),
};
let task = tokio::spawn(run_session_loop(config, Some((prepared, bootstrap_result))));
(task, ready_rx, outbound, runtime_ready)
@@ -534,6 +555,10 @@ struct SessionConfig {
session_id_updates: Option<watch::Sender<Option<String>>>,
ready_tx: watch::Sender<bool>,
runtime_ready: Arc<AtomicBool>,
/// Publishes whether the current session delivers configuration
/// authoritatively, so a stream-started runtime resumes polling when a
/// reconnect lands on a gateway that does not enable apply.
config_apply_updates: Option<watch::Sender<bool>>,
}
async fn run_session_loop(
@@ -593,10 +618,45 @@ async fn run_single_session(
None,
None,
)
.await?;
.await
.map_err(OpenSessionError::into_error)?;
run_prepared_session(config, prepared, None).await
}
type SessionError = Box<dyn std::error::Error + Send + Sync>;
enum OpenSessionError {
/// The gateway never received the hello; the preparer is returned unused.
Connect(SessionError, Option<StartupPolicyPreparer>),
Session(SessionError),
}
impl OpenSessionError {
fn into_error(self) -> SessionError {
match self {
Self::Connect(error, _) | Self::Session(error) => error,
}
}
}
impl From<SessionError> for OpenSessionError {
fn from(error: SessionError) -> Self {
Self::Session(error)
}
}
impl From<&'static str> for OpenSessionError {
fn from(error: &'static str) -> Self {
Self::Session(error.into())
}
}
impl From<String> for OpenSessionError {
fn from(error: String) -> Self {
Self::Session(error.into())
}
}
async fn open_session(
endpoint: String,
sandbox_id: String,
@@ -604,12 +664,18 @@ async fn open_session(
connection_epoch: u64,
image_policy_discovery: Option<ImagePolicyDiscovery>,
mut prepare_policy: Option<StartupPolicyPreparer>,
) -> Result<PreparedSupervisorSession, Box<dyn std::error::Error + Send + Sync>> {
) -> Result<PreparedSupervisorSession, OpenSessionError> {
// The same authenticated channel carries the long-lived control stream
// and all data-plane RelayStream calls.
let channel = grpc_client::connect_channel_pub(&endpoint)
.await
.map_err(|e| format!("connect failed: {e}"))?;
let channel = match grpc_client::connect_channel_pub(&endpoint).await {
Ok(channel) => channel,
Err(e) => {
return Err(OpenSessionError::Connect(
format!("connect failed: {e}").into(),
prepare_policy,
));
}
};
let mut client = OpenShellClient::new(channel.clone());
// Create the outbound message stream.
@@ -631,11 +697,18 @@ async fn open_session(
.await
.map_err(|_| "failed to queue hello")?;
// Open the bidirectional stream.
let response = client
.connect_supervisor(outbound)
.await
.map_err(|e| format!("connect_supervisor RPC failed: {e}"))?;
// Open the bidirectional stream. An unavailable gateway rejects the call
// before it reads the hello, so the preparer is still unused.
let response = match client.connect_supervisor(outbound).await {
Ok(response) => response,
Err(e) if e.code() == tonic::Code::Unavailable => {
return Err(OpenSessionError::Connect(
format!("connect_supervisor RPC failed: {e}").into(),
prepare_policy,
));
}
Err(e) => return Err(format!("connect_supervisor RPC failed: {e}").into()),
};
let mut inbound = response.into_inner();
// The gateway may ask the initial supervisor to prepare its selected
@@ -765,8 +838,22 @@ async fn run_prepared_session(
.map_err(|_| "failed to queue configuration bootstrap result")?;
}
let config_sequences = Arc::new(Mutex::new(ConfigSequenceWatermarks::default()));
if let Some(updates) = &config.config_apply_updates {
updates.send_replace(prepared.config_apply_enabled);
}
if !prepared.config_apply_enabled {
config.ready_tx.send_replace(true);
// Without streamed admission the gateway learns readiness only from
// this report. A reconnect after the runtime started repeats it.
if config.runtime_ready.load(Ordering::Acquire) {
tx.send(SupervisorMessage {
payload: Some(supervisor_message::Payload::RuntimeReady(
openshell_core::proto::SupervisorRuntimeReady {},
)),
})
.await
.map_err(|_| "failed to queue runtime readiness")?;
}
}
// Main loop: receive gateway messages + send heartbeats.
crates/openshell-supervisor/src/lib.rs
+77 -2
View File
@@ -829,6 +829,12 @@ async fn run_sandbox_with_backend(
let uses_stream_configuration = prepared_supervisor_session.as_ref().is_some_and(
openshell_supervisor_process::supervisor_session::PreparedSupervisorSession::uses_stream_configuration,
);
// A gateway that keeps polling authoritative must not see this session
// until the workload is running. Close it and reconnect after startup, as
// polling supervisors always have.
if !uses_stream_configuration {
prepared_supervisor_session = None;
}
if stream_bootstrap
.as_ref()
@@ -1125,6 +1131,7 @@ async fn run_sandbox_with_backend(
let (workspace_tx, workspace_rx) = tokio::sync::watch::channel(String::new());
let (config_apply_tx, config_apply_rx) = tokio::sync::mpsc::channel(16);
let mut config_apply_rx = Some(config_apply_rx);
let (config_apply_updates, config_apply_enabled) = tokio::sync::watch::channel(true);
let remote_network_source = remote_boundary.0.network_mediation_source();
let remote_host_gateway_ip = remote_boundary.0.host_gateway_ip();
@@ -1328,6 +1335,7 @@ async fn run_sandbox_with_backend(
.as_ref()
.and_then(|bootstrap| bootstrap.sandbox_config.clone())
.map(Into::into),
config_apply_enabled: Some(config_apply_enabled),
endpoint_observation_tx,
endpoint_status_rx,
endpoint_policy: poll_endpoint_policy,
@@ -1378,6 +1386,7 @@ async fn run_sandbox_with_backend(
ssh_socket_path.as_deref(),
config_apply_tx.clone(),
Some(supervisor_session_updates.clone()),
config_apply_updates,
)
.await?,
),
@@ -4024,6 +4033,9 @@ struct PolicyPollLoopContext {
/// runtime state, so this seeds equality tracking and disables fetch-based
/// reconciliation for the current protocol.
initial_stream_snapshot: Option<openshell_core::grpc_client::SettingsPollResult>,
/// Whether the current supervisor session delivers configuration
/// authoritatively. A stream-started runtime polls while it is false.
config_apply_enabled: Option<tokio::sync::watch::Receiver<bool>>,
/// Producer shared with network enforcement and policy installation.
endpoint_observation_tx: Option<openshell_core::endpoint_status::EndpointObservationSender>,
/// Single FIFO consumed by the endpoint status reporter.
@@ -5264,7 +5276,15 @@ async fn run_policy_poll_loop_with_client<C: PolicyGatewayClient>(
}
let initial_stream_snapshot = ctx.initial_stream_snapshot.take();
let stream_authoritative = initial_stream_snapshot.is_some();
let stream_started = initial_stream_snapshot.is_some();
// A stream-started runtime stays stream-authoritative only while its
// current session applies configuration. A reconnect to a gateway that
// does not enable apply, such as after a rollback to poll mode, resumes
// polling. Keep the last value if the session task ends.
let mut config_apply_updates = ctx.config_apply_enabled.take();
let mut config_apply_enabled = config_apply_updates
.as_ref()
.is_none_or(|updates| *updates.borrow());
let mut current_config_revision: u64 = initial_stream_snapshot
.as_ref()
.map_or(0, |snapshot| snapshot.config_revision);
@@ -5446,9 +5466,22 @@ async fn run_policy_poll_loop_with_client<C: PolicyGatewayClient>(
let interval = Duration::from_secs(ctx.interval_secs);
loop {
if stream_authoritative {
if let Some(updates) = config_apply_updates.as_ref() {
config_apply_enabled = *updates.borrow();
}
if stream_started && config_apply_enabled {
let delay = next_poll_delay(&ctx.extension_credentials, interval);
tokio::select! {
changed = async {
match config_apply_updates.as_mut() {
Some(updates) => updates.changed().await.is_ok(),
None => std::future::pending().await,
}
} => {
if !changed {
config_apply_updates = None;
}
}
request = receive_config_apply(&mut config_apply_rx) => {
let Some(request) = request else {
return Err(miette::miette!("stream configuration apply channel closed"));
@@ -9153,6 +9186,7 @@ network_policies:
transparent_tcp: TransparentTcpReloadState::default(),
config_apply_rx: None,
initial_stream_snapshot: None,
config_apply_enabled: None,
endpoint_observation_tx: None,
endpoint_status_rx: None,
endpoint_policy: None,
@@ -10926,6 +10960,47 @@ network_policies:
handle.abort();
}
#[tokio::test]
async fn stream_started_runtime_polls_when_session_stops_applying() {
let initial = settings_poll_result(
Some(proto_policy_fixture()),
1,
openshell_core::proto::PolicySource::Sandbox,
);
let engine =
Arc::new(OpaEngine::from_proto(&proto_policy_fixture()).expect("build OPA engine"));
let mut ctx = policy_poll_test_context(
engine,
LoadedPolicyOrigin::Gateway {
revision: Some(LoadedPolicyRevision::from_snapshot(&initial)),
has_last_valid_policy: true,
},
default_middleware_connector(),
);
ctx.initial_stream_snapshot = Some(initial);
let (_config_apply_tx, config_apply_rx) = tokio::sync::mpsc::channel(1);
ctx.config_apply_rx = Some(config_apply_rx);
let (apply_enabled_tx, apply_enabled_rx) = tokio::sync::watch::channel(true);
ctx.config_apply_enabled = Some(apply_enabled_rx);
let (client, _polls, _reports) = scripted_policy_gateway();
let poll_calls = Arc::clone(&client.poll_calls);
let handle = tokio::spawn(run_policy_poll_loop_with_client(ctx, client));
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(poll_calls.load(Ordering::SeqCst), 0);
// A reconnect lands on a gateway that keeps polling authoritative.
apply_enabled_tx.send_replace(false);
timeout(Duration::from_secs(5), async {
while poll_calls.load(Ordering::SeqCst) == 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("a session without apply must resume polling");
handle.abort();
}
#[tokio::test]
async fn streamed_tcp_expansion_reports_retained_active_policy() {
use openshell_core::proto::{ConfigApplyOutcome, PolicySource};