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

package/packages/pegboard-envoy/src/metrics.rs

@@ -0,0 +1,44 @@
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
+
+lazy_static::lazy_static! {
+	pub static ref CONNECTION_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_envoy_connection_total",
+		"Count of envoy connections opened.",
+		&["namespace_id", "pool_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref EVICTION_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_envoy_eviction_total",
+		"Count of envoy connections evicted.",
+		&["namespace_id", "pool_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref CONNECTION_ACTIVE: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_envoy_connection_active",
+		"Count of envoy connections currently active.",
+		&["namespace_id", "pool_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref RECEIVE_INIT_PACKET_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"pegboard_envoy_receive_init_packet_duration",
+		"Duration to receive the init packet for a envoy connection.",
+		&["namespace_id", "pool_name"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	pub static ref EVENT_MULTIPLEXER_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"pegboard_envoy_event_multiplexer_count",
+		"Number of active actor event multiplexers.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref INGESTED_EVENTS_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"pegboard_envoy_ingested_events_total",
+		"Count of actor events.",
+		*REGISTRY
+	).unwrap();
+}

package/packages/pegboard-envoy/src/ping_task.rs

@@ -0,0 +1,61 @@
+use anyhow::ensure;
+use gas::prelude::*;
+use hyper_tungstenite::tungstenite::Message;
+use rand::Rng;
+use rivet_envoy_protocol::{self as protocol, versioned};
+use std::sync::{Arc, atomic::Ordering};
+use std::time::Duration;
+use tokio::sync::watch;
+use vbare::OwnedVersionedData;
+
+use crate::{LifecycleResult, conn::Conn};
+
+#[tracing::instrument(name="ping_task", skip_all, fields(ray_id=?ctx.ray_id(), req_id=?ctx.req_id(), envoy_key=%conn.envoy_key, protocol_version=%conn.protocol_version))]
+pub async fn task(
+	ctx: StandaloneCtx,
+	conn: Arc<Conn>,
+	mut ping_abort_rx: watch::Receiver<()>,
+) -> Result<LifecycleResult> {
+	let update_ping_interval =
+		Duration::from_millis(ctx.config().pegboard().envoy_update_ping_interval());
+	let ping_timeout_ms = ctx.config().pegboard().envoy_ping_timeout();
+
+	loop {
+		// Jitter sleep to prevent stampeding herds
+		let jitter = Duration::from_millis(rand::thread_rng().gen_range(0..128));
+		tokio::select! {
+			_ = tokio::time::sleep(update_ping_interval + jitter) => {}
+			_ = ping_abort_rx.changed() => {
+				return Ok(LifecycleResult::Aborted);
+			}
+		}
+
+		// Check if the last ping is past the timeout threshold
+		let last_ping_ts = conn.last_ping_ts.load(Ordering::SeqCst);
+		let now = util::timestamp::now();
+		ensure!(
+			now - last_ping_ts <= ping_timeout_ms,
+			"envoy ws ping timed out"
+		);
+
+		// Update ping
+		ctx.op(pegboard::ops::envoy::update_ping::Input {
+			namespace_id: conn.namespace_id,
+			envoy_key: conn.envoy_key.clone(),
+			update_lb: !conn.is_serverless,
+			rtt: conn.last_rtt.load(Ordering::Relaxed),
+		})
+		.await?;
+
+		// Send ping to envoy
+		let ping_msg = versioned::ToEnvoy::wrap_latest(protocol::ToEnvoy::ToEnvoyPing(
+			protocol::ToEnvoyPing {
+				ts: util::timestamp::now(),
+			},
+		));
+		let ping_msg_serialized = ping_msg.serialize(conn.protocol_version)?;
+		conn.ws_handle
+			.send(Message::Binary(ping_msg_serialized.into()))
+			.await?;
+	}
+}

package/packages/pegboard-envoy/src/tunnel_to_ws_task.rs

@@ -0,0 +1,183 @@
+use anyhow::Result;
+use gas::prelude::*;
+use hyper_tungstenite::tungstenite::Message;
+use pegboard::pubsub_subjects::GatewayReceiverSubject;
+use rivet_envoy_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
+use std::sync::Arc;
+use tokio::sync::watch;
+use universalpubsub as ups;
+use universalpubsub::{NextOutput, PublishOpts, Subscriber};
+use vbare::OwnedVersionedData;
+
+use crate::{LifecycleResult, conn::Conn, metrics};
+
+#[tracing::instrument(name="tunnel_to_ws_task", skip_all, fields(ray_id=?ctx.ray_id(), req_id=?ctx.req_id(), envoy_key=%conn.envoy_key, protocol_version=%conn.protocol_version))]
+pub async fn task(
+	ctx: StandaloneCtx,
+	conn: Arc<Conn>,
+	mut tunnel_sub: Subscriber,
+	mut eviction_sub: Subscriber,
+	mut tunnel_to_ws_abort_rx: watch::Receiver<()>,
+) -> Result<LifecycleResult> {
+	loop {
+		match recv_msg(
+			&conn,
+			&mut tunnel_sub,
+			&mut eviction_sub,
+			&mut tunnel_to_ws_abort_rx,
+		)
+		.await?
+		{
+			Ok(msg) => {
+				let evicted = handle_message(&ctx, &conn, msg).await?;
+				if evicted {
+					return Ok(LifecycleResult::Evicted);
+				}
+			}
+			Err(lifecycle_res) => return Ok(lifecycle_res),
+		}
+	}
+}
+
+async fn recv_msg(
+	conn: &Conn,
+	tunnel_sub: &mut Subscriber,
+	eviction_sub: &mut Subscriber,
+	tunnel_to_ws_abort_rx: &mut watch::Receiver<()>,
+) -> Result<std::result::Result<ups::Message, LifecycleResult>> {
+	let tunnel_msg = tokio::select! {
+		res = tunnel_sub.next() => {
+			if let NextOutput::Message(tunnel_msg) = res? {
+				tunnel_msg
+			} else {
+				tracing::debug!("tunnel sub closed");
+				bail!("tunnel sub closed");
+			}
+		}
+		_ = eviction_sub.next() => {
+			tracing::debug!("envoy evicted");
+
+			metrics::EVICTION_TOTAL
+				.with_label_values(&[
+					conn.namespace_id.to_string().as_str(),
+					&conn.pool_name,
+					conn.protocol_version.to_string().as_str(),
+				])
+				.inc();
+
+			return Ok(Err(LifecycleResult::Evicted));
+		}
+		_ = tunnel_to_ws_abort_rx.changed() => {
+			tracing::debug!("task aborted");
+			return Ok(Err(LifecycleResult::Aborted));
+		}
+	};
+
+	tracing::debug!(
+		payload_len = tunnel_msg.payload.len(),
+		"received message from pubsub, forwarding to WebSocket"
+	);
+
+	Ok(Ok(tunnel_msg))
+}
+
+async fn handle_message(
+	ctx: &StandaloneCtx,
+	conn: &Conn,
+	tunnel_msg: ups::Message,
+) -> Result<bool> {
+	// Parse message
+	let msg = match versioned::ToEnvoyConn::deserialize_with_embedded_version(&tunnel_msg.payload) {
+		Result::Ok(x) => x,
+		Err(err) => {
+			tracing::error!(?err, "failed to parse tunnel message");
+			return Ok(false);
+		}
+	};
+
+	// Convert to ToEnvoy types
+	let to_client_msg = match msg {
+		protocol::ToEnvoyConn::ToEnvoyConnPing(ping) => {
+			// Publish pong to UPS
+			let gateway_reply_to = GatewayReceiverSubject::new(ping.gateway_id).to_string();
+			let msg_serialized = versioned::ToGateway::wrap_latest(
+				protocol::ToGateway::ToGatewayPong(protocol::ToGatewayPong {
+					request_id: ping.request_id,
+					ts: ping.ts,
+				}),
+			)
+			.serialize_with_embedded_version(PROTOCOL_VERSION)
+			.context("failed to serialize pong message for gateway")?;
+			ctx.ups()
+				.context("failed to get UPS instance for tunnel message")?
+				.publish(&gateway_reply_to, &msg_serialized, PublishOpts::one())
+				.await
+				.with_context(|| {
+					format!(
+						"failed to publish tunnel message to gateway reply topic: {}",
+						gateway_reply_to
+					)
+				})?;
+
+			// Not sent to envoy
+			return Ok(false);
+		}
+		protocol::ToEnvoyConn::ToEnvoyConnClose => return Ok(true),
+		protocol::ToEnvoyConn::ToEnvoyCommands(mut command_wrappers) => {
+			// TODO: Parallelize
+			for command_wrapper in &mut command_wrappers {
+				if let protocol::Command::CommandStartActor(start) = &mut command_wrapper.inner {
+					let actor_id = Id::parse(&command_wrapper.checkpoint.actor_id)?;
+					let actor_name = start.config.name.clone();
+					let ids = ctx
+						.op(pegboard::ops::actor::hibernating_request::list::Input { actor_id })
+						.await?;
+
+					// Dynamically populate hibernating request ids
+					start.hibernating_requests = ids
+						.into_iter()
+						.map(|x| protocol::HibernatingRequest {
+							gateway_id: x.gateway_id,
+							request_id: x.request_id,
+						})
+						.collect();
+
+					if start.preloaded_kv.is_none() {
+						let db = ctx.udb()?;
+						start.preloaded_kv = pegboard::actor_kv::preload::fetch_preloaded_kv(
+							&db,
+							ctx.config().pegboard(),
+							actor_id,
+							conn.namespace_id,
+							&actor_name,
+						)
+						.await?;
+					}
+				}
+			}
+
+			// NOTE: `command_wrappers` is mutated in this match arm, it is not the same as the
+			// ToEnvoyConn data
+			protocol::ToEnvoy::ToEnvoyCommands(command_wrappers)
+		}
+		protocol::ToEnvoyConn::ToEnvoyAckEvents(x) => protocol::ToEnvoy::ToEnvoyAckEvents(x),
+		protocol::ToEnvoyConn::ToEnvoyTunnelMessage(x) => {
+			let _ = conn
+				.authorized_tunnel_routes
+				.insert_async((x.message_id.gateway_id, x.message_id.request_id), ())
+				.await;
+			protocol::ToEnvoy::ToEnvoyTunnelMessage(x)
+		}
+	};
+
+	// Forward raw message to WebSocket
+	let serialized_msg =
+		versioned::ToEnvoy::wrap_latest(to_client_msg).serialize(conn.protocol_version)?;
+	let ws_msg = Message::Binary(serialized_msg.into());
+	conn.ws_handle
+		.send(ws_msg)
+		.await
+		.context("failed to send message to WebSocket")?;
+
+	Ok(false)
+}

package/packages/pegboard-envoy/src/utils.rs

@@ -0,0 +1,68 @@
+use gas::prelude::*;
+
+use crate::errors::WsError;
+
+#[derive(Clone)]
+pub struct UrlData {
+	pub protocol_version: u16,
+	pub namespace: String,
+	pub pool_name: String,
+	pub envoy_key: String,
+}
+
+impl UrlData {
+	pub fn parse_url(url: url::Url) -> Result<UrlData> {
+		// Read protocol version from query parameters
+		let protocol_version = url
+			.query_pairs()
+			.find_map(|(n, v)| (n == "protocol_version").then_some(v))
+			.context("missing `protocol_version` query parameter")?
+			.parse::<u16>()
+			.context("invalid `protocol_version` query parameter")?;
+
+		// Read namespace from query parameters
+		let namespace = url
+			.query_pairs()
+			.find_map(|(n, v)| (n == "namespace").then_some(v))
+			.context("missing `namespace` query parameter")?
+			.to_string();
+
+		// Read envoy pool name from query parameters
+		let pool_name = url
+			.query_pairs()
+			.find_map(|(n, v)| (n == "pool_name").then_some(v))
+			.context("missing `pool_name` query parameter")?
+			.to_string();
+
+		if pool_name.len() > 128 {
+			return Err(
+				WsError::InvalidInitialPacket("`pool_name` parameter too long (> 128)").build(),
+			);
+		}
+
+		// Read envoy key from query parameters
+		let envoy_key = url
+			.query_pairs()
+			.find_map(|(n, v)| (n == "envoy_key").then_some(v))
+			.context("missing `envoy_key` query parameter")?
+			.to_string();
+
+		if envoy_key.len() < 3 {
+			return Err(
+				WsError::InvalidInitialPacket("`envoy_key` parameter too short (< 3)").build(),
+			);
+		}
+		if envoy_key.len() > 128 {
+			return Err(
+				WsError::InvalidInitialPacket("`envoy_key` parameter too long (> 128)").build(),
+			);
+		}
+
+		Ok(UrlData {
+			protocol_version,
+			namespace,
+			pool_name,
+			envoy_key,
+		})
+	}
+}