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

package/packages/test-snapshot-gen/src/test_cluster.rs

@@ -0,0 +1,234 @@
+use std::collections::HashMap;
+
+use anyhow::{Context, Result};
+use epoxy::types;
+use epoxy_protocol::protocol::ReplicaId;
+use gas::prelude::{TestCtx as WorkflowTestCtx, *};
+use rivet_util::Id;
+use url::Url;
+
+struct ReplicaMetadata {
+	api_peer_port: u16,
+	guard_port: u16,
+}
+
+struct ReplicaContext {
+	wf_ctx: WorkflowTestCtx,
+	api_server_handle: tokio::task::JoinHandle<()>,
+}
+
+pub struct TestCluster {
+	test_id: Uuid,
+	leader_id: ReplicaId,
+	coordinator_workflow_id: Id,
+	replica_metadata: HashMap<ReplicaId, ReplicaMetadata>,
+	replica_contexts: HashMap<ReplicaId, ReplicaContext>,
+}
+
+impl TestCluster {
+	pub async fn new(replica_ids: &[ReplicaId]) -> Result<Self> {
+		let leader_id = replica_ids[0];
+
+		let mut cluster = TestCluster {
+			test_id: Uuid::new_v4(),
+			leader_id,
+			coordinator_workflow_id: Id::new_v1(leader_id as u16),
+			replica_contexts: HashMap::new(),
+			replica_metadata: HashMap::new(),
+		};
+
+		// Add metadata for all replicas before starting any.
+		for &replica_id in replica_ids {
+			let api_peer_port =
+				portpicker::pick_unused_port().context("failed to pick API peer port")?;
+			let guard_port = portpicker::pick_unused_port().context("failed to pick guard port")?;
+			cluster.replica_metadata.insert(
+				replica_id,
+				ReplicaMetadata {
+					api_peer_port,
+					guard_port,
+				},
+			);
+		}
+
+		// Start each replica.
+		for &replica_id in replica_ids {
+			cluster.start_replica(replica_id).await?;
+		}
+
+		// Start coordinator workflow on leader.
+		let mut config_sub = cluster
+			.get_ctx(leader_id)
+			.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>(("replica", leader_id))
+			.await?;
+
+		let coordinator_workflow_id =
+			setup_coordinator(&cluster.replica_contexts, leader_id).await?;
+		cluster.coordinator_workflow_id = coordinator_workflow_id;
+
+		// Wait for all replicas to become active.
+		tracing::info!("waiting for replicas to become ready");
+		loop {
+			let config_msg = config_sub.next().await?;
+			let all_active = config_msg
+				.config
+				.replicas
+				.iter()
+				.all(|r| r.status == types::ReplicaStatus::Active);
+			if all_active {
+				break;
+			}
+		}
+
+		Ok(cluster)
+	}
+
+	pub fn get_ctx(&self, replica_id: ReplicaId) -> &WorkflowTestCtx {
+		&self
+			.replica_contexts
+			.get(&replica_id)
+			.expect("replica not started")
+			.wf_ctx
+	}
+
+	pub fn leader_id(&self) -> ReplicaId {
+		self.leader_id
+	}
+
+	pub async fn shutdown(&mut self) -> Result<()> {
+		tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+
+		let mut ids = self.replica_contexts.keys().copied().collect::<Vec<_>>();
+		ids.sort_unstable();
+
+		for replica_id in ids {
+			let mut ctx = self.replica_contexts.remove(&replica_id).unwrap();
+			ctx.wf_ctx.shutdown().await?;
+			ctx.api_server_handle.abort();
+			let _ = (&mut ctx.api_server_handle).await;
+		}
+
+		Ok(())
+	}
+
+	async fn start_replica(&mut self, replica_id: ReplicaId) -> Result<()> {
+		tracing::info!(%replica_id, "starting replica");
+
+		let metadata = self.replica_metadata.get(&replica_id).unwrap();
+		let datacenters = self.build_datacenters()?;
+
+		let test_deps = rivet_test_deps::setup_single_datacenter(
+			self.test_id,
+			replica_id as u16,
+			datacenters,
+			metadata.api_peer_port,
+			metadata.guard_port,
+		)
+		.await?;
+
+		let reg = epoxy::registry()?
+			.merge(namespace::registry()?)?
+			.merge(pegboard::registry()?)?;
+		let test_ctx = WorkflowTestCtx::new_with_deps(reg, test_deps).await?;
+
+		let api_handle = setup_api_server(
+			test_ctx.config().clone(),
+			test_ctx.pools().clone(),
+			metadata.api_peer_port,
+		)
+		.await?;
+
+		// Start replica workflow.
+		let workflow_id = test_ctx
+			.workflow(epoxy::workflows::replica::Input {})
+			.tag("replica", replica_id)
+			.dispatch()
+			.await?;
+		tracing::info!(%workflow_id, %replica_id, "created epoxy replica");
+
+		self.replica_contexts.insert(
+			replica_id,
+			ReplicaContext {
+				wf_ctx: test_ctx,
+				api_server_handle: api_handle,
+			},
+		);
+
+		Ok(())
+	}
+
+	fn build_datacenters(
+		&self,
+	) -> Result<HashMap<String, rivet_config::config::topology::Datacenter>> {
+		let mut datacenters = HashMap::new();
+		let mut ids: Vec<_> = self.replica_metadata.keys().copied().collect();
+		ids.sort_unstable();
+
+		for &id in &ids {
+			let metadata = &self.replica_metadata[&id];
+			let name = format!("dc-{}", id);
+			datacenters.insert(
+				name.clone(),
+				rivet_config::config::topology::Datacenter {
+					name: format!("dc-{}", id),
+					datacenter_label: id as u16,
+					is_leader: id == self.leader_id,
+					peer_url: Url::parse(&format!("http://127.0.0.1:{}", metadata.api_peer_port))?,
+					public_url: Url::parse(&format!("http://127.0.0.1:{}", metadata.guard_port))?,
+					proxy_url: None,
+					valid_hosts: None,
+				},
+			);
+		}
+
+		Ok(datacenters)
+	}
+}
+
+async fn setup_coordinator(
+	replica_contexts: &HashMap<ReplicaId, ReplicaContext>,
+	leader_id: ReplicaId,
+) -> Result<Id> {
+	let leader_ctx = &replica_contexts
+		.get(&leader_id)
+		.expect("leader not in replica contexts")
+		.wf_ctx;
+
+	let workflow_id = leader_ctx
+		.workflow(epoxy::workflows::coordinator::Input {})
+		.tag("replica", leader_id)
+		.dispatch()
+		.await?;
+
+	let mut sub = leader_ctx
+		.subscribe::<epoxy::workflows::coordinator::ConfigChangeMessage>(("replica", leader_id))
+		.await?;
+	leader_ctx
+		.signal(epoxy::workflows::coordinator::Reconfigure {})
+		.to_workflow_id(workflow_id)
+		.send()
+		.await?;
+	sub.next().await?;
+
+	Ok(workflow_id)
+}
+
+async fn setup_api_server(
+	config: rivet_config::Config,
+	pools: rivet_pools::Pools,
+	port: u16,
+) -> Result<tokio::task::JoinHandle<()>> {
+	let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
+
+	let router = rivet_api_builder::create_router("api-peer-test", config, pools, |router| {
+		epoxy::http_routes::mount_routes(router)
+	})
+	.await?;
+
+	let listener = tokio::net::TcpListener::bind(addr).await?;
+	let handle = tokio::spawn(async move {
+		axum::serve(listener, router).await.unwrap();
+	});
+
+	Ok(handle)
+}

package/packages/tracing-reconfigure/src/lib.rs

@@ -18,7 +18,7 @@ pub async fn start(_config: rivet_config::Config, pools: rivet_pools::Pools) ->
 	let subject = "rivet.debug.tracing.config";
 	let mut sub = ups.subscribe(subject).await?;
 
-	tracing::debug!(subject = ?subject, "subscribed to tracing config updates");
+	tracing::debug!(%subject, "subscribed to tracing config updates");
 
 	// Process incoming messages
 	while let Ok(NextOutput::Message(msg)) = sub.next().await {

package/packages/tracing-utils/src/lib.rs

@@ -8,8 +8,6 @@ use std::{
 use futures_util::future;
 use tracing::{Instrument, instrument::Instrumented};
 
-use rivet_metrics::KeyValue;
-
 /// Attempts to create a new future to select over a list of futures.
 /// Non-panicking version of [futures_util::future::select_all](https://docs.rs/futures/0.3.15/futures/future/fn.select_all.html).
 ///
@@ -57,13 +55,10 @@ impl<T: Future> Future for CustomInstrumented<T> {
 			Poll::Ready(val) => {
 				if let Some(metadata) = metadata {
 					if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
-						metrics::INSTRUMENTED_FUTURE_DURATION.record(
-							this.start.elapsed().as_secs_f64(),
-							&[
-								KeyValue::new("location", format!("{file}:{line}")),
-								KeyValue::new("name", metadata.name()),
-							],
-						);
+						let location = format!("{file}:{line}");
+						metrics::INSTRUMENTED_FUTURE_DURATION
+							.with_label_values(&[location.as_str(), metadata.name()])
+							.observe(this.start.elapsed().as_secs_f64());
 					}
 				}
 				Poll::Ready(val)
@@ -74,18 +69,15 @@ impl<T: Future> Future for CustomInstrumented<T> {
 }
 
 mod metrics {
-	use rivet_metrics::{
-		MICRO_BUCKETS,
-		otel::{global::*, metrics::*},
-	};
+	use rivet_metrics::{MICRO_BUCKETS, REGISTRY, prometheus::*};
 
 	lazy_static::lazy_static! {
-		static ref METER: Meter = meter("rivet-util-core");
-
-		/// Expected attributes: "location", "name"
-		pub static ref INSTRUMENTED_FUTURE_DURATION: Histogram<f64> = METER.f64_histogram("rivet_instrumented_future_duration")
-			.with_description("Duration of a future.")
-			.with_boundaries(MICRO_BUCKETS.to_vec())
-			.build();
+		pub static ref INSTRUMENTED_FUTURE_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+			"instrumented_future_duration",
+			"Duration of a future.",
+			&["location", "name"],
+			MICRO_BUCKETS.to_vec(),
+			*REGISTRY
+		).unwrap();
 	}
 }

package/packages/types/src/actor/error.rs

@@ -0,0 +1,59 @@
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+
+/// Keep in sync with `website/src/content/docs/actors/troubleshooting.mdx` and
+/// `frontend/src/components/actors/actor-status-label.tsx`.
+#[derive(Debug, Clone, Serialize, Deserialize, Hash)]
+#[serde(rename_all = "snake_case")]
+pub enum RunnerPoolError {
+	/// Serverless: SSE returned non-200 status code (e.g., 404, 500)
+	ServerlessHttpError { status_code: u16, body: String },
+
+	/// Serverless: SSE stream ended unexpectedly before runner initialized
+	ServerlessStreamEndedEarly,
+
+	/// Serverless: SSE connection or network error
+	ServerlessConnectionError { message: String },
+
+	/// Serverless: Runner sent invalid payload
+	ServerlessInvalidSsePayload {
+		message: String,
+		raw_payload: Option<String>,
+	},
+
+	/// RivetKit was downgraded to a version after being upgraded to use envoys
+	Downgrade,
+
+	/// Internal error
+	InternalError,
+}
+
+/// Keep in sync with `website/src/content/docs/actors/troubleshooting.mdx` and
+/// `frontend/src/components/actors/actor-status-label.tsx`.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ActorError {
+	/// Runner pool-related errors
+	RunnerPoolError(RunnerPoolError),
+	/// No runners available matching the runner name
+	NoCapacity,
+	/// Runner was allocated but never started the actor (GC timeout)
+	RunnerNoResponse { runner_id: Id },
+	/// Runner connection was lost (no recent ping, network issue, or crash)
+	RunnerConnectionLost { runner_id: Id },
+	/// Runner was draining but actor didn't stop in time
+	RunnerDrainingTimeout { runner_id: Id },
+	/// Actor exited with an error and is now sleeping
+	Crashed { message: Option<String> },
+	/// Actor has an internal error
+	InternalError,
+
+	/// Actor cannot allocate due to the concurrent actor limit.
+	ConcurrentActorLimitReached,
+	/// No envoys available matching the pool name
+	NoEnvoys,
+	/// Envoy was allocated but never started the actor (GC timeout)
+	EnvoyNoResponse { envoy_key: String },
+	/// Envoy connection was lost (no recent ping, network issue, or crash)
+	EnvoyConnectionLost { envoy_key: String },
+}

package/packages/types/src/actor/mod.rs

@@ -0,0 +1,2 @@
+pub mod error;
+pub use error::*;

package/packages/types/src/actors.rs

@@ -30,6 +30,11 @@ pub struct Actor {
 	pub reschedule_ts: Option<i64>,
 	/// Denotes when the actor was destroyed.
 	pub destroy_ts: Option<i64>,
+
+	/// Error details if the actor failed to start.
+	#[serde(skip_serializing_if = "Option::is_none")]
+	#[schema(value_type = Option<Object>, additional_properties = true)]
+	pub error: Option<crate::actor::ActorError>,
 }
 
 #[derive(Debug, Copy, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash, ToSchema)]

package/packages/types/src/envoys.rs

@@ -0,0 +1,21 @@
+use gas::prelude::*;
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+#[serde(deny_unknown_fields)]
+pub struct Envoy {
+	pub envoy_key: String,
+	pub namespace_id: Id,
+	pub datacenter: String,
+	pub pool_name: String,
+	pub version: u32,
+	pub slots: u64,
+	pub create_ts: i64,
+	pub stop_ts: Option<i64>,
+	pub last_ping_ts: i64,
+	pub last_connected_ts: Option<i64>,
+	pub last_rtt: u32,
+	pub metadata: Option<serde_json::Map<String, serde_json::Value>>,
+}

package/packages/types/src/keys/backfill.rs

@@ -0,0 +1,5 @@
+use universaldb::prelude::*;
+
+pub fn subspace() -> universaldb::utils::Subspace {
+	universaldb::utils::Subspace::new(&(RIVET, BACKFILL))
+}

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

@@ -1,2 +1,3 @@
+pub mod backfill;
 pub mod namespace;
 pub mod pegboard;

package/packages/types/src/lib.rs

@@ -1,7 +1,8 @@
+pub mod actor;
 pub mod actors;
 pub mod datacenters;
+pub mod envoys;
 pub mod keys;
-pub mod msgs;
 pub mod namespaces;
 pub mod runner_configs;
 pub mod runners;

package/packages/types/src/runner_configs.rs

@@ -9,6 +9,8 @@ pub struct RunnerConfig {
 	pub kind: RunnerConfigKind,
 	#[serde(default, skip_serializing_if = "Option::is_none")]
 	pub metadata: Option<serde_json::Value>,
+	#[serde(default = "default_drain_on_version_upgrade")]
+	pub drain_on_version_upgrade: bool,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
@@ -20,40 +22,61 @@ pub enum RunnerConfigKind {
 		headers: HashMap<String, String>,
 		/// Seconds.
 		request_lifespan: u32,
+		max_concurrent_actors: u64,
+		/// Deprecated.
 		slots_per_runner: u32,
+		/// Deprecated.
 		min_runners: u32,
+		/// Deprecated.
 		max_runners: u32,
+		/// Deprecated.
 		runners_margin: u32,
+		/// Milliseconds between metadata polling. If not set, uses the global default.
+		#[serde(default, skip_serializing_if = "Option::is_none")]
+		metadata_poll_interval: Option<u64>,
 	},
 }
 
-impl From<RunnerConfig> for rivet_data::generated::namespace_runner_config_v2::RunnerConfig {
+fn default_drain_on_version_upgrade() -> bool {
+	false
+}
+
+impl From<RunnerConfig> for rivet_data::generated::namespace_runner_config_v5::RunnerConfig {
 	fn from(value: RunnerConfig) -> Self {
-		let RunnerConfig { kind, metadata } = value;
-		rivet_data::generated::namespace_runner_config_v2::RunnerConfig {
+		let RunnerConfig {
+			kind,
+			metadata,
+			drain_on_version_upgrade,
+		} = value;
+		rivet_data::generated::namespace_runner_config_v5::RunnerConfig {
 			metadata: metadata.and_then(|value| serde_json::to_string(&value).ok()),
+			drain_on_version_upgrade,
 			kind: match kind {
 				RunnerConfigKind::Normal {} => {
-					rivet_data::generated::namespace_runner_config_v2::RunnerConfigKind::Normal
+					rivet_data::generated::namespace_runner_config_v5::RunnerConfigKind::Normal
 				}
 				RunnerConfigKind::Serverless {
 					url,
 					headers,
 					request_lifespan,
+					max_concurrent_actors,
 					slots_per_runner,
 					min_runners,
 					max_runners,
 					runners_margin,
+					metadata_poll_interval,
 				} => {
-					rivet_data::generated::namespace_runner_config_v2::RunnerConfigKind::Serverless(
-						rivet_data::generated::namespace_runner_config_v2::Serverless {
+					rivet_data::generated::namespace_runner_config_v5::RunnerConfigKind::Serverless(
+						rivet_data::generated::namespace_runner_config_v5::Serverless {
 							url,
 							headers: headers.into(),
 							request_lifespan,
+							max_concurrent_actors,
 							slots_per_runner,
 							min_runners,
 							max_runners,
 							runners_margin,
+							metadata_poll_interval,
 						},
 					)
 				}
@@ -62,26 +85,32 @@ impl From<RunnerConfig> for rivet_data::generated::namespace_runner_config_v2::R
 	}
 }
 
-impl From<rivet_data::generated::namespace_runner_config_v2::RunnerConfig> for RunnerConfig {
-	fn from(value: rivet_data::generated::namespace_runner_config_v2::RunnerConfig) -> Self {
-		let rivet_data::generated::namespace_runner_config_v2::RunnerConfig { metadata, kind } =
-			value;
+impl From<rivet_data::generated::namespace_runner_config_v5::RunnerConfig> for RunnerConfig {
+	fn from(value: rivet_data::generated::namespace_runner_config_v5::RunnerConfig) -> Self {
+		let rivet_data::generated::namespace_runner_config_v5::RunnerConfig {
+			metadata,
+			kind,
+			drain_on_version_upgrade,
+		} = value;
 		RunnerConfig {
 			metadata: metadata.and_then(|raw| serde_json::from_str(&raw).ok()),
+			drain_on_version_upgrade,
 			kind: match kind {
-				rivet_data::generated::namespace_runner_config_v2::RunnerConfigKind::Normal => {
+				rivet_data::generated::namespace_runner_config_v5::RunnerConfigKind::Normal => {
 					RunnerConfigKind::Normal {}
 				}
-				rivet_data::generated::namespace_runner_config_v2::RunnerConfigKind::Serverless(
+				rivet_data::generated::namespace_runner_config_v5::RunnerConfigKind::Serverless(
 					o,
 				) => RunnerConfigKind::Serverless {
 					url: o.url,
 					headers: o.headers.into(),
 					request_lifespan: o.request_lifespan,
+					max_concurrent_actors: o.max_concurrent_actors,
 					slots_per_runner: o.slots_per_runner,
 					min_runners: o.min_runners,
 					max_runners: o.max_runners,
 					runners_margin: o.runners_margin,
+					metadata_poll_interval: o.metadata_poll_interval,
 				},
 			},
 		}
@@ -89,8 +118,8 @@ impl From<rivet_data::generated::namespace_runner_config_v2::RunnerConfig> for R
 }
 
 impl RunnerConfig {
-	/// If updates to this run config affects the autoscaler.
-	pub fn affects_autoscaler(&self) -> bool {
+	/// If updates to this run config affects the pool.
+	pub fn affects_pool(&self) -> bool {
 		matches!(self.kind, RunnerConfigKind::Serverless { .. })
 	}
 }

package/packages/universaldb/Cargo.toml

@@ -11,16 +11,20 @@ async-trait.workspace = true
 deadpool-postgres.workspace = true
 foundationdb-tuple.workspace = true
 futures-util.workspace = true
+hex.workspace = true
 lazy_static.workspace = true
 rand.workspace = true
 rivet-metrics.workspace = true
+rivet-postgres-util.workspace = true
 rivet-tracing-utils.workspace = true
 rocksdb.workspace = true
 serde.workspace = true
 thiserror.workspace = true
+tokio-postgres-rustls.workspace = true
 tokio-postgres.workspace = true
 tokio.workspace = true
 tracing.workspace = true
+url.workspace = true
 uuid.workspace = true
 
 [dev-dependencies]

package/packages/universaldb/src/database.rs

@@ -1,11 +1,18 @@
-use std::future::Future;
+use std::path::Path;
+use std::time::Instant;
+use std::{
+	future::Future,
+	sync::atomic::{AtomicUsize, Ordering},
+};
 
 use anyhow::{Context, Result, anyhow};
 use futures_util::FutureExt;
 use rivet_tracing_utils::CustomInstrumentExt;
+use tracing::Instrument;
 
 use crate::{
 	driver::{DatabaseDriverHandle, Erased},
+	metrics,
 	options::DatabaseOption,
 	transaction::{RetryableTransaction, Transaction},
 };
@@ -20,7 +27,7 @@ impl Database {
 		Database { driver }
 	}
 
-	/// Run a closure with automatic retry logic
+	/// Run a closure with automatic retry logic.
 	#[tracing::instrument(skip_all)]
 	pub async fn run<'a, F, Fut, T>(&'a self, closure: F) -> Result<T>
 	where
@@ -28,11 +35,32 @@ impl Database {
 		Fut: Future<Output = Result<T>> + Send,
 		T: Send + 'a + 'static,
 	{
+		self.txn("unnamed", closure).in_current_span().await
+	}
+
+	/// Run a closure with automatic retry logic and a name.
+	#[tracing::instrument(skip_all)]
+	pub async fn txn<'a, F, Fut, T>(&'a self, name: &'static str, closure: F) -> Result<T>
+	where
+		F: Fn(RetryableTransaction) -> Fut + Send + Sync,
+		Fut: Future<Output = Result<T>> + Send,
+		T: Send + 'a + 'static,
+	{
+		let start = Instant::now();
+		let attempts = AtomicUsize::new(0);
+		metrics::TRANSACTION_TOTAL.with_label_values(&[name]).inc();
+		metrics::TRANSACTION_PENDING
+			.with_label_values(&[name])
+			.inc();
+
 		let closure = &closure;
-		self.driver
+		let res = self
+			.driver
 			.run(Box::new(|tx| {
+				attempts.fetch_add(1, Ordering::AcqRel);
+
 				async move { closure(tx).await.map(|value| Box::new(value) as Erased) }
-					.custom_instrument(tracing::info_span!("run_attempt"))
+					.custom_instrument(tracing::info_span!("txn_attempt"))
 					.boxed()
 			}))
 			.await
@@ -41,7 +69,19 @@ impl Database {
 					.map(|x| *x)
 					.map_err(|_| anyhow!("failed to downcast `run` return type"))
 			})
-			.context("transaction failed")
+			.context("transaction failed");
+
+		metrics::TRANSACTION_ATTEMPTS
+			.with_label_values(&[name])
+			.observe(attempts.load(Ordering::Acquire) as f64);
+		metrics::TRANSACTION_PENDING
+			.with_label_values(&[name])
+			.dec();
+		metrics::TRANSACTION_DURATION
+			.with_label_values(&[name])
+			.observe(start.elapsed().as_secs_f64());
+
+		res
 	}
 
 	/// Creates a new txn instance.
@@ -53,4 +93,9 @@ impl Database {
 	pub fn set_option(&self, opt: DatabaseOption) -> Result<()> {
 		self.driver.set_option(opt)
 	}
+
+	/// Create a consistent point-in-time snapshot of the database at the given path.
+	pub fn checkpoint(&self, path: &Path) -> Result<()> {
+		self.driver.checkpoint(path)
+	}
 }