1use std::env;
2use std::path::{Path, PathBuf};
3use std::process::Stdio;
4use std::time::{Duration, Instant};
5
6use anyhow::{Context, Result};
7use reqwest::{Client, Url};
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10use tokio::process::{Child, Command};
11use tokio::task::JoinHandle;
12
13mod error;
14
15pub use error::EngineProcessError;
16
17const ENGINE_RUNTIME: &str = "engine";
18const RIVETKIT_RUNTIME: &str = "rivetkit";
19const ENGINE_VERSION_ENV: &str = "RIVETKIT_ENGINE_VERSION";
20const RELEASES_ENDPOINT_ENV: &str = "RIVETKIT_ENGINE_RELEASES_ENDPOINT";
21const RELEASES_ENDPOINT: &str = "https://releases.rivet.dev";
22const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(60);
23
24#[derive(Debug, Deserialize)]
25struct EngineHealthResponse {
26 status: Option<String>,
27 runtime: Option<String>,
28 version: Option<String>,
29}
30
31#[derive(Clone, Debug)]
32pub struct EngineResolverConfig {
33 pub endpoint: String,
34 pub explicit_binary_path: Option<PathBuf>,
35 pub bind_host: Option<String>,
36 pub bind_port: Option<u16>,
37 pub public_url: Option<String>,
38 pub auto_download: bool,
39 pub version: String,
40 pub releases_endpoint: String,
41}
42
43impl EngineResolverConfig {
44 pub fn from_parts(
45 endpoint: &str,
46 explicit_binary_path: Option<PathBuf>,
47 bind_host: Option<String>,
48 bind_port: Option<u16>,
49 auto_download: bool,
50 ) -> Self {
51 Self {
52 endpoint: endpoint.to_owned(),
53 explicit_binary_path,
54 bind_host,
55 bind_port,
56 public_url: None,
57 auto_download,
58 version: std::env::var(ENGINE_VERSION_ENV)
59 .unwrap_or_else(|_| env!("CARGO_PKG_VERSION").to_owned()),
60 releases_endpoint: std::env::var(RELEASES_ENDPOINT_ENV)
61 .unwrap_or_else(|_| RELEASES_ENDPOINT.to_owned()),
62 }
63 }
64}
65
66#[derive(Debug, PartialEq, Eq)]
67pub enum ResolvedEngine {
68 Existing,
69 Binary(PathBuf),
70}
71
72#[derive(Clone, Debug, Serialize, Deserialize)]
76pub struct EngineRuntimeStamp {
77 pub pid: u32,
78 pub endpoint: String,
79 pub bind_host: String,
80 pub public_url: Option<String>,
81}
82
83impl EngineRuntimeStamp {
84 pub fn binds_loopback_only(&self) -> bool {
87 if self.bind_host == "localhost" {
88 return true;
89 }
90 self.bind_host
91 .parse::<std::net::IpAddr>()
92 .map(|ip| ip.is_loopback())
93 .unwrap_or(false)
94 }
95}
96
97#[derive(Debug)]
99pub enum EngineStartup {
100 Spawned,
102 Reused { stamp: Option<EngineRuntimeStamp> },
106}
107
108#[derive(Debug)]
130pub struct EngineProcessManager {
131 watcher: Option<JoinHandle<()>>,
132 startup: EngineStartup,
133}
134
135impl EngineProcessManager {
136 pub async fn start_or_reuse(config: EngineResolverConfig) -> Result<Self> {
137 let resolved = resolve_engine_binary(&config).await?;
138 Self::start_resolved(resolved, &config).await
139 }
140
141 pub fn startup(&self) -> &EngineStartup {
143 &self.startup
144 }
145
146 async fn start_resolved(
147 resolved: ResolvedEngine,
148 config: &EngineResolverConfig,
149 ) -> Result<Self> {
150 let endpoint = &config.endpoint;
151 if matches!(resolved, ResolvedEngine::Existing) {
152 tracing::info!(
153 endpoint = %endpoint,
154 "reusing already-running engine process"
155 );
156 return Ok(Self {
157 watcher: None,
158 startup: EngineStartup::Reused {
159 stamp: read_engine_stamp(),
160 },
161 });
162 }
163
164 let ResolvedEngine::Binary(binary_path) = resolved else {
165 unreachable!("existing engine handled above");
166 };
167 if let Some(health) = probe_existing_engine(endpoint).await? {
168 tracing::info!(
169 endpoint = %endpoint,
170 status = ?health.status,
171 runtime = ?health.runtime,
172 version = ?health.version,
173 "reusing already-running engine process"
174 );
175 return Ok(Self {
176 watcher: None,
177 startup: EngineStartup::Reused {
178 stamp: read_engine_stamp(),
179 },
180 });
181 }
182
183 if !binary_path.exists() {
184 return Err(EngineProcessError::BinaryNotFound {
185 path: binary_path.display().to_string(),
186 }
187 .build());
188 }
189
190 let env = engine_env(config)?;
191 let config_path = write_engine_config(config)?;
192 let db_path = engine_db_path()?;
193 let logs_dir = storage_root()?
194 .join("var")
195 .join("logs")
196 .join("rivet-engine");
197 ensure_dir(&db_path).context("create engine db directory")?;
198 ensure_dir(&logs_dir).context("create engine logs directory")?;
199
200 let timestamp = log_timestamp();
201 let stdout_log_path = logs_dir.join(format!("engine-{timestamp}-stdout.log"));
202 let stderr_log_path = logs_dir.join(format!("engine-{timestamp}-stderr.log"));
203 let stdout_file = open_log_file(&stdout_log_path)
204 .with_context(|| format!("open engine stdout log `{}`", stdout_log_path.display()))?;
205 let stderr_file = open_log_file(&stderr_log_path)
206 .with_context(|| format!("open engine stderr log `{}`", stderr_log_path.display()))?;
207
208 let mut command = Command::new(&binary_path);
209 command.arg("start");
210 if let Some(config_path) = &config_path {
211 command.arg("--config").arg(config_path);
212 }
213 for (key, value) in &env {
214 command.env(key, value);
215 }
216 command
217 .stdin(Stdio::null())
218 .stdout(Stdio::from(stdout_file))
219 .stderr(Stdio::from(stderr_file));
220
221 #[cfg(unix)]
227 command.process_group(0);
228
229 let mut child = command
230 .spawn()
231 .with_context(|| format!("spawn engine binary `{}`", binary_path.display()))?;
232 let pid = child
233 .id()
234 .ok_or_else(|| EngineProcessError::MissingPid.build())?;
235
236 tracing::info!(
237 pid,
238 path = %binary_path.display(),
239 endpoint = %endpoint,
240 db_path = %db_path.display(),
241 "spawned engine process (intentionally orphaned, will outlive this process)"
242 );
243 tracing::info!(
244 stdout_log = %stdout_log_path.display(),
245 stderr_log = %stderr_log_path.display(),
246 "engine stdout/stderr piped to log files"
247 );
248
249 let health_url = engine_health_url(endpoint);
250 let health = match wait_for_engine_health(&health_url).await {
251 Ok(health) => health,
252 Err(error) => {
253 let error = match child.try_wait() {
254 Ok(Some(status)) => error.context(format!(
255 "engine process exited before becoming healthy with status {status}"
256 )),
257 Ok(None) => error,
258 Err(wait_error) => error.context(format!(
259 "failed to inspect engine process status: {wait_error:#}"
260 )),
261 };
262 if let Err(cleanup_error) = terminate_failed_spawn(&mut child).await {
263 tracing::warn!(
264 ?cleanup_error,
265 "failed to terminate engine process that never became healthy"
266 );
267 }
268 return Err(error);
269 }
270 };
271
272 tracing::info!(
273 pid,
274 status = ?health.status,
275 runtime = ?health.runtime,
276 version = ?health.version,
277 "engine process is healthy"
278 );
279
280 let stamp = EngineRuntimeStamp {
284 pid,
285 endpoint: endpoint.clone(),
286 bind_host: resolve_bind_host(config)?,
287 public_url: config.public_url.clone(),
288 };
289 if let Err(error) = write_engine_stamp(&stamp) {
290 tracing::warn!(?error, "failed to write engine runtime stamp");
291 }
292
293 Ok(Self {
294 watcher: Some(spawn_engine_watcher(child, pid)),
295 startup: EngineStartup::Spawned,
296 })
297 }
298}
299
300pub fn engine_db_path() -> Result<PathBuf> {
302 Ok(storage_root()?.join("var").join("engine").join("db"))
303}
304
305fn engine_stamp_path() -> Result<PathBuf> {
306 Ok(storage_root()?
307 .join("var")
308 .join("engine")
309 .join("runtime.json"))
310}
311
312fn resolve_bind_host(config: &EngineResolverConfig) -> Result<String> {
315 if let Some(bind_host) = &config.bind_host {
316 return Ok(bind_host.clone());
317 }
318 let endpoint = &config.endpoint;
319 let endpoint_url =
320 Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
321 Ok(endpoint_url
322 .host_str()
323 .ok_or_else(|| invalid_endpoint(endpoint, "missing host"))?
324 .to_owned())
325}
326
327fn write_engine_stamp(stamp: &EngineRuntimeStamp) -> Result<()> {
328 let path = engine_stamp_path()?;
329 if let Some(dir) = path.parent() {
330 ensure_dir(dir)?;
331 }
332 let contents = serde_json::to_vec_pretty(stamp).context("serialize engine stamp")?;
333 std::fs::write(&path, contents)
334 .with_context(|| format!("write engine stamp `{}`", path.display()))?;
335 Ok(())
336}
337
338fn read_engine_stamp() -> Option<EngineRuntimeStamp> {
344 let path = engine_stamp_path().ok()?;
345 let contents = std::fs::read(&path).ok()?;
346 let stamp: EngineRuntimeStamp = serde_json::from_slice(&contents).ok()?;
347 pid_is_alive(stamp.pid).then_some(stamp)
348}
349
350#[cfg(unix)]
351fn pid_is_alive(pid: u32) -> bool {
352 let result = unsafe { libc::kill(pid as libc::pid_t, 0) };
355 result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
356}
357
358#[cfg(not(unix))]
359fn pid_is_alive(_pid: u32) -> bool {
360 false
364}
365
366pub fn engine_env(config: &EngineResolverConfig) -> Result<Vec<(String, String)>> {
375 let endpoint = &config.endpoint;
376 let endpoint_url =
377 Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
378 let guard_host = resolve_bind_host(config)?;
379 let guard_port = endpoint_url
380 .port_or_known_default()
381 .ok_or_else(|| invalid_endpoint(endpoint, "missing port"))?;
382 let guard_port = config.bind_port.unwrap_or(guard_port);
383 let api_peer_port = guard_port
384 .checked_add(1)
385 .ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?;
386 let metrics_port = guard_port
387 .checked_add(10)
388 .ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?;
389
390 let db_path = engine_db_path()?;
391
392 Ok(vec![
393 (
394 "RIVET__AUTH__ADMIN_TOKEN".to_owned(),
395 env::var("RIVET__AUTH__ADMIN_TOKEN").unwrap_or_else(|_| "default".to_owned()),
396 ),
397 ("RIVET__GUARD__HOST".to_owned(), guard_host.clone()),
398 ("RIVET__GUARD__PORT".to_owned(), guard_port.to_string()),
399 ("RIVET__API_PEER__HOST".to_owned(), guard_host.clone()),
400 (
401 "RIVET__API_PEER__PORT".to_owned(),
402 api_peer_port.to_string(),
403 ),
404 ("RIVET__METRICS__HOST".to_owned(), guard_host),
405 ("RIVET__METRICS__PORT".to_owned(), metrics_port.to_string()),
406 (
407 "RIVET__FILE_SYSTEM__PATH".to_owned(),
408 db_path.to_string_lossy().into_owned(),
409 ),
410 (
411 "RIVET__PEGBOARD__RETRY_RESET_DURATION".to_owned(),
412 "100".to_owned(),
413 ),
414 (
415 "RIVET__PEGBOARD__BASE_RETRY_TIMEOUT".to_owned(),
416 "100".to_owned(),
417 ),
418 (
419 "RIVET__PEGBOARD__RESCHEDULE_BACKOFF_MAX_EXPONENT".to_owned(),
420 "1".to_owned(),
421 ),
422 (
423 "RIVET__PEGBOARD__RUNNER_ELIGIBLE_THRESHOLD".to_owned(),
424 "5000".to_owned(),
425 ),
426 (
427 "RIVET__PEGBOARD__RUNNER_LOST_THRESHOLD".to_owned(),
428 "7000".to_owned(),
429 ),
430 (
431 "RIVET__PEGBOARD__ENVOY_ELIGIBLE_THRESHOLD".to_owned(),
432 "5000".to_owned(),
433 ),
434 (
435 "RIVET__PEGBOARD__ENVOY_LOST_THRESHOLD".to_owned(),
436 "7000".to_owned(),
437 ),
438 (
439 "RIVET__PEGBOARD__MIN_METADATA_POLL_INTERVAL".to_owned(),
440 "1000".to_owned(),
441 ),
442 (
443 "RIVET__FEATURES__GUARD_GATEWAY_V3__MODE".to_owned(),
444 "on".to_owned(),
445 ),
446 (
447 "RIVET__FEATURES__GUARD_GATEWAY_V3__PERCENTAGE".to_owned(),
448 "100".to_owned(),
449 ),
450 (
451 "RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION".to_owned(),
452 "1".to_owned(),
453 ),
454 (
455 "RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION".to_owned(),
456 "1".to_owned(),
457 ),
458 (
459 "RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION".to_owned(),
460 "2".to_owned(),
461 ),
462 ])
463}
464
465fn write_engine_config(config: &EngineResolverConfig) -> Result<Option<PathBuf>> {
466 let Some(public_url) = &config.public_url else {
467 return Ok(None);
468 };
469
470 let public_url = Url::parse(public_url)
471 .with_context(|| format!("parse engine public URL `{public_url}`"))?;
472 let peer_url = peer_url_for_public_url(&public_url)?;
473 let dir = storage_root()?.join("var").join("engine");
474 ensure_dir(&dir)?;
475 let path = dir.join("config.json");
476 let config = serde_json::json!({
477 "topology": {
478 "datacenter_label": 1,
479 "datacenters": {
480 "default": {
481 "datacenter_label": 1,
482 "is_leader": true,
483 "public_url": public_url.as_str(),
484 "peer_url": peer_url,
485 }
486 }
487 }
488 });
489 std::fs::write(&path, serde_json::to_string_pretty(&config)?)
490 .with_context(|| format!("write engine config `{}`", path.display()))?;
491 Ok(Some(path))
492}
493
494fn peer_url_for_public_url(public_url: &Url) -> Result<String> {
495 let mut peer_url = public_url.clone();
496 let port = public_url
497 .port_or_known_default()
498 .ok_or_else(|| invalid_endpoint(public_url.as_str(), "missing port"))?
499 .checked_add(1)
500 .ok_or_else(|| invalid_endpoint(public_url.as_str(), "port is too large"))?;
501 if peer_url.set_port(Some(port)).is_err() {
502 return Err(invalid_endpoint(
503 public_url.as_str(),
504 "could not derive peer URL port",
505 ));
506 }
507 peer_url.set_path("");
508 peer_url.set_query(None);
509 peer_url.set_fragment(None);
510 Ok(peer_url.to_string().trim_end_matches('/').to_string())
511}
512
513pub async fn resolve_engine_binary(config: &EngineResolverConfig) -> Result<ResolvedEngine> {
514 if let Some(path) = config.explicit_binary_path.as_ref() {
515 return verify_binary_path(path);
516 }
517
518 if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
519 return verify_binary_path(&path);
520 }
521
522 if probe_existing_engine(&config.endpoint).await?.is_some() {
523 return Ok(ResolvedEngine::Existing);
524 }
525
526 let local_roots = local_engine_search_roots();
527 let cached = cached_engine_path(&config.version)?;
528 resolve_engine_binary_after_probe(config, false, &local_roots, cached).await
529}
530
531pub async fn resolve_engine_binary_path(config: &EngineResolverConfig) -> Result<PathBuf> {
538 if let Some(path) = config.explicit_binary_path.as_ref() {
539 verify_binary_path(path)?;
540 return Ok(path.clone());
541 }
542
543 if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
544 verify_binary_path(&path)?;
545 return Ok(path);
546 }
547
548 let local_roots = local_engine_search_roots();
549 let cached = cached_engine_path(&config.version)?;
550 match resolve_engine_binary_after_probe(config, false, &local_roots, cached).await? {
551 ResolvedEngine::Binary(path) => Ok(path),
552 ResolvedEngine::Existing => {
553 unreachable!("no-probe resolution never returns Existing")
554 }
555 }
556}
557
558async fn resolve_engine_binary_after_probe(
559 config: &EngineResolverConfig,
560 existing_engine: bool,
561 local_roots: &[PathBuf],
562 cached: PathBuf,
563) -> Result<ResolvedEngine> {
564 if existing_engine {
565 return Ok(ResolvedEngine::Existing);
566 }
567
568 if let Some(path) = find_local_engine_binary_in_roots(local_roots) {
569 return Ok(ResolvedEngine::Binary(path));
570 }
571
572 if cached.exists() {
573 return Ok(ResolvedEngine::Binary(cached));
574 }
575
576 if !config.auto_download {
577 return Err(EngineProcessError::BinaryUnavailable {
578 version: config.version.clone(),
579 }
580 .build());
581 }
582
583 download_engine_binary(config, &cached).await?;
584 Ok(ResolvedEngine::Binary(cached))
585}
586
587fn verify_binary_path(path: &Path) -> Result<ResolvedEngine> {
588 if !path.exists() {
589 return Err(EngineProcessError::BinaryNotFound {
590 path: path.display().to_string(),
591 }
592 .build());
593 }
594 Ok(ResolvedEngine::Binary(path.to_path_buf()))
595}
596
597fn local_engine_search_roots() -> Vec<PathBuf> {
598 Path::new(env!("CARGO_MANIFEST_DIR"))
599 .ancestors()
600 .map(Path::to_path_buf)
601 .collect()
602}
603
604fn find_local_engine_binary_in_roots(roots: &[PathBuf]) -> Option<PathBuf> {
605 for root in roots {
606 for profile in ["debug", "release"] {
607 let candidate = root
608 .join("target")
609 .join(profile)
610 .join(exe_name("rivet-engine"));
611 if candidate.exists() {
612 return Some(candidate);
613 }
614 }
615 }
616 None
617}
618
619fn cached_engine_path(version: &str) -> Result<PathBuf> {
620 Ok(storage_root()?
621 .join("engine")
622 .join(version)
623 .join(engine_artifact_name()))
624}
625
626async fn download_engine_binary(config: &EngineResolverConfig, destination: &Path) -> Result<()> {
627 let artifact = engine_artifact_name();
628 let base = config.releases_endpoint.trim_end_matches('/');
629 let artifact_url = format!("{base}/rivet/{}/engine/{artifact}", config.version);
630 let manifest_url = format!("{base}/rivet/{}/engine/SHA256SUMS", config.version);
631 let client = Client::builder()
632 .timeout(DOWNLOAD_TIMEOUT)
633 .build()
634 .context("build reqwest client for engine download")?;
635
636 let manifest = fetch_text(&client, &manifest_url).await?;
637 let expected = checksum_for_artifact(&manifest, &artifact).ok_or_else(|| {
638 EngineProcessError::DownloadFailed {
639 url: manifest_url.clone(),
640 reason: format!("manifest does not contain `{artifact}`"),
641 }
642 .build()
643 })?;
644
645 let bytes = fetch_bytes(&client, &artifact_url).await?;
646 let received = sha256_hex(&bytes);
647 if !received.eq_ignore_ascii_case(&expected) {
648 return Err(EngineProcessError::ChecksumMismatch {
649 artifact,
650 expected,
651 received,
652 }
653 .build());
654 }
655
656 let parent = destination
657 .parent()
658 .context("engine cache destination has no parent")?;
659 ensure_dir(parent)?;
660 std::fs::write(destination, bytes)
661 .with_context(|| format!("write engine binary `{}`", destination.display()))?;
662 make_executable(destination)?;
663 Ok(())
664}
665
666async fn fetch_text(client: &Client, url: &str) -> Result<String> {
667 let response = client.get(url).send().await.map_err(|error| {
668 EngineProcessError::DownloadFailed {
669 url: url.to_owned(),
670 reason: error.to_string(),
671 }
672 .build()
673 })?;
674 if !response.status().is_success() {
675 let status = response.status();
676 return Err(EngineProcessError::DownloadFailed {
677 url: url.to_owned(),
678 reason: format!("unexpected status {status}"),
679 }
680 .build());
681 }
682 response.text().await.map_err(|error| {
683 EngineProcessError::DownloadFailed {
684 url: url.to_owned(),
685 reason: error.to_string(),
686 }
687 .build()
688 })
689}
690
691async fn fetch_bytes(client: &Client, url: &str) -> Result<Vec<u8>> {
692 let response = client.get(url).send().await.map_err(|error| {
693 EngineProcessError::DownloadFailed {
694 url: url.to_owned(),
695 reason: error.to_string(),
696 }
697 .build()
698 })?;
699 if !response.status().is_success() {
700 let status = response.status();
701 return Err(EngineProcessError::DownloadFailed {
702 url: url.to_owned(),
703 reason: format!("unexpected status {status}"),
704 }
705 .build());
706 }
707 response
708 .bytes()
709 .await
710 .map(|bytes| bytes.to_vec())
711 .map_err(|error| {
712 EngineProcessError::DownloadFailed {
713 url: url.to_owned(),
714 reason: error.to_string(),
715 }
716 .build()
717 })
718}
719
720fn checksum_for_artifact(manifest: &str, artifact: &str) -> Option<String> {
721 manifest.lines().find_map(|line| {
722 let mut parts = line.split_whitespace();
723 let checksum = parts.next()?;
724 let name = parts.next()?.trim_start_matches('*');
725 (checksum.len() == 64 && name == artifact).then(|| checksum.to_owned())
726 })
727}
728
729fn sha256_hex(bytes: &[u8]) -> String {
730 let digest = Sha256::digest(bytes);
731 let mut out = String::with_capacity(digest.len() * 2);
732 for byte in digest {
733 use std::fmt::Write;
734 let _ = write!(&mut out, "{byte:02x}");
735 }
736 out
737}
738
739fn engine_artifact_name() -> String {
740 let arch = match std::env::consts::ARCH {
741 "x86_64" => "x86_64",
742 "aarch64" => "aarch64",
743 other => other,
744 };
745 let target = match std::env::consts::OS {
746 "linux" => format!("{arch}-unknown-linux-musl"),
747 "macos" => format!("{arch}-apple-darwin"),
748 "windows" => format!("{arch}-pc-windows-gnu.exe"),
749 other => format!("{arch}-{other}"),
750 };
751 format!("rivet-engine-{target}")
752}
753
754fn exe_name(base: &str) -> String {
755 if cfg!(windows) {
756 format!("{base}.exe")
757 } else {
758 base.to_owned()
759 }
760}
761
762fn make_executable(path: &Path) -> Result<()> {
763 #[cfg(unix)]
764 {
765 use std::os::unix::fs::PermissionsExt;
766 let mut permissions = std::fs::metadata(path)
767 .with_context(|| format!("read metadata for `{}`", path.display()))?
768 .permissions();
769 permissions.set_mode(0o755);
770 std::fs::set_permissions(path, permissions)
771 .with_context(|| format!("mark `{}` executable", path.display()))?;
772 }
773 #[cfg(not(unix))]
774 {
775 let _ = path;
776 }
777 Ok(())
778}
779
780impl Drop for EngineProcessManager {
781 fn drop(&mut self) {
782 if let Some(handle) = self.watcher.take() {
783 handle.abort();
789 tracing::debug!(
790 "aborted engine watcher; engine continues running (intentional orphan)"
791 );
792 }
793 }
794}
795
796fn spawn_engine_watcher(mut child: Child, pid: u32) -> JoinHandle<()> {
801 tokio::spawn(async move {
802 match child.wait().await {
803 Ok(status) if status.success() => {
804 tracing::warn!(
805 pid,
806 ?status,
807 "engine process exited cleanly while rivetkit was still running; \
808 rivetkit expected the engine to outlive it"
809 );
810 }
811 Ok(status) => {
812 tracing::error!(
813 pid,
814 ?status,
815 "engine process crashed while rivetkit was still running"
816 );
817 }
818 Err(error) => {
819 tracing::error!(
820 pid,
821 ?error,
822 "failed to wait on engine process; cannot detect crashes"
823 );
824 }
825 }
826 })
827}
828
829async fn probe_existing_engine(endpoint: &str) -> Result<Option<EngineHealthResponse>> {
836 let health_url = engine_health_url(endpoint);
837 let client = Client::builder()
838 .build()
839 .context("build reqwest client for engine probe")?;
840
841 let response = match client
842 .get(&health_url)
843 .timeout(Duration::from_secs(1))
844 .send()
845 .await
846 {
847 Ok(response) => response,
848 Err(_) => return Ok(None),
849 };
850
851 if !response.status().is_success() {
852 return Ok(None);
853 }
854
855 let health = response
856 .json::<EngineHealthResponse>()
857 .await
858 .context("decode existing engine health response")?;
859
860 match health.runtime.as_deref() {
861 Some(ENGINE_RUNTIME) => Ok(Some(health)),
862 Some(RIVETKIT_RUNTIME) => Err(EngineProcessError::PortOccupied {
863 endpoint: endpoint.to_owned(),
864 runtime: RIVETKIT_RUNTIME.to_owned(),
865 }
866 .build()),
867 Some(other) => Err(EngineProcessError::PortOccupied {
868 endpoint: endpoint.to_owned(),
869 runtime: other.to_owned(),
870 }
871 .build()),
872 None => Err(EngineProcessError::PortOccupied {
873 endpoint: endpoint.to_owned(),
874 runtime: "unknown".to_owned(),
875 }
876 .build()),
877 }
878}
879
880fn engine_health_url(endpoint: &str) -> String {
881 format!("{}/health", endpoint.trim_end_matches('/'))
882}
883
884fn storage_root() -> Result<PathBuf> {
885 if let Ok(path) = std::env::var("RIVETKIT_STORAGE_PATH") {
886 return Ok(PathBuf::from(path).join(".rivetkit"));
887 }
888 let home = std::env::var("HOME")
889 .map(PathBuf::from)
890 .or_else(|_| std::env::current_dir())
891 .context("locate home directory for engine storage path")?;
892 Ok(home.join(".rivetkit"))
893}
894
895fn ensure_dir(path: &Path) -> Result<()> {
896 std::fs::create_dir_all(path).with_context(|| format!("create directory `{}`", path.display()))
897}
898
899fn open_log_file(path: &Path) -> Result<std::fs::File> {
900 std::fs::OpenOptions::new()
901 .create(true)
902 .append(true)
903 .open(path)
904 .with_context(|| format!("open log file `{}`", path.display()))
905}
906
907fn log_timestamp() -> String {
908 let now = std::time::SystemTime::now()
909 .duration_since(std::time::UNIX_EPOCH)
910 .unwrap_or_default();
911 format!("{}", now.as_secs())
912}
913
914async fn wait_for_engine_health(health_url: &str) -> Result<EngineHealthResponse> {
915 const HEALTH_MAX_WAIT: Duration = Duration::from_secs(10);
916 const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
917 const HEALTH_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
918 const HEALTH_MAX_BACKOFF: Duration = Duration::from_secs(1);
919
920 let client = Client::builder()
921 .build()
922 .context("build reqwest client for engine health check")?;
923 let deadline = Instant::now() + HEALTH_MAX_WAIT;
924 let mut attempt = 0u32;
925 let mut backoff = HEALTH_INITIAL_BACKOFF;
926
927 loop {
928 attempt += 1;
929
930 let last_error = match client
931 .get(health_url)
932 .timeout(HEALTH_REQUEST_TIMEOUT)
933 .send()
934 .await
935 {
936 Ok(response) if response.status().is_success() => {
937 let health = response
938 .json::<EngineHealthResponse>()
939 .await
940 .context("decode engine health response")?;
941 return Ok(health);
942 }
943 Ok(response) => format!("unexpected status {}", response.status()),
944 Err(error) => error.to_string(),
945 };
946
947 if Instant::now() >= deadline {
948 return Err(EngineProcessError::HealthCheckFailed {
949 attempts: attempt,
950 reason: last_error,
951 }
952 .build());
953 }
954
955 tokio::time::sleep(backoff).await;
956 backoff = std::cmp::min(backoff * 2, HEALTH_MAX_BACKOFF);
957 }
958}
959
960async fn terminate_failed_spawn(child: &mut Child) -> Result<()> {
965 const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
966
967 if child
968 .try_wait()
969 .context("check engine process status")?
970 .is_some()
971 {
972 return Ok(());
973 }
974
975 child
976 .start_kill()
977 .context("kill half-started engine process")?;
978 match tokio::time::timeout(SHUTDOWN_TIMEOUT, child.wait()).await {
979 Ok(result) => {
980 let status = result.context("wait for half-started engine to exit")?;
981 tracing::info!(?status, "half-started engine process exited");
982 Ok(())
983 }
984 Err(_) => {
985 tracing::warn!("half-started engine process did not exit within timeout");
986 Ok(())
987 }
988 }
989}
990
991fn invalid_endpoint(endpoint: &str, reason: &str) -> anyhow::Error {
992 EngineProcessError::InvalidEndpoint {
993 endpoint: endpoint.to_owned(),
994 reason: reason.to_owned(),
995 }
996 .build()
997}
998
999#[cfg(test)]
1000mod tests {
1001 use std::collections::HashMap;
1002
1003 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1004 use tokio::net::TcpListener;
1005
1006 use super::*;
1007
1008 fn test_config(releases_endpoint: String, auto_download: bool) -> EngineResolverConfig {
1009 EngineResolverConfig {
1010 endpoint: "http://127.0.0.1:1".to_owned(),
1011 explicit_binary_path: None,
1012 bind_host: None,
1013 bind_port: None,
1014 public_url: None,
1015 auto_download,
1016 version: "test-version".to_owned(),
1017 releases_endpoint,
1018 }
1019 }
1020
1021 fn stamp_with_bind(bind_host: &str) -> EngineRuntimeStamp {
1022 EngineRuntimeStamp {
1023 pid: 1,
1024 endpoint: "http://127.0.0.1:6420".to_owned(),
1025 bind_host: bind_host.to_owned(),
1026 public_url: None,
1027 }
1028 }
1029
1030 #[test]
1031 fn binds_loopback_only_detects_loopback_addresses() {
1032 assert!(stamp_with_bind("127.0.0.1").binds_loopback_only());
1033 assert!(stamp_with_bind("localhost").binds_loopback_only());
1034 assert!(stamp_with_bind("::1").binds_loopback_only());
1035 }
1036
1037 #[test]
1038 fn binds_loopback_only_false_for_reachable_addresses() {
1039 assert!(!stamp_with_bind("0.0.0.0").binds_loopback_only());
1040 assert!(!stamp_with_bind("::").binds_loopback_only());
1041 assert!(!stamp_with_bind("192.168.1.5").binds_loopback_only());
1042 }
1043
1044 #[test]
1045 fn resolve_bind_host_prefers_override_then_endpoint() {
1046 let mut config = test_config(String::new(), false);
1047 config.endpoint = "http://127.0.0.1:6420".to_owned();
1048 assert_eq!(resolve_bind_host(&config).unwrap(), "127.0.0.1");
1049 config.bind_host = Some("0.0.0.0".to_owned());
1050 assert_eq!(resolve_bind_host(&config).unwrap(), "0.0.0.0");
1051 }
1052
1053 #[test]
1054 fn engine_env_uses_development_defaults() {
1055 let env = engine_env(&test_config(String::new(), false)).expect("build engine env");
1056 let env = env.into_iter().collect::<HashMap<_, _>>();
1057
1058 assert_eq!(env["RIVET__PEGBOARD__RETRY_RESET_DURATION"], "100");
1059 assert_eq!(env["RIVET__PEGBOARD__BASE_RETRY_TIMEOUT"], "100");
1060 assert_eq!(env["RIVET__PEGBOARD__RESCHEDULE_BACKOFF_MAX_EXPONENT"], "1");
1061 assert_eq!(env["RIVET__PEGBOARD__RUNNER_ELIGIBLE_THRESHOLD"], "5000");
1062 assert_eq!(env["RIVET__PEGBOARD__RUNNER_LOST_THRESHOLD"], "7000");
1063 assert_eq!(env["RIVET__PEGBOARD__ENVOY_ELIGIBLE_THRESHOLD"], "5000");
1064 assert_eq!(env["RIVET__PEGBOARD__ENVOY_LOST_THRESHOLD"], "7000");
1065 assert_eq!(env["RIVET__PEGBOARD__MIN_METADATA_POLL_INTERVAL"], "1000");
1066 assert_eq!(env["RIVET__AUTH__ADMIN_TOKEN"], "default");
1067 assert_eq!(env["RIVET__FEATURES__GUARD_GATEWAY_V3__MODE"], "on");
1068 assert_eq!(env["RIVET__FEATURES__GUARD_GATEWAY_V3__PERCENTAGE"], "100");
1069 assert_eq!(env["RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION"], "1");
1070 assert_eq!(env["RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION"], "1");
1071 assert_eq!(env["RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION"], "2");
1072 }
1073
1074 #[cfg(unix)]
1075 #[test]
1076 fn pid_is_alive_true_for_self_false_for_missing() {
1077 assert!(pid_is_alive(std::process::id()));
1078 assert!(!pid_is_alive(0x7FFF_FFFF));
1080 }
1081
1082 #[tokio::test]
1083 async fn resolver_prefers_existing_engine_before_filesystem_paths() {
1084 let temp = tempfile::tempdir().expect("create temp dir");
1085 let local = temp
1086 .path()
1087 .join("target")
1088 .join("debug")
1089 .join(exe_name("rivet-engine"));
1090 std::fs::create_dir_all(local.parent().expect("local parent")).expect("create local dir");
1091 std::fs::write(&local, b"local").expect("write local binary");
1092 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1093
1094 let resolved = resolve_engine_binary_after_probe(
1095 &test_config("http://127.0.0.1:1".to_owned(), false),
1096 true,
1097 &[temp.path().to_path_buf()],
1098 cached,
1099 )
1100 .await
1101 .expect("resolve engine");
1102
1103 assert_eq!(resolved, ResolvedEngine::Existing);
1104 }
1105
1106 #[tokio::test]
1107 async fn resolver_prefers_local_binary_before_cached_binary() {
1108 let temp = tempfile::tempdir().expect("create temp dir");
1109 let local = temp
1110 .path()
1111 .join("target")
1112 .join("debug")
1113 .join(exe_name("rivet-engine"));
1114 std::fs::create_dir_all(local.parent().expect("local parent")).expect("create local dir");
1115 std::fs::write(&local, b"local").expect("write local binary");
1116 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1117 std::fs::create_dir_all(cached.parent().expect("cached parent")).expect("create cache dir");
1118 std::fs::write(&cached, b"cached").expect("write cached binary");
1119
1120 let resolved = resolve_engine_binary_after_probe(
1121 &test_config("http://127.0.0.1:1".to_owned(), false),
1122 false,
1123 &[temp.path().to_path_buf()],
1124 cached,
1125 )
1126 .await
1127 .expect("resolve engine");
1128
1129 assert_eq!(resolved, ResolvedEngine::Binary(local));
1130 }
1131
1132 #[tokio::test]
1133 async fn resolver_reuses_cached_binary_without_download() {
1134 let temp = tempfile::tempdir().expect("create temp dir");
1135 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1136 std::fs::create_dir_all(cached.parent().expect("cached parent")).expect("create cache dir");
1137 std::fs::write(&cached, b"cached").expect("write cached binary");
1138
1139 let resolved = resolve_engine_binary_after_probe(
1140 &test_config("http://127.0.0.1:1".to_owned(), false),
1141 false,
1142 &[],
1143 cached.clone(),
1144 )
1145 .await
1146 .expect("resolve engine");
1147
1148 assert_eq!(resolved, ResolvedEngine::Binary(cached));
1149 }
1150
1151 #[tokio::test]
1152 async fn resolver_reports_actionable_error_without_binary_or_download() {
1153 let temp = tempfile::tempdir().expect("create temp dir");
1154 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1155
1156 let error = resolve_engine_binary_after_probe(
1157 &test_config("http://127.0.0.1:1".to_owned(), false),
1158 false,
1159 &[],
1160 cached,
1161 )
1162 .await
1163 .expect_err("missing binary should fail");
1164 let message = error.to_string();
1165
1166 assert!(message.contains("No usable engine binary was found"));
1167 assert!(message.contains("Build `rivet-engine`"));
1168 assert!(message.contains("RIVET_ENGINE_BINARY_PATH"));
1169 }
1170
1171 #[tokio::test]
1172 async fn resolver_download_checks_manifest_checksum() {
1173 let temp = tempfile::tempdir().expect("create temp dir");
1174 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1175 let artifact = engine_artifact_name();
1176 let expected = sha256_hex(b"different bytes");
1177 let manifest = format!("{expected} {artifact}\n");
1178 let releases_endpoint = spawn_download_server(HashMap::from([
1179 (
1180 format!("/rivet/test-version/engine/SHA256SUMS"),
1181 manifest.into_bytes(),
1182 ),
1183 (
1184 format!("/rivet/test-version/engine/{artifact}"),
1185 b"actual bytes".to_vec(),
1186 ),
1187 ]))
1188 .await;
1189
1190 let error = resolve_engine_binary_after_probe(
1191 &test_config(releases_endpoint, true),
1192 false,
1193 &[],
1194 cached,
1195 )
1196 .await
1197 .expect_err("checksum mismatch should fail");
1198
1199 assert!(
1200 error
1201 .to_string()
1202 .contains("Engine binary checksum mismatch")
1203 );
1204 }
1205
1206 async fn spawn_download_server(routes: HashMap<String, Vec<u8>>) -> String {
1207 let listener = TcpListener::bind("127.0.0.1:0")
1208 .await
1209 .expect("bind download server");
1210 let addr = listener.local_addr().expect("download server address");
1211 tokio::spawn(async move {
1212 for _ in 0..routes.len() {
1213 let (mut socket, _) = listener.accept().await.expect("accept download request");
1214 let mut buffer = [0_u8; 2048];
1215 let n = socket
1216 .read(&mut buffer)
1217 .await
1218 .expect("read download request");
1219 let request = String::from_utf8_lossy(&buffer[..n]);
1220 let path = request
1221 .split_whitespace()
1222 .nth(1)
1223 .expect("request path")
1224 .to_owned();
1225 let body = routes.get(&path).expect("route body");
1226 let header = format!(
1227 "HTTP/1.1 200 OK\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
1228 body.len()
1229 );
1230 socket
1231 .write_all(header.as_bytes())
1232 .await
1233 .expect("write response header");
1234 socket.write_all(body).await.expect("write response body");
1235 }
1236 });
1237 format!("http://{addr}")
1238 }
1239}