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