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";
22const 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#[derive(Clone, Debug, Serialize, Deserialize)]
84pub struct EngineRuntimeStamp {
85 pub pid: u32,
86 #[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 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#[derive(Debug)]
117pub enum EngineStartup {
118 Spawned,
120 Reused { stamp: Option<EngineRuntimeStamp> },
124}
125
126#[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 let _ = read_engine_stamp();
142 let resolved = resolve_engine_binary(&config).await?;
143 Self::start_resolved(resolved, &config).await
144 }
145
146 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 wait_for_engine_exit(&stamp, Duration::from_secs(20)).await?;
188 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 #[cfg(unix)]
249 command.process_group(0);
250 #[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 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 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
357async 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
378const RUNTIME_INFO_TIMEOUT: Duration = Duration::from_secs(15);
383const RUNTIME_INFO_MAX_BYTES: usize = 8192;
385
386async 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 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
452pub fn engine_data_path() -> Result<PathBuf> {
455 Ok(storage_root()?.join("var").join("engine").join("depot"))
456}
457
458pub 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
475fn 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
501fn 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
521fn 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 == ¤t,
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 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 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 false
574}
575
576pub 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 ("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 ("RIVET__STORAGE__BALANCING__ENABLED", "false".to_owned()),
610 ("RIVET__STORAGE__AUTO_PROVISION_PARTITIONS", "1".to_owned()),
611 (
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 ("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
639const LOG_MAX_DISK_BYTES: u64 = 4 * 1024 * 1024 * 1024;
641
642struct 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
673fn 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
720pub 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
808fn 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
980fn 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
1042async 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
1100fn 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 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
1198async 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
1249async 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}