diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 09d86f1a9..e7d527fe9 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -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] diff --git a/crates/openshell-server/src/config_delivery/mod.rs b/crates/openshell-server/src/config_delivery/mod.rs index 5e1f0af3c..db147e1c6 100644 --- a/crates/openshell-server/src/config_delivery/mod.rs +++ b/crates/openshell-server/src/config_delivery/mod.rs @@ -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, 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); } }); } diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 8794d8640..4dd78f762 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -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, 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, shutdown: oneshot::Sender<()>, config_slots: Option>, - 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, provider_readiness: ProviderReadinessEvidence, config_push: Option, @@ -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(®istry)); let (_tx, mut outbound) = register_push_session(®istry, "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] diff --git a/crates/openshell-supervisor-process/src/delegated.rs b/crates/openshell-supervisor-process/src/delegated.rs index 66f7a19ca..eac67f1a7 100644 --- a/crates/openshell-supervisor-process/src/delegated.rs +++ b/crates/openshell-supervisor-process/src/delegated.rs @@ -100,6 +100,7 @@ pub async fn start_prepared_supervisor_session( ssh_socket_path: Option<&str>, config_apply_tx: tokio::sync::mpsc::Sender, supervisor_session_updates: Option>>, + config_apply_updates: tokio::sync::watch::Sender, ) -> Result { 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, diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index 00219e091..aaef937aa 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -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, String> + Send + 'static, ) -> Result> { - 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 = 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, config_apply_tx: mpsc::Sender, session_id_updates: Option>>, + config_apply_updates: watch::Sender, ) -> ( tokio::task::JoinHandle<()>, watch::Receiver, @@ -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>>, ready_tx: watch::Sender, runtime_ready: Arc, + /// 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>, } 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; + +enum OpenSessionError { + /// The gateway never received the hello; the preparer is returned unused. + Connect(SessionError, Option), + Session(SessionError), +} + +impl OpenSessionError { + fn into_error(self) -> SessionError { + match self { + Self::Connect(error, _) | Self::Session(error) => error, + } + } +} + +impl From 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 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, mut prepare_policy: Option, -) -> Result> { +) -> Result { // 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. diff --git a/crates/openshell-supervisor/src/lib.rs b/crates/openshell-supervisor/src/lib.rs index 670053332..784d2f113 100644 --- a/crates/openshell-supervisor/src/lib.rs +++ b/crates/openshell-supervisor/src/lib.rs @@ -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, + /// Whether the current supervisor session delivers configuration + /// authoritatively. A stream-started runtime polls while it is false. + config_apply_enabled: Option>, /// Producer shared with network enforcement and policy installation. endpoint_observation_tx: Option, /// Single FIFO consumed by the endpoint status reporter. @@ -5264,7 +5276,15 @@ async fn run_policy_poll_loop_with_client( } 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( 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};