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

package/packages/namespace/src/keys/metric.rs

@@ -0,0 +1,301 @@
+use anyhow::Result;
+use gas::prelude::*;
+use universaldb::prelude::*;
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum Metric {
+	/// Seconds (actor name)
+	ActorAwake(String),
+	/// Count (actor name)
+	TotalActors(String),
+	/// Bytes (actor name)
+	KvStorageUsed(String),
+	/// Bytes (actor name)
+	KvRead(String),
+	/// Bytes (actor name)
+	KvWrite(String),
+	/// Count (actor name)
+	AlarmsSet(String),
+	/// Bytes (actor name, type)
+	GatewayIngress(String, String),
+	/// Bytes (actor name)
+	GatewayEgress(String, String),
+	/// Count (actor name, type)
+	Requests(String, String),
+	/// Count (actor name, type)
+	ActiveRequests(String, String),
+}
+
+impl Metric {
+	fn variant(&self) -> MetricVariant {
+		match self {
+			Metric::ActorAwake(_) => MetricVariant::ActorAwake,
+			Metric::TotalActors(_) => MetricVariant::TotalActors,
+			Metric::KvStorageUsed(_) => MetricVariant::KvStorageUsed,
+			Metric::KvRead(_) => MetricVariant::KvRead,
+			Metric::KvWrite(_) => MetricVariant::KvWrite,
+			Metric::AlarmsSet(_) => MetricVariant::AlarmsSet,
+			Metric::GatewayIngress(_, _) => MetricVariant::GatewayIngress,
+			Metric::GatewayEgress(_, _) => MetricVariant::GatewayEgress,
+			Metric::Requests(_, _) => MetricVariant::Requests,
+			Metric::ActiveRequests(_, _) => MetricVariant::ActiveRequests,
+		}
+	}
+}
+
+#[derive(strum::FromRepr)]
+enum MetricVariant {
+	ActorAwake = 0,
+	TotalActors = 1,
+	KvStorageUsed = 2,
+	KvRead = 3,
+	KvWrite = 4,
+	AlarmsSet = 5,
+	GatewayIngress = 6,
+	GatewayEgress = 7,
+	Requests = 8,
+	ActiveRequests = 9,
+}
+
+impl std::fmt::Display for MetricVariant {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		match self {
+			MetricVariant::ActorAwake => write!(f, "actor_awake"),
+			MetricVariant::TotalActors => write!(f, "total_actors"),
+			MetricVariant::KvStorageUsed => write!(f, "kv_storage_used"),
+			MetricVariant::KvRead => write!(f, "kv_read"),
+			MetricVariant::KvWrite => write!(f, "kv_write"),
+			MetricVariant::AlarmsSet => write!(f, "alarms_set"),
+			MetricVariant::GatewayIngress => write!(f, "gateway_ingress"),
+			MetricVariant::GatewayEgress => write!(f, "gateway_egress"),
+			MetricVariant::Requests => write!(f, "requests"),
+			MetricVariant::ActiveRequests => write!(f, "active_requests"),
+		}
+	}
+}
+
+#[derive(Debug)]
+pub struct MetricKey {
+	pub namespace_id: Id,
+	pub metric: Metric,
+}
+
+impl MetricKey {
+	pub fn new(namespace_id: Id, metric: Metric) -> Self {
+		MetricKey {
+			namespace_id,
+			metric,
+		}
+	}
+
+	pub fn subspace(namespace_id: Id) -> MetricSubspaceKey {
+		MetricSubspaceKey::new(namespace_id)
+	}
+}
+
+impl FormalKey for MetricKey {
+	// IMPORTANT: Uses LE bytes, not BE
+	/// Count.
+	type Value = i64;
+
+	fn deserialize(&self, raw: &[u8]) -> Result<Self::Value> {
+		Ok(i64::from_le_bytes(raw.try_into()?))
+	}
+
+	fn serialize(&self, value: Self::Value) -> Result<Vec<u8>> {
+		Ok(value.to_le_bytes().to_vec())
+	}
+}
+
+impl TuplePack for MetricKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let mut offset = VersionstampOffset::None { size: 0 };
+
+		let t = (METRIC, self.namespace_id, self.metric.variant() as usize);
+		offset += t.pack(w, tuple_depth)?;
+
+		offset += match &self.metric {
+			Metric::ActorAwake(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::TotalActors(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::KvStorageUsed(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::KvRead(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::KvWrite(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::AlarmsSet(actor_name) => actor_name.pack(w, tuple_depth)?,
+			Metric::GatewayIngress(actor_name, req_type) => {
+				(actor_name, req_type).pack(w, tuple_depth)?
+			}
+			Metric::GatewayEgress(actor_name, req_type) => {
+				(actor_name, req_type).pack(w, tuple_depth)?
+			}
+			Metric::Requests(actor_name, req_type) => {
+				(actor_name, req_type).pack(w, tuple_depth)?
+			}
+			Metric::ActiveRequests(actor_name, req_type) => {
+				(actor_name, req_type).pack(w, tuple_depth)?
+			}
+		};
+
+		std::result::Result::Ok(offset)
+	}
+}
+
+impl<'de> TupleUnpack<'de> for MetricKey {
+	fn unpack(input: &[u8], tuple_depth: TupleDepth) -> PackResult<(&[u8], Self)> {
+		let (input, (_, namespace_id, variant)) = <(usize, Id, usize)>::unpack(input, tuple_depth)?;
+		let variant = MetricVariant::from_repr(variant).ok_or_else(|| {
+			PackError::Message(format!("invalid metric variant `{variant}` in key").into())
+		})?;
+
+		let (input, v) = match variant {
+			MetricVariant::ActorAwake => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::ActorAwake(actor_name),
+					},
+				)
+			}
+			MetricVariant::TotalActors => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::TotalActors(actor_name),
+					},
+				)
+			}
+			MetricVariant::KvStorageUsed => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::KvStorageUsed(actor_name),
+					},
+				)
+			}
+			MetricVariant::KvRead => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::KvRead(actor_name),
+					},
+				)
+			}
+			MetricVariant::KvWrite => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::KvWrite(actor_name),
+					},
+				)
+			}
+			MetricVariant::AlarmsSet => {
+				let (input, actor_name) = String::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::AlarmsSet(actor_name),
+					},
+				)
+			}
+			MetricVariant::GatewayIngress => {
+				let (input, (actor_name, req_type)) =
+					<(String, String)>::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::GatewayIngress(actor_name, req_type),
+					},
+				)
+			}
+			MetricVariant::GatewayEgress => {
+				let (input, (actor_name, req_type)) =
+					<(String, String)>::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::GatewayEgress(actor_name, req_type),
+					},
+				)
+			}
+			MetricVariant::Requests => {
+				let (input, (actor_name, req_type)) =
+					<(String, String)>::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::Requests(actor_name, req_type),
+					},
+				)
+			}
+			MetricVariant::ActiveRequests => {
+				let (input, (actor_name, req_type)) =
+					<(String, String)>::unpack(input, tuple_depth)?;
+
+				(
+					input,
+					MetricKey {
+						namespace_id,
+						metric: Metric::ActiveRequests(actor_name, req_type),
+					},
+				)
+			}
+		};
+
+		Ok((input, v))
+	}
+}
+
+pub struct MetricSubspaceKey {
+	namespace_id: Id,
+}
+
+impl MetricSubspaceKey {
+	pub fn new(namespace_id: Id) -> Self {
+		MetricSubspaceKey { namespace_id }
+	}
+}
+
+impl TuplePack for MetricSubspaceKey {
+	fn pack<W: std::io::Write>(
+		&self,
+		w: &mut W,
+		tuple_depth: TupleDepth,
+	) -> std::io::Result<VersionstampOffset> {
+		let t = (METRIC, self.namespace_id);
+		t.pack(w, tuple_depth)
+	}
+}
+
+pub fn inc(tx: &universaldb::Transaction, namespace_id: Id, metric: Metric, by: i64) {
+	tx.atomic_op(
+		&MetricKey::new(namespace_id, metric),
+		&by.to_le_bytes(),
+		MutationType::Add,
+	);
+}

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

@@ -2,7 +2,7 @@ use anyhow::Result;
 use gas::prelude::*;
 use universaldb::prelude::*;
 
-pub mod runner_config;
+pub mod metric;
 
 pub fn subspace() -> universaldb::utils::Subspace {
 	universaldb::utils::Subspace::new(&(RIVET, NAMESPACE))

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

@@ -39,10 +39,7 @@ pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result<V
 							.send()
 							.await?;
 
-						let res = rivet_api_util::parse_response::<
-							rivet_api_types::namespaces::list::ListResponse,
-						>(res)
-						.await?;
+						let res = rivet_api_util::parse_response::<ListResponse>(res).await?;
 
 						for ns in res.namespaces {
 							let namespace_id = ns.namespace_id;
@@ -56,3 +53,9 @@ pub async fn namespace_get_global(ctx: &OperationCtx, input: &Input) -> Result<V
 			.await
 	}
 }
+
+// TODO: Cyclical dependency with rivet_api_types
+#[derive(Deserialize)]
+struct ListResponse {
+	namespaces: Vec<Namespace>,
+}

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

@@ -16,24 +16,40 @@ pub async fn namespace_get_local(ctx: &OperationCtx, input: &Input) -> Result<Ve
 		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();
+	ctx.cache()
+		.clone()
+		.request()
+		.fetch_all_json(
+			"namespace.get_local",
+			input.namespace_ids.clone(),
+			move |mut cache, namespace_ids| async move {
+				let namespace_ids = &namespace_ids;
+				let namespaces = ctx
+					.udb()?
+					.run(|tx| async move {
+						futures_util::stream::iter(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?;
+								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)
+				for ns in namespaces {
+					let namespace_id = ns.namespace_id;
+					cache.resolve(&&namespace_id, ns);
+				}
+
+				Ok(cache)
+			},
+		)
+		.await
 }
 
 pub(crate) async fn get_inner(

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

@@ -3,4 +3,3 @@ 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

@@ -38,10 +38,7 @@ pub async fn namespace_resolve_for_name_global(
 							.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 res = rivet_api_util::parse_response::<ListResponse>(res).await?;
 
 						let ns = res.namespaces.into_iter().next();
 
@@ -55,3 +52,9 @@ pub async fn namespace_resolve_for_name_global(
 			.map(|x| x.flatten())
 	}
 }
+
+// TODO: Cyclical dependency with rivet_api_types
+#[derive(Deserialize)]
+struct ListResponse {
+	namespaces: Vec<Namespace>,
+}

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

@@ -18,24 +18,44 @@ pub async fn namespace_resolve_for_name_local(
 		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"))
+	ctx.cache()
+		.clone()
+		.request()
+		.fetch_one_json(
+			"namespace.resolve_for_name_local",
+			input.name.clone(),
+			move |mut cache, key| {
+				async move {
+					let ns = 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?;
+
+					if let Some(ns) = ns {
+						cache.resolve(&key, ns);
+					}
+
+					Ok(cache)
+				}
+			},
+		)
 		.await
-		.map_err(Into::into)
 }

package/packages/namespace/src/workflows/namespace.rs

@@ -23,7 +23,7 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
 
 	if let Err(error) = validation_res {
 		ctx.msg(Failed { error })
-			.tag("namespace_id", input.namespace_id)
+			.topic(("namespace_id", input.namespace_id))
 			.send()
 			.await?;
 
@@ -42,7 +42,7 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
 
 	if let Err(error) = insert_res {
 		ctx.msg(Failed { error })
-			.tag("namespace_id", input.namespace_id)
+			.topic(("namespace_id", input.namespace_id))
 			.send()
 			.await?;
 
@@ -51,7 +51,7 @@ pub async fn namespace(ctx: &mut WorkflowCtx, input: &Input) -> Result<()> {
 	}
 
 	ctx.msg(CreateComplete {})
-		.tag("namespace_id", input.namespace_id)
+		.topic(("namespace_id", input.namespace_id))
 		.send()
 		.await?;
 

package/packages/pegboard/Cargo.toml

@@ -8,27 +8,49 @@ edition.workspace = true
 [dependencies]
 anyhow.workspace = true
 base64.workspace = true
+bytes.workspace = true
 epoxy.workspace = true
+epoxy-protocol.workspace = true
+foundationdb-tuple.workspace = true
+futures-util.workspace = true
 gas.workspace = true
 lazy_static.workspace = true
 namespace.workspace = true
 nix.workspace = true
 rand.workspace = true
+reqwest-eventsource.workspace = true
+reqwest.workspace = true
 rivet-api-types.workspace = true
 rivet-api-util.workspace = true
+rivet-config.workspace = true
 rivet-data.workspace = true
+rivet-envoy-protocol.workspace = true
 rivet-error.workspace = true
 rivet-metrics.workspace = true
+rivet-pools.workspace = true
 rivet-runner-protocol.workspace = true
+rivet-runtime.workspace = true
 rivet-types.workspace = true
 rivet-util.workspace = true
 serde_bare.workspace = true
 serde_json.workspace = true
 serde.workspace = true
 strum.workspace = true
+tokio.workspace = true
 tracing.workspace = true
 universaldb.workspace = true
 universalpubsub.workspace = true
+url.workspace = true
 utoipa.workspace = true
 uuid.workspace = true
 vbare.workspace = true
+
+[dev-dependencies]
+portpicker.workspace = true
+test-snapshot-gen.workspace = true
+rivet-config.workspace = true
+rivet-test-deps.workspace = true
+tokio.workspace = true
+tracing-subscriber.workspace = true
+url.workspace = true
+uuid.workspace = true

package/packages/pegboard/src/actor_kv/entry.rs

@@ -0,0 +1,47 @@
+use anyhow::{Context, Result, ensure};
+use rivet_envoy_protocol as ep;
+
+use crate::keys::actor_kv::KeyWrapper;
+
+pub struct EntryBuilder {
+	pub key: KeyWrapper,
+	metadata: Option<ep::KvMetadata>,
+	value: Vec<u8>,
+	next_idx: usize,
+}
+
+impl EntryBuilder {
+	pub fn new(key: KeyWrapper) -> Self {
+		EntryBuilder {
+			key,
+			metadata: None,
+			value: Vec::new(),
+			next_idx: 0,
+		}
+	}
+
+	pub fn append_metadata(&mut self, metadata: ep::KvMetadata) {
+		// We ignore setting the metadata again because it means the same key was given twice in the
+		// input keys for `get`. We don't perform automatic deduplication.
+		if self.metadata.is_none() {
+			self.metadata = Some(metadata);
+		}
+	}
+
+	pub fn append_chunk(&mut self, idx: usize, chunk: &[u8]) {
+		if idx >= self.next_idx {
+			self.value.extend(chunk);
+			self.next_idx = idx + 1;
+		}
+	}
+
+	pub fn build(self) -> Result<(ep::KvKey, ep::KvValue, ep::KvMetadata)> {
+		ensure!(!self.value.is_empty(), "empty value at key");
+
+		Ok((
+			self.key.0,
+			self.value,
+			self.metadata.context("no metadata for key")?,
+		))
+	}
+}

package/packages/pegboard/src/actor_kv/metrics.rs

@@ -0,0 +1,19 @@
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
+
+lazy_static::lazy_static! {
+	pub static ref ACTOR_KV_OPERATION_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"actor_kv_operation_duration_seconds",
+		"Duration of actor KV operations including UDB transaction.",
+		&["op"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+
+	pub static ref ACTOR_KV_KEYS_PER_OP: HistogramVec = register_histogram_vec_with_registry!(
+		"actor_kv_keys_per_operation",
+		"Number of keys per actor KV operation.",
+		&["op"],
+		vec![1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0],
+		*REGISTRY
+	).unwrap();
+}