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