feat: manage multiple ssh machines from one client (#3670)

1 files changed, 99 insertions, 0 deletions

diff --git a/src/client/endpoint/health.rs b/src/client/endpoint/health.rs
new file mode 100644
index 00000000..87c08952
--- /dev/null
+++ b/src/client/endpoint/health.rs
@@ -0,0 +1,99 @@
+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);
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(super) enum HealthAction {
+    None,
+    Ping,
+    Expired,
+}
+
+pub(super) struct EndpointHealth {
+    connected_at: Instant,
+    last_received: Instant,
+    ping_sent_at: Option<Instant>,
+    ready: bool,
+}
+
+impl EndpointHealth {
+    pub(super) fn new(now: Instant) -> Self {
+        Self {
+            connected_at: now,
+            last_received: now,
+            ping_sent_at: None,
+            ready: false,
+        }
+    }
+
+    pub(super) fn received(&mut self, now: Instant) {
+        self.last_received = now;
+        self.ping_sent_at = None;
+    }
+
+    pub(super) fn ready(&mut self) {
+        self.ready = true;
+    }
+
+    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);
+        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
+        {
+            HealthAction::Ping
+        } else {
+            HealthAction::None
+        }
+    }
+
+    pub(super) fn ping_sent(&mut self, now: Instant) {
+        self.ping_sent_at = Some(now);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn quiet_connection_is_probed_then_expires_without_a_reply() {
+        let now = Instant::now();
+        let mut health = EndpointHealth::new(now);
+        assert_eq!(health.action(now), HealthAction::None);
+        assert_eq!(health.action(now + HEARTBEAT_INTERVAL), HealthAction::Ping);
+        health.ping_sent(now + HEARTBEAT_INTERVAL);
+        assert_eq!(
+            health.action(now + HEARTBEAT_INTERVAL + HEARTBEAT_TIMEOUT),
+            HealthAction::Expired
+        );
+    }
+
+    #[test]
+    fn any_incoming_message_satisfies_an_outstanding_probe() {
+        let now = Instant::now();
+        let mut health = EndpointHealth::new(now);
+        health.ready();
+        health.ping_sent(now);
+        health.received(now + HEARTBEAT_TIMEOUT - Duration::from_millis(1));
+        assert_eq!(health.action(now + HEARTBEAT_TIMEOUT), HealthAction::None);
+    }
+
+    #[test]
+    fn heartbeats_do_not_hide_a_missing_initial_snapshot() {
+        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));
+        assert_eq!(
+            health.action(now + HEARTBEAT_TIMEOUT),
+            HealthAction::Expired
+        );
+    }
+}