@rivetkit/engine 1.0.0

package/packages/namespace/src/ops/get_global.rs

@@ -0,0 +1,58 @@
+use gas::prelude::*;
+use rivet_types::namespaces::Namespace;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_ids: Vec<Id>,
+}
+
+#[operation]
+pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result<Vec<Namespace>> {
+	if ctx.config().is_leader() {
+		ctx.op(super::get_local::Input {
+			namespace_ids: input.namespace_ids.clone(),
+		})
+		.await
+	} else {
+		let leader_dc = ctx.config().leader_dc()?;
+		let client = rivet_pools::reqwest::client().await?;
+
+		ctx.cache()
+			.clone()
+			.request()
+			.fetch_all_json("namespace.get_global", input.namespace_ids.clone(), {
+				let leader_dc = leader_dc.clone();
+				let client = client.clone();
+				move |mut cache, namespace_ids| {
+					let leader_dc = leader_dc.clone();
+					let client = client.clone();
+					async move {
+						let url = leader_dc.peer_url.join(&format!("/namespaces"))?;
+						let res = client
+							.get(url)
+							.query(
+								&namespace_ids
+									.iter()
+									.map(|ns_id| ("namespace_id", ns_id))
+									.collect::<Vec<_>>(),
+							)
+							.send()
+							.await?;
+
+						let res = rivet_api_util::parse_response::<
+							rivet_api_types::namespaces::list::ListResponse,
+						>(res)
+						.await?;
+
+						for ns in res.namespaces {
+							let namespace_id = ns.namespace_id;
+							cache.resolve(&&namespace_id, ns);
+						}
+
+						Ok(cache)
+					}
+				}
+			})
+			.await
+	}
+}

package/packages/namespace/src/ops/get_local.rs

@@ -0,0 +1,69 @@
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use rivet_types::namespaces::Namespace;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::{errors, keys};
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_ids: Vec<Id>,
+}
+
+#[operation]
+pub async fn namespace_get_local(ctx: &OperationCtx, input: &Input) -> Result<Vec<Namespace>> {
+	if !ctx.config().is_leader() {
+		return Err(errors::Namespace::NotLeader.build());
+	}
+
+	let namespaces = ctx
+		.udb()?
+		.run(|tx| async move {
+			futures_util::stream::iter(input.namespace_ids.clone())
+				.map(|namespace_id| {
+					let tx = tx.clone();
+
+					async move { get_inner(namespace_id, &tx).await }
+				})
+				.buffer_unordered(1024)
+				.try_filter_map(|x| std::future::ready(Ok(x)))
+				.try_collect::<Vec<_>>()
+				.await
+		})
+		.custom_instrument(tracing::info_span!("namespace_get_local_tx"))
+		.await?;
+
+	Ok(namespaces)
+}
+
+pub(crate) async fn get_inner(
+	namespace_id: Id,
+	tx: &universaldb::Transaction,
+) -> Result<Option<Namespace>> {
+	let tx = tx.with_subspace(keys::subspace());
+
+	let name_key = keys::NameKey::new(namespace_id);
+	let display_name_key = keys::DisplayNameKey::new(namespace_id);
+	let create_ts_key = keys::CreateTsKey::new(namespace_id);
+
+	let (name, display_name, create_ts) = tokio::try_join!(
+		tx.read_opt(&name_key, Serializable),
+		tx.read_opt(&display_name_key, Serializable),
+		tx.read_opt(&create_ts_key, Serializable),
+	)?;
+
+	// Namespace not found
+	let Some(name) = name else {
+		return Ok(None);
+	};
+
+	let display_name = display_name.context("key should exist")?;
+	let create_ts = create_ts.context("key should exist")?;
+
+	Ok(Some(Namespace {
+		namespace_id,
+		name,
+		display_name,
+		create_ts,
+	}))
+}

package/packages/namespace/src/ops/list.rs

@@ -0,0 +1,69 @@
+use anyhow::Result;
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_types::namespaces::Namespace;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::{errors, keys};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub limit: Option<usize>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Output {
+	pub namespaces: Vec<Namespace>,
+}
+
+#[operation]
+pub async fn namespace_list(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	if !ctx.config().is_leader() {
+		return Err(errors::Namespace::NotLeader.build());
+	}
+
+	let namespaces = ctx
+		.udb()?
+		.run(|tx| async move {
+			let mut namespaces = Vec::new();
+			let limit = input.limit.unwrap_or(1000); // Default limit to 1000
+
+			let mut stream = tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::Iterator,
+					..(&keys::subspace()).into()
+				},
+				Snapshot,
+			);
+
+			let mut seen_namespaces = std::collections::HashSet::new();
+
+			while let Some(entry) = stream.try_next().await? {
+				// Try to unpack as a NameKey
+				if let Ok(name_key) = keys::subspace().unpack::<keys::NameKey>(entry.key()) {
+					let namespace_id = name_key.namespace_id();
+
+					// Skip if we've already seen this namespace
+					if !seen_namespaces.insert(namespace_id) {
+						continue;
+					}
+
+					// Get the full namespace data
+					if let Some(namespace) = super::get_local::get_inner(namespace_id, &tx).await? {
+						namespaces.push(namespace);
+
+						if namespaces.len() >= limit {
+							break;
+						}
+					}
+				}
+			}
+
+			Ok(namespaces)
+		})
+		.custom_instrument(tracing::info_span!("namespace_list_tx"))
+		.await?;
+
+	Ok(Output { namespaces })
+}

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

@@ -0,0 +1,6 @@
+pub mod get_global;
+pub mod get_local;
+pub mod list;
+pub mod resolve_for_name_global;
+pub mod resolve_for_name_local;
+pub mod runner_config;

package/packages/namespace/src/ops/resolve_for_name_global.rs

@@ -0,0 +1,57 @@
+use gas::prelude::*;
+use rivet_types::namespaces::Namespace;
+
+#[derive(Debug)]
+pub struct Input {
+	pub name: String,
+}
+
+#[operation]
+pub async fn namespace_resolve_for_name_global(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<Option<Namespace>> {
+	if ctx.config().is_leader() {
+		ctx.op(crate::ops::resolve_for_name_local::Input {
+			name: input.name.clone(),
+		})
+		.await
+	} else {
+		let leader_dc = ctx.config().leader_dc()?;
+		let client = rivet_pools::reqwest::client().await?;
+
+		ctx.cache()
+			.clone()
+			.request()
+			.fetch_one_json("namespace.resolve_for_name_global", input.name.clone(), {
+				let leader_dc = leader_dc.clone();
+				let client = client.clone();
+				move |mut cache, key| {
+					let leader_dc = leader_dc.clone();
+					let client = client.clone();
+					async move {
+						let url = leader_dc.peer_url.join("/namespaces")?;
+						let res = client
+							.get(url)
+							.query(&[("name", &input.name)])
+							.send()
+							.custom_instrument(tracing::info_span!("namespaces_http_request"))
+							.await?;
+
+						let res = rivet_api_util::parse_response::<
+							rivet_api_types::namespaces::list::ListResponse,
+						>(res)
+						.await?;
+
+						let ns = res.namespaces.into_iter().next();
+
+						cache.resolve(&key, ns);
+
+						Ok(cache)
+					}
+				}
+			})
+			.await
+			.map(|x| x.flatten())
+	}
+}

package/packages/namespace/src/ops/resolve_for_name_local.rs

@@ -0,0 +1,41 @@
+use gas::prelude::*;
+use rivet_types::namespaces::Namespace;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::{errors, keys, ops::get_local::get_inner};
+
+#[derive(Debug)]
+pub struct Input {
+	pub name: String,
+}
+
+#[operation]
+pub async fn namespace_resolve_for_name_local(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<Option<Namespace>> {
+	if !ctx.config().is_leader() {
+		return Err(errors::Namespace::NotLeader.build());
+	}
+
+	ctx.udb()?
+		.run(|tx| {
+			let name = input.name.clone();
+			async move {
+				let tx = tx.with_subspace(keys::subspace());
+
+				let Some(namespace_id) = tx
+					.read_opt(&keys::ByNameKey::new(name.clone()), Serializable)
+					.await?
+				else {
+					// Namespace not found
+					return Ok(None);
+				};
+
+				get_inner(namespace_id, &tx).await
+			}
+		})
+		.custom_instrument(tracing::info_span!("namespace_resolve_for_name_local_tx"))
+		.await
+		.map_err(Into::into)
+}

package/packages/namespace/src/ops/runner_config/delete.rs

@@ -0,0 +1,53 @@
+use gas::prelude::*;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::{keys, utils::runner_config_variant};
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub name: String,
+}
+
+#[operation]
+pub async fn namespace_runner_config_delete(ctx: &OperationCtx, input: &Input) -> Result<()> {
+	let bump_autoscaler = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			// Read existing config to determine variant
+			let runner_config_key =
+				keys::runner_config::DataKey::new(input.namespace_id, input.name.clone());
+
+			let bump_autoscaler =
+				if let Some(config) = tx.read_opt(&runner_config_key, Serializable).await? {
+					tx.delete(&runner_config_key);
+
+					// Clear secondary idx
+					let variant = runner_config_variant(&config);
+					tx.delete(&keys::runner_config::ByVariantKey::new(
+						input.namespace_id,
+						variant,
+						input.name.clone(),
+					));
+
+					config.affects_autoscaler()
+				} else {
+					false
+				};
+
+			Ok(bump_autoscaler)
+		})
+		.custom_instrument(tracing::info_span!("runner_config_delete_tx"))
+		.await?;
+
+	// Bump autoscaler when a serverless config is modified
+	if bump_autoscaler {
+		ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
+			.send()
+			.await?;
+	}
+
+	Ok(())
+}

package/packages/namespace/src/ops/runner_config/get.rs

@@ -0,0 +1,94 @@
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use serde::{Deserialize, Serialize};
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub runners: Vec<(Id, String)>,
+	pub bypass_cache: bool,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct RunnerConfig {
+	pub namespace_id: Id,
+	pub name: String,
+	pub config: rivet_types::runner_configs::RunnerConfig,
+}
+
+#[operation]
+pub async fn namespace_runner_config_get(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<Vec<RunnerConfig>> {
+	if input.bypass_cache {
+		runner_config_get_inner(ctx, input.runners.clone()).await
+	} else {
+		ctx.cache()
+			.clone()
+			.request()
+			// Short TTL for faster updates
+			.ttl(5000)
+			.fetch_all_json("namespace.runner_config.get", input.runners.clone(), {
+				|mut cache, runners| async move {
+					let runner_configs = runner_config_get_inner(ctx, runners).await?;
+
+					for runner_config in runner_configs {
+						cache.resolve(
+							&(runner_config.namespace_id, runner_config.name.clone()),
+							runner_config,
+						);
+					}
+
+					Ok(cache)
+				}
+			})
+			.await
+	}
+}
+
+async fn runner_config_get_inner(
+	ctx: &OperationCtx,
+	runners: Vec<(Id, String)>,
+) -> Result<Vec<RunnerConfig>> {
+	ctx.udb()?
+		.run(|tx| {
+			let runners = runners.clone();
+			async move {
+				futures_util::stream::iter(runners)
+					.map(|(namespace_id, runner_name)| {
+						let tx = tx.clone();
+
+						async move {
+							let tx = tx.with_subspace(keys::subspace());
+
+							let runner_config_key = keys::runner_config::DataKey::new(
+								namespace_id,
+								runner_name.clone(),
+							);
+
+							let Some(runner_config) =
+								tx.read_opt(&runner_config_key, Serializable).await?
+							else {
+								// Runner config not found
+								return Ok(None);
+							};
+
+							Ok(Some(RunnerConfig {
+								namespace_id,
+								name: runner_name,
+								config: runner_config,
+							}))
+						}
+					})
+					.buffer_unordered(1024)
+					.try_filter_map(|x| std::future::ready(Ok(x)))
+					.try_collect::<Vec<_>>()
+					.await
+			}
+		})
+		.custom_instrument(tracing::info_span!("runner_config_get_tx"))
+		.await
+}

package/packages/namespace/src/ops/runner_config/list.rs

@@ -0,0 +1,94 @@
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use rivet_types::keys::namespace::runner_config::RunnerConfigVariant;
+use rivet_types::runner_configs::RunnerConfig;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub variant: Option<RunnerConfigVariant>,
+	pub after_name: Option<String>,
+	pub limit: usize,
+}
+
+// TODO: Needs to return default configs if they exist (currently no way to list from epoxy)
+#[operation]
+pub async fn namespace_runner_config_list(
+	ctx: &OperationCtx,
+	input: &Input,
+) -> Result<Vec<(String, RunnerConfig)>> {
+	let runner_configs = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			let (start, end) = if let Some(variant) = input.variant {
+				let (start, end) = keys::subspace()
+					.subspace(&keys::runner_config::ByVariantKey::subspace_with_variant(
+						input.namespace_id,
+						variant,
+					))
+					.range();
+
+				let start = if let Some(name) = &input.after_name {
+					tx.pack(&keys::runner_config::ByVariantKey::new(
+						input.namespace_id,
+						variant,
+						name.clone(),
+					))
+				} else {
+					start
+				};
+
+				(start, end)
+			} else {
+				let (start, end) = keys::subspace()
+					.subspace(&keys::runner_config::DataKey::subspace(input.namespace_id))
+					.range();
+
+				let start = if let Some(name) = &input.after_name {
+					tx.pack(&keys::runner_config::DataKey::new(
+						input.namespace_id,
+						name.clone(),
+					))
+				} else {
+					start
+				};
+
+				(start, end)
+			};
+
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					limit: Some(input.limit),
+					..(start, end).into()
+				},
+				Serializable,
+			)
+			.map(|res| match res {
+				Ok(entry) => {
+					if input.variant.is_some() {
+						let (key, config) =
+							tx.read_entry::<keys::runner_config::ByVariantKey>(&entry)?;
+						Ok((key.name, config))
+					} else {
+						let (key, config) =
+							tx.read_entry::<keys::runner_config::DataKey>(&entry)?;
+						Ok((key.name, config))
+					}
+				}
+				Err(err) => Err(err.into()),
+			})
+			.try_collect()
+			.await
+		})
+		.custom_instrument(tracing::info_span!("runner_config_list_tx"))
+		.await?;
+
+	Ok(runner_configs)
+}

package/packages/namespace/src/ops/runner_config/mod.rs

@@ -0,0 +1,4 @@
+pub mod delete;
+pub mod get;
+pub mod list;
+pub mod upsert;

package/packages/namespace/src/ops/runner_config/upsert.rs

@@ -0,0 +1,148 @@
+use gas::prelude::*;
+use rivet_types::runner_configs::{RunnerConfig, RunnerConfigKind};
+use universaldb::{options::MutationType, utils::IsolationLevel::*};
+
+use crate::{errors, keys, utils::runner_config_variant};
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub name: String,
+	pub config: RunnerConfig,
+}
+
+#[operation]
+pub async fn namespace_runner_config_upsert(ctx: &OperationCtx, input: &Input) -> Result<bool> {
+	let endpoint_config_changed = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			let runner_config_key =
+				keys::runner_config::DataKey::new(input.namespace_id, input.name.clone());
+
+			// Check if config changed (for serverless, compare URL and headers)
+			let endpoint_config_changed = if let Some(existing_config) =
+				tx.read_opt(&runner_config_key, Serializable).await?
+			{
+				// Delete previous index
+				tx.delete(&keys::runner_config::ByVariantKey::new(
+					input.namespace_id,
+					runner_config_variant(&existing_config),
+					input.name.clone(),
+				));
+
+				// Check if serverless endpoint config changed
+				match (&existing_config.kind, &input.config.kind) {
+					(
+						RunnerConfigKind::Serverless {
+							url: old_url,
+							headers: old_headers,
+							..
+						},
+						RunnerConfigKind::Serverless {
+							url: new_url,
+							headers: new_headers,
+							..
+						},
+					) => old_url != new_url || old_headers != new_headers,
+					_ => {
+						// Config type changed or is not serverless
+						true
+					}
+				}
+			} else {
+				// New config
+				true
+			};
+
+			// Write new config
+			tx.write(&runner_config_key, input.config.clone())?;
+			tx.write(
+				&keys::runner_config::ByVariantKey::new(
+					input.namespace_id,
+					runner_config_variant(&input.config),
+					input.name.clone(),
+				),
+				input.config.clone(),
+			)?;
+
+			match &input.config.kind {
+				RunnerConfigKind::Normal { .. } => {}
+				RunnerConfigKind::Serverless {
+					url,
+					headers,
+					slots_per_runner,
+					..
+				} => {
+					// Validate url
+					if let Err(err) = url::Url::parse(url) {
+						return Ok(Err(errors::RunnerConfig::Invalid {
+							reason: format!("invalid serverless url: {err}"),
+						}));
+					}
+
+					if headers.len() > 16 {
+						return Ok(Err(errors::RunnerConfig::Invalid {
+							reason: "too many headers (max 16)".to_string(),
+						}));
+					}
+
+					for (n, v) in headers {
+						if n.len() > 128 {
+							return Ok(Err(errors::RunnerConfig::Invalid {
+								reason: format!("invalid header name: too long (max 128)"),
+							}));
+						}
+						if let Err(err) = n.parse::<reqwest::header::HeaderName>() {
+							return Ok(Err(errors::RunnerConfig::Invalid {
+								reason: format!("invalid header name: {err}"),
+							}));
+						}
+						if v.len() > 4096 {
+							return Ok(Err(errors::RunnerConfig::Invalid {
+								reason: format!("invalid header value: too long (max 4096)"),
+							}));
+						}
+						if let Err(err) = v.parse::<reqwest::header::HeaderValue>() {
+							return Ok(Err(errors::RunnerConfig::Invalid {
+								reason: format!("invalid header value: {err}"),
+							}));
+						}
+					}
+
+					// Validate slots per runner
+					if *slots_per_runner == 0 {
+						return Ok(Err(errors::RunnerConfig::Invalid {
+							reason: "`slots_per_runner` cannot be 0".to_string(),
+						}));
+					}
+
+					// Sets desired count to 0 if it doesn't exist
+					let tx = tx.with_subspace(rivet_types::keys::pegboard::subspace());
+					tx.atomic_op(
+						&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::new(
+							input.namespace_id,
+							input.name.clone(),
+						),
+						&0i64.to_le_bytes(),
+						MutationType::Add,
+					);
+				}
+			}
+
+			Ok(Ok(endpoint_config_changed))
+		})
+		.custom_instrument(tracing::info_span!("runner_config_upsert_tx"))
+		.await?
+		.map_err(|err| err.build())?;
+
+	// Bump autoscaler
+	if input.config.affects_autoscaler() {
+		ctx.msg(rivet_types::msgs::pegboard::BumpServerlessAutoscaler {})
+			.send()
+			.await?;
+	}
+
+	Ok(endpoint_config_changed)
+}