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

package/packages/runtime/src/lib.rs

@@ -1,6 +1,6 @@
 use std::{env, future::Future, sync::Arc, time::Duration};
 
-use rivet_metrics::{KeyValue, init_otel_providers};
+use rivet_metrics::init_otel_providers;
 use tokio::sync::{Notify, OnceCell};
 mod metrics;
 mod traces;
@@ -53,60 +53,54 @@ fn build_tokio_runtime_builder() -> tokio::runtime::Builder {
 	rt_builder.enable_all();
 
 	rt_builder.on_thread_start(move || {
-		metrics::TOKIO_THREAD_COUNT.add(1, &[]);
+		metrics::TOKIO_THREAD_COUNT.inc();
 	});
 	rt_builder.on_thread_stop(move || {
-		metrics::TOKIO_THREAD_COUNT.add(-1, &[]);
+		metrics::TOKIO_THREAD_COUNT.dec();
 	});
 
-	rt_builder.on_task_spawn(move |_| metrics::TOKIO_TASK_TOTAL.add(1, &[]));
+	rt_builder.on_task_spawn(move |_| metrics::TOKIO_TASK_TOTAL.inc());
 
 	if env::var("TOKIO_RUNTIME_METRICS").is_ok() {
 		rt_builder.on_before_task_poll(|_| {
 			let metrics = tokio::runtime::Handle::current().metrics();
-			// let buckets = metrics.poll_time_histogram_num_buckets();
+			let buckets = metrics.poll_time_histogram_num_buckets();
 
-			metrics::TOKIO_GLOBAL_QUEUE_DEPTH.record(metrics.global_queue_depth() as u64, &[]);
-			metrics::TOKIO_ACTIVE_TASK_COUNT.record(metrics.num_alive_tasks() as u64, &[]);
+			metrics::TOKIO_GLOBAL_QUEUE_DEPTH.set(metrics.global_queue_depth() as i64);
+			metrics::TOKIO_ACTIVE_TASK_COUNT.set(metrics.num_alive_tasks() as i64);
 
 			for worker in 0..metrics.num_workers() {
-				metrics::TOKIO_WORKER_OVERFLOW_COUNT.record(
-					metrics.worker_overflow_count(worker),
-					&[KeyValue::new("worker", worker.to_string())],
-				);
-				metrics::TOKIO_WORKER_LOCAL_QUEUE_DEPTH.record(
-					metrics.worker_local_queue_depth(worker) as u64,
-					&[KeyValue::new("worker", worker.to_string())],
-				);
-
-				// TODO: Internal histogram data no longer accessable in OTEL
-				// needs to be turned into: 1. 20 gauges, or 2. store internal
-				// bucket counts manually
-
-				// use rivet_metrics::prometheus::core::Metric;
-				// // Has some sort of internal lock, must read data before loop
-				// let prom_buckets = {
-				// 	metrics::TOKIO_TASK_POLL_DURATION
-				// 		.metric()
-				// 		.get_histogram()
-				// 		.get_bucket()
-				// 		.iter()
-				// 		.map(|bucket| bucket.get_cumulative_count())
-				// 		.collect::<Vec<_>>()
-				// };
-
-				// for (bucket, prom_bucket_count) in (0..buckets).zip(prom_buckets) {
-				// 	let range = metrics.poll_time_histogram_bucket_range(bucket);
-				// 	let count = metrics.poll_time_histogram_bucket_count(worker, bucket);
-				// 	// Calculate difference in tokio's internal bucket counts and
-				// 	// prom's internal count
-				// 	let diff = count.saturating_sub(prom_bucket_count);
-
-				// 	// Match prom's count with tokio's for given bucket
-				// 	for _ in 0..diff {
-				// 		metrics::TOKIO_TASK_POLL_DURATION.record(range.start.as_secs_f64(), &[]);
-				// 	}
-				// }
+				metrics::TOKIO_WORKER_OVERFLOW_COUNT
+					.with_label_values(&[&worker.to_string()])
+					.set(metrics.worker_overflow_count(worker) as i64);
+				metrics::TOKIO_WORKER_LOCAL_QUEUE_DEPTH
+					.with_label_values(&[&worker.to_string()])
+					.set(metrics.worker_local_queue_depth(worker) as i64);
+
+				use rivet_metrics::prometheus::core::Metric;
+				// Has some sort of internal lock, must read data before loop
+				let prom_buckets = {
+					metrics::TOKIO_TASK_POLL_DURATION
+						.metric()
+						.get_histogram()
+						.get_bucket()
+						.iter()
+						.map(|bucket| bucket.cumulative_count())
+						.collect::<Vec<_>>()
+				};
+
+				for (bucket, prom_bucket_count) in (0..buckets).zip(prom_buckets) {
+					let range = metrics.poll_time_histogram_bucket_range(bucket);
+					let count = metrics.poll_time_histogram_bucket_count(worker, bucket);
+					// Calculate difference in tokio's internal bucket counts and
+					// prom's internal count
+					let diff = count.saturating_sub(prom_bucket_count);
+
+					// Match prom's count with tokio's for given bucket
+					for _ in 0..diff {
+						metrics::TOKIO_TASK_POLL_DURATION.observe(range.start.as_secs_f64());
+					}
+				}
 			}
 		});
 
@@ -130,8 +124,14 @@ fn build_tokio_runtime_builder() -> tokio::runtime::Builder {
 		rt_builder.thread_stack_size(8 * 1024 * 1024);
 	}
 
-	if let Ok(worker_threads) = env::var("TOKIO_WORKER_THREADS") {
-		rt_builder.worker_threads(worker_threads.parse().unwrap());
+	// Sets the thread count to available_parallelism - N
+	if let Ok(subtract_worker_threads) = env::var("TOKIO_WORKER_THREADS_SUBTRACT") {
+		rt_builder.worker_threads(
+			std::thread::available_parallelism()
+				.map_or(1, std::num::NonZeroUsize::get)
+				.saturating_sub(subtract_worker_threads.parse().unwrap())
+				.max(1),
+		);
 	}
 
 	if let Ok(max_blocking_threads) = env::var("TOKIO_MAX_BLOCKING_THREADS") {

package/packages/runtime/src/metrics.rs

@@ -1,34 +1,43 @@
-use rivet_metrics::{
-	TASK_POLL_BUCKETS,
-	otel::{global::*, metrics::*},
-};
+use rivet_metrics::{REGISTRY, TASK_POLL_BUCKETS, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-runtime");
-
 	// MARK: Tokio
-	pub static ref TOKIO_THREAD_COUNT: UpDownCounter<i64> = METER.i64_up_down_counter("rivet_tokio_thread_count")
-		.with_description("Total number of Tokio threads.")
-		.build();
-	pub static ref TOKIO_GLOBAL_QUEUE_DEPTH: Gauge<u64> = METER.u64_gauge("rivet_tokio_global_queue_depth")
-		.with_description("Number of pending tasks in the global queue.")
-		.build();
-	pub static ref TOKIO_TASK_TOTAL: Counter<u64> = METER.u64_counter("rivet_tokio_task_total")
-		.with_description("Total number of spawned tasks.")
-		.build();
-	pub static ref TOKIO_ACTIVE_TASK_COUNT: Gauge<u64> = METER.u64_gauge("rivet_tokio_active_task_count")
-		.with_description("Total number of active (running or sleeping) tasks.")
-		.build();
-	/// Expected attributes: "worker"
-	pub static ref TOKIO_WORKER_OVERFLOW_COUNT: Gauge<u64> = METER.u64_gauge("rivet_tokio_worker_overflow_count")
-		.with_description("Number of times the given worker thread saturated its local queue.")
-		.build();
-	/// Expected attributes: "worker"
-	pub static ref TOKIO_WORKER_LOCAL_QUEUE_DEPTH: Gauge<u64> = METER.u64_gauge("rivet_tokio_worker_local_queue_depth")
-		.with_description("Number of pending tasks in a worker's queue.")
-		.build();
-	pub static ref TOKIO_TASK_POLL_DURATION: Histogram<f64> = METER.f64_histogram("rivet_tokio_task_poll_duration")
-		.with_description("Duration to poll a task.")
-		.with_boundaries(TASK_POLL_BUCKETS.to_vec())
-		.build();
+	pub static ref TOKIO_THREAD_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"tokio_thread_count",
+		"Total number of Tokio threads.",
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_GLOBAL_QUEUE_DEPTH: IntGauge = register_int_gauge_with_registry!(
+		"tokio_global_queue_depth",
+		"Number of pending tasks in the global queue.",
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_TASK_TOTAL: IntCounter = register_int_counter_with_registry!(
+		"tokio_task_total",
+		"Total number of spawned tasks.",
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_ACTIVE_TASK_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"tokio_active_task_count",
+		"Total number of active (running or sleeping) tasks.",
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_WORKER_OVERFLOW_COUNT: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"tokio_worker_overflow_count",
+		"Number of times the given worker thread saturated its local queue.",
+		&["worker"],
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_WORKER_LOCAL_QUEUE_DEPTH: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"tokio_worker_local_queue_depth",
+		"Number of pending tasks in a worker's queue.",
+		&["worker"],
+		*REGISTRY
+	).unwrap();
+	pub static ref TOKIO_TASK_POLL_DURATION: Histogram = register_histogram_with_registry!(
+		"tokio_task_poll_duration",
+		"Duration to poll a task.",
+		TASK_POLL_BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
 }

package/packages/runtime/src/term_signal.rs

@@ -1,9 +1,7 @@
 use anyhow::Result;
 
-use tokio::{
-	sync::{OnceCell, watch},
-	task::JoinHandle,
-};
+use std::sync::OnceLock;
+use tokio::{sync::watch, task::JoinHandle};
 
 #[cfg(unix)]
 use tokio::signal::unix::{Signal, SignalKind, signal};
@@ -13,7 +11,7 @@ use tokio::signal::windows::ctrl_c as windows_ctrl_c;
 
 const FORCE_CLOSE_THRESHOLD: usize = 3;
 
-static HANDLER_CELL: OnceCell<(watch::Receiver<bool>, JoinHandle<()>)> = OnceCell::const_new();
+static HANDLER_CELL: OnceLock<(watch::Receiver<bool>, Option<JoinHandle<()>>)> = OnceLock::new();
 
 /// Cross-platform termination signal wrapper that handles:
 /// - Unix: SIGTERM and SIGINT
@@ -81,18 +79,30 @@ impl TermSignalHandler {
 pub struct TermSignal(watch::Receiver<bool>);
 
 impl TermSignal {
-	pub async fn new() -> Self {
+	pub fn get() -> Self {
 		let rx = HANDLER_CELL
 			.get_or_init(|| {
 				let term_signal = TermSignalHandler::new()
 					.expect("failed initializing termination signal handler");
 				let rx = term_signal.tx.subscribe();
 
-				let join_handle = tokio::spawn(term_signal.run());
-
-				std::future::ready((rx, join_handle))
+				let tokio_join_handle =
+					if std::env::var("RIVET_TEST_RUNTIME") == Ok("1".to_string()) {
+						std::thread::spawn(|| {
+							let runtime = tokio::runtime::Builder::new_current_thread()
+								.enable_all()
+								.build()
+								.expect("failed to start term signal handler runtime for tests");
+
+							runtime.block_on(term_signal.run());
+						});
+						None
+					} else {
+						Some(tokio::spawn(term_signal.run()))
+					};
+
+				(rx, tokio_join_handle)
 			})
-			.await
 			.0
 			.clone();
 
@@ -105,9 +115,12 @@ impl TermSignal {
 		*self.0.borrow()
 	}
 
+	// Stop background handler task
 	pub fn stop() {
-		if let Some((_, join_handle)) = HANDLER_CELL.get() {
-			join_handle.abort();
+		if let Some((_, tokio_join_handle)) = HANDLER_CELL.get() {
+			if let Some(tokio_join_handle) = tokio_join_handle {
+				tokio_join_handle.abort();
+			}
 		}
 	}
 }

package/packages/runtime/src/traces.rs

@@ -4,7 +4,7 @@ use console_subscriber;
 use opentelemetry::trace::{TraceContextExt, TracerProvider};
 use rivet_metrics::OtelProviderGuard;
 use std::sync::OnceLock;
-use tracing_opentelemetry::{MetricsLayer, OpenTelemetryLayer};
+use tracing_opentelemetry::OpenTelemetryLayer;
 use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, reload, util::SubscriberInitExt};
 
 type ReloadHandle = reload::Handle<EnvFilter, tracing_subscriber::Registry>;
@@ -22,24 +22,20 @@ pub fn init_tracing_subscriber(otel_providers: &Option<OtelProviderGuard>) {
 	let registry = tracing_subscriber::registry();
 
 	// Build and apply otel layers to the registry if otel is enabled
-	let (otel_trace_layer, otel_metric_layer) = match otel_providers {
+	let otel_trace_layer = match otel_providers {
 		Some(providers) => {
 			let tracer = providers.tracer_provider.tracer("tracing-otel-subscriber");
 
 			let otel_trace_layer = OpenTelemetryLayer::new(tracer)
 				.with_filter(build_filter_from_env_var("RUST_TRACE"));
 
-			let otel_metric_layer = MetricsLayer::new(providers.meter_provider.clone())
-				.with_filter(build_filter_from_env_var("RUST_TRACE"));
-
-			(Some(otel_trace_layer), Some(otel_metric_layer))
+			Some(otel_trace_layer)
 		}
-		None => (None, None),
+		None => None,
 	};
 
 	let registry = registry
 		.with(reload_layer)
-		.with(otel_metric_layer)
 		.with(otel_trace_layer)
 		.with(sentry::integrations::tracing::layer())
 		.with(SentryOtelLayer);
@@ -122,6 +118,7 @@ fn build_filter_from_spec(filter_spec: &str) -> anyhow::Result<EnvFilter> {
 		.add_directive("info".parse()?)
 		// Disable verbose logs
 		.add_directive("tokio_cron_scheduler=warn".parse()?)
+		.add_directive("async_nats=warn".parse()?)
 		.add_directive("tokio=warn".parse()?)
 		.add_directive("hyper=warn".parse()?)
 		.add_directive("h2=warn".parse()?);

package/packages/service-manager/src/lib.rs

@@ -10,6 +10,7 @@ use std::{
 
 use anyhow::{Context, Result, ensure};
 use futures_util::{StreamExt, stream::FuturesUnordered};
+use tokio::task::JoinHandle;
 
 #[derive(Clone)]
 pub struct Service {
@@ -128,6 +129,12 @@ impl RunConfigData {
 	}
 }
 
+struct ServiceTask {
+	name: String,
+	handle: JoinHandle<()>,
+	requires_graceful_shutdown: bool,
+}
+
 /// Runs services & waits for completion.
 ///
 /// Useful in order to allow for easily configuring an entrypoint where a custom set of services
@@ -135,14 +142,21 @@ impl RunConfigData {
 pub async fn start(
 	config: rivet_config::Config,
 	pools: rivet_pools::Pools,
-	services: Vec<Service>,
+	mut services: Vec<Service>,
 ) -> Result<()> {
+	services.push(Service::new(
+		"metrics",
+		ServiceKind::Core,
+		|config, _pools| rivet_metrics::run_standalone(config),
+		false,
+	));
+
 	// Spawn services
 	tracing::info!(services=?services.len(), "starting services");
 	let mut running_services = Vec::new();
 	let cron_schedule = tokio_cron_scheduler::JobScheduler::new().await?;
 
-	let mut term_signal = rivet_runtime::TermSignal::new().await;
+	let mut term_signal = rivet_runtime::TermSignal::get();
 	let shutting_down = Arc::new(AtomicBool::new(false));
 
 	for service in services {
@@ -153,8 +167,10 @@ pub async fn start(
 				let config = config.clone();
 				let pools = pools.clone();
 				let shutting_down = shutting_down.clone();
+				let task_name = format!("rivet::service::{}", service.name);
+
 				let join_handle = tokio::task::Builder::new()
-					.name(&format!("rivet::service::{}", service.name))
+					.name(&task_name)
 					.spawn(async move {
 						tracing::debug!(service=%service.name, "starting service");
 
@@ -184,14 +200,20 @@ pub async fn start(
 					})
 					.context("failed to spawn service")?;
 
-				running_services.push((service.requires_graceful_shutdown, join_handle));
+				running_services.push(ServiceTask {
+					name: task_name,
+					handle: join_handle,
+					requires_graceful_shutdown: service.requires_graceful_shutdown,
+				});
 			}
 			ServiceBehavior::Oneshot => {
 				let config = config.clone();
 				let pools = pools.clone();
 				let shutting_down = shutting_down.clone();
+				let task_name = format!("rivet::oneoff::{}", service.name);
+
 				let join_handle = tokio::task::Builder::new()
-					.name(&format!("rivet::oneoff::{}", service.name))
+					.name(&task_name)
 					.spawn(async move {
 						tracing::debug!(oneoff=%service.name, "starting oneoff");
 
@@ -217,7 +239,11 @@ pub async fn start(
 					})
 					.context("failed to spawn oneoff")?;
 
-				running_services.push((service.requires_graceful_shutdown, join_handle));
+				running_services.push(ServiceTask {
+					name: task_name,
+					handle: join_handle,
+					requires_graceful_shutdown: service.requires_graceful_shutdown,
+				});
 			}
 			ServiceBehavior::Cron(cron_config) => {
 				// Spawn immediate task
@@ -226,8 +252,10 @@ pub async fn start(
 					let config = config.clone();
 					let pools = pools.clone();
 					let shutting_down = shutting_down.clone();
+					let task_name = format!("rivet::cron_immediate::{}", service.name);
+
 					let join_handle = tokio::task::Builder::new()
-						.name(&format!("rivet::cron_immediate::{}", service.name))
+						.name(&task_name)
 						.spawn(async move {
 							tracing::debug!(cron=%service.name, "starting immediate cron");
 
@@ -255,7 +283,11 @@ pub async fn start(
 						})
 						.context("failed to spawn cron")?;
 
-					running_services.push((service.requires_graceful_shutdown, join_handle));
+					running_services.push(ServiceTask {
+						name: task_name,
+						handle: join_handle,
+						requires_graceful_shutdown: service.requires_graceful_shutdown,
+					});
 				}
 
 				// Spawn cron
@@ -263,6 +295,8 @@ pub async fn start(
 				let pools = pools.clone();
 				let service2 = service.clone();
 				let shutting_down = shutting_down.clone();
+				let task_name = format!("rivet::cron_dummy::{}", service.name);
+
 				cron_schedule
 					.add(tokio_cron_scheduler::Job::new_async_tz(
 						&cron_config.schedule,
@@ -303,10 +337,15 @@ pub async fn start(
 
 				// Add dummy task to prevent start command from stopping if theres a cron
 				let join_handle = tokio::task::Builder::new()
-					.name(&format!("rivet::cron_dummy::{}", service.name))
+					.name(&task_name)
 					.spawn(std::future::pending())
 					.context("failed creating dummy cron task")?;
-				running_services.push((false, join_handle));
+
+				running_services.push(ServiceTask {
+					name: task_name,
+					handle: join_handle,
+					requires_graceful_shutdown: false,
+				});
 			}
 		}
 	}
@@ -318,7 +357,7 @@ pub async fn start(
 		let join_fut = async {
 			let mut handle_futs = running_services
 				.iter_mut()
-				.filter_map(|(_, handle)| (!handle.is_finished()).then_some(handle))
+				.filter_map(|task| (!task.handle.is_finished()).then_some(&mut task.handle))
 				.collect::<FuturesUnordered<_>>();
 
 			while let Some(_) = handle_futs.next().await {}
@@ -330,15 +369,27 @@ pub async fn start(
 				break;
 			}
 			abort = term_signal.recv() => {
+				if !shutting_down.load(Ordering::SeqCst) {
+					// Spawn force exit task in case of a lingering task
+					let force_shutdown_duration = config.runtime.force_shutdown_duration();
+					tokio::spawn(async move {
+						tracing::info!(?force_shutdown_duration, "force shutdown timer started");
+						tokio::time::sleep(force_shutdown_duration).await;
+						tracing::warn!("force shutdown timeout reached, exiting process, this indicates a bug");
+						std::process::exit(1);
+					});
+				}
+
 				shutting_down.store(true, Ordering::SeqCst);
 
 				// Abort services that don't require graceful shutdown
-				running_services.retain(|(requires_graceful_shutdown, handle)| {
-					if !requires_graceful_shutdown {
-						handle.abort();
+				running_services.retain(|task| {
+					if !task.requires_graceful_shutdown {
+						tracing::debug!(name=%task.name, "aborting service");
+						task.handle.abort();
 					}
 
-					*requires_graceful_shutdown
+					task.requires_graceful_shutdown
 				});
 
 				if abort {

package/packages/test-deps/src/datacenter.rs

@@ -1,4 +1,5 @@
-use anyhow::*;
+use anyhow::Result;
+use std::collections::HashMap;
 use uuid::Uuid;
 
 use crate::TestDeps;
@@ -9,7 +10,7 @@ use rivet_test_deps_docker::{TestDatabase, TestPubSub};
 pub async fn setup_single_datacenter(
 	test_id: Uuid,
 	dc_label: u16,
-	datacenters: Vec<rivet_config::config::topology::Datacenter>,
+	datacenters: HashMap<String, rivet_config::config::topology::Datacenter>,
 	api_peer_port: u16,
 	guard_port: u16,
 ) -> Result<TestDeps> {
@@ -18,8 +19,9 @@ pub async fn setup_single_datacenter(
 
 	let dc = datacenters
 		.iter()
-		.find(|x| x.datacenter_label == dc_label)
+		.find(|(_, x)| x.datacenter_label == dc_label)
 		.expect("invalid dc label")
+		.1
 		.clone();
 
 	tracing::info!(
@@ -37,14 +39,13 @@ pub async fn setup_single_datacenter(
 		let was_started = docker_config.start().await?;
 		container_names.push(docker_config.container_name.clone());
 
-		// If Postgres was just started, wait for it to be ready
-		if was_started && test_database == TestDatabase::Postgres {
+		if was_started {
 			tracing::info!(
 				dc = dc.datacenter_label,
 				port = docker_config.port_mapping.0,
-				"waiting for Postgres to be ready"
+				"waiting for database to be ready"
 			);
-			TestDatabase::wait_for_postgres_ready(docker_config.port_mapping.0, 10).await?;
+			test_database.wait_for_ready(&docker_config).await?;
 		}
 	}
 
@@ -80,13 +81,26 @@ pub async fn setup_single_datacenter(
 
 	root.topology = Some(rivet_config::config::topology::Topology {
 		datacenter_label: dc.datacenter_label,
-		datacenters,
+		datacenters: rivet_config::config::topology::DatacentersRepr::Map(datacenters),
 	});
 
 	root.guard = Some(rivet_config::config::guard::Guard {
 		host: None,
 		port: Some(guard_port),
 		https: None,
+		..Default::default()
+	});
+	root.metrics = rivet_config::config::metrics::Metrics {
+		host: None,
+		port: Some(0),
+	};
+
+	// Use short timeouts for tests
+	root.pegboard = Some(rivet_config::config::pegboard::Pegboard {
+		actor_start_threshold: Some(3_000), // 3 seconds instead of 30
+		serverless_base_retry_timeout: Some(500), // 500ms instead of 2s
+		serverless_backoff_max_exponent: Some(2), // Max 2^2 = 4x base = 2s
+		..Default::default()
 	});
 
 	tracing::info!(