From 49f2fb2e9cf3ddb8279c006206828f7e71923d08 Mon Sep 17 00:00:00 2001
From: Dev Chiniwala <dev.chiniwala@gmail.com>
Date: Wed, 7 Oct 2026 20:28:48 +0530
Subject: [PATCH] fix(envoy-client): release global handle after shutdown

---
 engine/sdks/rust/envoy-client/src/envoy.rs    |  23 +++-
 .../tests/global_envoy_lifecycle.rs           | 102 ++++++++++++++++++
 2 files changed, 122 insertions(+), 3 deletions(-)
 create mode 100644 engine/sdks/rust/envoy-client/tests/global_envoy_lifecycle.rs

diff --git a/engine/sdks/rust/envoy-client/src/envoy.rs b/engine/sdks/rust/envoy-client/src/envoy.rs
index 9424be9a09..ca942ba427 100644
--- a/engine/sdks/rust/envoy-client/src/envoy.rs
+++ b/engine/sdks/rust/envoy-client/src/envoy.rs
@@ -44,9 +44,8 @@ use crate::tunnel::{
 };
 use crate::utils::{BufferMap, EnvoyShutdownError, SleepFuture, boxed_sleep, spawn_detached};
 
-/// Process-wide envoy slot. Holds the handle inside a mutex so a stopped
-/// handle (e.g. from a shutdown-during-build race in serverless mode) can be
-/// replaced on the next `start_envoy_sync` call.
+/// Process-wide envoy slot. The cached handle is cleared after cleanup so it
+/// does not keep callback state alive after shutdown.
 #[cfg(not(target_arch = "wasm32"))]
 static GLOBAL_ENVOY: OnceLock<Mutex<Option<EnvoyHandle>>> = OnceLock::new();
 
@@ -385,6 +384,22 @@ pub fn start_envoy_sync(config: EnvoyConfig) -> EnvoyHandle {
 	}
 }
 
+#[cfg(not(target_arch = "wasm32"))]
+fn clear_global_envoy(shared: &Arc<SharedContext>) {
+	let Some(slot) = GLOBAL_ENVOY.get() else {
+		return;
+	};
+
+	let mut guard = slot.lock();
+	let is_current = match guard.as_ref() {
+		Some(handle) => Arc::ptr_eq(&handle.shared, shared),
+		None => false,
+	};
+	if is_current {
+		*guard = None;
+	}
+}
+
 fn start_envoy_sync_inner(config: EnvoyConfig) -> EnvoyHandle {
 	let (envoy_tx, envoy_rx) = mpsc::unbounded_channel::<ToEnvoyMessage>();
 	let (start_tx, start_rx) = tokio::sync::watch::channel(());
@@ -644,6 +659,8 @@ async fn envoy_loop(
 	tracing::info!("envoy stopped");
 
 	ctx.shared.config.callbacks.on_shutdown();
+	#[cfg(not(target_arch = "wasm32"))]
+	clear_global_envoy(&ctx.shared);
 
 	// Latched signal: waiters on `EnvoyHandle::wait_stopped` observe this and
 	// any future callers of `wait_stopped` resolve immediately because watch
diff --git a/engine/sdks/rust/envoy-client/tests/global_envoy_lifecycle.rs b/engine/sdks/rust/envoy-client/tests/global_envoy_lifecycle.rs
new file mode 100644
index 0000000000..f78d93500b
--- /dev/null
+++ b/engine/sdks/rust/envoy-client/tests/global_envoy_lifecycle.rs
@@ -0,0 +1,102 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+
+use rivet_envoy_client::config::{
+	BoxFuture, EnvoyCallbacks, EnvoyConfig, HttpRequest, HttpResponse, WebSocketHandler,
+	WebSocketSender,
+};
+use rivet_envoy_client::envoy::start_envoy_sync;
+use rivet_envoy_client::handle::EnvoyHandle;
+use rivet_envoy_protocol as protocol;
+use tokio::sync::Notify;
+
+struct DropCallbacks {
+	dropped: Arc<Notify>,
+}
+
+impl Drop for DropCallbacks {
+	fn drop(&mut self) {
+		self.dropped.notify_one();
+	}
+}
+
+impl EnvoyCallbacks for DropCallbacks {
+	fn on_actor_start(
+		&self,
+		_handle: EnvoyHandle,
+		_actor_id: String,
+		_generation: u32,
+		_config: protocol::ActorConfig,
+		_preloaded_kv: Option<protocol::PreloadedKv>,
+	) -> BoxFuture<anyhow::Result<()>> {
+		Box::pin(async { Ok(()) })
+	}
+
+	fn on_shutdown(&self) {}
+
+	fn fetch(
+		&self,
+		_handle: EnvoyHandle,
+		_actor_id: String,
+		_gateway_id: protocol::GatewayId,
+		_request_id: protocol::RequestId,
+		_request: HttpRequest,
+	) -> BoxFuture<anyhow::Result<HttpResponse>> {
+		Box::pin(async { anyhow::bail!("fetch should not be called in lifecycle tests") })
+	}
+
+	fn websocket(
+		&self,
+		_handle: EnvoyHandle,
+		_actor_id: String,
+		_gateway_id: protocol::GatewayId,
+		_request_id: protocol::RequestId,
+		_request: HttpRequest,
+		_path: String,
+		_headers: HashMap<String, String>,
+		_is_hibernatable: bool,
+		_is_restoring_hibernatable: bool,
+		_sender: WebSocketSender,
+	) -> BoxFuture<anyhow::Result<WebSocketHandler>> {
+		Box::pin(async { anyhow::bail!("websocket should not be called in lifecycle tests") })
+	}
+
+	fn can_hibernate(
+		&self,
+		_actor_id: &str,
+		_gateway_id: &protocol::GatewayId,
+		_request_id: &protocol::RequestId,
+		_request: &HttpRequest,
+	) -> BoxFuture<anyhow::Result<bool>> {
+		Box::pin(async { Ok(false) })
+	}
+}
+
+#[tokio::test]
+async fn stopped_global_envoy_releases_callback_context() {
+	let dropped = Arc::new(Notify::new());
+	let callbacks: Arc<dyn EnvoyCallbacks> = Arc::new(DropCallbacks {
+		dropped: dropped.clone(),
+	});
+	let handle = start_envoy_sync(EnvoyConfig {
+		version: 1,
+		endpoint: "http://127.0.0.1:1".to_owned(),
+		token: None,
+		namespace: "test".to_owned(),
+		pool_name: "test".to_owned(),
+		prepopulate_actor_names: HashMap::new(),
+		metadata: None,
+		not_global: false,
+		debug_latency_ms: None,
+		callbacks,
+	});
+
+	handle.shutdown(true);
+	handle.wait_stopped().await;
+	drop(handle);
+
+	tokio::time::timeout(Duration::from_secs(5), dropped.notified())
+		.await
+		.expect("stopped global envoy retained its callback context");
+}