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