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__ENVOY_ELIGIBLE_THRESHOLD".to_owned(),
419 "5000".to_owned(),
420 ),
421 (
422 "RIVET__PEGBOARD__ENVOY_LOST_THRESHOLD".to_owned(),
423 "7000".to_owned(),
424 ),
425 (
426 "RIVET__PEGBOARD__MIN_METADATA_POLL_INTERVAL".to_owned(),
427 "1000".to_owned(),
428 ),
429 (
430 "RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION".to_owned(),
431 "1".to_owned(),
432 ),
433 (
434 "RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION".to_owned(),
435 "1".to_owned(),
436 ),
437 (
438 "RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION".to_owned(),
439 "2".to_owned(),
440 ),
441 ])
442}
443
444fn write_engine_config(config: &EngineResolverConfig) -> Result<Option<PathBuf>> {
445 let Some(public_url) = &config.public_url else {
446 return Ok(None);
447 };
448
449 let public_url = Url::parse(public_url)
450 .with_context(|| format!("parse engine public URL `{public_url}`"))?;
451 let peer_url = peer_url_for_public_url(&public_url)?;
452 let dir = storage_root()?.join("var").join("engine");
453 ensure_dir(&dir)?;
454 let path = dir.join("config.json");
455 let config = serde_json::json!({
456 "topology": {
457 "datacenter_label": 1,
458 "datacenters": {
459 "default": {
460 "datacenter_label": 1,
461 "is_leader": true,
462 "public_url": public_url.as_str(),
463 "peer_url": peer_url,
464 }
465 }
466 }
467 });
468 std::fs::write(&path, serde_json::to_string_pretty(&config)?)
469 .with_context(|| format!("write engine config `{}`", path.display()))?;
470 Ok(Some(path))
471}
472
473fn peer_url_for_public_url(public_url: &Url) -> Result<String> {
474 let mut peer_url = public_url.clone();
475 let port = public_url
476 .port_or_known_default()
477 .ok_or_else(|| invalid_endpoint(public_url.as_str(), "missing port"))?
478 .checked_add(1)
479 .ok_or_else(|| invalid_endpoint(public_url.as_str(), "port is too large"))?;
480 if peer_url.set_port(Some(port)).is_err() {
481 return Err(invalid_endpoint(
482 public_url.as_str(),
483 "could not derive peer URL port",
484 ));
485 }
486 peer_url.set_path("");
487 peer_url.set_query(None);
488 peer_url.set_fragment(None);
489 Ok(peer_url.to_string().trim_end_matches('/').to_string())
490}
491
492pub async fn resolve_engine_binary(config: &EngineResolverConfig) -> Result<ResolvedEngine> {
493 if let Some(path) = config.explicit_binary_path.as_ref() {
494 return verify_binary_path(path);
495 }
496
497 if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
498 return verify_binary_path(&path);
499 }
500
501 if probe_existing_engine(&config.endpoint).await?.is_some() {
502 return Ok(ResolvedEngine::Existing);
503 }
504
505 let local_roots = local_engine_search_roots();
506 let cached = cached_engine_path(&config.version)?;
507 resolve_engine_binary_after_probe(config, false, &local_roots, cached).await
508}
509
510pub async fn resolve_engine_binary_path(config: &EngineResolverConfig) -> Result<PathBuf> {
517 if let Some(path) = config.explicit_binary_path.as_ref() {
518 verify_binary_path(path)?;
519 return Ok(path.clone());
520 }
521
522 if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
523 verify_binary_path(&path)?;
524 return Ok(path);
525 }
526
527 let local_roots = local_engine_search_roots();
528 let cached = cached_engine_path(&config.version)?;
529 match resolve_engine_binary_after_probe(config, false, &local_roots, cached).await? {
530 ResolvedEngine::Binary(path) => Ok(path),
531 ResolvedEngine::Existing => {
532 unreachable!("no-probe resolution never returns Existing")
533 }
534 }
535}
536
537async fn resolve_engine_binary_after_probe(
538 config: &EngineResolverConfig,
539 existing_engine: bool,
540 local_roots: &[PathBuf],
541 cached: PathBuf,
542) -> Result<ResolvedEngine> {
543 if existing_engine {
544 return Ok(ResolvedEngine::Existing);
545 }
546
547 if let Some(path) = find_local_engine_binary_in_roots(local_roots) {
548 return Ok(ResolvedEngine::Binary(path));
549 }
550
551 if cached.exists() {
552 return Ok(ResolvedEngine::Binary(cached));
553 }
554
555 if !config.auto_download {
556 return Err(EngineProcessError::BinaryUnavailable {
557 version: config.version.clone(),
558 }
559 .build());
560 }
561
562 download_engine_binary(config, &cached).await?;
563 Ok(ResolvedEngine::Binary(cached))
564}
565
566fn verify_binary_path(path: &Path) -> Result<ResolvedEngine> {
567 if !path.exists() {
568 return Err(EngineProcessError::BinaryNotFound {
569 path: path.display().to_string(),
570 }
571 .build());
572 }
573 Ok(ResolvedEngine::Binary(path.to_path_buf()))
574}
575
576fn local_engine_search_roots() -> Vec<PathBuf> {
577 Path::new(env!("CARGO_MANIFEST_DIR"))
578 .ancestors()
579 .map(Path::to_path_buf)
580 .collect()
581}
582
583fn find_local_engine_binary_in_roots(roots: &[PathBuf]) -> Option<PathBuf> {
584 for root in roots {
585 for profile in ["debug", "release"] {
586 let candidate = root
587 .join("target")
588 .join(profile)
589 .join(exe_name("rivet-engine"));
590 if candidate.exists() {
591 return Some(candidate);
592 }
593 }
594 }
595 None
596}
597
598fn cached_engine_path(version: &str) -> Result<PathBuf> {
599 Ok(storage_root()?
600 .join("engine")
601 .join(version)
602 .join(engine_artifact_name()))
603}
604
605async fn download_engine_binary(config: &EngineResolverConfig, destination: &Path) -> Result<()> {
606 let artifact = engine_artifact_name();
607 let base = config.releases_endpoint.trim_end_matches('/');
608 let artifact_url = format!("{base}/rivet/{}/engine/{artifact}", config.version);
609 let manifest_url = format!("{base}/rivet/{}/engine/SHA256SUMS", config.version);
610 let client = Client::builder()
611 .timeout(DOWNLOAD_TIMEOUT)
612 .build()
613 .context("build reqwest client for engine download")?;
614
615 let manifest = fetch_text(&client, &manifest_url).await?;
616 let expected = checksum_for_artifact(&manifest, &artifact).ok_or_else(|| {
617 EngineProcessError::DownloadFailed {
618 url: manifest_url.clone(),
619 reason: format!("manifest does not contain `{artifact}`"),
620 }
621 .build()
622 })?;
623
624 let bytes = fetch_bytes(&client, &artifact_url).await?;
625 let received = sha256_hex(&bytes);
626 if !received.eq_ignore_ascii_case(&expected) {
627 return Err(EngineProcessError::ChecksumMismatch {
628 artifact,
629 expected,
630 received,
631 }
632 .build());
633 }
634
635 let parent = destination
636 .parent()
637 .context("engine cache destination has no parent")?;
638 ensure_dir(parent)?;
639 std::fs::write(destination, bytes)
640 .with_context(|| format!("write engine binary `{}`", destination.display()))?;
641 make_executable(destination)?;
642 Ok(())
643}
644
645async fn fetch_text(client: &Client, url: &str) -> Result<String> {
646 let response = client.get(url).send().await.map_err(|error| {
647 EngineProcessError::DownloadFailed {
648 url: url.to_owned(),
649 reason: error.to_string(),
650 }
651 .build()
652 })?;
653 if !response.status().is_success() {
654 let status = response.status();
655 return Err(EngineProcessError::DownloadFailed {
656 url: url.to_owned(),
657 reason: format!("unexpected status {status}"),
658 }
659 .build());
660 }
661 response.text().await.map_err(|error| {
662 EngineProcessError::DownloadFailed {
663 url: url.to_owned(),
664 reason: error.to_string(),
665 }
666 .build()
667 })
668}
669
670async fn fetch_bytes(client: &Client, url: &str) -> Result<Vec<u8>> {
671 let response = client.get(url).send().await.map_err(|error| {
672 EngineProcessError::DownloadFailed {
673 url: url.to_owned(),
674 reason: error.to_string(),
675 }
676 .build()
677 })?;
678 if !response.status().is_success() {
679 let status = response.status();
680 return Err(EngineProcessError::DownloadFailed {
681 url: url.to_owned(),
682 reason: format!("unexpected status {status}"),
683 }
684 .build());
685 }
686 response
687 .bytes()
688 .await
689 .map(|bytes| bytes.to_vec())
690 .map_err(|error| {
691 EngineProcessError::DownloadFailed {
692 url: url.to_owned(),
693 reason: error.to_string(),
694 }
695 .build()
696 })
697}
698
699fn checksum_for_artifact(manifest: &str, artifact: &str) -> Option<String> {
700 manifest.lines().find_map(|line| {
701 let mut parts = line.split_whitespace();
702 let checksum = parts.next()?;
703 let name = parts.next()?.trim_start_matches('*');
704 (checksum.len() == 64 && name == artifact).then(|| checksum.to_owned())
705 })
706}
707
708fn sha256_hex(bytes: &[u8]) -> String {
709 let digest = Sha256::digest(bytes);
710 let mut out = String::with_capacity(digest.len() * 2);
711 for byte in digest {
712 use std::fmt::Write;
713 let _ = write!(&mut out, "{byte:02x}");
714 }
715 out
716}
717
718fn engine_artifact_name() -> String {
719 let arch = match std::env::consts::ARCH {
720 "x86_64" => "x86_64",
721 "aarch64" => "aarch64",
722 other => other,
723 };
724 let target = match std::env::consts::OS {
725 "linux" => format!("{arch}-unknown-linux-musl"),
726 "macos" => format!("{arch}-apple-darwin"),
727 "windows" => format!("{arch}-pc-windows-gnu.exe"),
728 other => format!("{arch}-{other}"),
729 };
730 format!("rivet-engine-{target}")
731}
732
733fn exe_name(base: &str) -> String {
734 if cfg!(windows) {
735 format!("{base}.exe")
736 } else {
737 base.to_owned()
738 }
739}
740
741fn make_executable(path: &Path) -> Result<()> {
742 #[cfg(unix)]
743 {
744 use std::os::unix::fs::PermissionsExt;
745 let mut permissions = std::fs::metadata(path)
746 .with_context(|| format!("read metadata for `{}`", path.display()))?
747 .permissions();
748 permissions.set_mode(0o755);
749 std::fs::set_permissions(path, permissions)
750 .with_context(|| format!("mark `{}` executable", path.display()))?;
751 }
752 #[cfg(not(unix))]
753 {
754 let _ = path;
755 }
756 Ok(())
757}
758
759impl Drop for EngineProcessManager {
760 fn drop(&mut self) {
761 if let Some(handle) = self.watcher.take() {
762 handle.abort();
768 tracing::debug!(
769 "aborted engine watcher; engine continues running (intentional orphan)"
770 );
771 }
772 }
773}
774
775fn spawn_engine_watcher(mut child: Child, pid: u32) -> JoinHandle<()> {
780 tokio::spawn(async move {
781 match child.wait().await {
782 Ok(status) if status.success() => {
783 tracing::warn!(
784 pid,
785 ?status,
786 "engine process exited cleanly while rivetkit was still running; \
787 rivetkit expected the engine to outlive it"
788 );
789 }
790 Ok(status) => {
791 tracing::error!(
792 pid,
793 ?status,
794 "engine process crashed while rivetkit was still running"
795 );
796 }
797 Err(error) => {
798 tracing::error!(
799 pid,
800 ?error,
801 "failed to wait on engine process; cannot detect crashes"
802 );
803 }
804 }
805 })
806}
807
808async fn probe_existing_engine(endpoint: &str) -> Result<Option<EngineHealthResponse>> {
815 let health_url = engine_health_url(endpoint);
816 let client = Client::builder()
817 .build()
818 .context("build reqwest client for engine probe")?;
819
820 let response = match client
821 .get(&health_url)
822 .timeout(Duration::from_secs(1))
823 .send()
824 .await
825 {
826 Ok(response) => response,
827 Err(_) => return Ok(None),
828 };
829
830 if !response.status().is_success() {
831 return Ok(None);
832 }
833
834 let health = response
835 .json::<EngineHealthResponse>()
836 .await
837 .context("decode existing engine health response")?;
838
839 match health.runtime.as_deref() {
840 Some(ENGINE_RUNTIME) => Ok(Some(health)),
841 Some(RIVETKIT_RUNTIME) => Err(EngineProcessError::PortOccupied {
842 endpoint: endpoint.to_owned(),
843 runtime: RIVETKIT_RUNTIME.to_owned(),
844 }
845 .build()),
846 Some(other) => Err(EngineProcessError::PortOccupied {
847 endpoint: endpoint.to_owned(),
848 runtime: other.to_owned(),
849 }
850 .build()),
851 None => Err(EngineProcessError::PortOccupied {
852 endpoint: endpoint.to_owned(),
853 runtime: "unknown".to_owned(),
854 }
855 .build()),
856 }
857}
858
859fn engine_health_url(endpoint: &str) -> String {
860 format!("{}/health", endpoint.trim_end_matches('/'))
861}
862
863fn storage_root() -> Result<PathBuf> {
864 if let Ok(path) = std::env::var("RIVETKIT_STORAGE_PATH") {
865 return Ok(PathBuf::from(path).join(".rivetkit"));
866 }
867 let home = std::env::var("HOME")
868 .map(PathBuf::from)
869 .or_else(|_| std::env::current_dir())
870 .context("locate home directory for engine storage path")?;
871 Ok(home.join(".rivetkit"))
872}
873
874fn ensure_dir(path: &Path) -> Result<()> {
875 std::fs::create_dir_all(path).with_context(|| format!("create directory `{}`", path.display()))
876}
877
878fn open_log_file(path: &Path) -> Result<std::fs::File> {
879 std::fs::OpenOptions::new()
880 .create(true)
881 .append(true)
882 .open(path)
883 .with_context(|| format!("open log file `{}`", path.display()))
884}
885
886fn log_timestamp() -> String {
887 let now = std::time::SystemTime::now()
888 .duration_since(std::time::UNIX_EPOCH)
889 .unwrap_or_default();
890 format!("{}", now.as_secs())
891}
892
893async fn wait_for_engine_health(health_url: &str) -> Result<EngineHealthResponse> {
894 const HEALTH_MAX_WAIT: Duration = Duration::from_secs(10);
895 const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
896 const HEALTH_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
897 const HEALTH_MAX_BACKOFF: Duration = Duration::from_secs(1);
898
899 let client = Client::builder()
900 .build()
901 .context("build reqwest client for engine health check")?;
902 let deadline = Instant::now() + HEALTH_MAX_WAIT;
903 let mut attempt = 0u32;
904 let mut backoff = HEALTH_INITIAL_BACKOFF;
905
906 loop {
907 attempt += 1;
908
909 let last_error = match client
910 .get(health_url)
911 .timeout(HEALTH_REQUEST_TIMEOUT)
912 .send()
913 .await
914 {
915 Ok(response) if response.status().is_success() => {
916 let health = response
917 .json::<EngineHealthResponse>()
918 .await
919 .context("decode engine health response")?;
920 return Ok(health);
921 }
922 Ok(response) => format!("unexpected status {}", response.status()),
923 Err(error) => error.to_string(),
924 };
925
926 if Instant::now() >= deadline {
927 return Err(EngineProcessError::HealthCheckFailed {
928 attempts: attempt,
929 reason: last_error,
930 }
931 .build());
932 }
933
934 tokio::time::sleep(backoff).await;
935 backoff = std::cmp::min(backoff * 2, HEALTH_MAX_BACKOFF);
936 }
937}
938
939async fn terminate_failed_spawn(child: &mut Child) -> Result<()> {
944 const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
945
946 if child
947 .try_wait()
948 .context("check engine process status")?
949 .is_some()
950 {
951 return Ok(());
952 }
953
954 child
955 .start_kill()
956 .context("kill half-started engine process")?;
957 match tokio::time::timeout(SHUTDOWN_TIMEOUT, child.wait()).await {
958 Ok(result) => {
959 let status = result.context("wait for half-started engine to exit")?;
960 tracing::info!(?status, "half-started engine process exited");
961 Ok(())
962 }
963 Err(_) => {
964 tracing::warn!("half-started engine process did not exit within timeout");
965 Ok(())
966 }
967 }
968}
969
970fn invalid_endpoint(endpoint: &str, reason: &str) -> anyhow::Error {
971 EngineProcessError::InvalidEndpoint {
972 endpoint: endpoint.to_owned(),
973 reason: reason.to_owned(),
974 }
975 .build()
976}
977
978#[cfg(test)]
979mod tests {
980 use std::collections::HashMap;
981
982 use tokio::io::{AsyncReadExt, AsyncWriteExt};
983 use tokio::net::TcpListener;
984
985 use super::*;
986
987 fn test_config(releases_endpoint: String, auto_download: bool) -> EngineResolverConfig {
988 EngineResolverConfig {
989 endpoint: "http://127.0.0.1:1".to_owned(),
990 explicit_binary_path: None,
991 bind_host: None,
992 bind_port: None,
993 public_url: None,
994 auto_download,
995 version: "test-version".to_owned(),
996 releases_endpoint,
997 }
998 }
999
1000 fn stamp_with_bind(bind_host: &str) -> EngineRuntimeStamp {
1001 EngineRuntimeStamp {
1002 pid: 1,
1003 endpoint: "http://127.0.0.1:6420".to_owned(),
1004 bind_host: bind_host.to_owned(),
1005 public_url: None,
1006 }
1007 }
1008
1009 #[test]
1010 fn binds_loopback_only_detects_loopback_addresses() {
1011 assert!(stamp_with_bind("127.0.0.1").binds_loopback_only());
1012 assert!(stamp_with_bind("localhost").binds_loopback_only());
1013 assert!(stamp_with_bind("::1").binds_loopback_only());
1014 }
1015
1016 #[test]
1017 fn binds_loopback_only_false_for_reachable_addresses() {
1018 assert!(!stamp_with_bind("0.0.0.0").binds_loopback_only());
1019 assert!(!stamp_with_bind("::").binds_loopback_only());
1020 assert!(!stamp_with_bind("192.168.1.5").binds_loopback_only());
1021 }
1022
1023 #[test]
1024 fn resolve_bind_host_prefers_override_then_endpoint() {
1025 let mut config = test_config(String::new(), false);
1026 config.endpoint = "http://127.0.0.1:6420".to_owned();
1027 assert_eq!(resolve_bind_host(&config).unwrap(), "127.0.0.1");
1028 config.bind_host = Some("0.0.0.0".to_owned());
1029 assert_eq!(resolve_bind_host(&config).unwrap(), "0.0.0.0");
1030 }
1031
1032 #[test]
1033 fn engine_env_uses_development_recovery_thresholds() {
1034 let env = engine_env(&test_config(String::new(), false)).expect("build engine env");
1035 let env = env.into_iter().collect::<HashMap<_, _>>();
1036
1037 assert_eq!(env["RIVET__PEGBOARD__RETRY_RESET_DURATION"], "100");
1038 assert_eq!(env["RIVET__PEGBOARD__BASE_RETRY_TIMEOUT"], "100");
1039 assert_eq!(
1040 env["RIVET__PEGBOARD__RESCHEDULE_BACKOFF_MAX_EXPONENT"],
1041 "1"
1042 );
1043 assert_eq!(env["RIVET__PEGBOARD__ENVOY_ELIGIBLE_THRESHOLD"], "5000");
1044 assert_eq!(env["RIVET__PEGBOARD__ENVOY_LOST_THRESHOLD"], "7000");
1045 assert_eq!(env["RIVET__PEGBOARD__MIN_METADATA_POLL_INTERVAL"], "1000");
1046 assert_eq!(env["RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION"], "1");
1047 assert_eq!(env["RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION"], "1");
1048 assert_eq!(env["RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION"], "2");
1049 }
1050
1051 #[cfg(unix)]
1052 #[test]
1053 fn pid_is_alive_true_for_self_false_for_missing() {
1054 assert!(pid_is_alive(std::process::id()));
1055 assert!(!pid_is_alive(0x7FFF_FFFF));
1057 }
1058
1059 #[tokio::test]
1060 async fn resolver_prefers_existing_engine_before_filesystem_paths() {
1061 let temp = tempfile::tempdir().expect("create temp dir");
1062 let local = temp
1063 .path()
1064 .join("target")
1065 .join("debug")
1066 .join(exe_name("rivet-engine"));
1067 std::fs::create_dir_all(local.parent().expect("local parent")).expect("create local dir");
1068 std::fs::write(&local, b"local").expect("write local binary");
1069 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1070
1071 let resolved = resolve_engine_binary_after_probe(
1072 &test_config("http://127.0.0.1:1".to_owned(), false),
1073 true,
1074 &[temp.path().to_path_buf()],
1075 cached,
1076 )
1077 .await
1078 .expect("resolve engine");
1079
1080 assert_eq!(resolved, ResolvedEngine::Existing);
1081 }
1082
1083 #[tokio::test]
1084 async fn resolver_prefers_local_binary_before_cached_binary() {
1085 let temp = tempfile::tempdir().expect("create temp dir");
1086 let local = temp
1087 .path()
1088 .join("target")
1089 .join("debug")
1090 .join(exe_name("rivet-engine"));
1091 std::fs::create_dir_all(local.parent().expect("local parent")).expect("create local dir");
1092 std::fs::write(&local, b"local").expect("write local binary");
1093 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1094 std::fs::create_dir_all(cached.parent().expect("cached parent")).expect("create cache dir");
1095 std::fs::write(&cached, b"cached").expect("write cached binary");
1096
1097 let resolved = resolve_engine_binary_after_probe(
1098 &test_config("http://127.0.0.1:1".to_owned(), false),
1099 false,
1100 &[temp.path().to_path_buf()],
1101 cached,
1102 )
1103 .await
1104 .expect("resolve engine");
1105
1106 assert_eq!(resolved, ResolvedEngine::Binary(local));
1107 }
1108
1109 #[tokio::test]
1110 async fn resolver_reuses_cached_binary_without_download() {
1111 let temp = tempfile::tempdir().expect("create temp dir");
1112 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1113 std::fs::create_dir_all(cached.parent().expect("cached parent")).expect("create cache dir");
1114 std::fs::write(&cached, b"cached").expect("write cached binary");
1115
1116 let resolved = resolve_engine_binary_after_probe(
1117 &test_config("http://127.0.0.1:1".to_owned(), false),
1118 false,
1119 &[],
1120 cached.clone(),
1121 )
1122 .await
1123 .expect("resolve engine");
1124
1125 assert_eq!(resolved, ResolvedEngine::Binary(cached));
1126 }
1127
1128 #[tokio::test]
1129 async fn resolver_reports_actionable_error_without_binary_or_download() {
1130 let temp = tempfile::tempdir().expect("create temp dir");
1131 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1132
1133 let error = resolve_engine_binary_after_probe(
1134 &test_config("http://127.0.0.1:1".to_owned(), false),
1135 false,
1136 &[],
1137 cached,
1138 )
1139 .await
1140 .expect_err("missing binary should fail");
1141 let message = error.to_string();
1142
1143 assert!(message.contains("No usable engine binary was found"));
1144 assert!(message.contains("Build `rivet-engine`"));
1145 assert!(message.contains("RIVET_ENGINE_BINARY_PATH"));
1146 }
1147
1148 #[tokio::test]
1149 async fn resolver_download_checks_manifest_checksum() {
1150 let temp = tempfile::tempdir().expect("create temp dir");
1151 let cached = temp.path().join("cache").join(exe_name("rivet-engine"));
1152 let artifact = engine_artifact_name();
1153 let expected = sha256_hex(b"different bytes");
1154 let manifest = format!("{expected} {artifact}\n");
1155 let releases_endpoint = spawn_download_server(HashMap::from([
1156 (
1157 format!("/rivet/test-version/engine/SHA256SUMS"),
1158 manifest.into_bytes(),
1159 ),
1160 (
1161 format!("/rivet/test-version/engine/{artifact}"),
1162 b"actual bytes".to_vec(),
1163 ),
1164 ]))
1165 .await;
1166
1167 let error = resolve_engine_binary_after_probe(
1168 &test_config(releases_endpoint, true),
1169 false,
1170 &[],
1171 cached,
1172 )
1173 .await
1174 .expect_err("checksum mismatch should fail");
1175
1176 assert!(
1177 error
1178 .to_string()
1179 .contains("Engine binary checksum mismatch")
1180 );
1181 }
1182
1183 async fn spawn_download_server(routes: HashMap<String, Vec<u8>>) -> String {
1184 let listener = TcpListener::bind("127.0.0.1:0")
1185 .await
1186 .expect("bind download server");
1187 let addr = listener.local_addr().expect("download server address");
1188 tokio::spawn(async move {
1189 for _ in 0..routes.len() {
1190 let (mut socket, _) = listener.accept().await.expect("accept download request");
1191 let mut buffer = [0_u8; 2048];
1192 let n = socket
1193 .read(&mut buffer)
1194 .await
1195 .expect("read download request");
1196 let request = String::from_utf8_lossy(&buffer[..n]);
1197 let path = request
1198 .split_whitespace()
1199 .nth(1)
1200 .expect("request path")
1201 .to_owned();
1202 let body = routes.get(&path).expect("route body");
1203 let header = format!(
1204 "HTTP/1.1 200 OK\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
1205 body.len()
1206 );
1207 socket
1208 .write_all(header.as_bytes())
1209 .await
1210 .expect("write response header");
1211 socket.write_all(body).await.expect("write response body");
1212 }
1213 });
1214 format!("http://{addr}")
1215 }
1216}