@rivetkit/engine 1.0.0

package/packages/runtime/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "rivet-runtime"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+console-subscriber.workspace = true
+lazy_static.workspace = true
+opentelemetry_sdk.workspace = true
+opentelemetry-otlp.workspace = true
+opentelemetry-semantic-conventions.workspace = true
+opentelemetry-stdout.workspace = true
+opentelemetry.workspace = true
+reqwest.workspace = true
+rivet-env.workspace = true
+rivet-metrics.workspace = true
+sentry.workspace = true
+serde_json.workspace = true
+serde.workspace = true
+thiserror.workspace = true
+tokio.workspace = true
+tracing-core.workspace = true
+tracing-logfmt.workspace = true
+tracing-opentelemetry.workspace = true
+tracing-subscriber.workspace = true
+tracing.workspace = true

package/packages/runtime/src/lib.rs

@@ -0,0 +1,154 @@
+use std::{env, future::Future, sync::Arc, time::Duration};
+
+use rivet_metrics::{KeyValue, init_otel_providers};
+use tokio::sync::{Notify, OnceCell};
+mod metrics;
+mod traces;
+
+mod term_signal;
+
+pub use term_signal::TermSignal;
+pub use traces::reload_log_filter;
+
+static SHUTDOWN: OnceCell<Arc<Notify>> = OnceCell::const_new();
+
+/// Returns `None` if the runtime was shut down manually.
+pub fn run<F: Future>(f: F) -> Option<F::Output> {
+	// Build runtime
+	let mut rt_builder = build_tokio_runtime_builder();
+	let rt = rt_builder.build().expect("failed to build tokio runtime");
+	let output = rt.block_on(async move {
+		let notify = SHUTDOWN
+			.get_or_init(|| std::future::ready(Arc::new(Notify::new())))
+			.await
+			.clone();
+
+		// Must be called from within a tokio context
+		let providers = init_otel_providers();
+		traces::init_tracing_subscriber(&providers);
+
+		tokio::select! {
+			_ = notify.notified() => {
+				tracing::info!("shutting down runtime");
+				None
+			},
+			res = f => Some(res),
+		}
+	});
+
+	output
+}
+
+/// Shuts down the entire rivet runtime, if one is running.
+pub async fn shutdown() {
+	if let Some(shutdown) = SHUTDOWN.get() {
+		shutdown.notify_one();
+	} else {
+		tracing::error!("no runtime to shutdown");
+	};
+}
+
+fn build_tokio_runtime_builder() -> tokio::runtime::Builder {
+	let mut rt_builder = tokio::runtime::Builder::new_multi_thread();
+	rt_builder.enable_all();
+
+	rt_builder.on_thread_start(move || {
+		metrics::TOKIO_THREAD_COUNT.add(1, &[]);
+	});
+	rt_builder.on_thread_stop(move || {
+		metrics::TOKIO_THREAD_COUNT.add(-1, &[]);
+	});
+
+	rt_builder.on_task_spawn(move |_| metrics::TOKIO_TASK_TOTAL.add(1, &[]));
+
+	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();
+
+			metrics::TOKIO_GLOBAL_QUEUE_DEPTH.record(metrics.global_queue_depth() as u64, &[]);
+			metrics::TOKIO_ACTIVE_TASK_COUNT.record(metrics.num_alive_tasks() as u64, &[]);
+
+			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(), &[]);
+				// 	}
+				// }
+			}
+		});
+
+		rt_builder.metrics_poll_time_histogram_configuration(
+			tokio::runtime::HistogramConfiguration::log(
+				tokio::runtime::LogHistogram::builder()
+					.min_value(Duration::from_micros(20))
+					.max_value(Duration::from_millis(32))
+					.precision_exact(0)
+					.max_buckets(20)
+					.unwrap(),
+			),
+		);
+		rt_builder.enable_metrics_poll_time_histogram();
+	}
+
+	if let Ok(thread_stack_size) = env::var("TOKIO_THREAD_STACK_SIZE") {
+		rt_builder.thread_stack_size(thread_stack_size.parse().unwrap());
+	} else {
+		// async-nats requires a fat stack
+		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());
+	}
+
+	if let Ok(max_blocking_threads) = env::var("TOKIO_MAX_BLOCKING_THREADS") {
+		rt_builder.max_blocking_threads(max_blocking_threads.parse().unwrap());
+	}
+
+	if let Ok(global_queue_interval) = env::var("TOKIO_GLOBAL_QUEUE_INTERVAL") {
+		rt_builder.global_queue_interval(global_queue_interval.parse().unwrap());
+	}
+
+	if let Ok(event_interval) = env::var("TOKIO_EVENT_INTERVAL") {
+		rt_builder.event_interval(event_interval.parse().unwrap());
+	}
+
+	if let Ok(thread_keep_alive) = env::var("TOKIO_THREAD_KEEP_ALIVE") {
+		rt_builder.thread_keep_alive(Duration::from_millis(thread_keep_alive.parse().unwrap()));
+	}
+
+	rt_builder
+}

package/packages/runtime/src/metrics.rs

@@ -0,0 +1,34 @@
+use rivet_metrics::{
+	TASK_POLL_BUCKETS,
+	otel::{global::*, metrics::*},
+};
+
+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();
+}

package/packages/runtime/src/term_signal.rs

@@ -0,0 +1,113 @@
+use anyhow::Result;
+
+use tokio::{
+	sync::{OnceCell, watch},
+	task::JoinHandle,
+};
+
+#[cfg(unix)]
+use tokio::signal::unix::{Signal, SignalKind, signal};
+
+#[cfg(windows)]
+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();
+
+/// Cross-platform termination signal wrapper that handles:
+/// - Unix: SIGTERM and SIGINT
+/// - Windows: Ctrl+C
+struct TermSignalHandler {
+	count: usize,
+	tx: watch::Sender<bool>,
+
+	#[cfg(unix)]
+	sigterm: Signal,
+	#[cfg(unix)]
+	sigint: Signal,
+	#[cfg(windows)]
+	ctrl_c: tokio::signal::windows::CtrlC,
+}
+
+impl TermSignalHandler {
+	/// Returns existing termination signal handler or initializes it.
+	fn new() -> Result<Self> {
+		tracing::debug!("initialized termination signal handler");
+
+		Ok(Self {
+			count: 0,
+			tx: watch::channel(false).0,
+
+			#[cfg(unix)]
+			sigterm: signal(SignalKind::terminate())?,
+			#[cfg(unix)]
+			sigint: signal(SignalKind::interrupt())?,
+			#[cfg(windows)]
+			ctrl_c: windows_ctrl_c()?,
+		})
+	}
+
+	async fn run(mut self) {
+		loop {
+			#[cfg(unix)]
+			{
+				tokio::select! {
+					_ = self.sigterm.recv() => {}
+					_ = self.sigint.recv() => {}
+				}
+			}
+
+			#[cfg(windows)]
+			{
+				self.ctrl_c.recv().await;
+			}
+
+			self.count += 1;
+
+			if self.count == 1 {
+				tracing::info!("received SIGTERM");
+			} else {
+				tracing::warn!(count=%self.count, "received another SIGTERM");
+			}
+
+			if self.tx.send(self.count >= FORCE_CLOSE_THRESHOLD).is_err() {
+				tracing::debug!("no sigterm subscribers");
+			}
+		}
+	}
+}
+
+pub struct TermSignal(watch::Receiver<bool>);
+
+impl TermSignal {
+	pub async fn new() -> 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))
+			})
+			.await
+			.0
+			.clone();
+
+		TermSignal(rx)
+	}
+
+	/// Returns true if the user should abort any graceful attempt to shutdown and shutdown immediately.
+	pub async fn recv(&mut self) -> bool {
+		let _ = self.0.changed().await;
+		*self.0.borrow()
+	}
+
+	pub fn stop() {
+		if let Some((_, join_handle)) = HANDLER_CELL.get() {
+			join_handle.abort();
+		}
+	}
+}

package/packages/runtime/src/traces.rs

@@ -0,0 +1,158 @@
+// Based off of https://github.com/tokio-rs/tracing-opentelemetry/blob/v0.1.x/examples/opentelemetry-otlp.rs
+
+use console_subscriber;
+use opentelemetry::trace::{TraceContextExt, TracerProvider};
+use rivet_metrics::OtelProviderGuard;
+use std::sync::OnceLock;
+use tracing_opentelemetry::{MetricsLayer, OpenTelemetryLayer};
+use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, reload, util::SubscriberInitExt};
+
+type ReloadHandle = reload::Handle<EnvFilter, tracing_subscriber::Registry>;
+
+static RELOAD_HANDLE: OnceLock<ReloadHandle> = OnceLock::new();
+
+/// Initialize tracing-subscriber
+pub fn init_tracing_subscriber(otel_providers: &Option<OtelProviderGuard>) {
+	// Create reloadable env filter for RUST_LOG
+	let (reload_layer, reload_handle) = reload::Layer::new(build_filter_from_env_var("RUST_LOG"));
+
+	// Store handle globally for later reloading
+	let _ = RELOAD_HANDLE.set(reload_handle);
+
+	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 {
+		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))
+		}
+		None => (None, None),
+	};
+
+	let registry = registry
+		.with(reload_layer)
+		.with(otel_metric_layer)
+		.with(otel_trace_layer)
+		.with(sentry::integrations::tracing::layer())
+		.with(SentryOtelLayer);
+
+	// Check if tokio console is enabled
+	let enable_tokio_console = std::env::var("TOKIO_CONSOLE_ENABLE").map_or(false, |x| x == "1");
+
+	registry
+		.with(
+			// Add tokio console if its enabled
+			//
+			// This code is here because console layer depends
+			// on tracing_subscriber's weird layered registry type.
+			if enable_tokio_console {
+				Some(
+					console_subscriber::ConsoleLayer::builder()
+						.with_default_env()
+						.spawn(),
+				)
+			} else {
+				None
+			},
+		)
+		.with(
+			tracing_logfmt::builder()
+				.with_span_name(std::env::var("RUST_LOG_SPAN_NAME").map_or(false, |x| x == "1"))
+				.with_span_path(std::env::var("RUST_LOG_SPAN_PATH").map_or(false, |x| x == "1"))
+				.with_target(std::env::var("RUST_LOG_TARGET").map_or(false, |x| x == "1"))
+				.with_location(std::env::var("RUST_LOG_LOCATION").map_or(false, |x| x == "1"))
+				.with_module_path(std::env::var("RUST_LOG_MODULE_PATH").map_or(false, |x| x == "1"))
+				.with_ansi_color(std::env::var("RUST_LOG_ANSI_COLOR").map_or(false, |x| x == "1"))
+				.layer(),
+		)
+		.init()
+}
+
+/// Replaces sentry's trace id with otel's trace id
+struct SentryOtelLayer;
+
+impl<S> Layer<S> for SentryOtelLayer
+where
+	S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
+{
+	fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
+		if let Some(span) = ctx.event_span(event) {
+			// The OTel layer stores the OTel context in span extensions
+			let extensions = span.extensions();
+
+			if let Some(otel_data) = extensions.get::<tracing_opentelemetry::OtelData>() {
+				let span = otel_data.parent_cx.span();
+				let span_context = span.span_context();
+				let trace_id = if span_context.is_valid() {
+					Some(span_context.trace_id())
+				} else {
+					otel_data.builder.trace_id
+				};
+
+				if let (Some(trace_id), Some(span_id)) = (trace_id, otel_data.builder.span_id) {
+					sentry::configure_scope(|scope| {
+						scope.set_context(
+							"trace",
+							sentry::protocol::TraceContext {
+								trace_id: trace_id.to_bytes().into(),
+								span_id: span_id.to_bytes().into(),
+								..Default::default()
+							},
+						);
+					});
+				}
+			}
+		}
+	}
+}
+
+/// Build an EnvFilter from a filter specification string
+fn build_filter_from_spec(filter_spec: &str) -> anyhow::Result<EnvFilter> {
+	// Create env filter with defaults
+	let mut env_filter = EnvFilter::default()
+		// Default filter
+		.add_directive("info".parse()?)
+		// Disable verbose logs
+		.add_directive("tokio_cron_scheduler=warn".parse()?)
+		.add_directive("tokio=warn".parse()?)
+		.add_directive("hyper=warn".parse()?)
+		.add_directive("h2=warn".parse()?);
+
+	// Add user-provided directives
+	for s in filter_spec.split(',').filter(|x| !x.is_empty()) {
+		env_filter = env_filter.add_directive(s.parse()?);
+	}
+
+	Ok(env_filter)
+}
+
+/// Build an EnvFilter by reading from an environment variable
+fn build_filter_from_env_var(env_var_name: &str) -> EnvFilter {
+	let filter_spec = std::env::var(env_var_name).unwrap_or_default();
+	build_filter_from_spec(&filter_spec).expect("invalid env filter")
+}
+
+/// Reload the log filter with a new specification
+pub fn reload_log_filter(filter_spec: &str) -> anyhow::Result<()> {
+	let handle = RELOAD_HANDLE
+		.get()
+		.ok_or_else(|| anyhow::anyhow!("reload handle not initialized"))?;
+
+	// Build the new filter
+	let env_filter = build_filter_from_spec(filter_spec)?;
+
+	// Reload the filter
+	handle.reload(env_filter)?;
+
+	tracing::debug!(?filter_spec, "reloaded log filter");
+
+	Ok(())
+}

package/packages/service-manager/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "rivet-service-manager"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+chrono.workspace = true
+futures-util.workspace = true
+include_dir.workspace = true
+rivet-config.workspace = true
+rivet-metrics.workspace = true
+rivet-pools.workspace = true
+rivet-runtime.workspace = true
+tokio-cron-scheduler.workspace = true
+tokio.workspace = true
+tracing.workspace = true