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

package/packages/pegboard/src/workflows/runner_pool.rs

@@ -0,0 +1,298 @@
+use std::hash::{DefaultHasher, Hash, Hasher};
+
+use futures_util::FutureExt;
+use gas::prelude::*;
+use rivet_types::{keys, runner_configs::RunnerConfigKind};
+use universaldb::prelude::*;
+
+use super::{runner_pool_error_tracker, runner_pool_metadata_poller, serverless};
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub runner_name: String,
+}
+
+#[derive(Debug, Serialize, Deserialize, Default)]
+struct LifecycleState {
+	runners: Vec<RunnerState>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct RunnerState {
+	receiver_wf_id: Id,
+	details_hash: u64,
+}
+
+#[workflow]
+pub async fn pegboard_runner_pool(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
+	// Exit before starting sidecar workflows if there is no serverless runner config to manage.
+	if matches!(
+		ctx.v(4)
+			.activity(ReadDesiredInput {
+				namespace_id: input.namespace_id,
+				runner_name: input.runner_name.clone(),
+			})
+			.await?,
+		ReadDesiredOutput::Stop
+	) {
+		return Ok(());
+	}
+
+	ctx.v(2)
+		.workflow(runner_pool_error_tracker::Input {
+			namespace_id: input.namespace_id,
+			runner_name: input.runner_name.clone(),
+		})
+		.tag("namespace_id", input.namespace_id)
+		.tag("runner_name", &input.runner_name)
+		.unique()
+		.dispatch()
+		.await?;
+
+	ctx.v(3)
+		.workflow(runner_pool_metadata_poller::Input {
+			namespace_id: input.namespace_id,
+			runner_name: input.runner_name.clone(),
+		})
+		.tag("namespace_id", input.namespace_id)
+		.tag("runner_name", &input.runner_name)
+		.unique()
+		.dispatch()
+		.await?;
+
+	ctx.lupe()
+		.commit_interval(5)
+		.with_state(LifecycleState::default())
+		.run(|ctx, state| {
+			let input = input.clone();
+			async move {
+				// Get desired count -> drain and start counts
+				let ReadDesiredOutput::Desired {
+					desired_count,
+					details_hash,
+				} = ctx.activity(ReadDesiredInput {
+					namespace_id: input.namespace_id,
+					runner_name: input.runner_name.clone(),
+				})
+				.await?
+				else {
+					// Drain all
+					for runner in &state.runners {
+						ctx.signal(serverless::receiver::Drain {})
+							.to_workflow_id(runner.receiver_wf_id)
+							.send()
+							.await?;
+					}
+
+					return Ok(Loop::Break(()));
+				};
+
+				// Remove runners that have an outdated hash. This is done outside of the below draining mechanism
+				// because we drain specific runners, not just a number of runners
+				let (new, outdated) = std::mem::take(&mut state.runners)
+					.into_iter()
+					.partition::<Vec<_>, _>(|r| r.details_hash == details_hash);
+				state.runners = new;
+
+				for runner in outdated {
+					// TODO: Spawn sub wf to process these so this is not blocking the loop
+					ctx.signal(serverless::receiver::Drain {})
+						.to_workflow_id(runner.receiver_wf_id)
+						.send()
+						.await?;
+				}
+
+				let drain_count = state.runners.len().saturating_sub(desired_count);
+				let start_count = desired_count.saturating_sub(state.runners.len());
+
+				// Drain unnecessary runners
+				if drain_count != 0 {
+					// TODO: Implement smart logic of draining runners with the lowest allocated actors
+					let draining_runners =
+						state.runners.iter().take(drain_count).collect::<Vec<_>>();
+
+					// TODO: Spawn sub wf to process these so this is not blocking the loop
+					for runner in draining_runners {
+						ctx.signal(serverless::receiver::Drain {})
+							.to_workflow_id(runner.receiver_wf_id)
+							.send()
+							.await?;
+					}
+				}
+
+				// Dispatch new runner workflows
+				if start_count != 0 {
+					// TODO: Spawn sub wf to process these so this is not blocking the loop
+					for _ in 0..start_count {
+						let receiver_wf_id = ctx
+							.workflow(serverless::receiver::Input {
+								pool_wf_id: ctx.workflow_id(),
+								namespace_id: input.namespace_id,
+								runner_name: input.runner_name.clone(),
+							})
+							.tag("namespace_id", input.namespace_id)
+							.tag("runner_name", input.runner_name.clone())
+							.dispatch()
+							.await?;
+
+						state.runners.push(RunnerState {
+							receiver_wf_id,
+							details_hash,
+						});
+					}
+				}
+
+				// Wait for Bump or serverless signals until we tick again
+				for sig in ctx.listen_n::<Main>(256).await? {
+					match sig {
+						Main::OutboundConnDrainStarted(sig) => {
+							let (new, drain_started) = std::mem::take(&mut state.runners)
+								.into_iter()
+								.partition::<Vec<_>, _>(|r| r.receiver_wf_id != sig.receiver_wf_id);
+							state.runners = new;
+
+							for runner in drain_started {
+								// TODO: Spawn sub wf to process these so this is not blocking the loop
+								ctx.signal(serverless::receiver::Drain {})
+									.to_workflow_id(runner.receiver_wf_id)
+									.send()
+									.await?;
+							}
+						}
+						Main::Bump(bump) => {
+							if bump.endpoint_config_changed {
+								// Forward to metadata poller to trigger immediate metadata fetch
+								ctx.signal(runner_pool_metadata_poller::EndpointConfigChanged {})
+									.to_workflow::<runner_pool_metadata_poller::Workflow>()
+									.tag("namespace_id", input.namespace_id)
+									.tag("runner_name", &input.runner_name)
+									.send()
+									.await?;
+							}
+						}
+					}
+				}
+
+				Ok(Loop::Continue)
+			}
+			.boxed()
+		})
+		.await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct ReadDesiredInput {
+	namespace_id: Id,
+	runner_name: String,
+}
+
+// Should have `#[serde(rename_all = "snake_case")]` but doesn't
+#[derive(Debug, Serialize, Deserialize)]
+enum ReadDesiredOutput {
+	Desired {
+		desired_count: usize,
+		details_hash: u64,
+	},
+	Stop,
+}
+
+#[activity(ReadDesired)]
+async fn read_desired(ctx: &ActivityCtx, input: &ReadDesiredInput) -> Result<ReadDesiredOutput> {
+	let udb_pool = ctx.udb()?;
+	let (runner_config_res, desired_slots) = tokio::try_join!(
+		ctx.op(crate::ops::runner_config::get::Input {
+			runners: vec![(input.namespace_id, input.runner_name.clone())],
+			bypass_cache: false,
+		}),
+		udb_pool.run(|tx| async move {
+			let tx = tx.with_subspace(keys::pegboard::subspace());
+
+			let desired_slots = tx
+				.read_opt(
+					&keys::pegboard::ns::ServerlessDesiredSlotsKey {
+						namespace_id: input.namespace_id,
+						runner_name: input.runner_name.clone(),
+					},
+					Serializable,
+				)
+				.await?;
+
+			Ok(desired_slots.unwrap_or_default())
+		}),
+	)?;
+	let Some(runner_config) = runner_config_res.into_iter().next() else {
+		return Ok(ReadDesiredOutput::Stop);
+	};
+
+	let RunnerConfigKind::Serverless {
+		url,
+		headers,
+
+		slots_per_runner,
+		min_runners,
+		max_runners,
+		runners_margin,
+		..
+	} = runner_config.config.kind
+	else {
+		return Ok(ReadDesiredOutput::Stop);
+	};
+
+	let adjusted_desired_slots = if desired_slots < 0 {
+		tracing::error!(
+			namespace_id=%input.namespace_id,
+			runner_name=%input.runner_name,
+			?desired_slots,
+			"negative desired slots, scaling to 0"
+		);
+		0
+	} else {
+		desired_slots
+	};
+
+	// Won't overflow as these values are all in u32 range
+	let desired_count = (runners_margin
+		+ (adjusted_desired_slots as u32).div_ceil(slots_per_runner))
+	.max(min_runners)
+	.min(max_runners)
+	.min(
+		ctx.config()
+			.pegboard()
+			.pool_desired_max_override
+			.unwrap_or(u32::MAX),
+	)
+	.try_into()?;
+
+	// Compute consistent hash of serverless details
+	let mut hasher = DefaultHasher::new();
+	url.hash(&mut hasher);
+	let mut sorted_headers = headers.iter().collect::<Vec<_>>();
+	sorted_headers.sort();
+	sorted_headers.hash(&mut hasher);
+	let details_hash = hasher.finish();
+
+	Ok(ReadDesiredOutput::Desired {
+		desired_count,
+		details_hash,
+	})
+}
+
+#[signal("pegboard_runner_pool_bump")]
+#[derive(Debug, Default)]
+pub struct Bump {
+	#[serde(default)]
+	pub endpoint_config_changed: bool,
+}
+
+#[signal("pegboard_outbound_conn_drain_started")]
+pub struct OutboundConnDrainStarted {
+	pub receiver_wf_id: Id,
+}
+
+join_signal!(Main {
+	Bump,
+	OutboundConnDrainStarted,
+});

package/packages/pegboard/src/workflows/runner_pool_error_tracker.rs

@@ -0,0 +1,173 @@
+use std::time::Duration;
+
+use gas::prelude::*;
+use rivet_types::actor::RunnerPoolError;
+
+const SIGNAL_DEBOUNCE: Duration = Duration::from_millis(250);
+const SIGNAL_BATCH_SIZE: usize = 1024;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub runner_name: String,
+}
+
+#[derive(Debug, Serialize, Deserialize, Default, Clone)]
+pub struct State {
+	/// Persistent error state - set on error, cleared after consecutive successes.
+	/// Used to track errors during backoff periods when no new requests are made.
+	pub active_error: Option<ActiveError>,
+
+	/// Count of consecutive successes since last error.
+	/// Error is only cleared after reaching the configured threshold.
+	pub consecutive_successes: u32,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ActiveError {
+	pub timestamp: i64,
+	pub error: RunnerPoolError,
+}
+
+#[workflow]
+pub async fn pegboard_runner_pool_error_tracker(
+	ctx: &mut WorkflowCtx,
+	input: &Input,
+) -> Result<()> {
+	tracing::debug!(
+		namespace_id = %input.namespace_id,
+		runner_name = %input.runner_name,
+		"starting error tracker"
+	);
+
+	ctx.activity(InitStateInput {}).await?;
+
+	// Batch receive signals with debounce. This allows us to (a) not require polling if the pool
+	// is idle and has no signals and (b) avoid a hot loop by debouncing signal processing.
+	ctx.lupe()
+		// Txn sizes can quickly get large in this workflow, need to commit loop more often
+		.commit_interval(1)
+		.run(|ctx, _| {
+			Box::pin(async move {
+				// Sleep until we receive a signal
+				let signals_a = ctx.v(2).listen_n::<Main>(SIGNAL_BATCH_SIZE).await?;
+
+				// Debounce rest of signals if we haven't already reached the batch size
+				let remaining_signals = SIGNAL_BATCH_SIZE.saturating_sub(signals_a.len());
+				let signals_b = if remaining_signals > 0 {
+					ctx.listen_n_with_timeout::<Main>(SIGNAL_DEBOUNCE, remaining_signals)
+						.await?
+				} else {
+					Vec::new()
+				};
+
+				let signals_inner = signals_a
+					.into_iter()
+					.chain(signals_b.into_iter())
+					.map(|s| match s {
+						Main::ReportSuccess(x) => MainInner::ReportSuccess(x),
+						Main::ReportError(x) => MainInner::ReportError(x),
+						Main::Shutdown(x) => MainInner::Shutdown(x),
+					})
+					.collect();
+
+				// Process signals
+				let shutdown = ctx
+					.activity(ProcessSignalsInput {
+						signals: signals_inner,
+					})
+					.await?;
+
+				if shutdown {
+					Ok(Loop::Break(()))
+				} else {
+					Ok(Loop::Continue)
+				}
+			})
+		})
+		.await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct InitStateInput {}
+
+#[activity(InitState)]
+pub async fn init_state(ctx: &ActivityCtx, _input: &InitStateInput) -> Result<()> {
+	let mut state = ctx.state::<Option<State>>()?;
+	*state = Some(State::default());
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub struct ProcessSignalsInput {
+	pub signals: Vec<MainInner>,
+}
+
+/// Returns `true` if shutdown signal received.
+#[activity(ProcessSignals)]
+pub async fn process_signals(ctx: &ActivityCtx, input: &ProcessSignalsInput) -> Result<bool> {
+	let mut state = ctx.state::<State>()?;
+	let now = util::timestamp::now();
+
+	for signal in &input.signals {
+		match signal {
+			MainInner::ReportError(report) => {
+				state.active_error = Some(ActiveError {
+					timestamp: now,
+					error: report.error.clone(),
+				});
+				state.consecutive_successes = 0;
+			}
+			MainInner::ReportSuccess(_) => {
+				state.consecutive_successes += 1;
+
+				// Only clear error after threshold reached
+				let threshold = ctx
+					.config()
+					.pegboard()
+					.runner_pool_error_consecutive_successes_to_clear();
+				if state.consecutive_successes >= threshold {
+					if state.active_error.is_some() {
+						tracing::debug!("clearing active error after consecutive successes");
+					}
+					state.active_error = None;
+				}
+			}
+			MainInner::Shutdown(_) => {
+				return Ok(true);
+			}
+		}
+	}
+
+	Ok(false)
+}
+
+#[derive(Debug, Clone, Hash)]
+#[signal("pegboard_runner_pool_error_tracker_report_error")]
+pub struct ReportError {
+	pub error: RunnerPoolError,
+}
+
+#[derive(Debug, Clone, Hash)]
+#[signal("pegboard_runner_pool_error_tracker_report_success")]
+pub struct ReportSuccess {}
+
+#[derive(Debug, Clone, Hash)]
+#[signal("pegboard_runner_pool_error_tracker_shutdown")]
+pub struct Shutdown {}
+
+join_signal!(Main {
+	ReportError,
+	ReportSuccess,
+	Shutdown,
+});
+
+// HACK: Cannot implement `Hash` on `Main`
+#[derive(Debug, Serialize, Deserialize, Hash)]
+pub enum MainInner {
+	ReportError(ReportError),
+	ReportSuccess(ReportSuccess),
+	Shutdown(Shutdown),
+}

package/packages/pegboard/src/workflows/runner_pool_metadata_poller.rs

@@ -0,0 +1,237 @@
+use std::time::Duration;
+
+use futures_util::FutureExt;
+use gas::prelude::*;
+use rivet_types::runner_configs::RunnerConfigKind;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub runner_name: String,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, Default)]
+struct State {
+	last_retry_ts: i64,
+	retry_count: usize,
+}
+
+#[workflow]
+pub async fn pegboard_runner_pool_metadata_poller(
+	ctx: &mut WorkflowCtx,
+	input: &Input,
+) -> Result<()> {
+	ctx.repeat(|ctx| {
+		let input = input.clone();
+		async move {
+			let poll_result = ctx
+				.activity(PollMetadataInput {
+					namespace_id: input.namespace_id,
+					runner_name: input.runner_name.clone(),
+				})
+				.await?;
+
+			match poll_result {
+				PollMetadataOutput::Success {
+					poll_interval,
+					older_runner_workflow_ids,
+				} => {
+					// Send stop signals to older runners
+					if !older_runner_workflow_ids.is_empty() {
+						ctx.loope(older_runner_workflow_ids, |ctx, workflow_ids| {
+							async move {
+								let Some(workflow_id) = workflow_ids.pop() else {
+									return Ok(Loop::Break(()));
+								};
+
+								ctx.signal(crate::workflows::runner2::Stop {
+									reset_actor_rescheduling: false,
+								})
+								.to_workflow_id(workflow_id)
+								.send()
+								.await?;
+
+								Ok(Loop::Continue)
+							}
+							.boxed()
+						})
+						.await?;
+					}
+
+					// Wait at the end of the loop so the first poll happens immediately
+					// and the poll interval is read fresh from the config each iteration.
+					let _ = ctx
+						.listen_with_timeout::<EndpointConfigChanged>(Duration::from_millis(
+							poll_interval,
+						))
+						.await?;
+
+					Ok(Loop::Continue)
+				}
+				PollMetadataOutput::FetchError { poll_interval } => {
+					// Wait before retrying
+					let _ = ctx
+						.listen_with_timeout::<EndpointConfigChanged>(Duration::from_millis(
+							poll_interval,
+						))
+						.await?;
+
+					Ok(Loop::Continue)
+				}
+				PollMetadataOutput::Break => Ok(Loop::Break(())),
+			}
+		}
+		.boxed()
+	})
+	.await?;
+
+	Ok(())
+}
+
+#[derive(Debug, Serialize, Deserialize, Hash)]
+struct PollMetadataInput {
+	namespace_id: Id,
+	runner_name: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+enum PollMetadataOutput {
+	Success {
+		poll_interval: u64,
+		older_runner_workflow_ids: Vec<Id>,
+	},
+	FetchError {
+		poll_interval: u64,
+	},
+	Break,
+}
+
+/// Combined activity that fetches metadata, updates actor names, and drains older versions.
+#[activity(PollMetadata)]
+async fn poll_metadata(ctx: &ActivityCtx, input: &PollMetadataInput) -> Result<PollMetadataOutput> {
+	let mut state = ctx.state::<Option<State>>()?;
+	if state.is_none() {
+		*state = Some(State::default());
+	}
+	let state = state.as_mut().unwrap();
+
+	tracing::debug!(
+		namespace_id = %input.namespace_id,
+		runner_name = %input.runner_name,
+		"polling metadata"
+	);
+
+	// Get runner config
+	let runner_config_res = ctx
+		.op(crate::ops::runner_config::get::Input {
+			runners: vec![(input.namespace_id, input.runner_name.clone())],
+			bypass_cache: true,
+		})
+		.await?;
+
+	let Some(runner_config) = runner_config_res.into_iter().next() else {
+		tracing::debug!(
+			namespace_id = %input.namespace_id,
+			runner_name = %input.runner_name,
+			"runner config not found, stopping metadata poller"
+		);
+		return Ok(PollMetadataOutput::Break);
+	};
+
+	let RunnerConfigKind::Serverless {
+		url,
+		headers,
+		metadata_poll_interval,
+		..
+	} = runner_config.config.kind
+	else {
+		tracing::debug!(
+			namespace_id = %input.namespace_id,
+			runner_name = %input.runner_name,
+			"runner config is not serverless, stopping metadata poller"
+		);
+		return Ok(PollMetadataOutput::Break);
+	};
+
+	// Reset backoff
+	if state.last_retry_ts < util::timestamp::now() - util::duration::minutes(10) {
+		state.retry_count = 0;
+	}
+
+	// Calculate effective poll interval (in milliseconds)
+	let min_poll_interval = ctx.config().pegboard().min_metadata_poll_interval();
+	let default_poll_interval = ctx.config().pegboard().default_metadata_poll_interval();
+	let base_poll_interval = metadata_poll_interval
+		.unwrap_or(default_poll_interval)
+		.max(min_poll_interval);
+	let backoff = util::backoff::Backoff::new_at(
+		8,
+		None,
+		base_poll_interval as usize,
+		500,
+		state.retry_count,
+	);
+	let backoff_poll_interval = backoff.current_duration() as u64;
+
+	let result = ctx
+		.op(crate::ops::runner_config::refresh_metadata::Input {
+			namespace_id: input.namespace_id,
+			runner_name: input.runner_name.clone(),
+			url,
+			headers,
+		})
+		.await?;
+
+	let metadata = match result {
+		Ok(metadata) => metadata,
+		Err(error) => {
+			tracing::warn!(
+				namespace_id = %input.namespace_id,
+				runner_name = %input.runner_name,
+				?error,
+				"failed to fetch metadata, will retry"
+			);
+
+			state.last_retry_ts = util::timestamp::now();
+			state.retry_count += 1;
+
+			return Ok(PollMetadataOutput::FetchError {
+				poll_interval: backoff_poll_interval,
+			});
+		}
+	};
+
+	// Drain older runners if runner_version is set
+	let older_runner_workflow_ids = if let Some(version) = metadata.runner_version {
+		ctx.op(crate::ops::runner::drain::Input {
+			namespace_id: input.namespace_id,
+			name: input.runner_name.clone(),
+			version,
+			// Signals are sent by the workflow directly
+			send_runner_stop_signals: false,
+		})
+		.await?
+		.older_runner_workflow_ids
+	} else {
+		Vec::new()
+	};
+
+	// Drain older envoys if envoy_version is set
+	if let Some(version) = metadata.envoy_version {
+		ctx.op(crate::ops::envoy::drain::Input {
+			namespace_id: input.namespace_id,
+			pool_name: input.runner_name.clone(),
+			version,
+		})
+		.await?;
+	}
+
+	Ok(PollMetadataOutput::Success {
+		poll_interval: base_poll_interval,
+		older_runner_workflow_ids,
+	})
+}
+
+#[signal("pegboard_runner_pool_metadata_poller_endpoint_config_changed")]
+#[derive(Debug)]
+pub struct EndpointConfigChanged {}