@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/sdks/rust/envoy-client/src/commands.rs

@@ -0,0 +1,88 @@
+use rivet_envoy_protocol as protocol;
+
+use crate::actor::create_actor;
+use crate::connection::ws_send;
+use crate::envoy::{ActorEntry, EnvoyContext};
+
+pub const ACK_COMMANDS_INTERVAL_MS: u64 = 5 * 60 * 1000;
+
+pub async fn handle_commands(ctx: &mut EnvoyContext, commands: Vec<protocol::CommandWrapper>) {
+	tracing::info!(command_count = commands.len(), "received commands");
+
+	for command_wrapper in commands {
+		let checkpoint = command_wrapper.checkpoint;
+		match command_wrapper.inner {
+			protocol::Command::CommandStartActor(val) => {
+				let actor_name = val.config.name.clone();
+				let handle = create_actor(
+					ctx.shared.clone(),
+					checkpoint.actor_id.clone(),
+					checkpoint.generation,
+					val.config,
+					val.hibernating_requests,
+					val.preloaded_kv,
+				);
+
+				let generations = ctx
+					.actors
+					.entry(checkpoint.actor_id.clone())
+					.or_insert_with(std::collections::HashMap::new);
+				generations.insert(
+					checkpoint.generation,
+					ActorEntry {
+						handle,
+						name: actor_name,
+						event_history: Vec::new(),
+						last_command_idx: checkpoint.index,
+					},
+				);
+			}
+			protocol::Command::CommandStopActor(val) => {
+				let entry = ctx.get_actor_entry_mut(&checkpoint.actor_id, checkpoint.generation);
+
+				if let Some(entry) = entry {
+					entry.last_command_idx = checkpoint.index;
+					let _ = entry.handle.send(crate::actor::ToActor::Stop {
+						command_idx: checkpoint.index,
+						reason: val.reason,
+					});
+				} else {
+					tracing::warn!(
+						actor_id = %checkpoint.actor_id,
+						generation = checkpoint.generation,
+						"received stop actor command for unknown actor"
+					);
+				}
+			}
+		}
+	}
+}
+
+pub async fn send_command_ack(ctx: &EnvoyContext) {
+	let mut last_command_checkpoints: Vec<protocol::ActorCheckpoint> = Vec::new();
+
+	for (actor_id, generations) in &ctx.actors {
+		for (generation, entry) in generations {
+			if entry.last_command_idx < 0 {
+				continue;
+			}
+			last_command_checkpoints.push(protocol::ActorCheckpoint {
+				actor_id: actor_id.clone(),
+				generation: *generation,
+				index: entry.last_command_idx,
+			});
+		}
+	}
+
+	if last_command_checkpoints.is_empty() {
+		return;
+	}
+
+	ws_send(
+		&ctx.shared,
+		protocol::ToRivet::ToRivetAckCommands(protocol::ToRivetAckCommands {
+			last_command_checkpoints,
+		}),
+	)
+	.await;
+}

package/sdks/rust/envoy-client/src/config.rs

@@ -0,0 +1,159 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use rivet_envoy_protocol as protocol;
+use tokio::sync::mpsc;
+
+use crate::handle::EnvoyHandle;
+
+pub type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
+
+/// HTTP request/response types used by the envoy client.
+pub struct HttpRequest {
+	pub method: String,
+	pub path: String,
+	pub headers: HashMap<String, String>,
+	pub body: Option<Vec<u8>>,
+	/// If the request is streamed, body chunks arrive on this channel.
+	pub body_stream: Option<mpsc::UnboundedReceiver<Vec<u8>>>,
+}
+
+pub struct HttpResponse {
+	pub status: u16,
+	pub headers: HashMap<String, String>,
+	pub body: Option<Vec<u8>>,
+	/// If set, the response is streamed. The envoy client reads chunks and sends
+	/// `ToRivetResponseChunk` for each one.
+	pub body_stream: Option<mpsc::UnboundedReceiver<ResponseChunk>>,
+}
+
+/// A chunk in a streaming HTTP response.
+pub struct ResponseChunk {
+	pub data: Vec<u8>,
+	pub finish: bool,
+}
+
+pub struct EnvoyConfig {
+	pub version: u32,
+	pub endpoint: String,
+	pub token: Option<String>,
+	pub namespace: String,
+	pub pool_name: String,
+	pub prepopulate_actor_names: HashMap<String, PrepopulatedActor>,
+	pub metadata: Option<HashMap<String, String>>,
+	/// When `start_envoy` is called, create a new envoy every time instead of using a single global envoy
+	/// instance for the entire runtime.
+	pub not_global: bool,
+
+	/// Debug option to inject artificial latency (in ms) into WebSocket communication.
+	pub debug_latency_ms: Option<u64>,
+
+	pub callbacks: Arc<dyn EnvoyCallbacks>,
+}
+
+pub struct PrepopulatedActor {
+	pub metadata: String,
+}
+
+/// Callbacks that the consumer of the envoy client must implement.
+pub trait EnvoyCallbacks: Send + Sync + 'static {
+	fn on_actor_start(
+		&self,
+		handle: EnvoyHandle,
+		actor_id: String,
+		generation: u32,
+		config: protocol::ActorConfig,
+		preloaded_kv: Option<protocol::PreloadedKv>,
+	) -> BoxFuture<anyhow::Result<()>>;
+
+	fn on_actor_stop(
+		&self,
+		handle: EnvoyHandle,
+		actor_id: String,
+		generation: u32,
+		reason: protocol::StopActorReason,
+	) -> BoxFuture<anyhow::Result<()>>;
+
+	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>>;
+
+	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>>;
+
+	fn can_hibernate(
+		&self,
+		actor_id: &str,
+		gateway_id: &protocol::GatewayId,
+		request_id: &protocol::RequestId,
+		request: &HttpRequest,
+	) -> bool;
+}
+
+/// Handler returned by the websocket callback for receiving WebSocket events.
+pub struct WebSocketHandler {
+	pub on_message: Box<dyn Fn(WebSocketMessage) -> BoxFuture<()> + Send + Sync>,
+	pub on_close: Box<dyn Fn(u16, String) -> BoxFuture<()> + Send + Sync>,
+	pub on_open: Option<Box<dyn FnOnce(WebSocketSender) -> BoxFuture<()> + Send>>,
+}
+
+pub struct WebSocketMessage {
+	pub data: Vec<u8>,
+	pub binary: bool,
+	pub gateway_id: protocol::GatewayId,
+	pub request_id: protocol::RequestId,
+	pub message_index: u16,
+	/// Send data back on this WebSocket connection.
+	pub sender: WebSocketSender,
+}
+
+/// Allows sending messages back on a WebSocket connection from within the on_message callback.
+#[derive(Clone)]
+pub struct WebSocketSender {
+	pub(crate) tx: tokio::sync::mpsc::UnboundedSender<WsOutgoing>,
+}
+
+pub(crate) enum WsOutgoing {
+	Message {
+		data: Vec<u8>,
+		binary: bool,
+	},
+	Close {
+		code: Option<u16>,
+		reason: Option<String>,
+	},
+}
+
+impl WebSocketSender {
+	pub fn send(&self, data: Vec<u8>, binary: bool) {
+		let _ = self.tx.send(WsOutgoing::Message { data, binary });
+	}
+
+	pub fn send_text(&self, text: &str) {
+		self.send(text.as_bytes().to_vec(), false);
+	}
+
+	pub fn close(&self, code: Option<u16>, reason: Option<String>) {
+		let _ = self.tx.send(WsOutgoing::Close { code, reason });
+	}
+}

package/sdks/rust/envoy-client/src/connection.rs

@@ -0,0 +1,288 @@
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+
+use futures_util::{SinkExt, StreamExt};
+use rivet_envoy_protocol as protocol;
+use rivet_util_serde::HashableMap;
+use tokio::sync::mpsc;
+use tokio_tungstenite::tungstenite;
+use vbare::OwnedVersionedData;
+
+use crate::context::{SharedContext, WsTxMessage};
+use crate::envoy::ToEnvoyMessage;
+use crate::stringify::{stringify_to_envoy, stringify_to_rivet};
+use crate::utils::{BackoffOptions, calculate_backoff, parse_ws_close_reason};
+
+const STABLE_CONNECTION_MS: u64 = 60_000;
+
+pub fn start_connection(shared: Arc<SharedContext>) {
+	tokio::spawn(connection_loop(shared));
+}
+
+async fn connection_loop(shared: Arc<SharedContext>) {
+	let mut attempt = 0u32;
+
+	loop {
+		if shared.shutting_down.load(Ordering::Acquire) {
+			tracing::debug!("stopping reconnect loop because envoy is shutting down");
+			return;
+		}
+
+		let connected_at = std::time::Instant::now();
+
+		match single_connection(&shared).await {
+			Ok(close_reason) => {
+				if let Some(reason) = &close_reason {
+					if reason.group == "ws" && reason.error == "eviction" {
+						tracing::debug!("connection evicted");
+						let _ = shared
+							.envoy_tx
+							.send(ToEnvoyMessage::ConnClose { evict: true });
+						return;
+					}
+				}
+				let _ = shared
+					.envoy_tx
+					.send(ToEnvoyMessage::ConnClose { evict: false });
+			}
+			Err(error) => {
+				tracing::error!(?error, "connection failed");
+				let _ = shared
+					.envoy_tx
+					.send(ToEnvoyMessage::ConnClose { evict: false });
+			}
+		}
+
+		if connected_at.elapsed().as_millis() >= STABLE_CONNECTION_MS as u128 {
+			attempt = 0;
+		}
+
+		if shared.shutting_down.load(Ordering::Acquire) {
+			tracing::debug!("skipping reconnect because envoy is shutting down");
+			return;
+		}
+
+		let delay = calculate_backoff(attempt, &BackoffOptions::default());
+		tracing::info!(attempt, delay_ms = delay.as_millis() as u64, "reconnecting");
+		tokio::time::sleep(delay).await;
+		attempt += 1;
+	}
+}
+
+async fn single_connection(
+	shared: &Arc<SharedContext>,
+) -> anyhow::Result<Option<crate::utils::ParsedCloseReason>> {
+	let url = ws_url(shared);
+	let protocols = {
+		let mut p = vec!["rivet".to_string()];
+		if let Some(token) = &shared.config.token {
+			p.push(format!("rivet_token.{token}"));
+		}
+		p
+	};
+
+	// Initialize with a default CryptoProvider for rustls
+	let provider = rustls::crypto::ring::default_provider();
+	if provider.install_default().is_err() {
+		tracing::debug!("crypto provider already installed in this process");
+	}
+
+	let request = tungstenite::http::Request::builder()
+		.uri(&url)
+		.header("Sec-WebSocket-Protocol", protocols.join(", "))
+		.header("Connection", "Upgrade")
+		.header("Upgrade", "websocket")
+		.header(
+			"Sec-WebSocket-Key",
+			tungstenite::handshake::client::generate_key(),
+		)
+		.header("Sec-WebSocket-Version", "13")
+		.header("Host", extract_host(&url))
+		.body(())
+		.map_err(|e| anyhow::anyhow!("failed to build ws request: {e}"))?;
+
+	let (ws_stream, _) = tokio_tungstenite::connect_async(request).await?;
+	let (mut write, mut read) = ws_stream.split();
+
+	let (ws_tx, mut ws_rx) = mpsc::unbounded_channel::<WsTxMessage>();
+	{
+		let mut guard = shared.ws_tx.lock().await;
+		*guard = Some(ws_tx);
+	}
+
+	tracing::info!(
+		endpoint = %shared.config.endpoint,
+		namespace = %shared.config.namespace,
+		envoy_key = %shared.envoy_key,
+		has_token = shared.config.token.is_some(),
+		"websocket connected"
+	);
+
+	// Build prepopulate actor names map
+	let mut prepopulate_map = HashableMap::new();
+	for (name, actor) in &shared.config.prepopulate_actor_names {
+		prepopulate_map.insert(
+			name.clone(),
+			protocol::ActorName {
+				metadata: actor.metadata.clone(),
+			},
+		);
+	}
+
+	// Serialize metadata HashMap to JSON string for the protocol
+	let metadata_json = shared
+		.config
+		.metadata
+		.as_ref()
+		.map(|m| serde_json::to_string(m).unwrap_or_else(|_| "{}".to_string()));
+
+	// Send init
+	ws_send(
+		shared,
+		protocol::ToRivet::ToRivetInit(protocol::ToRivetInit {
+			envoy_key: shared.envoy_key.clone(),
+			version: shared.config.version,
+			prepopulate_actor_names: Some(prepopulate_map),
+			metadata: metadata_json,
+		}),
+	)
+	.await;
+
+	// Spawn write task
+	let write_handle = tokio::spawn(async move {
+		while let Some(msg) = ws_rx.recv().await {
+			match msg {
+				WsTxMessage::Send(data) => {
+					if let Err(e) = write.send(tungstenite::Message::Binary(data.into())).await {
+						tracing::error!(?e, "failed to send ws message");
+						break;
+					}
+				}
+				WsTxMessage::Close => {
+					let _ = write
+						.send(tungstenite::Message::Close(Some(
+							tungstenite::protocol::CloseFrame {
+								code: tungstenite::protocol::frame::coding::CloseCode::Normal,
+								reason: "envoy.shutdown".into(),
+							},
+						)))
+						.await;
+					break;
+				}
+			}
+		}
+	});
+
+	let mut result = None;
+
+	let debug_latency_ms = shared.config.debug_latency_ms;
+
+	while let Some(msg) = read.next().await {
+		match msg {
+			Ok(tungstenite::Message::Binary(data)) => {
+				crate::utils::inject_latency(debug_latency_ms).await;
+
+				let decoded = crate::protocol::versioned::ToEnvoy::deserialize(
+					&data,
+					protocol::PROTOCOL_VERSION,
+				)?;
+
+				if tracing::enabled!(tracing::Level::DEBUG) {
+					tracing::debug!(data = stringify_to_envoy(&decoded), "received message");
+				}
+
+				forward_to_envoy(shared, decoded).await;
+			}
+			Ok(tungstenite::Message::Close(frame)) => {
+				if let Some(frame) = frame {
+					let reason_str = frame.reason.to_string();
+					let code: u16 = frame.code.into();
+					tracing::info!(
+						code,
+						reason = %reason_str,
+						"websocket closed"
+					);
+					result = parse_ws_close_reason(&reason_str);
+				}
+				break;
+			}
+			Err(e) => {
+				tracing::error!(?e, "websocket error");
+				break;
+			}
+			_ => {}
+		}
+	}
+
+	// Clean up
+	{
+		let mut guard = shared.ws_tx.lock().await;
+		*guard = None;
+	}
+	write_handle.abort();
+
+	Ok(result)
+}
+
+async fn forward_to_envoy(shared: &SharedContext, message: protocol::ToEnvoy) {
+	match message {
+		protocol::ToEnvoy::ToEnvoyPing(ping) => {
+			ws_send(
+				shared,
+				protocol::ToRivet::ToRivetPong(protocol::ToRivetPong { ts: ping.ts }),
+			)
+			.await;
+		}
+		other => {
+			let _ = shared
+				.envoy_tx
+				.send(ToEnvoyMessage::ConnMessage { message: other });
+		}
+	}
+}
+
+/// Send a message over the WebSocket. Returns true if the message could not be sent.
+pub async fn ws_send(shared: &SharedContext, message: protocol::ToRivet) -> bool {
+	if tracing::enabled!(tracing::Level::DEBUG) {
+		tracing::debug!(data = stringify_to_rivet(&message), "sending message");
+	}
+
+	let guard = shared.ws_tx.lock().await;
+	let Some(tx) = guard.as_ref() else {
+		tracing::error!("websocket not available for sending");
+		return true;
+	};
+
+	let encoded = crate::protocol::versioned::ToRivet::wrap_latest(message)
+		.serialize(protocol::PROTOCOL_VERSION)
+		.expect("failed to encode message");
+	let _ = tx.send(WsTxMessage::Send(encoded));
+	false
+}
+
+fn ws_url(shared: &SharedContext) -> String {
+	let ws_endpoint = shared
+		.config
+		.endpoint
+		.replace("http://", "ws://")
+		.replace("https://", "wss://");
+	let base_url = ws_endpoint.trim_end_matches('/');
+
+	format!(
+		"{}/envoys/connect?protocol_version={}&namespace={}&envoy_key={}&pool_name={}",
+		base_url,
+		protocol::PROTOCOL_VERSION,
+		urlencoding::encode(&shared.config.namespace),
+		urlencoding::encode(&shared.envoy_key),
+		urlencoding::encode(&shared.config.pool_name),
+	)
+}
+
+fn extract_host(url: &str) -> String {
+	url.replace("ws://", "")
+		.replace("wss://", "")
+		.split('/')
+		.next()
+		.unwrap_or("localhost")
+		.to_string()
+}

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

@@ -0,0 +1,24 @@
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+use rivet_envoy_protocol as protocol;
+use tokio::sync::Mutex;
+use tokio::sync::mpsc;
+
+use crate::config::EnvoyConfig;
+use crate::envoy::ToEnvoyMessage;
+
+pub struct SharedContext {
+	pub config: EnvoyConfig,
+	pub envoy_key: String,
+	pub envoy_tx: mpsc::UnboundedSender<ToEnvoyMessage>,
+	pub ws_tx: Arc<Mutex<Option<mpsc::UnboundedSender<WsTxMessage>>>>,
+	pub protocol_metadata: Arc<Mutex<Option<protocol::ProtocolMetadata>>>,
+	pub shutting_down: AtomicBool,
+}
+
+#[derive(Debug)]
+pub enum WsTxMessage {
+	Send(Vec<u8>),
+	Close,
+}