feat(envoy-client): add ping health check

Author: MasterPtato

Cargo.lock

@@ -0,0 +1,108 @@
+//! Minimal TCP forwarder with a `freeze` mode used by network-fault tests.
+//!
+//! Toxiproxy can stall traffic but always relays a peer's TCP close to the other side, which
+//! defeats tests that need to observe one peer's behavior when it has no signal that the other
+//! peer hung up. `FreezeProxy` is a single-purpose forwarder that supports a true black-hole
+//! mode: while frozen, bytes are read from each peer and discarded, and an EOF from either peer
+//! is held instead of being forwarded.
+
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use anyhow::{Context, Result};
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
+use tokio::net::{TcpListener, TcpStream};
+
+pub struct FreezeProxy {
+	listen_addr: SocketAddr,
+	frozen: Arc<AtomicBool>,
+}
+
+impl FreezeProxy {
+	pub async fn start(upstream: SocketAddr) -> Result<Self> {
+		let listener = TcpListener::bind("127.0.0.1:0")
+			.await
+			.context("failed to bind FreezeProxy listener")?;
+		let listen_addr = listener
+			.local_addr()
+			.context("failed to read FreezeProxy listen addr")?;
+		let frozen = Arc::new(AtomicBool::new(false));
+
+		tokio::spawn({
+			let frozen = frozen.clone();
+			async move {
+				loop {
+					let (client, _) = match listener.accept().await {
+						Ok(pair) => pair,
+						Err(err) => {
+							tracing::warn!(?err, "freeze proxy accept failed");
+							return;
+						}
+					};
+					let _ = client.set_nodelay(true);
+					let frozen = frozen.clone();
+					tokio::spawn(async move {
+						let server = match TcpStream::connect(upstream).await {
+							Ok(stream) => stream,
+							Err(err) => {
+								tracing::warn!(?err, %upstream, "freeze proxy upstream connect failed");
+								return;
+							}
+						};
+						let _ = server.set_nodelay(true);
+						let (client_r, client_w) = client.into_split();
+						let (server_r, server_w) = server.into_split();
+						tokio::spawn(forward(client_r, server_w, frozen.clone()));
+						tokio::spawn(forward(server_r, client_w, frozen));
+					});
+				}
+			}
+		});
+
+		Ok(Self {
+			listen_addr,
+			frozen,
+		})
+	}
+
+	pub fn endpoint(&self) -> String {
+		format!("http://{}", self.listen_addr)
+	}
+
+	/// Stops shuttling bytes between the two peers and starts swallowing EOFs so neither peer
+	/// learns that the other has hung up. Bytes already in flight before this call may still
+	/// reach the other side.
+	pub fn freeze(&self) {
+		self.frozen.store(true, Ordering::SeqCst);
+	}
+}
+
+async fn forward(mut src: OwnedReadHalf, mut dst: OwnedWriteHalf, frozen: Arc<AtomicBool>) {
+	let mut buf = vec![0u8; 8192];
+	loop {
+		match src.read(&mut buf).await {
+			Ok(0) => {
+				if frozen.load(Ordering::SeqCst) {
+					// Hold the destination open: the peer's own send/recv buffer keeps it
+					// believing the connection is alive, with no FIN ever delivered.
+					std::future::pending::<()>().await;
+				}
+				return;
+			}
+			Ok(n) => {
+				if frozen.load(Ordering::SeqCst) {
+					// Drain-and-discard so the sender's TCP window stays open and it does not
+					// notice the link is dead via back-pressure.
+					continue;
+				}
+				if dst.write_all(&buf[..n]).await.is_err() {
+					return;
+				}
+			}
+			Err(_) => return,
+		}
+	}
+}

engine/packages/engine/tests/common/freeze_proxy.rs

@@ -3,6 +3,7 @@
 pub mod actors;
 pub mod api;
 pub mod ctx;
+pub mod freeze_proxy;
 pub mod test_envoy;
 pub mod test_helpers;
 pub mod test_runner;

engine/packages/engine/tests/common/mod.rs

@@ -254,6 +254,14 @@ impl Envoy {
 		}
 	}
 
+	pub async fn is_ping_healthy(&self) -> Option<bool> {
+		self.handle
+			.lock()
+			.await
+			.as_ref()
+			.map(|handle| handle.is_ping_healthy())
+	}
+
 	pub async fn shutdown(&self) {
 		if let Some(handle) = self.handle.lock().await.take() {
 			handle.shutdown_and_wait(false).await;
@@ -299,7 +307,10 @@ impl TestEnvoyCallbacks {
 
 impl rivet_test_envoy::EnvoyCallbacks for TestEnvoyCallbacks {
 	fn on_connect(&self, _handle: EnvoyHandle) {
-		let _ = self.inner.connection_tx.send(EnvoyConnectionEvent::Connected);
+		let _ = self
+			.inner
+			.connection_tx
+			.send(EnvoyConnectionEvent::Connected);
 	}
 
 	fn on_disconnect(&self, _handle: EnvoyHandle) {

engine/packages/engine/tests/common/test_envoy.rs

@@ -8,7 +8,9 @@ use super::super::common;
 #[test]
 fn envoy_reconnects_after_server_side_tcp_reset() {
 	common::run(
-		common::TestOpts::new(1).with_timeout(90).with_network_faults(),
+		common::TestOpts::new(1)
+			.with_timeout(90)
+			.with_network_faults(),
 		|ctx| async move {
 			let dc = ctx.leader_dc();
 			let (namespace, _) = common::setup_test_namespace(dc).await;
@@ -71,8 +73,9 @@ fn envoy_reconnects_after_server_side_tcp_reset() {
 			// The ping task writes every few seconds in the test config.
 			disconnect.wait().await;
 
-			let reconnect = envoy
-				.wait_for_next_connection_event(common::test_envoy::EnvoyConnectionEvent::Connected);
+			let reconnect = envoy.wait_for_next_connection_event(
+				common::test_envoy::EnvoyConnectionEvent::Connected,
+			);
 			envoy_proxy
 				.clear_toxics()
 				.await
@@ -90,6 +93,100 @@ fn envoy_reconnects_after_server_side_tcp_reset() {
 	);
 }
 
+#[test]
+fn engine_closes_envoy_ws_after_ping_timeout_while_envoy_remains_unaware() {
+	common::run(
+		common::TestOpts::new(1).with_timeout(120),
+		|ctx| async move {
+			let dc = ctx.leader_dc();
+			let (namespace, _) = common::setup_test_namespace(dc).await;
+
+			// Stand up our own forwarder so we can simulate a true network partition.
+			// Toxiproxy can stall traffic but always relays a peer's TCP close to the other
+			// side, which would let envoy-client notice the engine has hung up.
+			let freeze_proxy = common::freeze_proxy::FreezeProxy::start(
+				std::net::SocketAddr::from(([127, 0, 0, 1], dc.guard_port())),
+			)
+			.await
+			.expect("failed to start freeze proxy");
+
+			let envoy = common::setup_envoy(dc, &namespace, |builder| {
+				builder
+					.with_endpoint(freeze_proxy.endpoint())
+					.with_actor_behavior("network-fault-actor", |_| {
+						Box::new(
+							common::test_envoy::CustomActorBuilder::new()
+								.on_start(|_| {
+									Box::pin(async {
+										Ok(common::test_envoy::ActorStartResult::Running)
+									})
+								})
+								.build(),
+						)
+					})
+			})
+			.await;
+
+			let res = common::create_actor(
+				dc.guard_port(),
+				&namespace,
+				"network-fault-actor",
+				envoy.pool_name(),
+				rivet_types::actors::CrashPolicy::Sleep,
+			)
+			.await;
+			let actor_id = res.actor.actor_id.to_string();
+			wait_for_envoy_actor(&envoy, &actor_id).await;
+			wait_for_connectable(dc.guard_port(), &namespace, &actor_id).await;
+
+			let response = common::ping_actor_via_guard(dc, &actor_id).await;
+			assert_eq!(response["status"], "ok");
+
+			// Subscribe before injecting the fault so we can assert no event slips through.
+			let mut disconnect = envoy.wait_for_next_connection_event(
+				common::test_envoy::EnvoyConnectionEvent::Disconnected,
+			);
+			disconnect.assert_no_event();
+
+			// Black-hole the link in both directions. Bytes are read from both peers and
+			// discarded, and EOFs are swallowed so neither peer's TCP stack ever sees a FIN.
+			// The engine still keeps sending pings every few seconds (default 3s) but no pongs
+			// come back, so its application-level ping timeout (default 15s) will eventually
+			// fire and close the WebSocket. The envoy-client has no application-level
+			// liveness check of its own, so as long as its TCP socket stays open it continues
+			// to believe the connection is healthy.
+			freeze_proxy.freeze();
+
+			// Wait well past the engine's 15s ping timeout, then assert the engine did in fact
+			// tear down its side via the ping task error log...
+			tokio::time::sleep(std::time::Duration::from_secs(20)).await;
+			let logs = common::captured_logs_snapshot();
+			assert!(
+				logs.contains("ping task failed"),
+				"expected engine to log a ping timeout, got logs:\n{logs}"
+			);
+
+			// ...and that the envoy-client is still oblivious. The engine's close frame and
+			// TCP FIN never reach it because the freeze proxy is holding the link open from
+			// envoy-client's perspective.
+			disconnect.assert_no_event();
+
+			// Even though the envoy-client thinks the WebSocket is alive, its own ping-tracker
+			// must report unhealthy because no engine ping arrived in the last 20s. This is
+			// the signal the rivetkit `/health` endpoint uses to ask its host to recycle the
+			// container.
+			let healthy = envoy
+				.is_ping_healthy()
+				.await
+				.expect("envoy handle should exist");
+			assert!(
+				!healthy,
+				"envoy-client should report unhealthy after 20s without an engine ping"
+			);
+		},
+	);
+}
+
 async fn wait_for_envoy_actor(envoy: &common::test_envoy::TestEnvoy, actor_id: &str) {
 	common::wait_with_poll(
 		std::time::Duration::from_secs(5),
@@ -138,7 +235,9 @@ async fn ping_actor_via_gateway(guard_port: u16, actor_id: &str) -> serde_json::
 		.expect("failed to build reqwest client");
 
 	let response = client
-		.get(format!("http://127.0.0.1:{guard_port}/gateway/{actor_id}/ping"))
+		.get(format!(
+			"http://127.0.0.1:{guard_port}/gateway/{actor_id}/ping"
+		))
 		.send()
 		.await
 		.expect("failed to ping actor through gateway");

engine/packages/engine/tests/envoy/network_faults.rs

@@ -215,6 +215,19 @@ impl ToxiproxyProxy {
 		.await
 	}
 
+	pub async fn timeout_upstream(&self, timeout_ms: u64, toxicity: f32) -> Result<()> {
+		self.add_toxic(
+			"timeout-upstream",
+			"timeout",
+			ToxiproxyDirection::Upstream,
+			toxicity,
+			TimeoutAttributes {
+				timeout: timeout_ms,
+			},
+		)
+		.await
+	}
+
 	pub async fn bandwidth_downstream(&self, kbps: u64) -> Result<()> {
 		self.add_toxic(
 			"bandwidth-downstream",

engine/packages/test-deps-docker/src/toxiproxy.rs

@@ -1660,6 +1660,7 @@ mod tests {
 			)),
 			protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)),
 			shutting_down: std::sync::atomic::AtomicBool::new(false),
+			last_ping_ts: std::sync::atomic::AtomicI64::new(rivet_util::timestamp::now()),
 			stopped_tx: tokio::sync::watch::channel(true).0,
 		});
 		(shared, envoy_rx)

engine/sdks/rust/envoy-client/src/actor.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex as StdMutex;
-use std::sync::atomic::AtomicBool;
+use std::sync::atomic::{AtomicBool, AtomicI64};
 
 use crate::async_counter::AsyncCounter;
 use rivet_envoy_protocol as protocol;
@@ -32,6 +32,11 @@ pub struct SharedContext {
 	pub ws_tx: Arc<Mutex<Option<mpsc::UnboundedSender<WsTxMessage>>>>,
 	pub protocol_metadata: Arc<Mutex<Option<protocol::ProtocolMetadata>>>,
 	pub shutting_down: AtomicBool,
+	/// Epoch ms timestamp of the most recent ping packet received from the engine. Used by
+	/// `EnvoyHandle::is_ping_healthy` to surface a dead engine link to upstream health checks.
+	/// Initialized to the construction time so a freshly created envoy reports healthy until
+	/// its first ping arrives or the threshold elapses without one.
+	pub last_ping_ts: AtomicI64,
 	// Latched signal fired by `envoy_loop` after its cleanup block completes.
 	// Waiters observing `true` are guaranteed that the loop has exited and
 	// every pending KV/SQLite request has been resolved (with `EnvoyShutdownError`

engine/sdks/rust/envoy-client/src/context.rs

@@ -306,6 +306,7 @@ fn start_envoy_sync_inner(config: EnvoyConfig) -> EnvoyHandle {
 		ws_tx: Arc::new(tokio::sync::Mutex::new(None)),
 		protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)),
 		shutting_down: std::sync::atomic::AtomicBool::new(false),
+		last_ping_ts: std::sync::atomic::AtomicI64::new(rivet_util::timestamp::now()),
 		stopped_tx,
 	});