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

package/packages/pegboard/src/keys/runner.rs

@@ -183,6 +183,50 @@ impl<'de> TupleUnpack<'de> for TotalSlotsKey {
 	}
 }
 
+#[derive(Debug)]
+pub struct ProtocolVersionKey {
+	runner_id: Id,
+}
+
+impl ProtocolVersionKey {
+	pub fn new(runner_id: Id) -> Self {
+		ProtocolVersionKey { runner_id }
+	}
+}
+
+impl FormalKey for ProtocolVersionKey {
+	type Value = u16;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(u16::from_be_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_be_bytes().to_vec())
+	}
+}
+
+impl TuplePack for ProtocolVersionKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (RUNNER, DATA, self.runner_id, PROTOCOL_VERSION);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for ProtocolVersionKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, runner_id, _)) =
+			<(usize, usize, Id, usize)>::unpack(input, tuple_depth)?;
+		let v = ProtocolVersionKey { runner_id };
+
+		Ok((input, v))
+	}
+}
+
 #[derive(Debug)]
 pub struct ActorKey {
 	runner_id: Id,
@@ -816,3 +860,240 @@ impl<'de> TupleUnpack<'de> for MetadataChunkKey {
 		Ok((input, v))
 	}
 }
+
+#[derive(Debug)]
+pub struct ActorLastCommandIdxKey {
+	runner_id: Id,
+	actor_id: Id,
+	generation: u32,
+}
+
+impl ActorLastCommandIdxKey {
+	pub fn new(runner_id: Id, actor_id: Id, generation: u32) -> Self {
+		ActorLastCommandIdxKey {
+			runner_id,
+			actor_id,
+			generation,
+		}
+	}
+}
+
+impl FormalKey for ActorLastCommandIdxKey {
+	// Timestamp.
+	type Value = i64;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(i64::from_be_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_be_bytes().to_vec())
+	}
+}
+
+impl TuplePack for ActorLastCommandIdxKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (
+			RUNNER,
+			DATA,
+			self.runner_id,
+			ACTOR,
+			LAST_COMMAND_IDX,
+			self.actor_id,
+			self.generation,
+		);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for ActorLastCommandIdxKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, runner_id, _, _, actor_id, generation)) =
+			<(usize, usize, Id, usize, usize, Id, u32)>::unpack(input, tuple_depth)?;
+		let v = ActorLastCommandIdxKey {
+			runner_id,
+			actor_id,
+			generation,
+		};
+
+		Ok((input, v))
+	}
+}
+
+#[derive(Debug)]
+pub struct ActorCommandKey {
+	pub runner_id: Id,
+	pub actor_id: Id,
+	pub generation: u32,
+	pub index: i64,
+}
+
+impl ActorCommandKey {
+	pub fn new(runner_id: Id, actor_id: Id, generation: u32, index: i64) -> Self {
+		ActorCommandKey {
+			runner_id,
+			actor_id,
+			generation,
+			index,
+		}
+	}
+
+	pub fn subspace(runner_id: Id) -> ActorCommandSubspaceKey {
+		ActorCommandSubspaceKey::new(runner_id)
+	}
+
+	pub fn subspace_with_actor(
+		runner_id: Id,
+		actor_id: Id,
+		generation: u32,
+	) -> ActorCommandSubspaceKey {
+		ActorCommandSubspaceKey::new_with_actor(runner_id, actor_id, generation)
+	}
+
+	pub fn subspace_with_index(
+		runner_id: Id,
+		actor_id: Id,
+		generation: u32,
+		index: i64,
+	) -> ActorCommandSubspaceKey {
+		ActorCommandSubspaceKey::new_with_index(runner_id, actor_id, generation, index)
+	}
+}
+
+impl FormalKey for ActorCommandKey {
+	type Value = rivet_runner_protocol::mk2::ActorCommandKeyData;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		rivet_runner_protocol::versioned::ActorCommandKeyData::deserialize_with_embedded_version(
+			raw,
+		)
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		rivet_runner_protocol::versioned::ActorCommandKeyData::wrap_latest(value)
+			.serialize_with_embedded_version(rivet_runner_protocol::PROTOCOL_MK2_VERSION)
+	}
+}
+
+impl TuplePack for ActorCommandKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (
+			RUNNER,
+			DATA,
+			self.runner_id,
+			ACTOR,
+			COMMAND,
+			self.actor_id,
+			self.generation,
+			self.index,
+		);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for ActorCommandKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, runner_id, _, _, actor_id, generation, index)) =
+			<(usize, usize, Id, usize, usize, Id, u32, i64)>::unpack(input, tuple_depth)?;
+		let v = ActorCommandKey {
+			runner_id,
+			actor_id,
+			generation,
+			index,
+		};
+
+		Ok((input, v))
+	}
+}
+
+pub struct ActorCommandSubspaceKey {
+	runner_id: Id,
+	actor_id: Option<Id>,
+	generation: Option<u32>,
+	index: Option<i64>,
+}
+
+impl ActorCommandSubspaceKey {
+	pub fn new(runner_id: Id) -> Self {
+		ActorCommandSubspaceKey {
+			runner_id,
+			actor_id: None,
+			generation: None,
+			index: None,
+		}
+	}
+
+	pub fn new_with_actor(runner_id: Id, actor_id: Id, generation: u32) -> Self {
+		ActorCommandSubspaceKey {
+			runner_id,
+			actor_id: Some(actor_id),
+			generation: Some(generation),
+			index: None,
+		}
+	}
+
+	pub fn new_with_index(runner_id: Id, actor_id: Id, generation: u32, index: i64) -> Self {
+		ActorCommandSubspaceKey {
+			runner_id,
+			actor_id: Some(actor_id),
+			generation: Some(generation),
+			index: Some(index),
+		}
+	}
+}
+
+impl TuplePack for ActorCommandSubspaceKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
+		let t = (RUNNER, DATA, self.runner_id, ACTOR, COMMAND);
+		offset += t.pack(w, tuple_depth)?;
+
+		if let Some(actor_id) = &self.actor_id {
+			offset += actor_id.pack(w, tuple_depth)?;
+
+			if let Some(v) = &self.generation {
+				offset += v.pack(w, tuple_depth)?;
+
+				if let Some(index) = &self.index {
+					offset += index.pack(w, tuple_depth)?;
+				}
+			}
+		}
+
+		Ok(offset)
+	}
+}
+
+pub struct ActorDataSubspaceKey {
+	runner_id: Id,
+}
+
+impl ActorDataSubspaceKey {
+	pub fn new(runner_id: Id) -> Self {
+		ActorDataSubspaceKey { runner_id }
+	}
+}
+
+impl TuplePack for ActorDataSubspaceKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (RUNNER, DATA, self.runner_id, ACTOR);
+		t.pack(w, tuple_depth)
+	}
+}

package/packages/{namespace → pegboard}/src/keys/runner_config.rs

@@ -209,3 +209,56 @@ impl TuplePack for ByVariantSubspaceKey {
 		Ok(offset)
 	}
 }
+
+#[derive(Debug)]
+pub struct ProtocolVersionKey {
+	pub namespace_id: Id,
+	pub name: String,
+}
+
+impl ProtocolVersionKey {
+	pub fn new(namespace_id: Id, name: String) -> Self {
+		ProtocolVersionKey { namespace_id, name }
+	}
+}
+
+impl FormalKey for ProtocolVersionKey {
+	type Value = u16;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(u16::from_be_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_be_bytes().to_vec())
+	}
+}
+
+impl TuplePack for ProtocolVersionKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (
+			RUNNER,
+			CONFIG,
+			DATA,
+			self.namespace_id,
+			&self.name,
+			PROTOCOL_VERSION,
+		);
+		t.pack(w, tuple_depth)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for ProtocolVersionKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, _, _, namespace_id, name, _)) =
+			<(usize, usize, usize, Id, String, usize)>::unpack(input, tuple_depth)?;
+
+		let v = ProtocolVersionKey { namespace_id, name };
+
+		Ok((input, v))
+	}
+}

package/packages/pegboard/src/lib.rs

@@ -1,11 +1,12 @@
 use gas::prelude::*;
 
+pub mod actor_kv;
 pub mod errors;
 pub mod keys;
-mod metrics;
+pub mod metrics;
 pub mod ops;
 pub mod pubsub_subjects;
-mod utils;
+pub mod utils;
 pub mod workflows;
 
 pub fn registry() -> WorkflowResult<Registry> {
@@ -13,7 +14,19 @@ pub fn registry() -> WorkflowResult<Registry> {
 
 	let mut registry = Registry::new();
 	registry.register_workflow::<actor::Workflow>()?;
+	registry.register_workflow::<actor::metrics::Workflow>()?;
+	registry.register_workflow::<actor2::Workflow>()?;
+	registry.register_workflow::<actor2::metrics::Workflow>()?;
 	registry.register_workflow::<runner::Workflow>()?;
+	registry.register_workflow::<runner2::Workflow>()?;
+	registry.register_workflow::<runner_pool::Workflow>()?;
+	registry.register_workflow::<runner_pool_error_tracker::Workflow>()?;
+	registry.register_workflow::<runner_pool_metadata_poller::Workflow>()?;
+	registry.register_workflow::<serverless::receiver::Workflow>()?;
+	registry.register_workflow::<serverless::conn::Workflow>()?;
+	registry.register_workflow::<serverless::backfill::Workflow>()?;
+	registry.register_workflow::<metrics_aggregator::Workflow>()?;
+	registry.register_workflow::<actor_runner_name_selector_backfill::Workflow>()?;
 
 	Ok(registry)
 }

package/packages/pegboard/src/metrics.rs

@@ -1,19 +1,60 @@
-use rivet_metrics::{
-	MICRO_BUCKETS,
-	otel::{global::*, metrics::*},
-};
+use rivet_metrics::{MICRO_BUCKETS, REGISTRY, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-pegboard");
-
-	/// Expected attributes: "namespace_id", "runner_name"
-	pub static ref ACTOR_PENDING_ALLOCATION: Gauge<f64> = METER.f64_gauge("rivet_pegboard_actor_pending_allocation")
-		.with_description("Total actors waiting for availability.")
-		.build();
-
-	/// Expected attributes: "did_reserve"
-	pub static ref ACTOR_ALLOCATE_DURATION: Histogram<f64> = METER.f64_histogram("rivet_pegboard_actor_allocate_duration")
-		.with_description("Total duration to reserve resources for an actor.")
-		.with_boundaries(MICRO_BUCKETS.to_vec())
-		.build();
+	pub static ref ACTOR_PENDING_ALLOCATION: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_actor_pending_allocation",
+		"Total actors waiting for availability.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref ACTOR_ACTIVE: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_actor_active",
+		"Total actors currently allocated.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref SERVERLESS_DESIRED_SLOTS: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_serverless_desired_slots",
+		"Total amount of desired slots for serverless runners.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref ACTOR_ALLOCATE_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"pegboard_actor_allocate_duration",
+		"Total duration to allocate an actor.",
+		&["kind", "result"],
+		MICRO_BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	pub static ref RUNNER_VERSION_UPGRADE_DRAIN_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_runner_version_upgrade_drain_total",
+		"Count of runners drained due to version upgrade.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref ENVOY_VERSION_UPGRADE_DRAIN_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_envoy_version_upgrade_drain_total",
+		"Count of envoys drained due to version upgrade.",
+		&["namespace_id", "pool_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref SERVERLESS_OUTBOUND_REQ_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"pegboard_serverless_outbound_req_total",
+		"Count of serverless outbound requests made.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
+
+	pub static ref SERVERLESS_OUTBOUND_REQ_ACTIVE: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"pegboard_serverless_outbound_req_active",
+		"Count of serverless outbound requests currently active.",
+		&["namespace_id", "runner_name"],
+		*REGISTRY
+	).unwrap();
 }

package/packages/pegboard/src/ops/actor/create.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
 use gas::prelude::*;
 use rivet_api_util::{Method, request_remote_datacenter};
 use rivet_types::actors::{Actor, CrashPolicy};
@@ -30,63 +30,132 @@ pub struct Output {
 #[operation]
 pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result<Output> {
 	// Set up subscriptions before dispatching workflow
-	let mut create_sub = ctx
-		.subscribe::<crate::workflows::actor::CreateComplete>(("actor_id", input.actor_id))
-		.await?;
-	let mut fail_sub = ctx
-		.subscribe::<crate::workflows::actor::Failed>(("actor_id", input.actor_id))
-		.await?;
-	let mut destroy_sub = ctx
-		.subscribe::<crate::workflows::actor::DestroyStarted>(("actor_id", input.actor_id))
+	let (
+		mut create_sub,
+		mut fail_sub,
+		mut destroy_sub,
+		mut create_sub2,
+		mut fail_sub2,
+		mut destroy_sub2,
+		pool_res,
+	) = tokio::try_join!(
+		ctx.subscribe::<crate::workflows::actor::CreateComplete>(("actor_id", input.actor_id)),
+		ctx.subscribe::<crate::workflows::actor::Failed>(("actor_id", input.actor_id)),
+		ctx.subscribe::<crate::workflows::actor::DestroyStarted>(("actor_id", input.actor_id)),
+		ctx.subscribe::<crate::workflows::actor2::CreateComplete>(("actor_id", input.actor_id)),
+		ctx.subscribe::<crate::workflows::actor2::Failed>(("actor_id", input.actor_id)),
+		ctx.subscribe::<crate::workflows::actor2::DestroyStarted>(("actor_id", input.actor_id)),
+		ctx.op(crate::ops::runner_config::get::Input {
+			runners: vec![(input.namespace_id, input.runner_name_selector.clone())],
+			bypass_cache: false,
+		}),
+	)?;
+
+	let actor_v2 = pool_res
+		.into_iter()
+		.next()
+		.map(|p| p.protocol_version.is_some())
+		.unwrap_or_default();
+
+	if actor_v2 {
+		// Dispatch actor workflow
+		ctx.workflow(crate::workflows::actor2::Input {
+			actor_id: input.actor_id,
+			name: input.name.clone(),
+			pool_name: input.runner_name_selector.clone(),
+			key: input.key.clone(),
+			namespace_id: input.namespace_id,
+			input: input.input.clone(),
+			from_v1: false,
+		})
+		.tag("actor_id", input.actor_id)
+		.dispatch()
 		.await?;
 
-	// Dispatch actor workflow
-	ctx.workflow(crate::workflows::actor::Input {
-		actor_id: input.actor_id,
-		name: input.name.clone(),
-		key: input.key.clone(),
-		namespace_id: input.namespace_id,
-		runner_name_selector: input.runner_name_selector.clone(),
-		input: input.input.clone(),
-		crash_policy: input.crash_policy,
-	})
-	.tag("actor_id", input.actor_id)
-	.dispatch()
-	.await?;
+		// Wait for actor creation to complete, fail, or be destroyed
+		tokio::select! {
+			res = create_sub2.next() => { res?; },
+			res = fail_sub2.next() => {
+				let msg = res?;
+				let error = msg.into_body().error;
 
-	// Wait for actor creation to complete, fail, or be destroyed
-	tokio::select! {
-		res = create_sub.next() => { res?; },
-		res = fail_sub.next() => {
-			let msg = res?;
-			let error = msg.into_body().error;
-
-			// Check if this request needs to be forwarded
-			//
-			// We cannot forward if `datacenter_name` is specified because this actor is being
-			// restricted to the given datacenter.
-			if input.forward_request && input.datacenter_name.is_none() {
-				if let crate::errors::Actor::KeyReservedInDifferentDatacenter { datacenter_label } = &error {
-					// Forward the request to the correct datacenter
-					return forward_to_datacenter(
-						ctx,
-						*datacenter_label,
-						input.namespace_id,
-						input.name.clone(),
-						input.key.clone(),
-						input.runner_name_selector.clone(),
-						input.input.clone(),
-					input.crash_policy
-					).await;
+				// Check if this request needs to be forwarded
+				//
+				// We cannot forward if `datacenter_name` is specified because this actor is being
+				// restricted to the given datacenter.
+				if input.forward_request && input.datacenter_name.is_none() {
+					if let crate::errors::Actor::KeyReservedInDifferentDatacenter { datacenter_label } = &error {
+						// Forward the request to the correct datacenter
+						return forward_to_datacenter(
+							ctx,
+							*datacenter_label,
+							input.namespace_id,
+							input.name.clone(),
+							input.key.clone(),
+							input.runner_name_selector.clone(),
+							input.input.clone(),
+						input.crash_policy
+						).await;
+					}
 				}
-			}
 
-			// Otherwise, return the error as-is
-			return Err(error.build());
+				// Otherwise, return the error as-is
+				return Err(error.build());
+			}
+			res = destroy_sub2.next() => {
+				res?;
+				return Err(crate::errors::Actor::DestroyedDuringCreation.build());
+			}
 		}
-		res = destroy_sub.next() => {
-			res?;
-			return Err(crate::errors::Actor::DestroyedDuringCreation.build());
+	} else {
+		// Dispatch actor workflow
+		ctx.workflow(crate::workflows::actor::Input {
+			actor_id: input.actor_id,
+			name: input.name.clone(),
+			runner_name_selector: input.runner_name_selector.clone(),
+			key: input.key.clone(),
+			namespace_id: input.namespace_id,
+			crash_policy: input.crash_policy,
+			input: input.input.clone(),
+		})
+		.tag("actor_id", input.actor_id)
+		.dispatch()
+		.await?;
+
+		// Wait for actor creation to complete, fail, or be destroyed
+		tokio::select! {
+			res = create_sub.next() => { res?; },
+			res = fail_sub.next() => {
+				let msg = res?;
+				let error = msg.into_body().error;
+
+				// Check if this request needs to be forwarded
+				//
+				// We cannot forward if `datacenter_name` is specified because this actor is being
+				// restricted to the given datacenter.
+				if input.forward_request && input.datacenter_name.is_none() {
+					if let crate::errors::Actor::KeyReservedInDifferentDatacenter { datacenter_label } = &error {
+						// Forward the request to the correct datacenter
+						return forward_to_datacenter(
+							ctx,
+							*datacenter_label,
+							input.namespace_id,
+							input.name.clone(),
+							input.key.clone(),
+							input.runner_name_selector.clone(),
+							input.input.clone(),
+						input.crash_policy
+						).await;
+					}
+				}
+
+				// Otherwise, return the error as-is
+				return Err(error.build());
+			}
+			res = destroy_sub.next() => {
+				res?;
+				return Err(crate::errors::Actor::DestroyedDuringCreation.build());
+			}
 		}
 	}
 
@@ -94,6 +163,7 @@ pub async fn pegboard_actor_create(ctx: &OperationCtx, input: &Input) -> Result<
 	let actors_res = ctx
 		.op(crate::ops::actor::get::Input {
 			actor_ids: vec![input.actor_id],
+			fetch_error: false,
 		})
 		.await?;
 
@@ -121,7 +191,7 @@ async fn forward_to_datacenter(
 	let _target_dc = ctx
 		.config()
 		.dc_for_label(datacenter_label)
-		.ok_or_else(|| anyhow::anyhow!("datacenter not found for label {}", datacenter_label))?;
+		.with_context(|| format!("datacenter not found for label {}", datacenter_label))?;
 
 	// Get namespace name for the remote call
 	let namespace = ctx