From ee731a221e12d6d2e3358dfd487b2d87171640de Mon Sep 17 00:00:00 2001 From: JJ Liebig Date: Tue, 13 Oct 2026 18:00:54 +0200 Subject: [PATCH] fix: allow busy servers time to become ready (#5241) refs #4721 --- src/client/endpoint/health.rs | 68 +++++++++++++++++++++++----- src/client/endpoint/registry.rs | 6 ++- src/client/endpoint/supervisor.rs | 62 +++++++++++++++++++++++--- src/client/mod.rs | 1 + src/server/client_transport.rs | 74 +++++++++++++++++++++++++++++++ 5 files changed, 190 insertions(+), 21 deletions(-) diff --git a/src/client/endpoint/health.rs b/src/client/endpoint/health.rs index 87c08952f..230939c3e 100644 --- a/src/client/endpoint/health.rs +++ b/src/client/endpoint/health.rs @@ -2,6 +2,8 @@ use std::time::{Duration, Instant}; pub(super) const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); pub(super) const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(10); +// Match the remote setup allowance without extending established transport health checks. +const INITIAL_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(60); #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(super) enum HealthAction { @@ -12,7 +14,7 @@ pub(super) enum HealthAction { pub(super) struct EndpointHealth { connected_at: Instant, - last_received: Instant, + last_received: Option, ping_sent_at: Option, ready: bool, } @@ -21,14 +23,14 @@ impl EndpointHealth { pub(super) fn new(now: Instant) -> Self { Self { connected_at: now, - last_received: now, + last_received: None, ping_sent_at: None, ready: false, } } pub(super) fn received(&mut self, now: Instant) { - self.last_received = now; + self.last_received = Some(now); self.ping_sent_at = None; } @@ -37,15 +39,18 @@ impl EndpointHealth { } pub(super) fn action(&self, now: Instant) -> HealthAction { - let initial_snapshot_expired = - !self.ready && now.saturating_duration_since(self.connected_at) >= HEARTBEAT_TIMEOUT; - let probe_expired = self - .ping_sent_at - .is_some_and(|sent_at| now.saturating_duration_since(sent_at) >= HEARTBEAT_TIMEOUT); + let initial_snapshot_expired = !self.ready + && now.saturating_duration_since(self.connected_at) >= INITIAL_SNAPSHOT_TIMEOUT; + // Welcome precedes state-loop registration, which can postpone the first pong. + let probe_expired = (self.ready || self.last_received.is_some()) + && self + .ping_sent_at + .is_some_and(|sent_at| now.saturating_duration_since(sent_at) >= HEARTBEAT_TIMEOUT); if initial_snapshot_expired || probe_expired { HealthAction::Expired } else if self.ping_sent_at.is_none() - && now.saturating_duration_since(self.last_received) >= HEARTBEAT_INTERVAL + && now.saturating_duration_since(self.last_received.unwrap_or(self.connected_at)) + >= HEARTBEAT_INTERVAL { HealthAction::Ping } else { @@ -66,6 +71,7 @@ mod tests { fn quiet_connection_is_probed_then_expires_without_a_reply() { let now = Instant::now(); let mut health = EndpointHealth::new(now); + health.ready(); assert_eq!(health.action(now), HealthAction::None); assert_eq!(health.action(now + HEARTBEAT_INTERVAL), HealthAction::Ping); health.ping_sent(now + HEARTBEAT_INTERVAL); @@ -86,13 +92,51 @@ mod tests { } #[test] - fn heartbeats_do_not_hide_a_missing_initial_snapshot() { + fn delayed_registration_can_reply_after_the_normal_probe_deadline() { let now = Instant::now(); let mut health = EndpointHealth::new(now); health.ping_sent(now + HEARTBEAT_INTERVAL); - health.received(now + HEARTBEAT_INTERVAL + Duration::from_secs(1)); + let registered = now + Duration::from_secs(45); + assert_eq!(health.action(registered), HealthAction::None); + health.received(registered); + health.ready(); + assert_eq!(health.action(registered), HealthAction::None); + } + + #[test] + fn responding_startup_gets_time_for_its_snapshot_but_cannot_wait_forever() { + let now = Instant::now(); + let mut health = EndpointHealth::new(now); + for seconds in (5..60).step_by(5) { + let reply = now + Duration::from_secs(seconds); + health.received(reply); + assert_eq!(health.action(reply), HealthAction::None); + } assert_eq!( - health.action(now + HEARTBEAT_TIMEOUT), + health.action(now + Duration::from_secs(60)), + HealthAction::Expired + ); + } + + #[test] + fn startup_silence_is_bounded_and_confirmed_transport_keeps_normal_probe_timeout() { + let now = Instant::now(); + let mut health = EndpointHealth::new(now); + health.ping_sent(now + HEARTBEAT_INTERVAL); + assert_eq!( + health.action(now + Duration::from_secs(60)), + HealthAction::Expired + ); + + let mut health = EndpointHealth::new(now); + health.received(now + Duration::from_secs(6)); + health.ping_sent(now + Duration::from_secs(11)); + assert_eq!( + health.action(now + Duration::from_secs(20)), + HealthAction::None + ); + assert_eq!( + health.action(now + Duration::from_secs(21)), HealthAction::Expired ); } diff --git a/src/client/endpoint/registry.rs b/src/client/endpoint/registry.rs index 437f39f37..f575334d8 100644 --- a/src/client/endpoint/registry.rs +++ b/src/client/endpoint/registry.rs @@ -515,6 +515,7 @@ mod tests { false, ); let now = Instant::now(); + registry.mark_ready(&ssh_id, 2); registry.tick_health(now + super::super::health::HEARTBEAT_INTERVAL); assert!(matches!( sent.lock().unwrap().as_slice(), @@ -553,9 +554,10 @@ mod tests { ); let now = Instant::now(); registry.mark_ready(&ssh_id, 2); - registry.received(&ssh_id, 2, now + super::super::health::HEARTBEAT_INTERVAL); - registry.tick_health(now + super::super::health::HEARTBEAT_TIMEOUT); + registry.received(&ssh_id, 2, now + std::time::Duration::from_secs(80)); + registry.tick_health(now + std::time::Duration::from_secs(80)); assert!(registry.connection(&ssh_id).is_some()); + assert!(!registry.connection(&ssh_id).unwrap().surface_active); } #[test] diff --git a/src/client/endpoint/supervisor.rs b/src/client/endpoint/supervisor.rs index 471acd7d2..4be72784e 100644 --- a/src/client/endpoint/supervisor.rs +++ b/src/client/endpoint/supervisor.rs @@ -54,7 +54,7 @@ struct ReconnectState { next_attempt: Option, in_flight: bool, generation: Option, - online_since: Option, + ready_since: Option, } impl ReconnectState { @@ -65,7 +65,7 @@ impl ReconnectState { next_attempt: Some(now), in_flight: false, generation: None, - online_since: None, + ready_since: None, } } } @@ -208,23 +208,22 @@ impl EndpointSupervisors { if endpoint_id.is_local() { state.attempts = 0; } - state.online_since.get_or_insert(now); state.next_attempt = None; } ClientEndpointStatus::Attention => { - state.online_since = None; + state.ready_since = None; // Authentication or configuration may be repaired outside this client. state.next_attempt = (!endpoint_id.is_local()).then_some(now + Duration::from_secs(30)); } ClientEndpointStatus::Disabled => { - state.online_since = None; + state.ready_since = None; state.next_attempt = None; } ClientEndpointStatus::Connecting | ClientEndpointStatus::Reconnecting => { // A brief maintenance wake can complete a handshake without restoring the link. - if state.online_since.take().is_some_and(|connected| { - now.saturating_duration_since(connected) >= STABLE_CONNECTION_PERIOD + if state.ready_since.take().is_some_and(|ready| { + now.saturating_duration_since(ready) >= STABLE_CONNECTION_PERIOD }) { state.attempts = 0; } @@ -242,6 +241,21 @@ impl EndpointSupervisors { true } + pub(crate) fn mark_ready( + &mut self, + endpoint_id: &ClientEndpointId, + generation: u64, + now: Instant, + ) { + if let Some(state) = self + .endpoints + .get_mut(endpoint_id) + .filter(|state| state.generation == Some(generation)) + { + state.ready_since.get_or_insert(now); + } + } + /// A user asked for this machine now, so skip its pending backoff and restart the ladder. /// Returns false when no attempt is waiting: it is online, in flight, or not retried at all. pub(crate) fn retry_now(&mut self, endpoint_id: &ClientEndpointId, now: Instant) -> bool { @@ -481,6 +495,7 @@ mod tests { for attempt in 1..=5 { let connected = now + Duration::from_secs(attempt * 20); assert!(supervisors.record_status(&id, 2, ClientEndpointStatus::Online, connected)); + supervisors.mark_ready(&id, 2, connected); let failed = connected + Duration::from_secs(15); assert!(supervisors.disconnected(&id, 2, failed)); assert_eq!( @@ -490,14 +505,47 @@ mod tests { } let connected = now + Duration::from_secs(200); assert!(supervisors.record_status(&id, 2, ClientEndpointStatus::Online, connected)); + supervisors.mark_ready(&id, 1, connected); + let ready = connected + Duration::from_secs(45); + supervisors.mark_ready(&id, 2, ready); let failed = connected + Duration::from_secs(60); assert!(supervisors.disconnected(&id, 2, failed)); + assert_eq!( + supervisors.endpoints[&id].next_attempt, + Some(failed + Duration::from_secs(16)) + ); + + let connected = now + Duration::from_secs(300); + assert!(supervisors.record_status(&id, 2, ClientEndpointStatus::Online, connected)); + let ready = connected + Duration::from_secs(45); + supervisors.mark_ready(&id, 2, ready); + supervisors.mark_ready(&id, 2, ready + Duration::from_secs(30)); + let failed = ready + Duration::from_secs(60); + assert!(supervisors.disconnected(&id, 2, failed)); assert_eq!( supervisors.endpoints[&id].next_attempt, Some(failed + INITIAL_RETRY_DELAY) ); } + #[test] + fn a_handshake_without_readiness_does_not_reset_backoff() { + let now = Instant::now(); + let profile = profile(); + let id = ClientEndpointId::Ssh(profile.id.clone()); + let mut supervisors = EndpointSupervisors::new(&[profile], now); + let state = supervisors.endpoints.get_mut(&id).unwrap(); + state.generation = Some(2); + state.attempts = 9; + assert!(supervisors.record_status(&id, 2, ClientEndpointStatus::Online, now)); + let failed = now + Duration::from_secs(75); + assert!(supervisors.disconnected(&id, 2, failed)); + assert_eq!( + supervisors.endpoints[&id].next_attempt, + Some(failed + MAX_RETRY_DELAY) + ); + } + #[test] fn requested_retry_skips_backoff_without_duplicating_attempts() { let now = Instant::now(); diff --git a/src/client/mod.rs b/src/client/mod.rs index 51c181906..14482134a 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -2171,6 +2171,7 @@ async fn run_client_loop( } } write_stream.mark_ready(&endpoint_id, generation); + supervisors.mark_ready(&endpoint_id, generation, now); if endpoint_id.is_local() { if let Some(event) = take_ready_local_activation(&mut state, &write_stream) diff --git a/src/server/client_transport.rs b/src/server/client_transport.rs index 0536fe7e9..db985baca 100644 --- a/src/server/client_transport.rs +++ b/src/server/client_transport.rs @@ -2021,6 +2021,80 @@ mod tests { .expect("handshake thread result"); } + #[test] + fn endpoint_health_waits_for_registration_but_not_the_initial_snapshot() { + let (mut client, server, _path) = local_stream_pair("startup-health"); + let (events, mut pending) = mpsc::channel(1); + events + .blocking_send(ServerEvent::ClientWriterDrained { client_id: 99 }) + .unwrap(); + let server_events = events.clone(); + let server_thread = std::thread::spawn(move || { + handle_client_handshake( + server, + 43, + &server_events, + &Arc::new(AtomicBool::new(false)), + ) + }); + let mut reader = client.try_clone().unwrap(); + let (messages, received) = std::sync::mpsc::channel(); + let client_thread = std::thread::spawn(move || { + for _ in 0..3 { + let message: ServerMessage = + protocol::read_message(&mut reader, MAX_FRAME_SIZE).unwrap(); + messages.send(message).unwrap(); + } + }); + protocol::write_message(&mut client, &endpoint_hello(80, 29)).unwrap(); + let welcome = endpoint_welcome(received.recv_timeout(Duration::from_secs(2)).unwrap()); + assert!(welcome.error.is_none()); + let ping = ClientMessage::EndpointControl { + kind: crate::protocol::endpoint::HEALTH_PING_KIND.into(), + data: "startup".into(), + }; + protocol::write_message(&mut client, &ping).unwrap(); + assert!(matches!( + received.recv_timeout(Duration::from_millis(100)), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) + )); + assert!(matches!( + pending.blocking_recv(), + Some(ServerEvent::ClientWriterDrained { client_id: 99 }) + )); + let Some(ServerEvent::ClientShellConnected { writer, .. }) = pending.blocking_recv() else { + panic!("registration must precede other client events"); + }; + assert!(matches!( + received.recv_timeout(Duration::from_secs(2)).unwrap(), + ServerMessage::EndpointControl { kind, .. } + if kind == crate::protocol::endpoint::HEALTH_PONG_KIND + )); + + // Registered transports answer health checks while the state loop is still busy. + events + .blocking_send(ServerEvent::ClientWriterDrained { client_id: 99 }) + .unwrap(); + protocol::write_message(&mut client, &ping).unwrap(); + assert!(matches!( + received.recv_timeout(Duration::from_secs(2)).unwrap(), + ServerMessage::EndpointControl { kind, .. } + if kind == crate::protocol::endpoint::HEALTH_PONG_KIND + )); + protocol::write_message(&mut client, &ClientMessage::Detach).unwrap(); + assert!(matches!( + pending.blocking_recv(), + Some(ServerEvent::ClientWriterDrained { client_id: 99 }) + )); + assert!(matches!( + pending.blocking_recv(), + Some(ServerEvent::ClientDetach { client_id: 43 }) + )); + server_thread.join().unwrap().unwrap(); + client_thread.join().unwrap(); + drop(writer); + } + #[test] fn dedicated_client_shell_handshake_rejects_empty_surface() { let (mut client_stream, server_stream, _path) =