Skip to main content

rivetkit_engine_process/
lib.rs

1use std::process::Stdio;
2use std::time::{Duration, Instant};
3use std::{
4	env,
5	path::{Path, PathBuf},
6};
7
8use anyhow::{Context, Result};
9use reqwest::{Client, Url};
10use serde::{Deserialize, Serialize};
11use sha2::{Digest, Sha256};
12use tokio::io::AsyncReadExt;
13use tokio::process::{Child, Command};
14use tokio::sync::oneshot;
15use tokio::task::JoinHandle;
16
17mod error;
18
19pub use error::EngineProcessError;
20
21const ENGINE_RUNTIME: &str = "engine";
22/// Engine architecture this RivetKit requires. v2 engines report the same runtime and
23/// version but no architecture, so the field alone tells them apart.
24const ENGINE_ARCHITECTURE: &str = "v3";
25const RIVETKIT_RUNTIME: &str = "rivetkit";
26const ENGINE_VERSION_ENV: &str = "RIVETKIT_ENGINE_VERSION";
27const RELEASES_ENDPOINT_ENV: &str = "RIVETKIT_ENGINE_RELEASES_ENDPOINT";
28const RELEASES_ENDPOINT: &str = "https://releases.rivet.dev";
29const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(60);
30
31#[derive(Debug, Deserialize)]
32struct EngineHealthResponse {
33	status: Option<String>,
34	runtime: Option<String>,
35	version: Option<String>,
36	architecture: Option<String>,
37}
38
39#[derive(Clone, Debug)]
40pub struct EngineResolverConfig {
41	pub endpoint: String,
42	pub explicit_binary_path: Option<PathBuf>,
43	pub bind_host: Option<String>,
44	pub bind_port: Option<u16>,
45	pub public_url: Option<String>,
46	pub auto_download: bool,
47	pub version: String,
48	pub releases_endpoint: String,
49}
50
51impl EngineResolverConfig {
52	pub fn from_parts(
53		endpoint: &str,
54		explicit_binary_path: Option<PathBuf>,
55		bind_host: Option<String>,
56		bind_port: Option<u16>,
57		auto_download: bool,
58	) -> Self {
59		Self {
60			endpoint: endpoint.to_owned(),
61			explicit_binary_path,
62			bind_host,
63			bind_port,
64			public_url: None,
65			auto_download,
66			version: std::env::var(ENGINE_VERSION_ENV)
67				.unwrap_or_else(|_| env!("CARGO_PKG_VERSION").to_owned()),
68			releases_endpoint: std::env::var(RELEASES_ENDPOINT_ENV)
69				.unwrap_or_else(|_| RELEASES_ENDPOINT.to_owned()),
70		}
71	}
72}
73
74#[derive(Debug, PartialEq, Eq)]
75pub enum ResolvedEngine {
76	Existing,
77	Binary(PathBuf),
78}
79
80/// Effective runtime parameters of a spawned engine, stamped to disk so a later
81/// process that reattaches can tell how the running engine is configured without
82/// probing it over the network.
83#[derive(Clone, Debug, Serialize, Deserialize)]
84pub struct EngineRuntimeStamp {
85	pub pid: u32,
86	/// Linux process identity, independent of PID reuse and VM restarts.
87	/// Older records and other platforms retain conservative PID-only checks.
88	#[serde(default, skip_serializing_if = "Option::is_none")]
89	pub process_identity: Option<ProcessIdentity>,
90	pub endpoint: String,
91	pub bind_host: String,
92	pub public_url: Option<String>,
93}
94
95#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
96pub struct ProcessIdentity {
97	pub boot_id: String,
98	pub start_time_ticks: u64,
99}
100
101impl EngineRuntimeStamp {
102	/// True if the engine bound a loopback address only, so peers on other hosts
103	/// (for example a container reaching the host) cannot connect to it.
104	pub fn binds_loopback_only(&self) -> bool {
105		if self.bind_host == "localhost" {
106			return true;
107		}
108		self.bind_host
109			.parse::<std::net::IpAddr>()
110			.map(|ip| ip.is_loopback())
111			.unwrap_or(false)
112	}
113}
114
115/// How `start_or_reuse` satisfied the request.
116#[derive(Debug)]
117pub enum EngineStartup {
118	/// A new engine was spawned with the requested configuration.
119	Spawned,
120	/// An already-running engine was reused. `stamp` is its recorded runtime
121	/// parameters, or `None` when no live stamp was found (an engine from an
122	/// older CLI, started by hand, or a stale stamp whose process is gone).
123	Reused { stamp: Option<EngineRuntimeStamp> },
124}
125
126/// Manages the rivet-engine subprocess.
127///
128/// A spawned engine belongs to this host and is stopped on shutdown. An engine
129/// found at the endpoint is only reused, never stopped by this manager.
130#[derive(Debug)]
131pub struct EngineProcessManager {
132	watcher: Option<JoinHandle<()>>,
133	stop: Option<oneshot::Sender<()>>,
134	startup: EngineStartup,
135}
136
137impl EngineProcessManager {
138	pub async fn start_or_reuse(config: EngineResolverConfig) -> Result<Self> {
139		// Remove a dead process's record before deciding whether to attach or spawn.
140		// The endpoint probe below is still authoritative for a live engine.
141		let _ = read_engine_stamp();
142		let resolved = resolve_engine_binary(&config).await?;
143		Self::start_resolved(resolved, &config).await
144	}
145
146	/// How the engine was obtained: freshly spawned or reused.
147	pub fn startup(&self) -> &EngineStartup {
148		&self.startup
149	}
150
151	pub async fn shutdown(mut self) {
152		if let Some(stop) = self.stop.take() {
153			let _ = stop.send(());
154		}
155		if let Some(watcher) = self.watcher.take() {
156			let _ = watcher.await;
157		}
158	}
159
160	async fn start_resolved(
161		resolved: ResolvedEngine,
162		config: &EngineResolverConfig,
163	) -> Result<Self> {
164		let endpoint = &config.endpoint;
165		if matches!(resolved, ResolvedEngine::Existing) {
166			tracing::info!(
167				endpoint = %endpoint,
168				"reusing already-running engine process"
169			);
170			return Ok(Self {
171				watcher: None,
172				stop: None,
173				startup: EngineStartup::Reused {
174					stamp: read_engine_stamp(),
175				},
176			});
177		}
178
179		let ResolvedEngine::Binary(binary_path) = resolved else {
180			unreachable!("existing engine handled above");
181		};
182		let mut health = probe_existing_engine(endpoint).await?;
183		if health.is_none() {
184			if let Some(stamp) = read_engine_stamp().filter(|stamp| stamp.endpoint == *endpoint) {
185				// Shutdown closes the listener before releasing the storage locks.
186				// A watcher may restart its app while that Engine is still draining.
187				wait_for_engine_exit(&stamp, Duration::from_secs(20)).await?;
188				// Another app may have completed startup while we waited.
189				health = probe_existing_engine(endpoint).await?;
190			}
191		}
192		if let Some(health) = health {
193			tracing::info!(
194				endpoint = %endpoint,
195				status = ?health.status,
196				runtime = ?health.runtime,
197				version = ?health.version,
198				"reusing already-running engine process"
199			);
200			return Ok(Self {
201				watcher: None,
202				stop: None,
203				startup: EngineStartup::Reused {
204					stamp: read_engine_stamp(),
205				},
206			});
207		}
208
209		if !binary_path.exists() {
210			return Err(EngineProcessError::BinaryNotFound {
211				path: binary_path.display().to_string(),
212			}
213			.build());
214		}
215
216		ensure_compatible_binary(&binary_path).await?;
217
218		let env = engine_env(config)?;
219		let config_path = write_engine_config(config)?;
220		let data_path = engine_data_path()?;
221		let logs_dir = storage_root()?
222			.join("var")
223			.join("logs")
224			.join("rivet-engine");
225		ensure_dir(&logs_dir).context("create engine logs directory")?;
226		ensure_dir(&engine_log_directory()?).context("create engine log shard directory")?;
227
228		let timestamp = log_timestamp();
229		let stdout_log_path = logs_dir.join(format!("engine-{timestamp}-stdout.log"));
230		let stderr_log_path = logs_dir.join(format!("engine-{timestamp}-stderr.log"));
231		let stdout_file = open_log_file(&stdout_log_path)
232			.with_context(|| format!("open engine stdout log `{}`", stdout_log_path.display()))?;
233		let stderr_file = open_log_file(&stderr_log_path)
234			.with_context(|| format!("open engine stderr log `{}`", stderr_log_path.display()))?;
235
236		let mut command = Command::new(&binary_path);
237		command.arg("--config").arg(&config_path);
238		for (key, value) in &env {
239			command.env(key, value);
240		}
241		command
242			.stdin(Stdio::null())
243			.stdout(Stdio::from(stdout_file))
244			.stderr(Stdio::from(stderr_file));
245
246		// Do not send terminal signals to the engine twice: the host drains its
247		// actors first, then asks the watcher to stop the engine.
248		#[cfg(unix)]
249		command.process_group(0);
250		// A file watcher can terminate its child without running Rust destructors.
251		// Linux still sends the engine SIGTERM when that host process disappears.
252		#[cfg(target_os = "linux")]
253		unsafe {
254			use std::os::unix::process::CommandExt;
255			let parent_pid = libc::getpid();
256			command.as_std_mut().pre_exec(move || {
257				if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) != 0 {
258					return Err(std::io::Error::last_os_error());
259				}
260				if libc::getppid() != parent_pid {
261					libc::raise(libc::SIGTERM);
262				}
263				Ok(())
264			});
265		}
266
267		let mut child = command
268			.spawn()
269			.with_context(|| format!("spawn engine binary `{}`", binary_path.display()))?;
270		let pid = child
271			.id()
272			.ok_or_else(|| EngineProcessError::MissingPid.build())?;
273
274		tracing::info!(
275			pid,
276			path = %binary_path.display(),
277			endpoint = %endpoint,
278			data_path = %data_path.display(),
279			"spawned engine process"
280		);
281		tracing::info!(
282			stdout_log = %stdout_log_path.display(),
283			stderr_log = %stderr_log_path.display(),
284			"engine stdout/stderr piped to log files"
285		);
286
287		let health_url = engine_health_url(endpoint);
288		let health = match wait_for_engine_health(&health_url).await {
289			Ok(health) => health,
290			Err(error) => {
291				let error = match child.try_wait() {
292					Ok(Some(status)) => error.context(format!(
293						"engine process exited before becoming healthy with status {status}"
294					)),
295					Ok(None) => error,
296					Err(wait_error) => error.context(format!(
297						"failed to inspect engine process status: {wait_error:#}"
298					)),
299				};
300				if let Err(cleanup_error) = terminate_failed_spawn(&mut child).await {
301					tracing::warn!(
302						?cleanup_error,
303						"failed to terminate engine process that never became healthy"
304					);
305				}
306				return Err(error);
307			}
308		};
309
310		tracing::info!(
311			pid,
312			status = ?health.status,
313			runtime = ?health.runtime,
314			version = ?health.version,
315			"engine process is healthy"
316		);
317
318		// An explicit or cached binary may be an older engine. Never leave one running.
319		let ready = match ensure_compatible_engine(endpoint, &health) {
320			Ok(()) => wait_for_engine_ready(endpoint).await,
321			Err(error) => Err(error),
322		};
323		if let Err(error) = ready {
324			if let Err(cleanup_error) = terminate_failed_spawn(&mut child).await {
325				tracing::warn!(
326					?cleanup_error,
327					"failed to terminate engine process that never became ready"
328				);
329			}
330			return Err(error);
331		}
332		tracing::info!(pid, "engine process is ready");
333
334		// Record the effective runtime parameters so a later reattaching process
335		// can tell how this engine is bound. Best-effort: a failed write only
336		// costs that later process its reuse diagnostics.
337		let stamp = EngineRuntimeStamp {
338			pid,
339			process_identity: read_process_identity(pid),
340			endpoint: endpoint.clone(),
341			bind_host: resolve_bind_host(config)?,
342			public_url: config.public_url.clone(),
343		};
344		if let Err(error) = write_engine_stamp(&stamp) {
345			tracing::warn!(?error, "failed to write engine runtime stamp");
346		}
347
348		let (watcher, stop) = spawn_engine_watcher(child, pid);
349		Ok(Self {
350			watcher: Some(watcher),
351			stop: Some(stop),
352			startup: EngineStartup::Spawned,
353		})
354	}
355}
356
357/// Waits without signalling a previous Engine that may still own storage locks.
358async fn wait_for_engine_exit(stamp: &EngineRuntimeStamp, timeout: Duration) -> Result<()> {
359	let pid = stamp.pid;
360	tracing::info!(
361		pid,
362		"waiting for the previous local engine to finish stopping"
363	);
364	tokio::time::timeout(timeout, async {
365		while stamp_process_is_alive(stamp) {
366			tokio::time::sleep(Duration::from_millis(100)).await;
367		}
368	})
369	.await
370	.with_context(|| {
371		format!(
372			"previous local Engine process {pid} is still stopping; refusing to start another Engine against its storage"
373		)
374	})?;
375	Ok(())
376}
377
378/// How long the `--runtime-info` probe may run before the binary is treated as incompatible.
379// A freshly linked native binary can spend several seconds in the OS loader
380// before reaching main (notably macOS first-launch validation). Keep the probe
381// bounded without rejecting a compatible Engine just because its launch is cold.
382const RUNTIME_INFO_TIMEOUT: Duration = Duration::from_secs(15);
383/// The largest `--runtime-info` response read from a probed binary.
384const RUNTIME_INFO_MAX_BYTES: usize = 8192;
385
386/// Probe the executable before creating local data or starting services. Package versions
387/// do not distinguish the unreleased v3 architecture from released v2 binaries.
388async fn ensure_compatible_binary(path: &Path) -> Result<()> {
389	let incompatible = |reason: &str| {
390		EngineProcessError::IncompatibleBinary {
391			path: path.display().to_string(),
392			reason: reason.to_owned(),
393			required: ENGINE_ARCHITECTURE.to_owned(),
394		}
395		.build()
396	};
397	let mut child = Command::new(path)
398		.arg("--runtime-info")
399		.stdin(Stdio::null())
400		.stdout(Stdio::piped())
401		.stderr(Stdio::null())
402		.kill_on_drop(true)
403		.spawn()
404		.with_context(|| format!("probe engine binary `{}`", path.display()))?;
405	let stdout = child.stdout.take().context("probe stdout was not piped")?;
406	let probe = async {
407		let read = async {
408			let mut bytes = Vec::new();
409			let limit = RUNTIME_INFO_MAX_BYTES as u64 + 1;
410			stdout.take(limit).read_to_end(&mut bytes).await?;
411			if bytes.len() > RUNTIME_INFO_MAX_BYTES {
412				return Err(std::io::Error::other(
413					"runtime identity response exceeds its size limit",
414				));
415			}
416			Ok(bytes)
417		};
418		tokio::try_join!(child.wait(), read)
419	};
420	let (status, bytes) = match tokio::time::timeout(RUNTIME_INFO_TIMEOUT, probe).await {
421		Ok(Ok(output)) => output,
422		failure => {
423			// Kill and reap before reporting failure; kill_on_drop also covers caller cancellation.
424			child
425				.kill()
426				.await
427				.context("terminate failed runtime identity probe")?;
428			return Err(incompatible(if failure.is_err() {
429				"runtime identity probe timed out"
430			} else {
431				"runtime identity probe output could not be read"
432			}));
433		}
434	};
435	if !status.success() {
436		return Err(incompatible(
437			"runtime identity probe is unsupported or failed",
438		));
439	}
440	let info: EngineHealthResponse = serde_json::from_slice(&bytes)
441		.map_err(|_| incompatible("runtime identity response is invalid"))?;
442	if info.runtime.as_deref() != Some(ENGINE_RUNTIME)
443		|| info.architecture.as_deref() != Some(ENGINE_ARCHITECTURE)
444	{
445		return Err(incompatible(
446			"runtime identity does not advertise the required architecture",
447		));
448	}
449	Ok(())
450}
451
452/// Path to the local engine's Depot data under the shared storage root. It holds the log
453/// shard and the object store, and is kept apart from the `db` directory older engines used.
454pub fn engine_data_path() -> Result<PathBuf> {
455	Ok(storage_root()?.join("var").join("engine").join("depot"))
456}
457
458/// The local engine's log shard directory, which must exist before the engine starts. The
459/// engine creates its object store directory itself.
460pub fn engine_log_directory() -> Result<PathBuf> {
461	Ok(engine_data_path()?.join("log"))
462}
463
464fn admin_token() -> String {
465	env::var("RIVET__AUTH__ADMIN_TOKEN").unwrap_or_else(|_| "default".to_owned())
466}
467
468fn engine_stamp_path() -> Result<PathBuf> {
469	Ok(storage_root()?
470		.join("var")
471		.join("engine")
472		.join("runtime.json"))
473}
474
475/// The address the engine actually binds: the explicit `bind_host` override, or
476/// the endpoint host when unset. Must match `engine_env`'s guard host.
477fn resolve_bind_host(config: &EngineResolverConfig) -> Result<String> {
478	if let Some(bind_host) = &config.bind_host {
479		return Ok(bind_host.clone());
480	}
481	let endpoint = &config.endpoint;
482	let endpoint_url =
483		Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
484	Ok(endpoint_url
485		.host_str()
486		.ok_or_else(|| invalid_endpoint(endpoint, "missing host"))?
487		.to_owned())
488}
489
490fn write_engine_stamp(stamp: &EngineRuntimeStamp) -> Result<()> {
491	let path = engine_stamp_path()?;
492	if let Some(dir) = path.parent() {
493		ensure_dir(dir)?;
494	}
495	let contents = serde_json::to_vec_pretty(stamp).context("serialize engine stamp")?;
496	std::fs::write(&path, contents)
497		.with_context(|| format!("write engine stamp `{}`", path.display()))?;
498	Ok(())
499}
500
501/// Reads the runtime stamp of the currently running engine, if its recorded
502/// process is still alive. A stale stamp (process gone) reads as `None` so
503/// callers treat the binding as unknown rather than trusting old data. On
504/// platforms without a liveness check (see `pid_is_alive`) the stamp is always
505/// treated as unknown for the same reason.
506fn read_engine_stamp() -> Option<EngineRuntimeStamp> {
507	let path = engine_stamp_path().ok()?;
508	read_engine_stamp_from(&path)
509}
510
511fn read_engine_stamp_from(path: &Path) -> Option<EngineRuntimeStamp> {
512	let contents = std::fs::read(path).ok()?;
513	let stamp: EngineRuntimeStamp = serde_json::from_slice(&contents).ok()?;
514	if !stamp_process_is_alive(&stamp) {
515		let _ = std::fs::remove_file(path);
516		return None;
517	}
518	Some(stamp)
519}
520
521/// Missing identity (legacy records, unavailable procfs, or other platforms)
522/// must not bypass the protection against a real Engine still draining.
523fn stamp_process_is_alive(stamp: &EngineRuntimeStamp) -> bool {
524	if !pid_is_alive(stamp.pid) {
525		return false;
526	}
527	match (&stamp.process_identity, read_process_identity(stamp.pid)) {
528		(Some(recorded), Some(current)) => recorded == &current,
529		_ => true,
530	}
531}
532
533#[cfg(target_os = "linux")]
534fn read_process_identity(pid: u32) -> Option<ProcessIdentity> {
535	let boot_id = std::fs::read_to_string("/proc/sys/kernel/random/boot_id").ok()?;
536	let boot_id = boot_id.trim();
537	if boot_id.is_empty() {
538		return None;
539	}
540	let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
541	Some(ProcessIdentity {
542		boot_id: boot_id.to_owned(),
543		start_time_ticks: process_start_time_ticks(&stat)?,
544	})
545}
546
547#[cfg(target_os = "linux")]
548fn process_start_time_ticks(stat: &str) -> Option<u64> {
549	// Field 2 (comm) may contain spaces and parentheses. Start time is field
550	// 22, or offset 19 after the final closing parenthesis of comm.
551	let (_, fields) = stat.rsplit_once(')')?;
552	fields.split_whitespace().nth(19)?.parse().ok()
553}
554
555#[cfg(not(target_os = "linux"))]
556fn read_process_identity(_pid: u32) -> Option<ProcessIdentity> {
557	None
558}
559
560#[cfg(unix)]
561fn pid_is_alive(pid: u32) -> bool {
562	// `kill(pid, 0)` sends no signal but performs the existence and permission
563	// checks: success or EPERM means the process exists, ESRCH means it is gone.
564	let result = unsafe { libc::kill(pid as libc::pid_t, 0) };
565	result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
566}
567
568#[cfg(not(unix))]
569fn pid_is_alive(_pid: u32) -> bool {
570	// No portable liveness check here, so fail safe: report the process as gone
571	// so a stale stamp is never trusted. The binding check then falls back to
572	// its "unknown, warn" branch instead of silently trusting old data.
573	false
574}
575
576/// Computes the environment variables that configure a rivet-engine process
577/// for the given endpoint. The engine requires the log directory to exist; see
578/// [`engine_log_directory`].
579///
580/// Shared by the spawn path and by callers that exec the engine binary
581/// directly (for example the CLI `engine` proxy) so both operate on the same
582/// data, guard, api-peer, and metrics ports. These callers run one local
583/// development engine: it is the only Depot storage worker, keeps its cold
584/// data in a local directory, provisions its own tablet on first start, and uses
585/// shorter recovery and shutdown thresholds than the production defaults.
586pub fn engine_env(config: &EngineResolverConfig) -> Result<Vec<(String, String)>> {
587	let ports = EnginePorts::resolve(config)?;
588	let log_directory = engine_log_directory()?;
589	let object_store = engine_data_path()?.join("objects");
590
591	let env = [
592		("RIVET__AUTH__ADMIN_TOKEN", admin_token()),
593		("RIVET__GUARD__HOST", ports.guard_host.clone()),
594		("RIVET__GUARD__PORT", ports.guard.to_string()),
595		// The peer listener only serves this engine's own internal requests.
596		("RIVET__API_PEER__HOST", LOOPBACK.to_owned()),
597		("RIVET__API_PEER__PORT", ports.api_peer.to_string()),
598		("RIVET__METRICS__HOST", ports.guard_host),
599		("RIVET__METRICS__PORT", ports.metrics.to_string()),
600		(
601			"RIVET__STORAGE__LOG_DIRECTORY",
602			log_directory.to_string_lossy().into_owned(),
603		),
604		(
605			"RIVET__OBJECT_STORAGE__PRIMARY__FILE_SYSTEM__PATH",
606			object_store.to_string_lossy().into_owned(),
607		),
608		// One storage worker holds every tablet, so there is nothing to place or split.
609		("RIVET__STORAGE__BALANCING__ENABLED", "false".to_owned()),
610		("RIVET__STORAGE__AUTO_PROVISION_PARTITIONS", "1".to_owned()),
611		// Bound disk use instead of sizing from the host's free space. The cold store is
612		// already on local disk, so a disk cache in front of it only duplicates data.
613		(
614			"RIVET__STORAGE__LOG_SHARD__MAX_DISK_BYTES",
615			LOG_MAX_DISK_BYTES.to_string(),
616		),
617		(
618			"RIVET__STORAGE__COLD_STORE__DISK_CACHE_BYTES",
619			"0".to_owned(),
620		),
621		// Storage shutdown must outlast the cold-write age, and the forced shutdown must
622		// outlast both stop budgets; these keep a local engine's stop within seconds.
623		("RIVET__STORAGE__WRITE_TO_COLD_AFTER_MS", "5000".to_owned()),
624		("RIVET__STORAGE__STOP_TIMEOUT_MS", "6000".to_owned()),
625		("RIVET__STORAGE__STOP_TEARDOWN_MS", "2000".to_owned()),
626		("RIVET__PEGBOARD__NODE_LOST_THRESHOLD", "7000".to_owned()),
627		("RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION", "1".to_owned()),
628		("RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION", "1".to_owned()),
629		("RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION", "10".to_owned()),
630	];
631	Ok(env
632		.into_iter()
633		.map(|(key, value)| (key.to_owned(), value))
634		.collect())
635}
636
637const LOOPBACK: &str = "127.0.0.1";
638
639/// Log shard disk cap of the local engine.
640const LOG_MAX_DISK_BYTES: u64 = 4 * 1024 * 1024 * 1024;
641
642/// Listener addresses derived from the endpoint: guard on its port, the peer listener one
643/// above it and metrics ten above it.
644struct EnginePorts {
645	guard_host: String,
646	guard: u16,
647	api_peer: u16,
648	metrics: u16,
649}
650
651impl EnginePorts {
652	fn resolve(config: &EngineResolverConfig) -> Result<Self> {
653		let endpoint = &config.endpoint;
654		let endpoint_url =
655			Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
656		let guard_port = endpoint_url
657			.port_or_known_default()
658			.ok_or_else(|| invalid_endpoint(endpoint, "missing port"))?;
659		let guard = config.bind_port.unwrap_or(guard_port);
660		Ok(Self {
661			guard_host: resolve_bind_host(config)?,
662			guard,
663			api_peer: guard
664				.checked_add(1)
665				.ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?,
666			metrics: guard
667				.checked_add(10)
668				.ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?,
669		})
670	}
671}
672
673/// Writes the topology, which the engine cannot take from environment variables. The public
674/// URL is the explicit override or the endpoint; the peer URL is the loopback peer listener.
675fn write_engine_config(config: &EngineResolverConfig) -> Result<PathBuf> {
676	let public_url = config.public_url.as_deref().unwrap_or(&config.endpoint);
677	let public_url = Url::parse(public_url)
678		.with_context(|| format!("parse engine public URL `{public_url}`"))?;
679	let ports = EnginePorts::resolve(config)?;
680	let peer_url = format!("http://{LOOPBACK}:{}", ports.api_peer);
681	let dir = storage_root()?.join("var").join("engine");
682	ensure_dir(&dir)?;
683	let path = dir.join("config.json");
684	let config = serde_json::json!({
685		"topology": {
686			"datacenter_label": 1,
687			"datacenters": {
688				"default": {
689					"datacenter_label": 1,
690					"is_leader": true,
691					"public_url": public_url.as_str(),
692					"peer_url": peer_url,
693				}
694			}
695		}
696	});
697	std::fs::write(&path, serde_json::to_string_pretty(&config)?)
698		.with_context(|| format!("write engine config `{}`", path.display()))?;
699	Ok(path)
700}
701
702pub async fn resolve_engine_binary(config: &EngineResolverConfig) -> Result<ResolvedEngine> {
703	if let Some(path) = config.explicit_binary_path.as_ref() {
704		return verify_binary_path(path);
705	}
706
707	if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
708		return verify_binary_path(&path);
709	}
710
711	if probe_existing_engine(&config.endpoint).await?.is_some() {
712		return Ok(ResolvedEngine::Existing);
713	}
714
715	let local_roots = local_engine_search_roots();
716	let cached = cached_engine_path(&config.version)?;
717	resolve_engine_binary_after_probe(config, false, &local_roots, cached).await
718}
719
720/// Resolves the engine binary path without probing for an already-running
721/// engine.
722///
723/// Used by callers that exec the binary directly (for example the CLI `engine`
724/// proxy), where a running engine on the endpoint should not short-circuit
725/// resolution to `Existing` and leave the caller without a path to run.
726pub async fn resolve_engine_binary_path(config: &EngineResolverConfig) -> Result<PathBuf> {
727	if let Some(path) = config.explicit_binary_path.as_ref() {
728		verify_binary_path(path)?;
729		return Ok(path.clone());
730	}
731
732	if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
733		verify_binary_path(&path)?;
734		return Ok(path);
735	}
736
737	let local_roots = local_engine_search_roots();
738	let cached = cached_engine_path(&config.version)?;
739	match resolve_engine_binary_after_probe(config, false, &local_roots, cached).await? {
740		ResolvedEngine::Binary(path) => Ok(path),
741		ResolvedEngine::Existing => {
742			unreachable!("no-probe resolution never returns Existing")
743		}
744	}
745}
746
747async fn resolve_engine_binary_after_probe(
748	config: &EngineResolverConfig,
749	existing_engine: bool,
750	local_roots: &[PathBuf],
751	cached: PathBuf,
752) -> Result<ResolvedEngine> {
753	if existing_engine {
754		return Ok(ResolvedEngine::Existing);
755	}
756
757	if let Some(path) = find_local_engine_binary_in_roots(local_roots) {
758		return Ok(ResolvedEngine::Binary(path));
759	}
760
761	if cached.exists() {
762		return Ok(ResolvedEngine::Binary(cached));
763	}
764
765	if !config.auto_download {
766		return Err(EngineProcessError::BinaryUnavailable {
767			version: config.version.clone(),
768		}
769		.build());
770	}
771
772	download_engine_binary(config, &cached).await?;
773	Ok(ResolvedEngine::Binary(cached))
774}
775
776fn verify_binary_path(path: &Path) -> Result<ResolvedEngine> {
777	if !path.exists() {
778		return Err(EngineProcessError::BinaryNotFound {
779			path: path.display().to_string(),
780		}
781		.build());
782	}
783	Ok(ResolvedEngine::Binary(path.to_path_buf()))
784}
785
786fn local_engine_search_roots() -> Vec<PathBuf> {
787	Path::new(env!("CARGO_MANIFEST_DIR"))
788		.ancestors()
789		.map(Path::to_path_buf)
790		.collect()
791}
792
793fn find_local_engine_binary_in_roots(roots: &[PathBuf]) -> Option<PathBuf> {
794	for root in roots {
795		for profile in ["debug", "release"] {
796			let candidate = root
797				.join("target")
798				.join(profile)
799				.join(exe_name("rivet-engine"));
800			if candidate.exists() {
801				return Some(candidate);
802			}
803		}
804	}
805	None
806}
807
808/// Cached release binaries live under the architecture, so a v2 engine cached by an older
809/// RivetKit under the same version number is never picked up.
810fn cached_engine_path(version: &str) -> Result<PathBuf> {
811	Ok(storage_root()?
812		.join("engine")
813		.join(ENGINE_ARCHITECTURE)
814		.join(version)
815		.join(engine_artifact_name()))
816}
817
818async fn download_engine_binary(config: &EngineResolverConfig, destination: &Path) -> Result<()> {
819	let artifact = engine_artifact_name();
820	let base = config.releases_endpoint.trim_end_matches('/');
821	let artifact_url = format!("{base}/rivet/{}/engine/{artifact}", config.version);
822	let manifest_url = format!("{base}/rivet/{}/engine/SHA256SUMS", config.version);
823	let client = Client::builder()
824		.timeout(DOWNLOAD_TIMEOUT)
825		.build()
826		.context("build reqwest client for engine download")?;
827
828	let manifest = fetch_text(&client, &manifest_url).await?;
829	let expected = checksum_for_artifact(&manifest, &artifact).ok_or_else(|| {
830		EngineProcessError::DownloadFailed {
831			url: manifest_url.clone(),
832			reason: format!("manifest does not contain `{artifact}`"),
833		}
834		.build()
835	})?;
836
837	let bytes = fetch_bytes(&client, &artifact_url).await?;
838	let received = sha256_hex(&bytes);
839	if !received.eq_ignore_ascii_case(&expected) {
840		return Err(EngineProcessError::ChecksumMismatch {
841			artifact,
842			expected,
843			received,
844		}
845		.build());
846	}
847
848	let parent = destination
849		.parent()
850		.context("engine cache destination has no parent")?;
851	ensure_dir(parent)?;
852	std::fs::write(destination, bytes)
853		.with_context(|| format!("write engine binary `{}`", destination.display()))?;
854	make_executable(destination)?;
855	Ok(())
856}
857
858async fn fetch_text(client: &Client, url: &str) -> Result<String> {
859	let response = client.get(url).send().await.map_err(|error| {
860		EngineProcessError::DownloadFailed {
861			url: url.to_owned(),
862			reason: error.to_string(),
863		}
864		.build()
865	})?;
866	if !response.status().is_success() {
867		let status = response.status();
868		return Err(EngineProcessError::DownloadFailed {
869			url: url.to_owned(),
870			reason: format!("unexpected status {status}"),
871		}
872		.build());
873	}
874	response.text().await.map_err(|error| {
875		EngineProcessError::DownloadFailed {
876			url: url.to_owned(),
877			reason: error.to_string(),
878		}
879		.build()
880	})
881}
882
883async fn fetch_bytes(client: &Client, url: &str) -> Result<Vec<u8>> {
884	let response = client.get(url).send().await.map_err(|error| {
885		EngineProcessError::DownloadFailed {
886			url: url.to_owned(),
887			reason: error.to_string(),
888		}
889		.build()
890	})?;
891	if !response.status().is_success() {
892		let status = response.status();
893		return Err(EngineProcessError::DownloadFailed {
894			url: url.to_owned(),
895			reason: format!("unexpected status {status}"),
896		}
897		.build());
898	}
899	response
900		.bytes()
901		.await
902		.map(|bytes| bytes.to_vec())
903		.map_err(|error| {
904			EngineProcessError::DownloadFailed {
905				url: url.to_owned(),
906				reason: error.to_string(),
907			}
908			.build()
909		})
910}
911
912fn checksum_for_artifact(manifest: &str, artifact: &str) -> Option<String> {
913	manifest.lines().find_map(|line| {
914		let mut parts = line.split_whitespace();
915		let checksum = parts.next()?;
916		let name = parts.next()?.trim_start_matches('*');
917		(checksum.len() == 64 && name == artifact).then(|| checksum.to_owned())
918	})
919}
920
921fn sha256_hex(bytes: &[u8]) -> String {
922	let digest = Sha256::digest(bytes);
923	let mut out = String::with_capacity(digest.len() * 2);
924	for byte in digest {
925		use std::fmt::Write;
926		let _ = write!(&mut out, "{byte:02x}");
927	}
928	out
929}
930
931fn engine_artifact_name() -> String {
932	let arch = match std::env::consts::ARCH {
933		"x86_64" => "x86_64",
934		"aarch64" => "aarch64",
935		other => other,
936	};
937	let target = match std::env::consts::OS {
938		"linux" => format!("{arch}-unknown-linux-musl"),
939		"macos" => format!("{arch}-apple-darwin"),
940		"windows" => format!("{arch}-pc-windows-gnu.exe"),
941		other => format!("{arch}-{other}"),
942	};
943	format!("rivet-engine-{target}")
944}
945
946fn exe_name(base: &str) -> String {
947	if cfg!(windows) {
948		format!("{base}.exe")
949	} else {
950		base.to_owned()
951	}
952}
953
954fn make_executable(path: &Path) -> Result<()> {
955	#[cfg(unix)]
956	{
957		use std::os::unix::fs::PermissionsExt;
958		let mut permissions = std::fs::metadata(path)
959			.with_context(|| format!("read metadata for `{}`", path.display()))?
960			.permissions();
961		permissions.set_mode(0o755);
962		std::fs::set_permissions(path, permissions)
963			.with_context(|| format!("mark `{}` executable", path.display()))?;
964	}
965	#[cfg(not(unix))]
966	{
967		let _ = path;
968	}
969	Ok(())
970}
971
972impl Drop for EngineProcessManager {
973	fn drop(&mut self) {
974		if let Some(stop) = self.stop.take() {
975			let _ = stop.send(());
976		}
977	}
978}
979
980/// Spawns a background task that owns the `Child` and awaits `wait()` so we
981/// log a clear message if the engine dies while rivetkit is still up. Taking
982/// the `Child` into the task also reaps it via `waitpid` on exit, so a
983/// crashed engine never lingers as a zombie in our process table.
984fn spawn_engine_watcher(mut child: Child, pid: u32) -> (JoinHandle<()>, oneshot::Sender<()>) {
985	let (stop, stop_rx) = oneshot::channel();
986	let watcher = tokio::spawn(async move {
987		let exit = tokio::select! {
988			exit = child.wait() => Some(exit),
989			_ = stop_rx => None,
990		};
991		if exit.is_none() {
992			if matches!(child.try_wait(), Ok(None)) {
993				#[cfg(unix)]
994				unsafe {
995					libc::kill(pid as libc::pid_t, libc::SIGTERM);
996				}
997				#[cfg(not(unix))]
998				let _ = child.start_kill();
999			}
1000			match tokio::time::timeout(Duration::from_secs(15), child.wait()).await {
1001				Ok(Ok(_)) => {}
1002				result => {
1003					tracing::warn!(
1004						pid,
1005						?result,
1006						"engine did not stop after SIGTERM; killing it"
1007					);
1008					let _ = child.start_kill();
1009					let _ = child.wait().await;
1010				}
1011			}
1012			return;
1013		}
1014		match exit.expect("checked above") {
1015			Ok(status) if status.success() => {
1016				tracing::warn!(
1017					pid,
1018					?status,
1019					"engine process exited cleanly while rivetkit was still running; \
1020					 rivetkit expected the engine to outlive it"
1021				);
1022			}
1023			Ok(status) => {
1024				tracing::error!(
1025					pid,
1026					?status,
1027					"engine process crashed while rivetkit was still running"
1028				);
1029			}
1030			Err(error) => {
1031				tracing::error!(
1032					pid,
1033					?error,
1034					"failed to wait on engine process; cannot detect crashes"
1035				);
1036			}
1037		}
1038	});
1039	(watcher, stop)
1040}
1041
1042/// Probes the configured endpoint for an already-running, healthy engine.
1043///
1044/// Returns `Ok(Some(health))` if the endpoint is serving a `runtime: "engine"`
1045/// health response that we can reattach to. Returns `Ok(None)` if the port is
1046/// free. Returns `Err(...)` if the port is occupied by a non-engine process
1047/// (for example a stale rivetkit) which would conflict with a fresh spawn.
1048async fn probe_existing_engine(endpoint: &str) -> Result<Option<EngineHealthResponse>> {
1049	let health_url = engine_health_url(endpoint);
1050	let client = Client::builder()
1051		.build()
1052		.context("build reqwest client for engine probe")?;
1053
1054	let response = match client
1055		.get(&health_url)
1056		.timeout(Duration::from_secs(1))
1057		.send()
1058		.await
1059	{
1060		Ok(response) => response,
1061		Err(_) => return Ok(None),
1062	};
1063
1064	if !response.status().is_success() {
1065		return Ok(None);
1066	}
1067
1068	let health = response
1069		.json::<EngineHealthResponse>()
1070		.await
1071		.context("decode existing engine health response")?;
1072
1073	match health.runtime.as_deref() {
1074		Some(ENGINE_RUNTIME) => {
1075			ensure_compatible_engine(endpoint, &health)?;
1076			Ok(Some(health))
1077		}
1078		Some(RIVETKIT_RUNTIME) => Err(EngineProcessError::PortOccupied {
1079			endpoint: endpoint.to_owned(),
1080			runtime: RIVETKIT_RUNTIME.to_owned(),
1081		}
1082		.build()),
1083		Some(other) => Err(EngineProcessError::PortOccupied {
1084			endpoint: endpoint.to_owned(),
1085			runtime: other.to_owned(),
1086		}
1087		.build()),
1088		None => Err(EngineProcessError::PortOccupied {
1089			endpoint: endpoint.to_owned(),
1090			runtime: "unknown".to_owned(),
1091		}
1092		.build()),
1093	}
1094}
1095
1096fn engine_health_url(endpoint: &str) -> String {
1097	format!("{}/health", endpoint.trim_end_matches('/'))
1098}
1099
1100/// Refuses any engine that does not report the required architecture, so a v2 engine
1101/// listening on the endpoint is never reused or left serving.
1102fn ensure_compatible_engine(endpoint: &str, health: &EngineHealthResponse) -> Result<()> {
1103	if health.architecture.as_deref() == Some(ENGINE_ARCHITECTURE) {
1104		return Ok(());
1105	}
1106	Err(EngineProcessError::IncompatibleEngine {
1107		endpoint: endpoint.to_owned(),
1108		version: health
1109			.version
1110			.clone()
1111			.unwrap_or_else(|| "unknown".to_owned()),
1112		architecture: health
1113			.architecture
1114			.clone()
1115			.unwrap_or_else(|| "v2".to_owned()),
1116		required: ENGINE_ARCHITECTURE.to_owned(),
1117	}
1118	.build())
1119}
1120
1121fn storage_root() -> Result<PathBuf> {
1122	if let Ok(path) = std::env::var("RIVETKIT_STORAGE_PATH") {
1123		return Ok(PathBuf::from(path).join(".rivetkit"));
1124	}
1125	let home = std::env::var("HOME")
1126		.map(PathBuf::from)
1127		.or_else(|_| std::env::current_dir())
1128		.context("locate home directory for engine storage path")?;
1129	Ok(home.join(".rivetkit"))
1130}
1131
1132fn ensure_dir(path: &Path) -> Result<()> {
1133	std::fs::create_dir_all(path).with_context(|| format!("create directory `{}`", path.display()))
1134}
1135
1136fn open_log_file(path: &Path) -> Result<std::fs::File> {
1137	std::fs::OpenOptions::new()
1138		.create(true)
1139		.append(true)
1140		.open(path)
1141		.with_context(|| format!("open log file `{}`", path.display()))
1142}
1143
1144fn log_timestamp() -> String {
1145	let now = std::time::SystemTime::now()
1146		.duration_since(std::time::UNIX_EPOCH)
1147		.unwrap_or_default();
1148	format!("{}", now.as_secs())
1149}
1150
1151async fn wait_for_engine_health(health_url: &str) -> Result<EngineHealthResponse> {
1152	// A first start creates the log shard before the HTTP listener binds.
1153	const HEALTH_MAX_WAIT: Duration = Duration::from_secs(60);
1154	const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
1155	const HEALTH_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
1156	const HEALTH_MAX_BACKOFF: Duration = Duration::from_secs(1);
1157
1158	let client = Client::builder()
1159		.build()
1160		.context("build reqwest client for engine health check")?;
1161	let deadline = Instant::now() + HEALTH_MAX_WAIT;
1162	let mut attempt = 0u32;
1163	let mut backoff = HEALTH_INITIAL_BACKOFF;
1164
1165	loop {
1166		attempt += 1;
1167
1168		let last_error = match client
1169			.get(health_url)
1170			.timeout(HEALTH_REQUEST_TIMEOUT)
1171			.send()
1172			.await
1173		{
1174			Ok(response) if response.status().is_success() => {
1175				let health = response
1176					.json::<EngineHealthResponse>()
1177					.await
1178					.context("decode engine health response")?;
1179				return Ok(health);
1180			}
1181			Ok(response) => format!("unexpected status {}", response.status()),
1182			Err(error) => error.to_string(),
1183		};
1184
1185		if Instant::now() >= deadline {
1186			return Err(EngineProcessError::HealthCheckFailed {
1187				attempts: attempt,
1188				reason: last_error,
1189			}
1190			.build());
1191		}
1192
1193		tokio::time::sleep(backoff).await;
1194		backoff = std::cmp::min(backoff * 2, HEALTH_MAX_BACKOFF);
1195	}
1196}
1197
1198/// Waits until the engine routes requests to Depot. Guard answers health before the
1199/// first tablet is provisioned and opened, and actor requests fail until then, so this
1200/// reads an actor listing, which routes through Depot, until it succeeds.
1201async fn wait_for_engine_ready(endpoint: &str) -> Result<()> {
1202	const READY_MAX_WAIT: Duration = Duration::from_secs(60);
1203	const READY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
1204	const READY_BACKOFF: Duration = Duration::from_millis(250);
1205
1206	let url = format!(
1207		"{}/actors?namespace=default&name=rivetkit-readiness-probe&key=probe",
1208		endpoint.trim_end_matches('/')
1209	);
1210	let token = admin_token();
1211	let client = Client::builder()
1212		.build()
1213		.context("build reqwest client for engine readiness check")?;
1214	let deadline = Instant::now() + READY_MAX_WAIT;
1215	let mut attempt = 0u32;
1216
1217	loop {
1218		attempt += 1;
1219
1220		let last_error = match client
1221			.get(&url)
1222			.bearer_auth(&token)
1223			.timeout(READY_REQUEST_TIMEOUT)
1224			.send()
1225			.await
1226		{
1227			Ok(response) if response.status().is_success() => return Ok(()),
1228			Ok(response) => {
1229				let status = response.status();
1230				let body = response.text().await.unwrap_or_default();
1231				format!("unexpected status {status}: {body}")
1232			}
1233			Err(error) => error.to_string(),
1234		};
1235		tracing::debug!(attempt, reason = %last_error, "waiting for engine readiness");
1236
1237		if Instant::now() >= deadline {
1238			return Err(EngineProcessError::NotReady {
1239				attempts: attempt,
1240				reason: last_error,
1241			}
1242			.build());
1243		}
1244
1245		tokio::time::sleep(READY_BACKOFF).await;
1246	}
1247}
1248
1249/// Cleanup path for a spawn that never reached `healthy`. We *do* kill here
1250/// because the half-started engine has no useful state to preserve and
1251/// leaving it running would conflict with a retry. This is the only place
1252/// allowed to terminate the engine.
1253async fn terminate_failed_spawn(child: &mut Child) -> Result<()> {
1254	const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
1255
1256	if child
1257		.try_wait()
1258		.context("check engine process status")?
1259		.is_some()
1260	{
1261		return Ok(());
1262	}
1263
1264	child
1265		.start_kill()
1266		.context("kill half-started engine process")?;
1267	match tokio::time::timeout(SHUTDOWN_TIMEOUT, child.wait()).await {
1268		Ok(result) => {
1269			let status = result.context("wait for half-started engine to exit")?;
1270			tracing::info!(?status, "half-started engine process exited");
1271			Ok(())
1272		}
1273		Err(_) => {
1274			tracing::warn!("half-started engine process did not exit within timeout");
1275			Ok(())
1276		}
1277	}
1278}
1279
1280fn invalid_endpoint(endpoint: &str, reason: &str) -> anyhow::Error {
1281	EngineProcessError::InvalidEndpoint {
1282		endpoint: endpoint.to_owned(),
1283		reason: reason.to_owned(),
1284	}
1285	.build()
1286}