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

package/packages/pegboard/src/ops/serverless_metadata/fetch.rs

@@ -0,0 +1,223 @@
+use std::collections::HashMap;
+use std::time::Duration;
+
+use anyhow::Result;
+use gas::prelude::*;
+use reqwest::header::{HeaderMap as ReqwestHeaderMap, HeaderName, HeaderValue};
+use rivet_envoy_protocol::PROTOCOL_VERSION;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+const RESPONSE_BODY_MAX_LEN: usize = 1024;
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
+
+#[derive(Debug)]
+pub struct Input {
+	pub url: String,
+	pub headers: HashMap<String, String>,
+}
+
+#[derive(Deserialize, Serialize, ToSchema, Clone, Debug, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+#[schema(as = RunnerConfigsServerlessMetadataError)]
+pub enum ServerlessMetadataError {
+	InvalidRequest {},
+	RequestFailed {},
+	RequestTimedOut {},
+	NonSuccessStatus { status_code: u16, body: String },
+	InvalidResponseJson { body: String },
+	InvalidResponseSchema { runtime: String, version: String },
+}
+
+#[derive(Debug, Clone)]
+pub struct Output {
+	pub runtime: String,
+	pub version: String,
+	pub envoy_protocol_version: Option<u16>,
+	pub actor_names: Vec<ActorNameMetadata>,
+	pub runner_version: Option<u32>,
+	pub envoy_version: Option<u32>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Hash)]
+pub struct ActorNameMetadata {
+	pub name: String,
+	pub metadata: serde_json::Map<String, serde_json::Value>,
+}
+
+#[derive(Deserialize)]
+struct ServerlessMetadataEnvoy {
+	version: Option<u32>,
+}
+
+#[derive(Deserialize)]
+struct ServerlessMetadataRunner {
+	version: Option<u32>,
+}
+
+#[derive(Deserialize)]
+struct ServerlessMetadataPayload {
+	runtime: String,
+	version: String,
+	#[serde(rename = "envoyProtocolVersion")]
+	envoy_protocol_version: Option<u16>,
+	#[serde(rename = "actorNames", default)]
+	actor_names: HashMap<String, ActorName>,
+	envoy: Option<ServerlessMetadataEnvoy>,
+	runner: Option<ServerlessMetadataRunner>,
+}
+
+#[derive(Deserialize)]
+struct ActorName {
+	metadata: Option<serde_json::Value>,
+}
+
+fn truncate_response_body(body: &str) -> String {
+	let mut chars = body.chars();
+	let mut truncated: String = chars.by_ref().take(RESPONSE_BODY_MAX_LEN).collect();
+	if chars.next().is_some() {
+		truncated.push_str("...[truncated]");
+	}
+
+	truncated
+}
+
+/// Fetches metadata from a serverless runner at the given URL.
+///
+/// Returns metadata including runtime, version, and actor names if available.
+#[operation]
+#[tracing::instrument(skip_all)]
+pub async fn pegboard_serverless_metadata_fetch(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<std::result::Result<Output, ServerlessMetadataError>> {
+	tracing::debug!(url = ?input.url, "fetching serverless runner metadata");
+
+	let trimmed_url = input.url.trim();
+	if trimmed_url.is_empty() {
+		return Ok(Err(ServerlessMetadataError::InvalidRequest {}));
+	}
+
+	let metadata_url = format!("{}/metadata", trimmed_url.trim_end_matches('/'));
+
+	if reqwest::Url::parse(&metadata_url).is_err() {
+		return Ok(Err(ServerlessMetadataError::InvalidRequest {}));
+	}
+
+	let mut header_map = ReqwestHeaderMap::new();
+	for (name, value) in &input.headers {
+		let header_name = match HeaderName::from_bytes(name.trim().as_bytes()) {
+			Ok(n) => n,
+			Err(_) => return Ok(Err(ServerlessMetadataError::InvalidRequest {})),
+		};
+
+		let header_value = match HeaderValue::from_str(value.trim()) {
+			Ok(v) => v,
+			Err(_) => return Ok(Err(ServerlessMetadataError::InvalidRequest {})),
+		};
+
+		header_map.insert(header_name, header_value);
+	}
+
+	let client = match rivet_pools::reqwest::client().await {
+		Ok(c) => c,
+		Err(_) => return Ok(Err(ServerlessMetadataError::RequestFailed {})),
+	};
+
+	tracing::debug!("sending metadata request");
+	let response = match client
+		.get(&metadata_url)
+		.headers(header_map)
+		.timeout(REQUEST_TIMEOUT)
+		.send()
+		.custom_instrument(tracing::info_span!("fetch_metadata_request"))
+		.await
+	{
+		Ok(r) => r,
+		Err(err) => {
+			return Ok(Err(if err.is_timeout() {
+				ServerlessMetadataError::RequestTimedOut {}
+			} else {
+				ServerlessMetadataError::RequestFailed {}
+			}));
+		}
+	};
+
+	let status = response.status();
+	tracing::debug!(?status, "received metadata response");
+	let body_raw = response
+		.text()
+		.await
+		.unwrap_or_else(|_| String::from("<failed to read body>"));
+	let body_for_user = truncate_response_body(&body_raw);
+
+	if !status.is_success() {
+		return Ok(Err(ServerlessMetadataError::NonSuccessStatus {
+			status_code: status.as_u16(),
+			body: body_for_user,
+		}));
+	}
+
+	let payload = match serde_json::from_str::<ServerlessMetadataPayload>(&body_raw) {
+		Ok(p) => p,
+		Err(_) => {
+			return Ok(Err(ServerlessMetadataError::InvalidResponseJson {
+				body: body_for_user,
+			}));
+		}
+	};
+
+	let ServerlessMetadataPayload {
+		runtime,
+		version,
+		envoy_protocol_version,
+		actor_names,
+		envoy,
+		runner,
+	} = payload;
+
+	let runner_version = runner.and_then(|r| r.version);
+	let envoy_version = envoy.and_then(|e| e.version);
+
+	tracing::debug!(
+		?runtime,
+		?version,
+		actor_names_count = actor_names.len(),
+		?envoy_version,
+		"parsed metadata payload"
+	);
+
+	let trimmed_version = version.trim();
+	if runtime != "rivetkit" || trimmed_version.is_empty() {
+		return Ok(Err(ServerlessMetadataError::InvalidResponseSchema {
+			runtime,
+			version,
+		}));
+	}
+
+	if let Some(envoy_protocol_version) = envoy_protocol_version {
+		if envoy_protocol_version < 1 || envoy_protocol_version > PROTOCOL_VERSION {
+			return Ok(Err(ServerlessMetadataError::InvalidResponseJson {
+				body: body_for_user,
+			}));
+		}
+	}
+
+	// Convert actor names, filtering out non-object metadata
+	let actor_names: Vec<ActorNameMetadata> = actor_names
+		.into_iter()
+		.filter_map(|(name, data)| {
+			let metadata = data.metadata?.as_object()?.clone();
+			Some(ActorNameMetadata { name, metadata })
+		})
+		.collect();
+
+	Ok(Ok(Output {
+		runtime,
+		version: trimmed_version.to_owned(),
+		envoy_protocol_version,
+		actor_names,
+		runner_version,
+		envoy_version,
+	}))
+}

package/packages/pegboard/src/ops/serverless_metadata/mod.rs

@@ -0,0 +1 @@
+pub mod fetch;

package/packages/pegboard/src/pubsub_subjects.rs

@@ -78,3 +78,55 @@ impl std::fmt::Display for GatewayReceiverSubject {
 		)
 	}
 }
+
+pub struct EnvoyReceiverSubject {
+	namespace_id: Id,
+	envoy_key: String,
+}
+
+impl EnvoyReceiverSubject {
+	pub fn new(namespace_id: Id, envoy_key: String) -> Self {
+		Self {
+			namespace_id,
+			envoy_key,
+		}
+	}
+}
+
+impl std::fmt::Display for EnvoyReceiverSubject {
+	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+		write!(f, "pegboard.envoy.{}.{}", self.namespace_id, self.envoy_key)
+	}
+}
+
+pub struct EnvoyEvictionSubject {
+	namespace_id: Id,
+	envoy_key: String,
+}
+
+impl EnvoyEvictionSubject {
+	pub fn new(namespace_id: Id, envoy_key: String) -> Self {
+		Self {
+			namespace_id,
+			envoy_key,
+		}
+	}
+}
+
+impl std::fmt::Display for EnvoyEvictionSubject {
+	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+		write!(
+			f,
+			"pegboard.envoy.eviction.{}.{}",
+			self.namespace_id, self.envoy_key
+		)
+	}
+}
+
+pub struct ServerlessOutboundSubject;
+
+impl std::fmt::Display for ServerlessOutboundSubject {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "pegboard.serverless.outbound",)
+	}
+}

package/packages/pegboard/src/utils.rs

@@ -1,6 +1,12 @@
+use gas::prelude::*;
+use rivet_cache::Cache;
 use rivet_runner_protocol as protocol;
+use rivet_types::{
+	keys::namespace::runner_config::RunnerConfigVariant,
+	runner_configs::{RunnerConfig, RunnerConfigKind},
+};
 
-pub fn event_actor_id(event: &protocol::Event) -> &str {
+pub fn event_actor_id_mk1(event: &protocol::Event) -> &str {
 	match event {
 		protocol::Event::EventActorIntent(protocol::EventActorIntent { actor_id, .. }) => actor_id,
 		protocol::Event::EventActorStateUpdate(protocol::EventActorStateUpdate {
@@ -13,7 +19,7 @@ pub fn event_actor_id(event: &protocol::Event) -> &str {
 	}
 }
 
-pub fn event_generation(event: &protocol::Event) -> u32 {
+pub fn event_generation_mk1(event: &protocol::Event) -> u32 {
 	match event {
 		protocol::Event::EventActorIntent(protocol::EventActorIntent { generation, .. }) => {
 			*generation
@@ -27,3 +33,31 @@ pub fn event_generation(event: &protocol::Event) -> u32 {
 		}) => *generation,
 	}
 }
+
+pub fn runner_config_variant(runner_config: &RunnerConfig) -> RunnerConfigVariant {
+	match runner_config.kind {
+		RunnerConfigKind::Normal { .. } => RunnerConfigVariant::Normal,
+		RunnerConfigKind::Serverless { .. } => RunnerConfigVariant::Serverless,
+	}
+}
+
+pub async fn purge_runner_config_caches(
+	cache: &Cache,
+	namespace_id: Id,
+	runner_name: &str,
+) -> Result<()> {
+	let key = (namespace_id, runner_name.to_string());
+
+	cache
+		.clone()
+		.request()
+		.purge("namespace.runner_config.get", vec![key.clone()])
+		.await?;
+	cache
+		.clone()
+		.request()
+		.purge("runner.list_runner_config_enabled_dcs", vec![key])
+		.await?;
+
+	Ok(())
+}

package/packages/pegboard/src/workflows/actor/destroy.rs

@@ -1,5 +1,6 @@
 use gas::prelude::*;
 use rivet_data::converted::ActorByKeyKeyData;
+use rivet_runner_protocol::PROTOCOL_MK1_VERSION;
 use universaldb::options::MutationType;
 use universaldb::utils::IsolationLevel::*;
 
@@ -19,7 +20,7 @@ pub(crate) struct Input {
 #[workflow]
 pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
 	ctx.msg(DestroyStarted {})
-		.tag("actor_id", input.actor_id)
+		.topic(("actor_id", input.actor_id))
 		.send()
 		.await?;
 
@@ -29,12 +30,34 @@ pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input)
 		})
 		.await?;
 
-	// If a slot was allocated at the time of actor destruction then bump the serverless autoscaler so it can scale down
+	// If a slot was allocated at the time of actor destruction then bump the runner pool so it can scale down
 	// if needed
 	if res.allocated_serverless_slot {
-		ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
-			.send()
+		ctx.removed::<Message<super::BumpServerlessAutoscalerStub>>()
 			.await?;
+
+		let bump_res = ctx
+			.v(2)
+			.signal(crate::workflows::runner_pool::Bump::default())
+			.to_workflow::<crate::workflows::runner_pool::Workflow>()
+			.tag("namespace_id", input.namespace_id)
+			.tag("runner_name", res.runner_name_selector.clone())
+			.send()
+			.await;
+
+		if let Some(WorkflowError::WorkflowNotFound) = bump_res
+			.as_ref()
+			.err()
+			.and_then(|x| x.chain().find_map(|x| x.downcast_ref::<WorkflowError>()))
+		{
+			tracing::warn!(
+				namespace_id=%input.namespace_id,
+				runner_name=%res.runner_name_selector,
+				"serverless pool workflow not found, runner config likely deleted"
+			);
+		} else {
+			bump_res?;
+		}
 	}
 
 	// Clear KV
@@ -44,7 +67,7 @@ pub(crate) async fn pegboard_actor_destroy(ctx: &mut WorkflowCtx, input: &Input)
 	.await?;
 
 	ctx.msg(DestroyComplete {})
-		.tag("actor_id", input.actor_id)
+		.topic(("actor_id", input.actor_id))
 		.send()
 		.await?;
 
@@ -59,6 +82,7 @@ struct UpdateStateAndDbInput {
 #[derive(Debug, Serialize, Deserialize, Hash)]
 struct UpdateStateAndDbOutput {
 	allocated_serverless_slot: bool,
+	runner_name_selector: String,
 }
 
 #[activity(UpdateStateAndDb)]
@@ -69,53 +93,45 @@ async fn update_state_and_db(
 	let mut state = ctx.state::<State>()?;
 	let destroy_ts = util::timestamp::now();
 
+	let runner_id = state.runner_id;
+	let namespace_id = state.namespace_id;
+	let runner_name_selector = &state.runner_name_selector;
+	let allocated_serverless_slot = state.allocated_serverless_slot;
+	let name = &state.name;
+	let create_ts = state.create_ts;
+	let key = &state.key;
 	ctx.udb()?
 		.run(|tx| {
-			let state = (*state).clone();
-
 			async move {
 				let tx = tx.with_subspace(keys::subspace());
 
 				tx.write(&keys::actor::DestroyTsKey::new(input.actor_id), destroy_ts)?;
 
-				if let Some(runner_id) = state.runner_id {
-					clear_slot(
-						input.actor_id,
-						state.namespace_id,
-						&state.runner_name_selector,
-						runner_id,
-						state.for_serverless,
-						&tx,
-					)
-					.await?;
-				} else if state.allocated_serverless_slot {
-					// Clear the serverless slot even if we do not have a runner id. This happens when the
-					// actor is destroyed while pending allocation
-					tx.atomic_op(
-						&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new(
-							state.namespace_id,
-							state.runner_name_selector.clone(),
-						),
-						&(-1i64).to_le_bytes(),
-						MutationType::Add,
-					);
-				}
+				clear_slot(
+					input.actor_id,
+					namespace_id,
+					runner_name_selector,
+					runner_id,
+					allocated_serverless_slot,
+					&tx,
+				)
+				.await?;
 
 				// Update namespace indexes
 				tx.delete(&keys::ns::ActiveActorKey::new(
-					state.namespace_id,
-					state.name.clone(),
-					state.create_ts,
+					namespace_id,
+					name.clone(),
+					create_ts,
 					input.actor_id,
 				));
 
-				if let Some(k) = &state.key {
+				if let Some(key) = &key {
 					tx.write(
 						&keys::ns::ActorByKeyKey::new(
-							state.namespace_id,
-							state.name.clone(),
-							k.clone(),
-							state.create_ts,
+							namespace_id,
+							name.clone(),
+							key.clone(),
+							create_ts,
 							input.actor_id,
 						),
 						ActorByKeyKeyData {
@@ -125,6 +141,14 @@ async fn update_state_and_db(
 					)?;
 				}
 
+				// Update metrics
+				namespace::keys::metric::inc(
+					&tx.with_subspace(namespace::keys::subspace()),
+					namespace_id,
+					namespace::keys::metric::Metric::TotalActors(name.clone()),
+					-1,
+				);
+
 				Ok(())
 			}
 		})
@@ -139,6 +163,7 @@ async fn update_state_and_db(
 
 	Ok(UpdateStateAndDbOutput {
 		allocated_serverless_slot: old_allocated_serverless_slot,
+		runner_name_selector: state.runner_name_selector.clone(),
 	})
 }
 
@@ -149,20 +174,21 @@ struct ClearKvInput {
 
 #[derive(Debug, Serialize, Deserialize, Hash)]
 struct ClearKvOutput {
+	// Simply an estimate, not accurate under 3MiB
 	final_size: i64,
 }
 
 #[activity(ClearKv)]
 async fn clear_kv(ctx: &ActivityCtx, input: &ClearKvInput) -> Result<ClearKvOutput> {
-	// Matches `delete_all` from actor_kv (can't import because of cyclical dep)
 	let final_size = ctx
 		.udb()?
 		.run(|tx| async move {
-			let subspace = keys::actor_kv_subspace().subspace(&input.actor_id);
+			let subspace = keys::actor_kv::subspace(input.actor_id);
 
 			let (start, end) = subspace.range();
 			let final_size = tx.get_estimated_range_size_bytes(&start, &end).await?;
 
+			// Matches `delete_all` from actor kv
 			tx.clear_subspace_range(&subspace);
 
 			Ok(final_size)
@@ -177,79 +203,89 @@ pub(crate) async fn clear_slot(
 	actor_id: Id,
 	namespace_id: Id,
 	runner_name_selector: &str,
-	runner_id: Id,
-	for_serverless: bool,
+	runner_id: Option<Id>,
+	allocated_serverless_slot: bool,
 	tx: &universaldb::Transaction,
 ) -> Result<()> {
 	let tx = tx.with_subspace(keys::subspace());
 
-	tx.delete(&keys::actor::RunnerIdKey::new(actor_id));
-
-	// This is cleared when the state changes as well as when the actor is destroyed to ensure
-	// consistency during rescheduling and forced deletion.
-	tx.delete(&keys::runner::ActorKey::new(runner_id, actor_id));
-
-	let runner_workflow_id_key = keys::runner::WorkflowIdKey::new(runner_id);
-	let runner_version_key = keys::runner::VersionKey::new(runner_id);
-	let runner_remaining_slots_key = keys::runner::RemainingSlotsKey::new(runner_id);
-	let runner_total_slots_key = keys::runner::TotalSlotsKey::new(runner_id);
-	let runner_last_ping_ts_key = keys::runner::LastPingTsKey::new(runner_id);
-
-	let (
-		runner_workflow_id,
-		runner_version,
-		runner_remaining_slots,
-		runner_total_slots,
-		runner_last_ping_ts,
-	) = tokio::try_join!(
-		tx.read(&runner_workflow_id_key, Serializable),
-		tx.read(&runner_version_key, Serializable),
-		tx.read(&runner_remaining_slots_key, Serializable),
-		tx.read(&runner_total_slots_key, Serializable),
-		tx.read(&runner_last_ping_ts_key, Serializable),
-	)?;
-
-	let old_runner_remaining_millislots = (runner_remaining_slots * 1000) / runner_total_slots;
-	let new_runner_remaining_slots = runner_remaining_slots + 1;
-
-	// Write new remaining slots
-	tx.write(&runner_remaining_slots_key, new_runner_remaining_slots)?;
-
-	let old_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new(
-		namespace_id,
-		runner_name_selector.to_string(),
-		runner_version,
-		old_runner_remaining_millislots,
-		runner_last_ping_ts,
-		runner_id,
-	);
-
-	// Only update allocation idx if it existed before
-	if tx.exists(&old_runner_alloc_key, Serializable).await? {
-		// Clear old key
-		tx.delete(&old_runner_alloc_key);
-
-		let new_remaining_millislots = (new_runner_remaining_slots * 1000) / runner_total_slots;
-		let new_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new(
+	// Only clear slot if we have a runner id
+	if let Some(runner_id) = runner_id {
+		tx.delete(&keys::actor::RunnerIdKey::new(actor_id));
+
+		// This is cleared when the state changes as well as when the actor is destroyed to ensure
+		// consistency during rescheduling and forced deletion.
+		tx.delete(&keys::runner::ActorKey::new(runner_id, actor_id));
+
+		let runner_workflow_id_key = keys::runner::WorkflowIdKey::new(runner_id);
+		let runner_version_key = keys::runner::VersionKey::new(runner_id);
+		let runner_remaining_slots_key = keys::runner::RemainingSlotsKey::new(runner_id);
+		let runner_total_slots_key = keys::runner::TotalSlotsKey::new(runner_id);
+		let runner_last_ping_ts_key = keys::runner::LastPingTsKey::new(runner_id);
+		let runner_protocol_version_key = keys::runner::ProtocolVersionKey::new(runner_id);
+
+		let (
+			runner_workflow_id,
+			runner_version,
+			runner_remaining_slots,
+			runner_total_slots,
+			runner_last_ping_ts,
+			runner_protocol_version,
+		) = tokio::try_join!(
+			tx.read(&runner_workflow_id_key, Serializable),
+			tx.read(&runner_version_key, Serializable),
+			tx.read(&runner_remaining_slots_key, Serializable),
+			tx.read(&runner_total_slots_key, Serializable),
+			tx.read(&runner_last_ping_ts_key, Serializable),
+			tx.read_opt(&runner_protocol_version_key, Serializable),
+		)?;
+
+		let old_runner_remaining_millislots = (runner_remaining_slots * 1000) / runner_total_slots;
+		let new_runner_remaining_slots = runner_remaining_slots + 1;
+
+		// Write new remaining slots
+		tx.write(&runner_remaining_slots_key, new_runner_remaining_slots)?;
+
+		let old_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new(
 			namespace_id,
 			runner_name_selector.to_string(),
 			runner_version,
-			new_remaining_millislots,
+			old_runner_remaining_millislots,
 			runner_last_ping_ts,
 			runner_id,
 		);
 
-		tx.write(
-			&new_runner_alloc_key,
-			rivet_data::converted::RunnerAllocIdxKeyData {
-				workflow_id: runner_workflow_id,
-				remaining_slots: new_runner_remaining_slots,
-				total_slots: runner_total_slots,
-			},
-		)?;
+		// Only update allocation idx if it existed before
+		if tx.exists(&old_runner_alloc_key, Serializable).await? {
+			// Clear old key
+			tx.delete(&old_runner_alloc_key);
+
+			let new_remaining_millislots = (new_runner_remaining_slots * 1000) / runner_total_slots;
+			let new_runner_alloc_key = keys::ns::RunnerAllocIdxKey::new(
+				namespace_id,
+				runner_name_selector.to_string(),
+				runner_version,
+				new_remaining_millislots,
+				runner_last_ping_ts,
+				runner_id,
+			);
+
+			tx.write(
+				&new_runner_alloc_key,
+				rivet_data::converted::RunnerAllocIdxKeyData {
+					workflow_id: runner_workflow_id,
+					remaining_slots: new_runner_remaining_slots,
+					total_slots: runner_total_slots,
+					// We default here because its not important for mk1 protocol runners
+					protocol_version: runner_protocol_version.unwrap_or(PROTOCOL_MK1_VERSION),
+				},
+			)?;
+		}
 	}
 
-	if for_serverless {
+	if allocated_serverless_slot {
+		// Clear the serverless slot even if we do not have a runner id. This happens when the
+		// actor is destroyed while pending allocation
 		tx.atomic_op(
 			&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new(
 				namespace_id,