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

package/packages/pegboard-runner/src/lib.rs

@@ -6,28 +6,27 @@ use http_body_util::Full;
 use hyper::{Response, StatusCode};
 use pegboard::ops::runner::update_alloc_idx::Action;
 use rivet_guard_core::{
-	WebSocketHandle, custom_serve::CustomServeTrait, proxy_service::ResponseBody,
-	request_context::RequestContext,
+	ResponseBody, WebSocketHandle, custom_serve::CustomServeTrait, request_context::RequestContext,
 };
-use rivet_runner_protocol as protocol;
 use std::time::Duration;
 use tokio::sync::watch;
 use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
 use universalpubsub::PublishOpts;
 
+mod actor_event_demuxer;
 mod conn;
 mod errors;
+mod metrics;
 mod ping_task;
 mod tunnel_to_ws_task;
 mod utils;
 mod ws_to_tunnel_task;
 
-const UPDATE_PING_INTERVAL: Duration = Duration::from_secs(3);
-
 #[derive(Debug)]
 enum LifecycleResult {
 	Closed,
 	Aborted,
+	Evicted,
 }
 
 pub struct PegboardRunnerWsCustomServe {
@@ -44,11 +43,11 @@ impl PegboardRunnerWsCustomServe {
 
 #[async_trait]
 impl CustomServeTrait for PegboardRunnerWsCustomServe {
+	#[tracing::instrument(skip_all)]
 	async fn handle_request(
 		&self,
 		_req: hyper::Request<http_body_util::Full<hyper::body::Bytes>>,
-		_request_context: &mut RequestContext,
-		_request_id: protocol::RequestId,
+		_req_ctx: &mut RequestContext,
 	) -> Result<Response<ResponseBody>> {
 		// Pegboard runner ws doesn't handle regular HTTP requests
 		// Return a simple status response
@@ -62,28 +61,28 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 		Ok(response)
 	}
 
+	#[tracing::instrument(skip_all)]
 	async fn handle_websocket(
 		&self,
+		req_ctx: &mut RequestContext,
 		ws_handle: WebSocketHandle,
-		_headers: &hyper::HeaderMap,
-		path: &str,
-		_request_context: &mut RequestContext,
-		_unique_request_id: protocol::RequestId,
 		_after_hibernation: bool,
 	) -> Result<Option<CloseFrame>> {
+		let ctx = self.ctx.with_ray(req_ctx.ray_id(), req_ctx.req_id())?;
+
 		// Get UPS
-		let ups = self.ctx.ups().context("failed to get UPS instance")?;
+		let ups = ctx.ups().context("failed to get UPS instance")?;
 
 		// Parse URL to extract parameters
-		let url = url::Url::parse(&format!("ws://placeholder/{path}"))
+		let url = url::Url::parse(&format!("ws://placeholder/{}", req_ctx.path()))
 			.context("failed to parse WebSocket URL")?;
 		let url_data = utils::UrlData::parse_url(url)
 			.map_err(|err| errors::WsError::InvalidUrl(err.to_string()).build())?;
 
-		tracing::debug!(?path, "tunnel ws connection established");
+		tracing::debug!(path=%req_ctx.path(), "tunnel ws connection established");
 
 		// Create connection
-		let conn = conn::init_conn(&self.ctx, ws_handle.clone(), url_data)
+		let conn = conn::init_conn(&ctx, ws_handle.clone(), url_data)
 			.await
 			.context("failed to initialize runner connection")?;
 
@@ -133,12 +132,20 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 			},
 		)?;
 
+		metrics::CONNECTION_ACTIVE
+			.with_label_values(&[
+				conn.namespace_id.to_string().as_str(),
+				&conn.runner_name,
+				conn.protocol_version.to_string().as_str(),
+			])
+			.inc();
+
 		let (tunnel_to_ws_abort_tx, tunnel_to_ws_abort_rx) = watch::channel(());
 		let (ws_to_tunnel_abort_tx, ws_to_tunnel_abort_rx) = watch::channel(());
 		let (ping_abort_tx, ping_abort_rx) = watch::channel(());
 
 		let tunnel_to_ws = tokio::spawn(tunnel_to_ws_task::task(
-			self.ctx.clone(),
+			ctx.clone(),
 			conn.clone(),
 			sub,
 			eviction_sub,
@@ -146,7 +153,7 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 		));
 
 		let ws_to_tunnel = tokio::spawn(ws_to_tunnel_task::task(
-			self.ctx.clone(),
+			ctx.clone(),
 			conn.clone(),
 			ws_handle.recv(),
 			eviction_sub2,
@@ -154,10 +161,13 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 		));
 
 		// Update pings
+		let update_ping_interval =
+			Duration::from_millis(ctx.config().pegboard().runner_update_ping_interval_ms());
 		let ping = tokio::spawn(ping_task::task(
-			self.ctx.clone(),
+			ctx.clone(),
 			conn.clone(),
 			ping_abort_rx,
+			update_ping_interval,
 		));
 		let tunnel_to_ws_abort_tx2 = tunnel_to_ws_abort_tx.clone();
 		let ws_to_tunnel_abort_tx2 = ws_to_tunnel_abort_tx.clone();
@@ -213,35 +223,51 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 		);
 
 		// Determine single result from all tasks
-		let lifecycle_res = match (tunnel_to_ws_res, ws_to_tunnel_res, ping_res) {
+		let mut lifecycle_res = match (tunnel_to_ws_res, ws_to_tunnel_res, ping_res) {
 			// Prefer error
 			(Err(err), _, _) => Err(err),
 			(_, Err(err), _) => Err(err),
 			(_, _, Err(err)) => Err(err),
-			// Prefer non aborted result if both succeed
+			// Prefer non aborted result
 			(Ok(res), Ok(LifecycleResult::Aborted), _) => Ok(res),
 			(Ok(LifecycleResult::Aborted), Ok(res), _) => Ok(res),
 			// Unlikely case
 			(res, _, _) => res,
 		};
 
-		// Make runner immediately ineligible when it disconnects
-		let update_alloc_res = self
-			.ctx
-			.op(pegboard::ops::runner::update_alloc_idx::Input {
-				runners: vec![pegboard::ops::runner::update_alloc_idx::Runner {
-					runner_id: conn.runner_id,
-					action: Action::ClearIdx,
-				}],
-			})
-			.await;
-		if let Err(err) = update_alloc_res {
-			tracing::error!(
-				runner_id=?conn.runner_id,
-				?err,
-				"critical: failed to evict runner from allocation index during disconnect"
-			);
+		if let Ok(LifecycleResult::Evicted) = &lifecycle_res {
+			lifecycle_res = Err(errors::WsError::Eviction.build());
 		}
+		// Clear alloc idx if not evicted
+		else {
+			// Make runner immediately ineligible when it disconnects
+			let update_alloc_res = self
+				.ctx
+				.op(pegboard::ops::runner::update_alloc_idx::Input {
+					runners: vec![pegboard::ops::runner::update_alloc_idx::Runner {
+						runner_id: conn.runner_id,
+						action: Action::ClearIdx,
+					}],
+				})
+				.await;
+			if let Err(err) = update_alloc_res {
+				tracing::error!(
+					runner_id=?conn.runner_id,
+					?err,
+					"failed to evict runner from allocation index during disconnect"
+				);
+			}
+		}
+
+		tracing::debug!(%topic, "runner websocket closed");
+
+		metrics::CONNECTION_ACTIVE
+			.with_label_values(&[
+				conn.namespace_id.to_string().as_str(),
+				&conn.runner_name,
+				conn.protocol_version.to_string().as_str(),
+			])
+			.dec();
 
 		// This will determine the close frame sent back to the runner websocket
 		lifecycle_res.map(|_| None)

package/packages/pegboard-runner/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_runner_connection_total",
+		"Count of runner connections opened.",
+		&["namespace_id", "runner_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref EVICTION_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_runner_eviction_total",
+		"Count of runner connections evicted.",
+		&["namespace_id", "runner_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref CONNECTION_ACTIVE: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_runner_connection_active",
+		"Count of runner connections currently active.",
+		&["namespace_id", "runner_name", "protocol_version"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref RECEIVE_INIT_PACKET_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"pegboard_runner_receive_init_packet_duration",
+		"Duration to receive the init packet for a runner connection.",
+		&["namespace_id", "runner_name"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	pub static ref EVENT_MULTIPLEXER_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"pegboard_runner_event_multiplexer_count",
+		"Number of active actor event multiplexers.",
+		*REGISTRY
+	).unwrap();
+
+	pub static ref INGESTED_EVENTS_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"pegboard_runner_ingested_events_total",
+		"Count of actor events.",
+		*REGISTRY
+	).unwrap();
+}

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

@@ -1,35 +1,75 @@
+use anyhow::ensure;
 use gas::prelude::*;
+use hyper_tungstenite::tungstenite::Message;
 use pegboard::ops::runner::update_alloc_idx::{Action, RunnerEligibility};
+use rand::Rng;
+use rivet_runner_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, UPDATE_PING_INTERVAL, conn::Conn};
+use crate::{LifecycleResult, conn::Conn};
 
 /// Updates the ping of all runners requesting a ping update at once.
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(name="ping_task", skip_all, fields(ray_id=?ctx.ray_id(), req_id=?ctx.req_id(), runner_id=?conn.runner_id, workflow_id=?conn.workflow_id, protocol_version=%conn.protocol_version))]
 pub async fn task(
 	ctx: StandaloneCtx,
 	conn: Arc<Conn>,
 	mut ping_abort_rx: watch::Receiver<()>,
+	update_ping_interval: Duration,
 ) -> Result<LifecycleResult> {
+	let ping_timeout_ms = ctx.config().pegboard().runner_ping_timeout_ms();
+
 	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) => {}
+			_ = tokio::time::sleep(update_ping_interval + jitter) => {}
 			_ = ping_abort_rx.changed() => {
 				return Ok(LifecycleResult::Aborted);
 			}
 		}
 
+		// Check if the last pong is past the timeout threshold (mk2 only)
+		if protocol::is_mk2(conn.protocol_version) {
+			let last_ping_ts = conn.last_ping_ts.load(Ordering::Relaxed);
+			let now = util::timestamp::now();
+			ensure!(
+				now - last_ping_ts <= ping_timeout_ms,
+				"runner ws ping timed out"
+			);
+		}
+
 		update_runner_ping(&ctx, &conn).await?;
+
+		// Send ping to runner
+		if protocol::is_mk2(conn.protocol_version) {
+			let ping_msg = versioned::ToClientMk2::wrap_latest(
+				protocol::mk2::ToClient::ToClientPing(protocol::mk2::ToClientPing {
+					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?;
+		}
 	}
 }
 
 async fn update_runner_ping(ctx: &StandaloneCtx, conn: &Conn) -> Result<()> {
-	let Some(wf) = ctx
-		.workflow::<pegboard::workflows::runner::Input>(conn.workflow_id)
-		.get()
-		.await?
-	else {
+	let wf = if protocol::is_mk2(conn.protocol_version) {
+		ctx.workflow::<pegboard::workflows::runner2::Input>(conn.workflow_id)
+			.get()
+			.await?
+	} else {
+		ctx.workflow::<pegboard::workflows::runner::Input>(conn.workflow_id)
+			.get()
+			.await?
+	};
+
+	let Some(wf) = wf else {
 		tracing::error!(?conn.runner_id, "workflow does not exist");
 		return Ok(());
 	};
@@ -50,15 +90,22 @@ async fn update_runner_ping(ctx: &StandaloneCtx, conn: &Conn) -> Result<()> {
 		})
 		.await?;
 
-	// If runner became eligible again, then pull any pending actors
+	// If runner became eligible again, have the workflow check pending actor queue
 	for notif in res.notifications {
 		if let RunnerEligibility::ReEligible = notif.eligibility {
 			tracing::debug!(runner_id=?notif.runner_id, "runner has become eligible again");
 
-			ctx.signal(pegboard::workflows::runner::CheckQueue {})
-				.to_workflow_id(notif.workflow_id)
-				.send()
-				.await?;
+			if protocol::is_mk2(conn.protocol_version) {
+				ctx.signal(pegboard::workflows::runner2::CheckQueue {})
+					.to_workflow_id(notif.workflow_id)
+					.send()
+					.await?;
+			} else {
+				ctx.signal(pegboard::workflows::runner::CheckQueue {})
+					.to_workflow_id(notif.workflow_id)
+					.send()
+					.await?;
+			}
 		}
 	}