feat(rivetkit): expose runtime diagnostics

Author: NathanFlurry

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

@@ -1652,6 +1652,7 @@ mod tests {
 			envoy_key: "test-envoy".to_string(),
 			envoy_tx,
 			actors: Arc::new(std::sync::Mutex::new(HashMap::new())),
+			actors_notify: Arc::new(tokio::sync::Notify::new()),
 			live_tunnel_requests: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			pending_hibernation_restores: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			ws_tx: Arc::new(tokio::sync::Mutex::new(

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

@@ -5,8 +5,9 @@ use std::sync::atomic::AtomicBool;
 
 use crate::async_counter::AsyncCounter;
 use rivet_envoy_protocol as protocol;
-use tokio::sync::Mutex;
 use tokio::sync::mpsc;
+use tokio::sync::Mutex;
+use tokio::sync::Notify;
 use tokio::sync::watch;
 
 use crate::actor::ToActor;
@@ -24,6 +25,7 @@ pub struct SharedContext {
 	pub envoy_key: String,
 	pub envoy_tx: mpsc::UnboundedSender<ToEnvoyMessage>,
 	pub actors: Arc<StdMutex<HashMap<String, HashMap<u32, SharedActorEntry>>>>,
+	pub actors_notify: Arc<Notify>,
 	pub live_tunnel_requests: Arc<StdMutex<HashMap<[u8; 8], String>>>,
 	pub pending_hibernation_restores:
 		Arc<StdMutex<HashMap<String, Vec<HibernatingWebSocketMetadata>>>>,

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

@@ -183,6 +183,8 @@ impl EnvoyContext {
 				},
 			);
 
+		self.shared.actors_notify.notify_waiters();
+
 		if let Some(messages) = self.buffered_actor_messages.remove(&buffered_actor_id) {
 			for message in messages {
 				match message {
@@ -216,6 +218,7 @@ impl EnvoyContext {
 				shared.remove(actor_id);
 			}
 		}
+		self.shared.actors_notify.notify_waiters();
 	}
 
 	pub fn get_actor(&self, actor_id: &str, generation: Option<u32>) -> Option<&ActorEntry> {
@@ -297,6 +300,7 @@ fn start_envoy_sync_inner(config: EnvoyConfig) -> EnvoyHandle {
 		envoy_key,
 		envoy_tx: envoy_tx.clone(),
 		actors: Arc::new(std::sync::Mutex::new(HashMap::new())),
+		actors_notify: Arc::new(tokio::sync::Notify::new()),
 		live_tunnel_requests: Arc::new(std::sync::Mutex::new(HashMap::new())),
 		pending_hibernation_restores: Arc::new(std::sync::Mutex::new(HashMap::new())),
 		ws_tx: Arc::new(tokio::sync::Mutex::new(None)),

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

@@ -164,6 +164,7 @@ mod tests {
 			envoy_key: "test-envoy".to_string(),
 			envoy_tx,
 			actors: Arc::new(std::sync::Mutex::new(HashMap::new())),
+			actors_notify: Arc::new(tokio::sync::Notify::new()),
 			live_tunnel_requests: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			pending_hibernation_restores: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			ws_tx: Arc::new(tokio::sync::Mutex::new(

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

@@ -17,6 +17,12 @@ pub struct EnvoyHandle {
 	pub(crate) started_rx: tokio::sync::watch::Receiver<()>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ServerlessActorStart {
+	pub actor_id: String,
+	pub generation: u32,
+}
+
 impl EnvoyHandle {
 	#[doc(hidden)]
 	pub fn from_shared(shared: Arc<SharedContext>) -> Self {
@@ -85,6 +91,23 @@ impl EnvoyHandle {
 		&self.shared.config.namespace
 	}
 
+	pub fn active_actor_count(&self) -> usize {
+		let guard = self
+			.shared
+			.actors
+			.lock()
+			.expect("shared actor registry poisoned");
+		guard
+			.values()
+			.map(|generations| {
+				generations
+					.values()
+					.filter(|actor| !actor.handle.is_closed())
+					.count()
+			})
+			.sum()
+	}
+
 	pub fn pool_name(&self) -> &str {
 		&self.shared.config.pool_name
 	}
@@ -138,6 +161,40 @@ impl EnvoyHandle {
 		rx.await.ok().flatten()
 	}
 
+	pub async fn wait_actor_registered_then_stopped(&self, actor_id: &str, generation: u32) {
+		let mut registered = false;
+		loop {
+			let notified = self.shared.actors_notify.notified();
+			if self.is_stopped() {
+				return;
+			}
+
+			let actor_is_registered = {
+				let guard = self
+					.shared
+					.actors
+					.lock()
+					.expect("shared actor registry poisoned");
+				guard
+					.get(actor_id)
+					.and_then(|generations| generations.get(&generation))
+					.is_some()
+			};
+
+			if registered && !actor_is_registered {
+				return;
+			}
+			if actor_is_registered {
+				registered = true;
+			}
+
+			tokio::select! {
+				_ = notified => {}
+				_ = self.wait_stopped() => return,
+			}
+		}
+	}
+
 	pub fn http_request_counter(
 		&self,
 		actor_id: &str,
@@ -484,6 +541,35 @@ impl EnvoyHandle {
 	/// Inject a serverless start payload into the envoy.
 	/// The payload is a u16 LE protocol version followed by a serialized ToEnvoy message.
 	pub async fn start_serverless_actor(&self, payload: &[u8]) -> anyhow::Result<()> {
+		let (message, _) = decode_serverless_actor_start_payload(payload)?;
+
+		// Wait for envoy to be started before injecting
+		self.started().await?;
+
+		tracing::debug!(
+			data = crate::stringify::stringify_to_envoy(&message),
+			"received serverless start"
+		);
+		self.shared
+			.envoy_tx
+			.send(ToEnvoyMessage::ConnMessage { message })
+			.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
+
+		Ok(())
+	}
+
+	pub fn decode_serverless_actor_start(
+		&self,
+		payload: &[u8],
+	) -> anyhow::Result<ServerlessActorStart> {
+		let (_, actor_start) = decode_serverless_actor_start_payload(payload)?;
+		Ok(actor_start)
+	}
+}
+
+fn decode_serverless_actor_start_payload(
+	payload: &[u8],
+) -> anyhow::Result<(protocol::ToEnvoy, ServerlessActorStart)> {
 		use vbare::OwnedVersionedData;
 
 		if payload.len() < 2 {
@@ -524,21 +610,15 @@ impl EnvoyHandle {
 			anyhow::bail!("invalid serverless payload: expected CommandStartActor");
 		}
 
-		// Wait for envoy to be started before injecting
-		self.started().await?;
-
-		tracing::debug!(
-			data = crate::stringify::stringify_to_envoy(&message),
-			"received serverless start"
-		);
-		self.shared
-			.envoy_tx
-			.send(ToEnvoyMessage::ConnMessage { message })
-			.map_err(|_| anyhow::anyhow!("envoy channel closed"))?;
+		let actor_start = ServerlessActorStart {
+			actor_id: commands[0].checkpoint.actor_id.clone(),
+			generation: commands[0].checkpoint.generation,
+		};
 
-		Ok(())
+		Ok((message, actor_start))
 	}
 
+impl EnvoyHandle {
 	async fn send_kv_request(
 		&self,
 		actor_id: String,

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

@@ -463,6 +463,7 @@ mod tests {
 			envoy_key: "test-envoy".to_string(),
 			envoy_tx,
 			actors: Arc::new(std::sync::Mutex::new(HashMap::new())),
+			actors_notify: Arc::new(tokio::sync::Notify::new()),
 			live_tunnel_requests: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			pending_hibernation_restores: Arc::new(std::sync::Mutex::new(HashMap::new())),
 			ws_tx: Arc::new(tokio::sync::Mutex::new(