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

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

@@ -1,22 +1,41 @@
+// runner wf see how signal fail handling
 use base64::Engine;
 use base64::prelude::BASE64_STANDARD;
 use futures_util::StreamExt;
 use futures_util::TryStreamExt;
 use gas::prelude::*;
-use rivet_metrics::KeyValue;
-use rivet_runner_protocol as protocol;
-use rivet_types::{
-	actors::CrashPolicy, keys::namespace::runner_config::RunnerConfigVariant,
-	runner_configs::RunnerConfigKind,
+use rand::prelude::SliceRandom;
+use rivet_runner_protocol::{
+	self as protocol, PROTOCOL_MK1_VERSION, PROTOCOL_MK2_VERSION, versioned,
 };
+use rivet_types::actors::CrashPolicy;
+use rivet_types::runner_configs::RunnerConfigKind;
 use std::time::Instant;
-use universaldb::options::{ConflictRangeType, MutationType, StreamingMode};
-use universaldb::utils::{FormalKey, IsolationLevel::*};
+use universaldb::prelude::*;
+use universalpubsub::PublishOpts;
+use vbare::OwnedVersionedData;
+
+use super::FailureReason;
 
 use crate::{keys, metrics};
 
 use super::{Allocate, Destroy, Input, PendingAllocation, State, destroy};
 
+#[derive(Debug, Deserialize, Serialize)]
+pub struct LifecycleRunnerState {
+	pub last_event_idx: i64,
+	pub last_event_ack_idx: i64,
+}
+
+impl Default for LifecycleRunnerState {
+	fn default() -> Self {
+		LifecycleRunnerState {
+			last_event_idx: -1,
+			last_event_ack_idx: -1,
+		}
+	}
+}
+
 // TODO: Rewrite this as a series of nested structs/enums for better transparency of current state (likely
 // requires actor wf v2)
 #[derive(Deserialize, Serialize)]
@@ -26,6 +45,8 @@ pub struct LifecycleState {
 	// Set when currently running (not rescheduling or sleeping)
 	pub runner_id: Option<Id>,
 	pub runner_workflow_id: Option<Id>,
+	pub runner_protocol_version: Option<u16>,
+	pub runner_state: Option<LifecycleRunnerState>,
 
 	pub sleeping: bool,
 	#[serde(default)]
@@ -46,11 +67,18 @@ pub struct LifecycleState {
 }
 
 impl LifecycleState {
-	pub fn new(runner_id: Id, runner_workflow_id: Id, actor_start_threshold: i64) -> Self {
+	pub fn new(
+		runner_id: Id,
+		runner_workflow_id: Id,
+		runner_protocol_version: u16,
+		actor_start_threshold: i64,
+	) -> Self {
 		LifecycleState {
 			generation: 0,
 			runner_id: Some(runner_id),
 			runner_workflow_id: Some(runner_workflow_id),
+			runner_protocol_version: Some(runner_protocol_version),
+			runner_state: Some(LifecycleRunnerState::default()),
 			sleeping: false,
 			stopping: false,
 			going_away: false,
@@ -66,6 +94,8 @@ impl LifecycleState {
 			generation: 0,
 			runner_id: None,
 			runner_workflow_id: None,
+			runner_protocol_version: None,
+			runner_state: None,
 			sleeping: true,
 			stopping: false,
 			going_away: false,
@@ -80,6 +110,8 @@ impl LifecycleState {
 #[derive(Serialize, Deserialize)]
 pub struct LifecycleResult {
 	pub generation: u32,
+	#[serde(default)]
+	pub migrate_to_v2: bool,
 }
 
 #[derive(Serialize, Deserialize, Clone, Default)]
@@ -102,14 +134,26 @@ async fn update_runner(ctx: &ActivityCtx, input: &UpdateRunnerInput) -> Result<(
 
 	state.sleep_ts = None;
 	state.pending_allocation_ts = None;
+	state.failure_reason = None;
 	state.runner_id = Some(input.runner_id);
 	state.runner_workflow_id = Some(input.runner_workflow_id);
 
+	ctx.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			// Set actor as not sleeping
+			tx.delete(&keys::actor::SleepTsKey::new(input.actor_id));
+
+			Ok(())
+		})
+		.await?;
+
 	Ok(())
 }
 
 #[derive(Debug, Serialize, Deserialize, Hash)]
-struct AllocateActorInput {
+struct AllocateActorInputV1 {
 	actor_id: Id,
 	generation: u32,
 	force_allocate: bool,
@@ -117,7 +161,7 @@ struct AllocateActorInput {
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
-pub enum AllocateActorOutput {
+pub enum AllocateActorOutputV1 {
 	Allocated {
 		runner_id: Id,
 		runner_workflow_id: Id,
@@ -128,12 +172,73 @@ pub enum AllocateActorOutput {
 	Sleep,
 }
 
-// If no availability, returns the timestamp of the actor's queue key
 #[activity(AllocateActor)]
 async fn allocate_actor(
 	ctx: &ActivityCtx,
-	input: &AllocateActorInput,
-) -> Result<AllocateActorOutput> {
+	input: &AllocateActorInputV1,
+) -> Result<AllocateActorOutputV1> {
+	bail!("allocate actor v1 should never be called again")
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct AllocateActorInputV2 {
+	actor_id: Id,
+	generation: u32,
+	force_allocate: bool,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct AllocateActorOutputV2 {
+	status: AllocateActorStatus,
+	serverless: bool,
+}
+
+impl From<AllocateActorOutputV1> for AllocateActorOutputV2 {
+	fn from(value: AllocateActorOutputV1) -> Self {
+		Self {
+			serverless: false,
+			status: match value {
+				AllocateActorOutputV1::Allocated {
+					runner_id,
+					runner_workflow_id,
+				} => AllocateActorStatus::Allocated {
+					runner_id,
+					runner_workflow_id,
+					runner_protocol_version: None,
+				},
+				AllocateActorOutputV1::Pending {
+					pending_allocation_ts,
+				} => AllocateActorStatus::Pending {
+					pending_allocation_ts,
+				},
+				AllocateActorOutputV1::Sleep => AllocateActorStatus::Sleep,
+			},
+		}
+	}
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+enum AllocateActorStatus {
+	Allocated {
+		runner_id: Id,
+		runner_workflow_id: Id,
+		#[serde(default)]
+		runner_protocol_version: Option<u16>,
+	},
+	Pending {
+		pending_allocation_ts: i64,
+	},
+	Sleep,
+	MigrateToV2,
+}
+
+// If no availability, returns the timestamp of the actor's queue key
+#[activity(AllocateActorV2)]
+async fn allocate_actor_v2(
+	ctx: &ActivityCtx,
+	input: &AllocateActorInputV2,
+) -> Result<AllocateActorOutputV2> {
 	let start_instant = Instant::now();
 
 	let mut state = ctx.state::<State>()?;
@@ -141,45 +246,61 @@ async fn allocate_actor(
 	let crash_policy = state.crash_policy;
 	let runner_name_selector = &state.runner_name_selector;
 
-	// Check if valid serverless config exists for the current ns + runner name
-	let runner_config_res = ctx
-		.op(namespace::ops::runner_config::get::Input {
+	let runner_eligible_threshold = ctx.config().pegboard().runner_eligible_threshold();
+	let actor_allocation_candidate_sample_size = ctx
+		.config()
+		.pegboard()
+		.actor_allocation_candidate_sample_size();
+
+	let pool_res = ctx
+		.op(crate::ops::runner_config::get::Input {
 			runners: vec![(namespace_id, runner_name_selector.clone())],
 			bypass_cache: false,
 		})
 		.await?;
-	let has_valid_serverless = runner_config_res
-		.first()
-		.map(|runner| match &runner.config.kind {
-			RunnerConfigKind::Serverless { max_runners, .. } => *max_runners != 0,
-			_ => false,
-		})
-		.unwrap_or_default();
-
-	let runner_eligible_threshold = ctx.config().pegboard().runner_eligible_threshold();
+	let pool = pool_res.into_iter().next();
+	let for_serverless = pool
+		.as_ref()
+		.map(|pool| matches!(pool.config.kind, RunnerConfigKind::Serverless { .. }))
+		.unwrap_or(false);
+
+	// Protocol version is set or this is a serverless pool
+	if pool.and_then(|p| p.protocol_version).is_some() {
+		return Ok(AllocateActorOutputV2 {
+			status: AllocateActorStatus::MigrateToV2,
+			serverless: false,
+		});
+	}
 
 	// NOTE: This txn should closely resemble the one found in the allocate_pending_actors activity of the
 	// client wf
-	let (for_serverless, res) = ctx
+	let res = ctx
 		.udb()?
 		.run(|tx| async move {
 			let ping_threshold_ts = util::timestamp::now() - runner_eligible_threshold;
 
-			// Check if runner is an serverless runner
-			let for_serverless = tx
-				.with_subspace(namespace::keys::subspace())
-				.exists(
-					&namespace::keys::runner_config::ByVariantKey::new(
-						namespace_id,
-						RunnerConfigVariant::Serverless,
-						runner_name_selector.clone(),
-					),
-					Serializable,
-				)
-				.await?;
-
 			let tx = tx.with_subspace(keys::subspace());
 
+			// Check if a queue exists
+			let pending_actor_subspace = keys::subspace().subspace(
+				&keys::ns::PendingActorByRunnerNameSelectorKey::subspace(
+					namespace_id,
+					runner_name_selector.clone(),
+				),
+			);
+
+			let mut queue_stream = tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::Exact,
+					limit: Some(1),
+					..(&pending_actor_subspace).into()
+				},
+				// NOTE: This is not Serializable because we don't want to conflict with other
+				// inserts/clears to this range
+				Snapshot,
+			);
+			let queue_exists = queue_stream.next().await.is_some();
+
 			if for_serverless {
 				tx.atomic_op(
 					&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new(
@@ -191,28 +312,6 @@ async fn allocate_actor(
 				);
 			}
 
-			// Check if a queue exists
-			let pending_actor_subspace = keys::subspace().subspace(
-				&keys::ns::PendingActorByRunnerNameSelectorKey::subspace(
-					namespace_id,
-					runner_name_selector.clone(),
-				),
-			);
-			let queue_exists = tx
-				.get_ranges_keyvalues(
-					universaldb::RangeOption {
-						mode: StreamingMode::Exact,
-						limit: Some(1),
-						..(&pending_actor_subspace).into()
-					},
-					// NOTE: This is not Serializable because we don't want to conflict with other
-					// inserts/clears to this range
-					Snapshot,
-				)
-				.next()
-				.await
-				.is_some();
-
 			if !queue_exists {
 				let runner_alloc_subspace =
 					keys::subspace().subspace(&keys::ns::RunnerAllocIdxKey::subspace(
@@ -231,7 +330,9 @@ async fn allocate_actor(
 				);
 
 				let mut highest_version = None;
+				let mut candidates = Vec::with_capacity(actor_allocation_candidate_sample_size);
 
+				// Select valid runner candidates for allocation
 				loop {
 					let Some(entry) = stream.try_next().await? else {
 						break;
@@ -255,11 +356,26 @@ async fn allocate_actor(
 						break;
 					}
 
-					// Scan by last ping
+					// Ignore runners without valid ping
 					if old_runner_alloc_key.last_ping_ts < ping_threshold_ts {
 						continue;
 					}
 
+					candidates.push((old_runner_alloc_key, old_runner_alloc_key_data));
+
+					// Max candidate size reached
+					if candidates.len() >= actor_allocation_candidate_sample_size {
+						break;
+					}
+				}
+
+				if !candidates.is_empty() {
+					// Select a candidate at random, weighted by remaining slots
+					let (old_runner_alloc_key, old_runner_alloc_key_data) = candidates
+						.choose_weighted(&mut rand::thread_rng(), |(key, _)| {
+							key.remaining_millislots
+						})?;
+
 					// Add read conflict only for this key
 					tx.add_conflict_key(&old_runner_alloc_key, ConflictRangeType::Read)?;
 
@@ -285,6 +401,7 @@ async fn allocate_actor(
 							workflow_id: old_runner_alloc_key_data.workflow_id,
 							remaining_slots: new_remaining_slots,
 							total_slots: old_runner_alloc_key_data.total_slots,
+							protocol_version: old_runner_alloc_key_data.protocol_version,
 						},
 					)?;
 
@@ -312,22 +429,26 @@ async fn allocate_actor(
 					// Set actor as not sleeping
 					tx.delete(&keys::actor::SleepTsKey::new(input.actor_id));
 
-					return Ok((
-						for_serverless,
-						AllocateActorOutput::Allocated {
+					return Ok(AllocateActorOutputV2 {
+						serverless: for_serverless,
+						status: AllocateActorStatus::Allocated {
 							runner_id: old_runner_alloc_key.runner_id,
 							runner_workflow_id: old_runner_alloc_key_data.workflow_id,
+							runner_protocol_version: Some(
+								old_runner_alloc_key_data.protocol_version,
+							),
 						},
-					));
+					});
 				}
 			}
 
 			// At this point in the txn there is no availability
 
-			match (crash_policy, input.force_allocate, has_valid_serverless) {
-				(CrashPolicy::Sleep, false, false) => {
-					Ok((for_serverless, AllocateActorOutput::Sleep))
-				}
+			match (crash_policy, input.force_allocate, for_serverless) {
+				(CrashPolicy::Sleep, false, false) => Ok(AllocateActorOutputV2 {
+					serverless: false,
+					status: AllocateActorStatus::Sleep,
+				}),
 				// Write the actor to the alloc queue to wait
 				_ => {
 					let pending_allocation_ts = util::timestamp::now();
@@ -345,12 +466,12 @@ async fn allocate_actor(
 						input.generation,
 					)?;
 
-					Ok((
-						for_serverless,
-						AllocateActorOutput::Pending {
+					Ok(AllocateActorOutputV2 {
+						serverless: for_serverless,
+						status: AllocateActorStatus::Pending {
 							pending_allocation_ts,
 						},
-					))
+					})
 				}
 			}
 		})
@@ -358,38 +479,58 @@ async fn allocate_actor(
 		.await?;
 
 	let dt = start_instant.elapsed().as_secs_f64();
-	metrics::ACTOR_ALLOCATE_DURATION.record(
-		dt,
-		&[KeyValue::new(
-			"did_reserve",
-			matches!(res, AllocateActorOutput::Allocated { .. }).to_string(),
-		)],
-	);
-
-	state.for_serverless = for_serverless;
-	state.allocated_serverless_slot = for_serverless;
-
-	match &res {
-		AllocateActorOutput::Allocated {
+	metrics::ACTOR_ALLOCATE_DURATION
+		.with_label_values(&[
+			if res.serverless {
+				"serverless"
+			} else {
+				"serverful"
+			},
+			match res.status {
+				AllocateActorStatus::Allocated { .. } => "allocated",
+				AllocateActorStatus::Pending { .. } => "pending",
+				AllocateActorStatus::Sleep { .. } => "sleep",
+				AllocateActorStatus::MigrateToV2 => bail!("should not be migrate_to_v2"),
+			},
+		])
+		.observe(dt);
+
+	state.for_serverless = res.serverless;
+	state.allocated_serverless_slot = res.serverless;
+	state.reschedule_ts = None;
+
+	match &res.status {
+		AllocateActorStatus::Allocated {
 			runner_id,
 			runner_workflow_id,
+			..
 		} => {
 			state.sleep_ts = None;
 			state.pending_allocation_ts = None;
+			state.failure_reason = None;
 			state.runner_id = Some(*runner_id);
 			state.runner_workflow_id = Some(*runner_workflow_id);
 		}
-		AllocateActorOutput::Pending {
+		AllocateActorStatus::Pending {
 			pending_allocation_ts,
+			..
 		} => {
-			tracing::warn!(
+			tracing::debug!(
 				actor_id=?input.actor_id,
 				"failed to allocate (no availability), waiting for allocation",
 			);
 
 			state.pending_allocation_ts = Some(*pending_allocation_ts);
+			if state.failure_reason.is_none() {
+				state.failure_reason = Some(super::FailureReason::NoCapacity);
+			}
+		}
+		AllocateActorStatus::Sleep => {
+			if state.failure_reason.is_none() {
+				state.failure_reason = Some(super::FailureReason::NoCapacity);
+			}
 		}
-		AllocateActorOutput::Sleep => {}
+		AllocateActorStatus::MigrateToV2 => bail!("should not be migrate_to_v2"),
 	}
 
 	Ok(res)
@@ -406,6 +547,8 @@ pub async fn set_not_connectable(ctx: &ActivityCtx, input: &SetNotConnectableInp
 
 	ctx.udb()?
 		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
 			let connectable_key = keys::actor::ConnectableKey::new(input.actor_id);
 			tx.clear(&keys::subspace().pack(&connectable_key));
 
@@ -432,10 +575,10 @@ pub struct DeallocateOutput {
 #[activity(Deallocate)]
 pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<DeallocateOutput> {
 	let mut state = ctx.state::<State>()?;
-	let runner_name_selector = &state.runner_name_selector;
 	let namespace_id = state.namespace_id;
+	let runner_name_selector = &state.runner_name_selector;
 	let runner_id = state.runner_id;
-	let for_serverless = state.for_serverless;
+	let allocated_serverless_slot = state.allocated_serverless_slot;
 
 	ctx.udb()?
 		.run(|tx| async move {
@@ -443,18 +586,15 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<De
 
 			tx.delete(&keys::actor::ConnectableKey::new(input.actor_id));
 
-			// Only clear slot if we have a runner id
-			if let Some(runner_id) = runner_id {
-				destroy::clear_slot(
-					input.actor_id,
-					namespace_id,
-					runner_name_selector,
-					runner_id,
-					for_serverless,
-					&tx,
-				)
-				.await?;
-			}
+			destroy::clear_slot(
+				input.actor_id,
+				namespace_id,
+				runner_name_selector,
+				runner_id,
+				allocated_serverless_slot,
+				&tx,
+			)
+			.await?;
 
 			Ok(())
 		})
@@ -464,6 +604,7 @@ pub async fn deallocate(ctx: &ActivityCtx, input: &DeallocateInput) -> Result<De
 	state.connectable_ts = None;
 	state.runner_id = None;
 	state.runner_workflow_id = None;
+	state.runner_state = None;
 	// Slot was cleared by the above txn
 	state.allocated_serverless_slot = false;
 
@@ -479,6 +620,9 @@ pub enum AllocationOverride {
 	None,
 	/// Forces actors with CrashPolicy::Sleep to pend instead of sleep.
 	DontSleep { pending_timeout: Option<i64> },
+	/// If an allocation results in pending, it will be put to sleep if it is not allocated after this
+	/// timeout.
+	PendingTimeout { pending_timeout: i64 },
 }
 
 #[derive(Debug)]
@@ -486,9 +630,11 @@ pub enum SpawnActorOutput {
 	Allocated {
 		runner_id: Id,
 		runner_workflow_id: Id,
+		runner_protocol_version: u16,
 	},
 	Sleep,
 	Destroy,
+	MigrateToV2,
 }
 
 /// Wrapper around `allocate_actor` that handles pending state.
@@ -499,76 +645,179 @@ pub async fn spawn_actor(
 	allocation_override: AllocationOverride,
 ) -> Result<SpawnActorOutput> {
 	// Attempt allocation
-	let allocate_res = ctx
-		.activity(AllocateActorInput {
-			actor_id: input.actor_id,
-			generation,
-			force_allocate: matches!(&allocation_override, AllocationOverride::DontSleep { .. }),
-		})
-		.await?;
+	let allocate_res: AllocateActorOutputV2 = match ctx.check_version(2).await? {
+		1 => ctx
+			.activity(AllocateActorInputV1 {
+				actor_id: input.actor_id,
+				generation,
+				force_allocate: matches!(
+					&allocation_override,
+					AllocationOverride::DontSleep { .. }
+				),
+			})
+			.await?
+			.into(),
+		_latest => {
+			ctx.v(2)
+				.activity(AllocateActorInputV2 {
+					actor_id: input.actor_id,
+					generation,
+					force_allocate: matches!(
+						&allocation_override,
+						AllocationOverride::DontSleep { .. }
+					),
+				})
+				.await?
+		}
+	};
 
-	match allocate_res {
-		AllocateActorOutput::Allocated {
+	match allocate_res.status {
+		AllocateActorStatus::Allocated {
 			runner_id,
 			runner_workflow_id,
+			runner_protocol_version,
 		} => {
-			// Bump the autoscaler so it can scale up
-			ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
-				.send()
+			let runner_protocol_version = runner_protocol_version.unwrap_or(PROTOCOL_MK1_VERSION);
+
+			ctx.removed::<Message<super::BumpServerlessAutoscalerStub>>()
 				.await?;
 
-			ctx.signal(crate::workflows::runner::Command {
-				inner: protocol::Command::CommandStartActor(protocol::CommandStartActor {
-					actor_id: input.actor_id.to_string(),
+			// Bump the pool so it can scale up
+			if allocate_res.serverless {
+				let 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", input.runner_name_selector.clone())
+					.send()
+					.await;
+
+				if let Some(WorkflowError::WorkflowNotFound) = 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=%input.runner_name_selector,
+						"serverless pool workflow not found, respective runner config likely deleted"
+					);
+				} else {
+					res?;
+				}
+			}
+
+			if protocol::is_mk2(runner_protocol_version) {
+				ctx.activity(InsertAndSendCommandsInput {
+					actor_id: input.actor_id,
 					generation,
-					config: protocol::ActorConfig {
-						name: input.name.clone(),
-						key: input.key.clone(),
-						// HACK: We should not use dynamic timestamp here, but we don't validate if signal data
-						// changes (like activity inputs) so this is fine for now.
-						create_ts: util::timestamp::now(),
-						input: input
-							.input
-							.as_ref()
-							.map(|x| BASE64_STANDARD.decode(x))
-							.transpose()?,
-					},
-					// Empty because request ids are ephemeral. This is intercepted by guard and
-					// populated before it reaches the runner
-					hibernating_requests: Vec::new(),
-				}),
-			})
-			.to_workflow_id(runner_workflow_id)
-			.send()
-			.await?;
+					runner_id,
+					commands: vec![protocol::mk2::Command::CommandStartActor(
+						protocol::mk2::CommandStartActor {
+							config: protocol::mk2::ActorConfig {
+								name: input.name.clone(),
+								key: input.key.clone(),
+								// HACK: We should not use dynamic timestamp here, but we don't validate if signal data
+								// changes (like activity inputs) so this is fine for now.
+								create_ts: util::timestamp::now(),
+								input: input
+									.input
+									.as_ref()
+									.and_then(|x| BASE64_STANDARD.decode(x).ok()),
+							},
+							// Empty because request ids are ephemeral. This is intercepted by guard and
+							// populated before it reaches the runner
+							hibernating_requests: Vec::new(),
+						},
+					)],
+				})
+				.await?;
+			} else {
+				ctx.signal(crate::workflows::runner::Command {
+					inner: protocol::Command::CommandStartActor(protocol::CommandStartActor {
+						actor_id: input.actor_id.to_string(),
+						generation,
+						config: protocol::ActorConfig {
+							name: input.name.clone(),
+							key: input.key.clone(),
+							// HACK: We should not use dynamic timestamp here, but we don't validate if signal data
+							// changes (like activity inputs) so this is fine for now.
+							create_ts: util::timestamp::now(),
+							input: input
+								.input
+								.as_ref()
+								.map(|x| BASE64_STANDARD.decode(x))
+								.transpose()?,
+						},
+						// Empty because request ids are ephemeral. This is intercepted by guard and
+						// populated before it reaches the runner
+						hibernating_requests: Vec::new(),
+					}),
+				})
+				.to_workflow_id(runner_workflow_id)
+				.send()
+				.await?;
+			}
 
 			Ok(SpawnActorOutput::Allocated {
 				runner_id,
 				runner_workflow_id,
+				runner_protocol_version,
 			})
 		}
-		AllocateActorOutput::Pending {
+		AllocateActorStatus::Pending {
 			pending_allocation_ts,
 		} => {
-			// Bump the autoscaler so it can scale up
-			ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
-				.send()
+			ctx.removed::<Message<super::BumpServerlessAutoscalerStub>>()
 				.await?;
 
-			let signal = if let AllocationOverride::DontSleep {
-				pending_timeout: Some(timeout),
-			} = allocation_override
-			{
-				ctx.listen_with_timeout::<PendingAllocation>(timeout)
-					.await?
-			} else {
-				Some(ctx.listen::<PendingAllocation>().await?)
+			// Bump the pool so it can scale up
+			if allocate_res.serverless {
+				let 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", input.runner_name_selector.clone())
+					.send()
+					.await;
+
+				if let Some(WorkflowError::WorkflowNotFound) = 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=%input.runner_name_selector,
+						"serverless pool workflow not found, respective runner config likely deleted"
+					);
+				} else {
+					res?;
+				}
+			}
+
+			let signal = match allocation_override {
+				AllocationOverride::DontSleep {
+					pending_timeout: Some(timeout),
+				}
+				| AllocationOverride::PendingTimeout {
+					pending_timeout: timeout,
+				} => {
+					ctx.listen_with_timeout::<PendingAllocation>(timeout)
+						.await?
+				}
+				_ => Some(ctx.listen::<PendingAllocation>().await?),
 			};
 
 			// If allocation fails, the allocate txn already inserted this actor into the queue. Now we wait for
 			// an `Allocate` signal
 			match signal {
 				Some(PendingAllocation::Allocate(sig)) => {
+					let runner_protocol_version =
+						sig.runner_protocol_version.unwrap_or(PROTOCOL_MK1_VERSION);
+
 					ctx.activity(UpdateRunnerInput {
 						actor_id: input.actor_id,
 						runner_id: sig.runner_id,
@@ -576,32 +825,61 @@ pub async fn spawn_actor(
 					})
 					.await?;
 
-					ctx.signal(crate::workflows::runner::Command {
-						inner: protocol::Command::CommandStartActor(protocol::CommandStartActor {
-							actor_id: input.actor_id.to_string(),
+					if protocol::is_mk2(runner_protocol_version) {
+						ctx.activity(InsertAndSendCommandsInput {
+							actor_id: input.actor_id,
 							generation,
-							config: protocol::ActorConfig {
-								name: input.name.clone(),
-								key: input.key.clone(),
-								create_ts: util::timestamp::now(),
-								input: input
-									.input
-									.as_ref()
-									.map(|x| BASE64_STANDARD.decode(x))
-									.transpose()?,
-							},
-							// Empty because request ids are ephemeral. This is intercepted by guard and
-							// populated before it reaches the runner
-							hibernating_requests: Vec::new(),
-						}),
-					})
-					.to_workflow_id(sig.runner_workflow_id)
-					.send()
-					.await?;
+							runner_id: sig.runner_id,
+							commands: vec![protocol::mk2::Command::CommandStartActor(
+								protocol::mk2::CommandStartActor {
+									config: protocol::mk2::ActorConfig {
+										name: input.name.clone(),
+										key: input.key.clone(),
+										create_ts: util::timestamp::now(),
+										input: input
+											.input
+											.as_ref()
+											.map(|x| BASE64_STANDARD.decode(x))
+											.transpose()?,
+									},
+									// Empty because request ids are ephemeral. This is intercepted by guard and
+									// populated before it reaches the runner
+									hibernating_requests: Vec::new(),
+								},
+							)],
+						})
+						.await?;
+					} else {
+						ctx.signal(crate::workflows::runner::Command {
+							inner: protocol::Command::CommandStartActor(
+								protocol::CommandStartActor {
+									actor_id: input.actor_id.to_string(),
+									generation,
+									config: protocol::ActorConfig {
+										name: input.name.clone(),
+										key: input.key.clone(),
+										create_ts: util::timestamp::now(),
+										input: input
+											.input
+											.as_ref()
+											.map(|x| BASE64_STANDARD.decode(x))
+											.transpose()?,
+									},
+									// Empty because request ids are ephemeral. This is intercepted by guard and
+									// populated before it reaches the runner
+									hibernating_requests: Vec::new(),
+								},
+							),
+						})
+						.to_workflow_id(sig.runner_workflow_id)
+						.send()
+						.await?;
+					}
 
 					Ok(SpawnActorOutput::Allocated {
 						runner_id: sig.runner_id,
 						runner_workflow_id: sig.runner_workflow_id,
+						runner_protocol_version,
 					})
 				}
 				Some(PendingAllocation::Destroy(_)) => {
@@ -628,6 +906,31 @@ pub async fn spawn_actor(
 						})
 						.await?;
 					}
+					// Bump the pool so it can scale down
+					else if allocate_res.serverless {
+						let 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", input.runner_name_selector.clone())
+							.send()
+							.await;
+
+						if let Some(WorkflowError::WorkflowNotFound) = 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=%input.runner_name_selector,
+								"serverless pool workflow not found, respective runner config likely deleted"
+							);
+						} else {
+							res?;
+						}
+					}
 
 					Ok(SpawnActorOutput::Destroy)
 				}
@@ -647,6 +950,8 @@ pub async fn spawn_actor(
 					// wait for the allocated signal to prevent a race condition.
 					if !cleared {
 						let sig = ctx.listen::<Allocate>().await?;
+						let runner_protocol_version =
+							sig.runner_protocol_version.unwrap_or(PROTOCOL_MK1_VERSION);
 
 						ctx.activity(UpdateRunnerInput {
 							actor_id: input.actor_id,
@@ -655,42 +960,95 @@ pub async fn spawn_actor(
 						})
 						.await?;
 
-						ctx.signal(crate::workflows::runner::Command {
-							inner: protocol::Command::CommandStartActor(
-								protocol::CommandStartActor {
-									actor_id: input.actor_id.to_string(),
-									generation,
-									config: protocol::ActorConfig {
-										name: input.name.clone(),
-										key: input.key.clone(),
-										create_ts: util::timestamp::now(),
-										input: input
-											.input
-											.as_ref()
-											.map(|x| BASE64_STANDARD.decode(x))
-											.transpose()?,
+						if protocol::is_mk2(runner_protocol_version) {
+							ctx.activity(InsertAndSendCommandsInput {
+								actor_id: input.actor_id,
+								generation,
+								runner_id: sig.runner_id,
+								commands: vec![protocol::mk2::Command::CommandStartActor(
+									protocol::mk2::CommandStartActor {
+										config: protocol::mk2::ActorConfig {
+											name: input.name.clone(),
+											key: input.key.clone(),
+											create_ts: util::timestamp::now(),
+											input: input
+												.input
+												.as_ref()
+												.map(|x| BASE64_STANDARD.decode(x))
+												.transpose()?,
+										},
+										// Empty because request ids are ephemeral. This is intercepted by guard and
+										// populated before it reaches the runner
+										hibernating_requests: Vec::new(),
 									},
-									// Empty because request ids are ephemeral. This is intercepted by guard and
-									// populated before it reaches the runner
-									hibernating_requests: Vec::new(),
-								},
-							),
-						})
-						.to_workflow_id(sig.runner_workflow_id)
-						.send()
-						.await?;
+								)],
+							})
+							.await?;
+						} else {
+							ctx.signal(crate::workflows::runner::Command {
+								inner: protocol::Command::CommandStartActor(
+									protocol::CommandStartActor {
+										actor_id: input.actor_id.to_string(),
+										generation,
+										config: protocol::ActorConfig {
+											name: input.name.clone(),
+											key: input.key.clone(),
+											create_ts: util::timestamp::now(),
+											input: input
+												.input
+												.as_ref()
+												.map(|x| BASE64_STANDARD.decode(x))
+												.transpose()?,
+										},
+										// Empty because request ids are ephemeral. This is intercepted by guard and
+										// populated before it reaches the runner
+										hibernating_requests: Vec::new(),
+									},
+								),
+							})
+							.to_workflow_id(sig.runner_workflow_id)
+							.send()
+							.await?;
+						}
 
 						Ok(SpawnActorOutput::Allocated {
 							runner_id: sig.runner_id,
 							runner_workflow_id: sig.runner_workflow_id,
+							runner_protocol_version,
 						})
 					} else {
+						// Bump the pool so it can scale down
+						if allocate_res.serverless {
+							let 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", input.runner_name_selector.clone())
+								.send()
+								.await;
+
+							if let Some(WorkflowError::WorkflowNotFound) =
+								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=%input.runner_name_selector,
+									"serverless pool workflow not found, respective runner config likely deleted"
+								);
+							} else {
+								res?;
+							}
+						}
+
 						Ok(SpawnActorOutput::Sleep)
 					}
 				}
 			}
 		}
-		AllocateActorOutput::Sleep => Ok(SpawnActorOutput::Sleep),
+		AllocateActorStatus::Sleep => Ok(SpawnActorOutput::Sleep),
+		AllocateActorStatus::MigrateToV2 => Ok(SpawnActorOutput::MigrateToV2),
 	}
 }
 
@@ -700,6 +1058,7 @@ pub async fn reschedule_actor(
 	ctx: &mut WorkflowCtx,
 	input: &Input,
 	state: &mut LifecycleState,
+	metrics_workflow_id: Id,
 	allocation_override: AllocationOverride,
 ) -> Result<SpawnActorOutput> {
 	tracing::debug!(actor_id=?input.actor_id, "rescheduling actor");
@@ -747,15 +1106,26 @@ pub async fn reschedule_actor(
 	if let SpawnActorOutput::Allocated {
 		runner_id,
 		runner_workflow_id,
+		runner_protocol_version,
 	} = &spawn_res
 	{
 		state.generation = next_generation;
 		state.runner_id = Some(*runner_id);
 		state.runner_workflow_id = Some(*runner_workflow_id);
+		state.runner_protocol_version = Some(*runner_protocol_version);
 
 		// Reset gc timeout once allocated
 		state.gc_timeout_ts =
 			Some(util::timestamp::now() + ctx.config().pegboard().actor_start_threshold());
+
+		// Resume periodic metrics workflow
+		ctx.v(2)
+			.signal(super::metrics::Resume {
+				ts: util::timestamp::now(),
+			})
+			.to_workflow_id(metrics_workflow_id)
+			.send()
+			.await?;
 	}
 
 	Ok(spawn_res)
@@ -774,27 +1144,52 @@ pub async fn clear_pending_allocation(
 	ctx: &ActivityCtx,
 	input: &ClearPendingAllocationInput,
 ) -> Result<bool> {
+	let mut state = ctx.state::<State>()?;
+
+	let allocated_serverless_slot = state.allocated_serverless_slot;
+
 	// Clear self from alloc queue
 	let cleared = ctx
 		.udb()?
 		.run(|tx| async move {
-			let pending_alloc_key =
-				keys::subspace().pack(&keys::ns::PendingActorByRunnerNameSelectorKey::new(
-					input.namespace_id,
-					input.runner_name_selector.clone(),
-					input.pending_allocation_ts,
-					input.actor_id,
-				));
-
-			let exists = tx.get(&pending_alloc_key, Serializable).await?.is_some();
+			let tx = tx.with_subspace(keys::subspace());
 
-			tx.clear(&pending_alloc_key);
+			let pending_alloc_key = keys::ns::PendingActorByRunnerNameSelectorKey::new(
+				input.namespace_id,
+				input.runner_name_selector.clone(),
+				input.pending_allocation_ts,
+				input.actor_id,
+			);
+			let exists = tx.exists(&pending_alloc_key, Serializable).await?;
+
+			if exists {
+				tx.delete(&pending_alloc_key);
+
+				// If the pending actor key still exists, we must clear its desired slot because after this
+				// activity the actor will go to sleep or be destroyed. We don't clear the slot if the key
+				// doesn't exist because the actor may either be allocated or destroyed.
+				if allocated_serverless_slot {
+					tx.atomic_op(
+						&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new(
+							input.namespace_id,
+							input.runner_name_selector.clone(),
+						),
+						&(-1i64).to_le_bytes(),
+						MutationType::Add,
+					);
+				}
+			}
 
 			Ok(exists)
 		})
 		.custom_instrument(tracing::info_span!("actor_clear_pending_alloc_tx"))
 		.await?;
 
+	// Only mark allocated_serverless_slot as false if it was allocated before and cleared now
+	if allocated_serverless_slot && cleared {
+		state.allocated_serverless_slot = false;
+	}
+
 	Ok(cleared)
 }
 
@@ -836,14 +1231,17 @@ pub struct SetStartedInput {
 #[activity(SetStarted)]
 pub async fn set_started(ctx: &ActivityCtx, input: &SetStartedInput) -> Result<()> {
 	let mut state = ctx.state::<State>()?;
+	let now = util::timestamp::now();
 
 	if state.start_ts.is_none() {
-		state.start_ts = Some(util::timestamp::now());
+		state.start_ts = Some(now);
 	}
-	state.connectable_ts = Some(util::timestamp::now());
+	state.connectable_ts = Some(now);
 
 	ctx.udb()?
 		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
 			let connectable_key = keys::actor::ConnectableKey::new(input.actor_id);
 			tx.set(
 				&keys::subspace().pack(&connectable_key),
@@ -866,9 +1264,9 @@ pub struct SetSleepingInput {
 #[activity(SetSleeping)]
 pub async fn set_sleeping(ctx: &ActivityCtx, input: &SetSleepingInput) -> Result<()> {
 	let mut state = ctx.state::<State>()?;
-	let sleep_ts = util::timestamp::now();
+	let now = util::timestamp::now();
 
-	state.sleep_ts = Some(sleep_ts);
+	state.sleep_ts = Some(now);
 	state.connectable_ts = None;
 
 	ctx.udb()?
@@ -877,8 +1275,7 @@ pub async fn set_sleeping(ctx: &ActivityCtx, input: &SetSleepingInput) -> Result
 
 			// Make not connectable
 			tx.delete(&keys::actor::ConnectableKey::new(input.actor_id));
-
-			tx.write(&keys::actor::SleepTsKey::new(input.actor_id), sleep_ts)?;
+			tx.write(&keys::actor::SleepTsKey::new(input.actor_id), now)?;
 
 			Ok(())
 		})
@@ -907,3 +1304,179 @@ fn reschedule_backoff(
 ) -> util::backoff::Backoff {
 	util::backoff::Backoff::new_at(max_exponent, None, base_retry_timeout, 500, retry_count)
 }
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct InsertAndSendCommandsInput {
+	pub actor_id: Id,
+	pub generation: u32,
+	pub runner_id: Id,
+	pub commands: Vec<protocol::mk2::Command>,
+}
+
+#[activity(InsertAndSendCommands)]
+pub async fn insert_and_send_commands(
+	ctx: &ActivityCtx,
+	input: &InsertAndSendCommandsInput,
+) -> Result<()> {
+	let mut state = ctx.state::<State>()?;
+
+	let runner_state = state.runner_state.get_or_insert_default();
+	let old_last_command_idx = runner_state.last_command_idx;
+	runner_state.last_command_idx += input.commands.len() as i64;
+
+	// This does not have to be part of its own activity because the txn is idempotent
+	let last_command_idx = runner_state.last_command_idx;
+	ctx.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			tx.write(
+				&keys::runner::ActorLastCommandIdxKey::new(
+					input.runner_id,
+					input.actor_id,
+					input.generation,
+				),
+				last_command_idx,
+			)?;
+
+			for (i, command) in input.commands.iter().enumerate() {
+				tx.write(
+					&keys::runner::ActorCommandKey::new(
+						input.runner_id,
+						input.actor_id,
+						input.generation,
+						old_last_command_idx + i as i64 + 1,
+					),
+					match command {
+						protocol::mk2::Command::CommandStartActor(x) => {
+							protocol::mk2::ActorCommandKeyData::CommandStartActor(x.clone())
+						}
+						protocol::mk2::Command::CommandStopActor => {
+							protocol::mk2::ActorCommandKeyData::CommandStopActor
+						}
+					},
+				)?;
+			}
+
+			Ok(())
+		})
+		.await?;
+
+	let receiver_subject =
+		crate::pubsub_subjects::RunnerReceiverSubject::new(input.runner_id).to_string();
+
+	let message_serialized =
+		versioned::ToRunnerMk2::wrap_latest(protocol::mk2::ToRunner::ToClientCommands(
+			input
+				.commands
+				.iter()
+				.enumerate()
+				.map(|(i, command)| protocol::mk2::CommandWrapper {
+					checkpoint: protocol::mk2::ActorCheckpoint {
+						actor_id: input.actor_id.to_string(),
+						generation: input.generation,
+						index: old_last_command_idx + i as i64 + 1,
+					},
+					inner: command.clone(),
+				})
+				.collect(),
+		))
+		.serialize_with_embedded_version(PROTOCOL_MK2_VERSION)?;
+
+	ctx.ups()?
+		.publish(&receiver_subject, &message_serialized, PublishOpts::one())
+		.await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct SendMessagesToRunnerInput {
+	pub runner_id: Id,
+	pub messages: Vec<protocol::mk2::ToRunner>,
+}
+
+#[activity(SendMessagesToRunner)]
+pub async fn send_messages_to_runner(
+	ctx: &ActivityCtx,
+	input: &SendMessagesToRunnerInput,
+) -> Result<()> {
+	let receiver_subject =
+		crate::pubsub_subjects::RunnerReceiverSubject::new(input.runner_id).to_string();
+
+	for message in &input.messages {
+		let message_serialized = versioned::ToRunnerMk2::wrap_latest(message.clone())
+			.serialize_with_embedded_version(PROTOCOL_MK2_VERSION)?;
+
+		ctx.ups()?
+			.publish(&receiver_subject, &message_serialized, PublishOpts::one())
+			.await?;
+	}
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct CheckRunnersStubInput {}
+
+#[activity(CheckRunnersStub)]
+pub async fn check_runners(ctx: &ActivityCtx, input: &CheckRunnersStubInput) -> Result<()> {
+	unreachable!();
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct SetFailureReasonInput {
+	pub failure_reason: FailureReason,
+}
+
+/// Sets the failure reason on the actor workflow state.
+#[activity(SetFailureReason)]
+pub async fn set_failure_reason(ctx: &ActivityCtx, input: &SetFailureReasonInput) -> Result<()> {
+	let mut state = ctx.state::<State>()?;
+
+	// Runner-related errors are never overwritten, as they represent the root cause of the failure
+	// and should not be masked by subsequent errors like `Crashed`.
+	if let Some(existing) = &state.failure_reason
+		&& existing.is_runner_failure()
+	{
+		tracing::debug!(
+			?existing,
+			new=?input.failure_reason,
+			"preserving existing runner failure error"
+		);
+		return Ok(());
+	}
+
+	state.failure_reason = Some(input.failure_reason.clone());
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct RecordEventMetricsInput {
+	pub namespace_id: Id,
+	pub name: String,
+	pub alarms_set: usize,
+}
+
+#[activity(RecordEventMetrics)]
+pub async fn record_event_metrics(
+	ctx: &ActivityCtx,
+	input: &RecordEventMetricsInput,
+) -> Result<()> {
+	ctx.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				input.namespace_id,
+				namespace::keys::metric::Metric::AlarmsSet(input.name.clone()),
+				input.alarms_set as i64,
+			);
+
+			Ok(())
+		})
+		.await?;
+
+	Ok(())
+}