@rivetkit/engine 1.0.0

package/packages/pegboard/src/ops/actor/hibernating_request/upsert.rs

@@ -0,0 +1,51 @@
+use gas::prelude::*;
+use rivet_runner_protocol as protocol;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug, Default)]
+pub struct Input {
+	pub actor_id: Id,
+	pub gateway_id: protocol::GatewayId,
+	pub request_id: protocol::RequestId,
+}
+
+#[operation]
+pub async fn pegboard_actor_hibernating_request_upsert(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<()> {
+	ctx.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			let last_ping_ts_key =
+				keys::hibernating_request::LastPingTsKey::new(input.gateway_id, input.request_id);
+
+			if let Some(last_ping_ts) = tx.read_opt(&last_ping_ts_key, Serializable).await? {
+				tx.delete(&keys::actor::HibernatingRequestKey::new(
+					input.actor_id,
+					last_ping_ts,
+					input.gateway_id,
+					input.request_id,
+				));
+			}
+
+			let now = util::timestamp::now();
+			tx.write(&last_ping_ts_key, now)?;
+			tx.write(
+				&keys::actor::HibernatingRequestKey::new(
+					input.actor_id,
+					now,
+					input.gateway_id,
+					input.request_id,
+				),
+				(),
+			)?;
+
+			Ok(())
+		})
+		.custom_instrument(tracing::info_span!("hibernating_request_upsert_tx"))
+		.await
+}

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

@@ -0,0 +1,207 @@
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_types::actors::Actor;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug, Default)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub name: String,
+	pub key: Option<String>,
+	pub include_destroyed: bool,
+	pub created_before: Option<i64>,
+	pub limit: usize,
+}
+
+#[derive(Debug)]
+pub struct Output {
+	pub actors: Vec<Actor>,
+}
+
+#[operation]
+pub async fn pegboard_actor_list_for_ns(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	let actors_with_wf_ids = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+			let mut results = Vec::new();
+
+			if let Some(key) = &input.key {
+				let actor_subspace = keys::subspace().subspace(&keys::ns::ActorByKeyKey::subspace(
+					input.namespace_id,
+					input.name.clone(),
+					key.clone(),
+				));
+				let (start, end) = actor_subspace.range();
+
+				let end = if let Some(created_before) = input.created_before {
+					universaldb::utils::end_of_key_range(&tx.pack(
+						&keys::ns::ActorByKeyKey::subspace_with_create_ts(
+							input.namespace_id,
+							input.name.clone(),
+							key.clone(),
+							created_before,
+						),
+					))
+				} else {
+					end
+				};
+
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: StreamingMode::Iterator,
+						reverse: true,
+						..(start, end).into()
+					},
+					// NOTE: Does not have to be serializable because we are listing, stale data does not matter
+					Snapshot,
+				);
+
+				while let Some(entry) = stream.try_next().await? {
+					let (idx_key, data) = tx.read_entry::<keys::ns::ActorByKeyKey>(&entry)?;
+
+					if !data.is_destroyed || input.include_destroyed {
+						results.push((idx_key.actor_id, data.workflow_id));
+
+						if results.len() >= input.limit {
+							break;
+						}
+					}
+				}
+			} else if input.include_destroyed {
+				let actor_subspace = keys::subspace().subspace(&keys::ns::AllActorKey::subspace(
+					input.namespace_id,
+					input.name.clone(),
+				));
+				let (start, end) = actor_subspace.range();
+
+				let end = if let Some(created_before) = input.created_before {
+					universaldb::utils::end_of_key_range(&tx.pack(
+						&keys::ns::AllActorKey::subspace_with_create_ts(
+							input.namespace_id,
+							input.name.clone(),
+							created_before,
+						),
+					))
+				} else {
+					end
+				};
+
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: StreamingMode::Iterator,
+						reverse: true,
+						..(start, end).into()
+					},
+					// NOTE: Does not have to be serializable because we are listing, stale data does not matter
+					Snapshot,
+				);
+
+				while let Some(entry) = stream.try_next().await? {
+					let (idx_key, workflow_id) = tx.read_entry::<keys::ns::AllActorKey>(&entry)?;
+
+					results.push((idx_key.actor_id, workflow_id));
+
+					if results.len() >= input.limit {
+						break;
+					}
+				}
+			} else {
+				let actor_subspace = keys::subspace().subspace(
+					&keys::ns::ActiveActorKey::subspace(input.namespace_id, input.name.clone()),
+				);
+				let (start, end) = actor_subspace.range();
+
+				let end = if let Some(created_before) = input.created_before {
+					universaldb::utils::end_of_key_range(&tx.pack(
+						&keys::ns::ActiveActorKey::subspace_with_create_ts(
+							input.namespace_id,
+							input.name.clone(),
+							created_before,
+						),
+					))
+				} else {
+					end
+				};
+
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: StreamingMode::Iterator,
+						reverse: true,
+						..(start, end).into()
+					},
+					// NOTE: Does not have to be serializable because we are listing, stale data does not matter
+					Snapshot,
+				);
+
+				while let Some(entry) = stream.try_next().await? {
+					let (idx_key, workflow_id) =
+						tx.read_entry::<keys::ns::ActiveActorKey>(&entry)?;
+
+					results.push((idx_key.actor_id, workflow_id));
+
+					if results.len() >= input.limit {
+						break;
+					}
+				}
+			}
+
+			Ok(results)
+		})
+		.custom_instrument(tracing::info_span!("actor_list_tx"))
+		.await?;
+
+	let wfs = ctx
+		.get_workflows(
+			actors_with_wf_ids
+				.iter()
+				.map(|(_, workflow_id)| *workflow_id)
+				.collect(),
+		)
+		.await?;
+
+	let dc_name = ctx.config().dc_name()?.to_string();
+	let mut actors = Vec::with_capacity(wfs.len());
+
+	for (actor_id, workflow_id) in actors_with_wf_ids {
+		let Some(wf) = wfs.iter().find(|wf| wf.workflow_id == workflow_id) else {
+			// Actor not found
+			continue;
+		};
+
+		let actor_state = match wf.parse_state::<Option<crate::workflows::actor::State>>() {
+			Ok(Some(s)) => s,
+			Ok(None) => {
+				// Actor did not initialize state yet
+				continue;
+			}
+			Err(err) => {
+				tracing::error!(?actor_id, ?workflow_id, ?err, "failed to parse wf state");
+				continue;
+			}
+		};
+
+		actors.push(Actor {
+			actor_id,
+			name: actor_state.name.clone(),
+			key: actor_state.key,
+			namespace_id: actor_state.namespace_id,
+			datacenter: dc_name.to_string(),
+			runner_name_selector: actor_state.runner_name_selector,
+			crash_policy: actor_state.crash_policy,
+
+			create_ts: actor_state.create_ts,
+			pending_allocation_ts: actor_state.pending_allocation_ts,
+			start_ts: actor_state.start_ts,
+			sleep_ts: actor_state.sleep_ts,
+			connectable_ts: actor_state.connectable_ts,
+			reschedule_ts: actor_state.reschedule_ts,
+			destroy_ts: actor_state.destroy_ts,
+		});
+	}
+
+	Ok(Output { actors })
+}

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

@@ -0,0 +1,62 @@
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use rivet_data::converted::ActorNameKeyData;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub after_name: Option<String>,
+	pub limit: usize,
+}
+
+#[derive(Debug)]
+pub struct Output {
+	pub names: Vec<(String, ActorNameKeyData)>,
+}
+
+#[operation]
+pub async fn pegboard_actor_list_names(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	let names = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			let actor_name_subspace =
+				keys::subspace().subspace(&keys::ns::ActorNameKey::subspace(input.namespace_id));
+			let (start, end) = actor_name_subspace.range();
+
+			let start = if let Some(name) = &input.after_name {
+				tx.pack(&keys::ns::ActorNameKey::new(
+					input.namespace_id,
+					name.clone(),
+				))
+			} else {
+				start
+			};
+
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					limit: Some(input.limit),
+					..(start, end).into()
+				},
+				// NOTE: This is not Serializable to prevent contention with inserting new names
+				Snapshot,
+			)
+			.map(|res| {
+				let (key, metadata) = tx.read_entry::<keys::ns::ActorNameKey>(&res?)?;
+
+				Ok((key.name, metadata))
+			})
+			.try_collect::<Vec<(_, _)>>()
+			.await
+		})
+		.custom_instrument(tracing::info_span!("actor_list_names_tx"))
+		.await?;
+
+	Ok(Output { names })
+}

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

@@ -0,0 +1,9 @@
+pub mod create;
+pub mod get;
+pub mod get_for_gateway;
+pub mod get_for_key;
+pub mod get_reservation_for_key;
+pub mod get_runner;
+pub mod hibernating_request;
+pub mod list_for_ns;
+pub mod list_names;

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

@@ -0,0 +1,2 @@
+pub mod actor;
+pub mod runner;

package/packages/pegboard/src/ops/runner/find_dc_with_runner.rs

@@ -0,0 +1,222 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use futures_util::{FutureExt, StreamExt, TryFutureExt, stream::FuturesUnordered};
+use gas::prelude::*;
+use rivet_api_types::{runner_configs::list as runner_configs_list, runners::list as runners_list};
+use rivet_api_util::{Method, request_remote_datacenter};
+use rivet_types::runner_configs::RunnerConfigKind;
+use serde::de::DeserializeOwned;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub runner_name: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Output {
+	pub dc_label: Option<u16>,
+}
+
+/// Finds a datacenter that contains a given runner name in a namespace.
+///
+/// This is core to determining which datacenter actors should run in, since actors can only run in
+/// datacenters with supported runners.
+#[operation]
+pub async fn find_dc_with_runner(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	// TODO: We should figure out how to pre-emptively validate this cache so we don't have
+	// "stutters" where every 15s we have a high request duration
+	ctx.cache()
+		.clone()
+		.request()
+		.ttl(15_000)
+		.fetch_one_json(
+			"runner.find_dc_with_runner",
+			(input.namespace_id, input.runner_name.clone()),
+			{
+				move |mut cache, key| async move {
+					let dc_id = find_dc_with_runner_inner(ctx, input).await?;
+
+					if let Some(dc_id) = dc_id {
+						cache.resolve(&key, dc_id);
+					}
+
+					Ok(cache)
+				}
+			},
+		)
+		.await
+		.map(|dc_label| Output { dc_label })
+}
+
+async fn find_dc_with_runner_inner(ctx: &OperationCtx, input: &Input) -> Result<Option<u16>> {
+	// Check if this DC has any non-draining runners
+	let res = ctx
+		.op(super::list_for_ns::Input {
+			namespace_id: input.namespace_id,
+			name: Some(input.runner_name.clone()),
+			include_stopped: false,
+			created_before: None,
+			limit: 16,
+		})
+		.await?;
+	if res
+		.runners
+		.iter()
+		.filter(|runner| runner.drain_ts.is_none())
+		.count()
+		!= 0
+	{
+		return Ok(Some(ctx.config().dc_label()));
+	}
+
+	// Check if a serverless runner config with a max runners > 0 exists
+	let res = ctx
+		.op(namespace::ops::runner_config::get::Input {
+			runners: vec![(input.namespace_id, input.runner_name.clone())],
+			bypass_cache: false,
+		})
+		.await?;
+	if let Some(runner) = res.first() {
+		match &runner.config.kind {
+			RunnerConfigKind::Serverless { max_runners, .. } => {
+				if *max_runners != 0 {
+					return Ok(Some(ctx.config().dc_label()));
+				}
+			}
+			_ => {}
+		}
+	}
+
+	// Get namespace
+	let namespace = ctx
+		.op(namespace::ops::get_global::Input {
+			namespace_ids: vec![input.namespace_id],
+		})
+		.await?
+		.into_iter()
+		.next()
+		.ok_or_else(|| namespace::errors::Namespace::NotFound.build())?;
+
+	// Fanout two requests to all datacenters: runner list, and runner config list (with specific name)
+	let runners_fut =
+		race_request_to_datacenters::<runners_list::ListQuery, runners_list::ListResponse, _>(
+			ctx,
+			"/runners",
+			runners_list::ListQuery {
+				namespace: namespace.name.clone(),
+				name: Some(input.runner_name.clone()),
+				runner_ids: None,
+				include_stopped: Some(false),
+				limit: Some(16),
+				cursor: None,
+			},
+			// Check for non draining runners
+			|res| {
+				res.runners
+					.iter()
+					.filter(|runner| runner.drain_ts.is_none())
+					.count() != 0
+			},
+		)
+		.map(|res| res.map(|x| x.map(|x| x.0)))
+		.boxed();
+
+	let runner_configs_fut = race_request_to_datacenters::<
+		runner_configs_list::ListQuery,
+		runner_configs_list::ListResponse,
+		_,
+	>(
+		ctx,
+		"/runner-configs",
+		runner_configs_list::ListQuery {
+			namespace: namespace.name.clone(),
+			variant: None,
+			runner_names: Some(input.runner_name.clone()),
+			limit: None,
+			cursor: None,
+		},
+		// Check for configs with a max runners > 0
+		|res| {
+			res.runner_configs
+				.iter()
+				.filter(|(_, rc)| match rc.kind {
+					RunnerConfigKind::Serverless { max_runners, .. } => max_runners != 0,
+					_ => false,
+				})
+				.count() != 0
+		},
+	)
+	.map(|res| res.map(|x| x.map(|x| x.0)))
+	.boxed();
+
+	let mut futs = [runners_fut, runner_configs_fut]
+		.into_iter()
+		.collect::<FuturesUnordered<_>>();
+
+	Ok(futs.next().await.transpose()?.flatten())
+}
+
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
+
+/// Helper fn that will send a request to all datacenters and return the response of the first
+/// datacenter that matches `filter`.
+pub async fn race_request_to_datacenters<Q, R, F>(
+	ctx: &OperationCtx,
+	endpoint: &str,
+	query: Q,
+	filter: F,
+) -> Result<Option<(u16, R)>>
+where
+	R: DeserializeOwned + Send + 'static,
+	Q: Serialize + Clone + Send + 'static,
+	F: Fn(&R) -> bool,
+{
+	// Create futures for all dcs except the current
+	let dcs = &ctx.config().topology().datacenters;
+	let mut responses = dcs
+		.iter()
+		.filter(|dc| dc.datacenter_label != ctx.config().dc_label())
+		.map(|dc| {
+			let query = query.clone();
+			async move {
+				tokio::time::timeout(
+					REQUEST_TIMEOUT,
+					// Remote datacenter
+					request_remote_datacenter::<R>(
+						ctx.config(),
+						dc.datacenter_label,
+						&endpoint,
+						Method::GET,
+						Some(&query),
+						Option::<&()>::None,
+					)
+					.map_ok(|x| (dc.datacenter_label, x)),
+				)
+				.await
+			}
+		})
+		.collect::<FuturesUnordered<_>>();
+
+	// Collect responses until we reach quorum or all futures complete
+	while let Some(out) = responses.next().await {
+		match out {
+			std::result::Result::Ok(result) => match result {
+				std::result::Result::Ok((dc_label, response)) => {
+					if filter(&response) {
+						return Ok(Some((dc_label, response)));
+					}
+				}
+				std::result::Result::Err(err) => {
+					tracing::warn!(?err, "received error from replica");
+				}
+			},
+			std::result::Result::Err(err) => {
+				tracing::warn!(?err, "received timeout from replica");
+			}
+		}
+	}
+
+	Ok(None)
+}

package/packages/pegboard/src/ops/runner/get.rs

@@ -0,0 +1,143 @@
+use anyhow::Result;
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_types::runners::Runner;
+use universaldb::options::StreamingMode;
+use universaldb::utils::{FormalChunkedKey, IsolationLevel::*};
+
+use crate::keys;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub runner_ids: Vec<Id>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Output {
+	pub runners: Vec<Runner>,
+}
+
+#[operation]
+pub async fn pegboard_runner_get(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	let dc_name = ctx.config().dc_name()?.to_string();
+
+	let runners = ctx
+		.udb()?
+		.run(|tx| {
+			let dc_name = dc_name.to_string();
+			async move {
+				let mut runners = Vec::new();
+
+				for runner_id in input.runner_ids.clone() {
+					if let Some(runner) = get_inner(&dc_name, &tx, runner_id).await? {
+						runners.push(runner);
+					}
+				}
+
+				Ok(runners)
+			}
+		})
+		.custom_instrument(tracing::info_span!("runner_get_tx"))
+		.await?;
+
+	Ok(Output { runners })
+}
+
+pub(crate) async fn get_inner(
+	dc_name: &str,
+	tx: &universaldb::Transaction,
+	runner_id: Id,
+) -> Result<Option<Runner>> {
+	let tx = tx.with_subspace(keys::subspace());
+
+	// TODO: Make this part of the below try join to reduce round trip count
+	// Check if runner exists by looking for workflow ID
+	if !tx
+		.exists(&keys::runner::WorkflowIdKey::new(runner_id), Serializable)
+		.await?
+	{
+		return Ok(None);
+	}
+
+	let namespace_id_key = keys::runner::NamespaceIdKey::new(runner_id);
+	let name_key = keys::runner::NameKey::new(runner_id);
+	let key_key = keys::runner::KeyKey::new(runner_id);
+	let version_key = keys::runner::VersionKey::new(runner_id);
+	let total_slots_key = keys::runner::TotalSlotsKey::new(runner_id);
+	let remaining_slots_key = keys::runner::RemainingSlotsKey::new(runner_id);
+	let create_ts_key = keys::runner::CreateTsKey::new(runner_id);
+	let connected_ts_key = keys::runner::ConnectedTsKey::new(runner_id);
+	let drain_ts_key = keys::runner::DrainTsKey::new(runner_id);
+	let stop_ts_key = keys::runner::StopTsKey::new(runner_id);
+	let last_ping_ts_key = keys::runner::LastPingTsKey::new(runner_id);
+	let last_rtt_key = keys::runner::LastRttKey::new(runner_id);
+	let metadata_key = keys::runner::MetadataKey::new(runner_id);
+	let metadata_subspace = keys::subspace().subspace(&metadata_key);
+
+	let (
+		namespace_id,
+		name,
+		key,
+		version,
+		total_slots,
+		remaining_slots,
+		create_ts,
+		connected_ts,
+		drain_ts,
+		stop_ts,
+		last_ping_ts,
+		last_rtt,
+		metadata_chunks,
+	) = tokio::try_join!(
+		// NOTE: These are not Serializable because this op is meant for basic information (i.e. data for the
+		// API)
+		tx.read(&namespace_id_key, Snapshot),
+		tx.read(&name_key, Snapshot),
+		tx.read(&key_key, Snapshot),
+		tx.read(&version_key, Snapshot),
+		tx.read(&total_slots_key, Snapshot),
+		tx.read(&remaining_slots_key, Snapshot),
+		tx.read(&create_ts_key, Snapshot),
+		tx.read_opt(&connected_ts_key, Snapshot),
+		tx.read_opt(&drain_ts_key, Snapshot),
+		tx.read_opt(&stop_ts_key, Snapshot),
+		tx.read_opt(&last_ping_ts_key, Snapshot),
+		tx.read_opt(&last_rtt_key, Snapshot),
+		async {
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					..(&metadata_subspace).into()
+				},
+				Snapshot,
+			)
+			.try_collect::<Vec<_>>()
+			.await
+			.map_err(Into::into)
+		},
+	)?;
+
+	let metadata = if metadata_chunks.is_empty() {
+		None
+	} else {
+		Some(metadata_key.combine(metadata_chunks)?.metadata)
+	};
+
+	std::result::Result::Ok(Some(Runner {
+		runner_id,
+		namespace_id,
+		datacenter: dc_name.to_string(),
+		name,
+		key,
+		version,
+		total_slots,
+		remaining_slots,
+		create_ts,
+		last_connected_ts: connected_ts,
+		drain_ts,
+		stop_ts,
+		last_ping_ts: last_ping_ts.unwrap_or_default(),
+		last_rtt: last_rtt.unwrap_or_default(),
+		metadata,
+	}))
+}