use std::process::Stdio;
use std::time::{Duration, Instant};
use std::{
env,
path::{Path, PathBuf},
};
use anyhow::{Context, Result};
use reqwest::{Client, Url};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::io::AsyncReadExt;
use tokio::process::{Child, Command};
use tokio::task::JoinHandle;
mod error;
pub use error::EngineProcessError;
const ENGINE_RUNTIME: &str = "engine";
const ENGINE_ARCHITECTURE: &str = "v3";
const RIVETKIT_RUNTIME: &str = "rivetkit";
const ENGINE_VERSION_ENV: &str = "RIVETKIT_ENGINE_VERSION";
const RELEASES_ENDPOINT_ENV: &str = "RIVETKIT_ENGINE_RELEASES_ENDPOINT";
const RELEASES_ENDPOINT: &str = "https://releases.rivet.dev";
const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug, Deserialize)]
struct EngineHealthResponse {
status: Option<String>,
runtime: Option<String>,
version: Option<String>,
architecture: Option<String>,
}
#[derive(Clone, Debug)]
pub struct EngineResolverConfig {
pub endpoint: String,
pub explicit_binary_path: Option<PathBuf>,
pub bind_host: Option<String>,
pub bind_port: Option<u16>,
pub public_url: Option<String>,
pub auto_download: bool,
pub version: String,
pub releases_endpoint: String,
}
impl EngineResolverConfig {
pub fn from_parts(
endpoint: &str,
explicit_binary_path: Option<PathBuf>,
bind_host: Option<String>,
bind_port: Option<u16>,
auto_download: bool,
) -> Self {
Self {
endpoint: endpoint.to_owned(),
explicit_binary_path,
bind_host,
bind_port,
public_url: None,
auto_download,
version: std::env::var(ENGINE_VERSION_ENV)
.unwrap_or_else(|_| env!("CARGO_PKG_VERSION").to_owned()),
releases_endpoint: std::env::var(RELEASES_ENDPOINT_ENV)
.unwrap_or_else(|_| RELEASES_ENDPOINT.to_owned()),
}
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum ResolvedEngine {
Existing,
Binary(PathBuf),
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EngineRuntimeStamp {
pub pid: u32,
pub endpoint: String,
pub bind_host: String,
pub public_url: Option<String>,
}
impl EngineRuntimeStamp {
pub fn binds_loopback_only(&self) -> bool {
if self.bind_host == "localhost" {
return true;
}
self.bind_host
.parse::<std::net::IpAddr>()
.map(|ip| ip.is_loopback())
.unwrap_or(false)
}
}
#[derive(Debug)]
pub enum EngineStartup {
Spawned,
Reused { stamp: Option<EngineRuntimeStamp> },
}
#[derive(Debug)]
pub struct EngineProcessManager {
watcher: Option<JoinHandle<()>>,
startup: EngineStartup,
}
impl EngineProcessManager {
pub async fn start_or_reuse(config: EngineResolverConfig) -> Result<Self> {
let resolved = resolve_engine_binary(&config).await?;
Self::start_resolved(resolved, &config).await
}
pub fn startup(&self) -> &EngineStartup {
&self.startup
}
async fn start_resolved(
resolved: ResolvedEngine,
config: &EngineResolverConfig,
) -> Result<Self> {
let endpoint = &config.endpoint;
if matches!(resolved, ResolvedEngine::Existing) {
tracing::info!(
endpoint = %endpoint,
"reusing already-running engine process"
);
return Ok(Self {
watcher: None,
startup: EngineStartup::Reused {
stamp: read_engine_stamp(),
},
});
}
let ResolvedEngine::Binary(binary_path) = resolved else {
unreachable!("existing engine handled above");
};
if let Some(health) = probe_existing_engine(endpoint).await? {
tracing::info!(
endpoint = %endpoint,
status = ?health.status,
runtime = ?health.runtime,
version = ?health.version,
"reusing already-running engine process"
);
return Ok(Self {
watcher: None,
startup: EngineStartup::Reused {
stamp: read_engine_stamp(),
},
});
}
if !binary_path.exists() {
return Err(EngineProcessError::BinaryNotFound {
path: binary_path.display().to_string(),
}
.build());
}
ensure_compatible_binary(&binary_path).await?;
let env = engine_env(config)?;
let config_path = write_engine_config(config)?;
let data_path = engine_data_path()?;
let logs_dir = storage_root()?
.join("var")
.join("logs")
.join("rivet-engine");
ensure_dir(&logs_dir).context("create engine logs directory")?;
ensure_dir(&engine_log_directory()?).context("create engine log shard directory")?;
let timestamp = log_timestamp();
let stdout_log_path = logs_dir.join(format!("engine-{timestamp}-stdout.log"));
let stderr_log_path = logs_dir.join(format!("engine-{timestamp}-stderr.log"));
let stdout_file = open_log_file(&stdout_log_path)
.with_context(|| format!("open engine stdout log `{}`", stdout_log_path.display()))?;
let stderr_file = open_log_file(&stderr_log_path)
.with_context(|| format!("open engine stderr log `{}`", stderr_log_path.display()))?;
let mut command = Command::new(&binary_path);
command.arg("--config").arg(&config_path);
for (key, value) in &env {
command.env(key, value);
}
command
.stdin(Stdio::null())
.stdout(Stdio::from(stdout_file))
.stderr(Stdio::from(stderr_file));
#[cfg(unix)]
command.process_group(0);
let mut child = command
.spawn()
.with_context(|| format!("spawn engine binary `{}`", binary_path.display()))?;
let pid = child
.id()
.ok_or_else(|| EngineProcessError::MissingPid.build())?;
tracing::info!(
pid,
path = %binary_path.display(),
endpoint = %endpoint,
data_path = %data_path.display(),
"spawned engine process (intentionally orphaned, will outlive this process)"
);
tracing::info!(
stdout_log = %stdout_log_path.display(),
stderr_log = %stderr_log_path.display(),
"engine stdout/stderr piped to log files"
);
let health_url = engine_health_url(endpoint);
let health = match wait_for_engine_health(&health_url).await {
Ok(health) => health,
Err(error) => {
let error = match child.try_wait() {
Ok(Some(status)) => error.context(format!(
"engine process exited before becoming healthy with status {status}"
)),
Ok(None) => error,
Err(wait_error) => error.context(format!(
"failed to inspect engine process status: {wait_error:#}"
)),
};
if let Err(cleanup_error) = terminate_failed_spawn(&mut child).await {
tracing::warn!(
?cleanup_error,
"failed to terminate engine process that never became healthy"
);
}
return Err(error);
}
};
tracing::info!(
pid,
status = ?health.status,
runtime = ?health.runtime,
version = ?health.version,
"engine process is healthy"
);
let ready = match ensure_compatible_engine(endpoint, &health) {
Ok(()) => wait_for_engine_ready(endpoint).await,
Err(error) => Err(error),
};
if let Err(error) = ready {
if let Err(cleanup_error) = terminate_failed_spawn(&mut child).await {
tracing::warn!(
?cleanup_error,
"failed to terminate engine process that never became ready"
);
}
return Err(error);
}
tracing::info!(pid, "engine process is ready");
let stamp = EngineRuntimeStamp {
pid,
endpoint: endpoint.clone(),
bind_host: resolve_bind_host(config)?,
public_url: config.public_url.clone(),
};
if let Err(error) = write_engine_stamp(&stamp) {
tracing::warn!(?error, "failed to write engine runtime stamp");
}
Ok(Self {
watcher: Some(spawn_engine_watcher(child, pid)),
startup: EngineStartup::Spawned,
})
}
}
const RUNTIME_INFO_TIMEOUT: Duration = Duration::from_secs(5);
const RUNTIME_INFO_MAX_BYTES: usize = 8192;
async fn ensure_compatible_binary(path: &Path) -> Result<()> {
let incompatible = |reason: &str| {
EngineProcessError::IncompatibleBinary {
path: path.display().to_string(),
reason: reason.to_owned(),
required: ENGINE_ARCHITECTURE.to_owned(),
}
.build()
};
let mut child = Command::new(path)
.arg("--runtime-info")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.kill_on_drop(true)
.spawn()
.with_context(|| format!("probe engine binary `{}`", path.display()))?;
let stdout = child.stdout.take().context("probe stdout was not piped")?;
let probe = async {
let read = async {
let mut bytes = Vec::new();
let limit = RUNTIME_INFO_MAX_BYTES as u64 + 1;
stdout.take(limit).read_to_end(&mut bytes).await?;
if bytes.len() > RUNTIME_INFO_MAX_BYTES {
return Err(std::io::Error::other(
"runtime identity response exceeds its size limit",
));
}
Ok(bytes)
};
tokio::try_join!(child.wait(), read)
};
let (status, bytes) = match tokio::time::timeout(RUNTIME_INFO_TIMEOUT, probe).await {
Ok(Ok(output)) => output,
failure => {
child
.kill()
.await
.context("terminate failed runtime identity probe")?;
return Err(incompatible(if failure.is_err() {
"runtime identity probe timed out"
} else {
"runtime identity probe output could not be read"
}));
}
};
if !status.success() {
return Err(incompatible(
"runtime identity probe is unsupported or failed",
));
}
let info: EngineHealthResponse = serde_json::from_slice(&bytes)
.map_err(|_| incompatible("runtime identity response is invalid"))?;
if info.runtime.as_deref() != Some(ENGINE_RUNTIME)
|| info.architecture.as_deref() != Some(ENGINE_ARCHITECTURE)
{
return Err(incompatible(
"runtime identity does not advertise the required architecture",
));
}
Ok(())
}
pub fn engine_data_path() -> Result<PathBuf> {
Ok(storage_root()?.join("var").join("engine").join("depot"))
}
pub fn engine_log_directory() -> Result<PathBuf> {
Ok(engine_data_path()?.join("log"))
}
fn admin_token() -> String {
env::var("RIVET__AUTH__ADMIN_TOKEN").unwrap_or_else(|_| "default".to_owned())
}
fn engine_stamp_path() -> Result<PathBuf> {
Ok(storage_root()?
.join("var")
.join("engine")
.join("runtime.json"))
}
fn resolve_bind_host(config: &EngineResolverConfig) -> Result<String> {
if let Some(bind_host) = &config.bind_host {
return Ok(bind_host.clone());
}
let endpoint = &config.endpoint;
let endpoint_url =
Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
Ok(endpoint_url
.host_str()
.ok_or_else(|| invalid_endpoint(endpoint, "missing host"))?
.to_owned())
}
fn write_engine_stamp(stamp: &EngineRuntimeStamp) -> Result<()> {
let path = engine_stamp_path()?;
if let Some(dir) = path.parent() {
ensure_dir(dir)?;
}
let contents = serde_json::to_vec_pretty(stamp).context("serialize engine stamp")?;
std::fs::write(&path, contents)
.with_context(|| format!("write engine stamp `{}`", path.display()))?;
Ok(())
}
fn read_engine_stamp() -> Option<EngineRuntimeStamp> {
let path = engine_stamp_path().ok()?;
let contents = std::fs::read(&path).ok()?;
let stamp: EngineRuntimeStamp = serde_json::from_slice(&contents).ok()?;
pid_is_alive(stamp.pid).then_some(stamp)
}
#[cfg(unix)]
fn pid_is_alive(pid: u32) -> bool {
let result = unsafe { libc::kill(pid as libc::pid_t, 0) };
result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
#[cfg(not(unix))]
fn pid_is_alive(_pid: u32) -> bool {
false
}
pub fn engine_env(config: &EngineResolverConfig) -> Result<Vec<(String, String)>> {
let ports = EnginePorts::resolve(config)?;
let log_directory = engine_log_directory()?;
let object_store = engine_data_path()?.join("objects");
let env = [
("RIVET__AUTH__ADMIN_TOKEN", admin_token()),
("RIVET__GUARD__HOST", ports.guard_host.clone()),
("RIVET__GUARD__PORT", ports.guard.to_string()),
("RIVET__API_PEER__HOST", LOOPBACK.to_owned()),
("RIVET__API_PEER__PORT", ports.api_peer.to_string()),
("RIVET__METRICS__HOST", ports.guard_host),
("RIVET__METRICS__PORT", ports.metrics.to_string()),
(
"RIVET__STORAGE__LOG_DIRECTORY",
log_directory.to_string_lossy().into_owned(),
),
(
"RIVET__OBJECT_STORAGE__PRIMARY__FILE_SYSTEM__PATH",
object_store.to_string_lossy().into_owned(),
),
("RIVET__STORAGE__BALANCING__ENABLED", "false".to_owned()),
("RIVET__STORAGE__AUTO_PROVISION_PARTITIONS", "1".to_owned()),
(
"RIVET__STORAGE__LOG_SHARD__MAX_DISK_BYTES",
LOG_MAX_DISK_BYTES.to_string(),
),
(
"RIVET__STORAGE__COLD_STORE__DISK_CACHE_BYTES",
"0".to_owned(),
),
("RIVET__STORAGE__WRITE_TO_COLD_AFTER_MS", "5000".to_owned()),
("RIVET__STORAGE__STOP_TIMEOUT_MS", "6000".to_owned()),
("RIVET__STORAGE__STOP_TEARDOWN_MS", "2000".to_owned()),
("RIVET__PEGBOARD__NODE_LOST_THRESHOLD", "7000".to_owned()),
("RIVET__RUNTIME__WORKER_SHUTDOWN_DURATION", "1".to_owned()),
("RIVET__RUNTIME__GUARD_SHUTDOWN_DURATION", "1".to_owned()),
("RIVET__RUNTIME__FORCE_SHUTDOWN_DURATION", "10".to_owned()),
];
Ok(env
.into_iter()
.map(|(key, value)| (key.to_owned(), value))
.collect())
}
const LOOPBACK: &str = "127.0.0.1";
const LOG_MAX_DISK_BYTES: u64 = 4 * 1024 * 1024 * 1024;
struct EnginePorts {
guard_host: String,
guard: u16,
api_peer: u16,
metrics: u16,
}
impl EnginePorts {
fn resolve(config: &EngineResolverConfig) -> Result<Self> {
let endpoint = &config.endpoint;
let endpoint_url =
Url::parse(endpoint).with_context(|| format!("parse engine endpoint `{endpoint}`"))?;
let guard_port = endpoint_url
.port_or_known_default()
.ok_or_else(|| invalid_endpoint(endpoint, "missing port"))?;
let guard = config.bind_port.unwrap_or(guard_port);
Ok(Self {
guard_host: resolve_bind_host(config)?,
guard,
api_peer: guard
.checked_add(1)
.ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?,
metrics: guard
.checked_add(10)
.ok_or_else(|| invalid_endpoint(endpoint, "port is too large"))?,
})
}
}
fn write_engine_config(config: &EngineResolverConfig) -> Result<PathBuf> {
let public_url = config.public_url.as_deref().unwrap_or(&config.endpoint);
let public_url = Url::parse(public_url)
.with_context(|| format!("parse engine public URL `{public_url}`"))?;
let ports = EnginePorts::resolve(config)?;
let peer_url = format!("http://{LOOPBACK}:{}", ports.api_peer);
let dir = storage_root()?.join("var").join("engine");
ensure_dir(&dir)?;
let path = dir.join("config.json");
let config = serde_json::json!({
"topology": {
"datacenter_label": 1,
"datacenters": {
"default": {
"datacenter_label": 1,
"is_leader": true,
"public_url": public_url.as_str(),
"peer_url": peer_url,
}
}
}
});
std::fs::write(&path, serde_json::to_string_pretty(&config)?)
.with_context(|| format!("write engine config `{}`", path.display()))?;
Ok(path)
}
pub async fn resolve_engine_binary(config: &EngineResolverConfig) -> Result<ResolvedEngine> {
if let Some(path) = config.explicit_binary_path.as_ref() {
return verify_binary_path(path);
}
if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
return verify_binary_path(&path);
}
if probe_existing_engine(&config.endpoint).await?.is_some() {
return Ok(ResolvedEngine::Existing);
}
let local_roots = local_engine_search_roots();
let cached = cached_engine_path(&config.version)?;
resolve_engine_binary_after_probe(config, false, &local_roots, cached).await
}
pub async fn resolve_engine_binary_path(config: &EngineResolverConfig) -> Result<PathBuf> {
if let Some(path) = config.explicit_binary_path.as_ref() {
verify_binary_path(path)?;
return Ok(path.clone());
}
if let Some(path) = std::env::var_os("RIVET_ENGINE_BINARY_PATH").map(PathBuf::from) {
verify_binary_path(&path)?;
return Ok(path);
}
let local_roots = local_engine_search_roots();
let cached = cached_engine_path(&config.version)?;
match resolve_engine_binary_after_probe(config, false, &local_roots, cached).await? {
ResolvedEngine::Binary(path) => Ok(path),
ResolvedEngine::Existing => {
unreachable!("no-probe resolution never returns Existing")
}
}
}
async fn resolve_engine_binary_after_probe(
config: &EngineResolverConfig,
existing_engine: bool,
local_roots: &[PathBuf],
cached: PathBuf,
) -> Result<ResolvedEngine> {
if existing_engine {
return Ok(ResolvedEngine::Existing);
}
if let Some(path) = find_local_engine_binary_in_roots(local_roots) {
return Ok(ResolvedEngine::Binary(path));
}
if cached.exists() {
return Ok(ResolvedEngine::Binary(cached));
}
if !config.auto_download {
return Err(EngineProcessError::BinaryUnavailable {
version: config.version.clone(),
}
.build());
}
download_engine_binary(config, &cached).await?;
Ok(ResolvedEngine::Binary(cached))
}
fn verify_binary_path(path: &Path) -> Result<ResolvedEngine> {
if !path.exists() {
return Err(EngineProcessError::BinaryNotFound {
path: path.display().to_string(),
}
.build());
}
Ok(ResolvedEngine::Binary(path.to_path_buf()))
}
fn local_engine_search_roots() -> Vec<PathBuf> {
Path::new(env!("CARGO_MANIFEST_DIR"))
.ancestors()
.map(Path::to_path_buf)
.collect()
}
fn find_local_engine_binary_in_roots(roots: &[PathBuf]) -> Option<PathBuf> {
for root in roots {
for profile in ["debug", "release"] {
let candidate = root
.join("target")
.join(profile)
.join(exe_name("rivet-engine"));
if candidate.exists() {
return Some(candidate);
}
}
}
None
}
fn cached_engine_path(version: &str) -> Result<PathBuf> {
Ok(storage_root()?
.join("engine")
.join(ENGINE_ARCHITECTURE)
.join(version)
.join(engine_artifact_name()))
}
async fn download_engine_binary(config: &EngineResolverConfig, destination: &Path) -> Result<()> {
let artifact = engine_artifact_name();
let base = config.releases_endpoint.trim_end_matches('/');
let artifact_url = format!("{base}/rivet/{}/engine/{artifact}", config.version);
let manifest_url = format!("{base}/rivet/{}/engine/SHA256SUMS", config.version);
let client = Client::builder()
.timeout(DOWNLOAD_TIMEOUT)
.build()
.context("build reqwest client for engine download")?;
let manifest = fetch_text(&client, &manifest_url).await?;
let expected = checksum_for_artifact(&manifest, &artifact).ok_or_else(|| {
EngineProcessError::DownloadFailed {
url: manifest_url.clone(),
reason: format!("manifest does not contain `{artifact}`"),
}
.build()
})?;
let bytes = fetch_bytes(&client, &artifact_url).await?;
let received = sha256_hex(&bytes);
if !received.eq_ignore_ascii_case(&expected) {
return Err(EngineProcessError::ChecksumMismatch {
artifact,
expected,
received,
}
.build());
}
let parent = destination
.parent()
.context("engine cache destination has no parent")?;
ensure_dir(parent)?;
std::fs::write(destination, bytes)
.with_context(|| format!("write engine binary `{}`", destination.display()))?;
make_executable(destination)?;
Ok(())
}
async fn fetch_text(client: &Client, url: &str) -> Result<String> {
let response = client.get(url).send().await.map_err(|error| {
EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: error.to_string(),
}
.build()
})?;
if !response.status().is_success() {
let status = response.status();
return Err(EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: format!("unexpected status {status}"),
}
.build());
}
response.text().await.map_err(|error| {
EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: error.to_string(),
}
.build()
})
}
async fn fetch_bytes(client: &Client, url: &str) -> Result<Vec<u8>> {
let response = client.get(url).send().await.map_err(|error| {
EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: error.to_string(),
}
.build()
})?;
if !response.status().is_success() {
let status = response.status();
return Err(EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: format!("unexpected status {status}"),
}
.build());
}
response
.bytes()
.await
.map(|bytes| bytes.to_vec())
.map_err(|error| {
EngineProcessError::DownloadFailed {
url: url.to_owned(),
reason: error.to_string(),
}
.build()
})
}
fn checksum_for_artifact(manifest: &str, artifact: &str) -> Option<String> {
manifest.lines().find_map(|line| {
let mut parts = line.split_whitespace();
let checksum = parts.next()?;
let name = parts.next()?.trim_start_matches('*');
(checksum.len() == 64 && name == artifact).then(|| checksum.to_owned())
})
}
fn sha256_hex(bytes: &[u8]) -> String {
let digest = Sha256::digest(bytes);
let mut out = String::with_capacity(digest.len() * 2);
for byte in digest {
use std::fmt::Write;
let _ = write!(&mut out, "{byte:02x}");
}
out
}
fn engine_artifact_name() -> String {
let arch = match std::env::consts::ARCH {
"x86_64" => "x86_64",
"aarch64" => "aarch64",
other => other,
};
let target = match std::env::consts::OS {
"linux" => format!("{arch}-unknown-linux-musl"),
"macos" => format!("{arch}-apple-darwin"),
"windows" => format!("{arch}-pc-windows-gnu.exe"),
other => format!("{arch}-{other}"),
};
format!("rivet-engine-{target}")
}
fn exe_name(base: &str) -> String {
if cfg!(windows) {
format!("{base}.exe")
} else {
base.to_owned()
}
}
fn make_executable(path: &Path) -> Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut permissions = std::fs::metadata(path)
.with_context(|| format!("read metadata for `{}`", path.display()))?
.permissions();
permissions.set_mode(0o755);
std::fs::set_permissions(path, permissions)
.with_context(|| format!("mark `{}` executable", path.display()))?;
}
#[cfg(not(unix))]
{
let _ = path;
}
Ok(())
}
impl Drop for EngineProcessManager {
fn drop(&mut self) {
if let Some(handle) = self.watcher.take() {
handle.abort();
tracing::debug!(
"aborted engine watcher; engine continues running (intentional orphan)"
);
}
}
}
fn spawn_engine_watcher(mut child: Child, pid: u32) -> JoinHandle<()> {
tokio::spawn(async move {
match child.wait().await {
Ok(status) if status.success() => {
tracing::warn!(
pid,
?status,
"engine process exited cleanly while rivetkit was still running; \
rivetkit expected the engine to outlive it"
);
}
Ok(status) => {
tracing::error!(
pid,
?status,
"engine process crashed while rivetkit was still running"
);
}
Err(error) => {
tracing::error!(
pid,
?error,
"failed to wait on engine process; cannot detect crashes"
);
}
}
})
}
async fn probe_existing_engine(endpoint: &str) -> Result<Option<EngineHealthResponse>> {
let health_url = engine_health_url(endpoint);
let client = Client::builder()
.build()
.context("build reqwest client for engine probe")?;
let response = match client
.get(&health_url)
.timeout(Duration::from_secs(1))
.send()
.await
{
Ok(response) => response,
Err(_) => return Ok(None),
};
if !response.status().is_success() {
return Ok(None);
}
let health = response
.json::<EngineHealthResponse>()
.await
.context("decode existing engine health response")?;
match health.runtime.as_deref() {
Some(ENGINE_RUNTIME) => {
ensure_compatible_engine(endpoint, &health)?;
Ok(Some(health))
}
Some(RIVETKIT_RUNTIME) => Err(EngineProcessError::PortOccupied {
endpoint: endpoint.to_owned(),
runtime: RIVETKIT_RUNTIME.to_owned(),
}
.build()),
Some(other) => Err(EngineProcessError::PortOccupied {
endpoint: endpoint.to_owned(),
runtime: other.to_owned(),
}
.build()),
None => Err(EngineProcessError::PortOccupied {
endpoint: endpoint.to_owned(),
runtime: "unknown".to_owned(),
}
.build()),
}
}
fn engine_health_url(endpoint: &str) -> String {
format!("{}/health", endpoint.trim_end_matches('/'))
}
fn ensure_compatible_engine(endpoint: &str, health: &EngineHealthResponse) -> Result<()> {
if health.architecture.as_deref() == Some(ENGINE_ARCHITECTURE) {
return Ok(());
}
Err(EngineProcessError::IncompatibleEngine {
endpoint: endpoint.to_owned(),
version: health
.version
.clone()
.unwrap_or_else(|| "unknown".to_owned()),
architecture: health
.architecture
.clone()
.unwrap_or_else(|| "v2".to_owned()),
required: ENGINE_ARCHITECTURE.to_owned(),
}
.build())
}
fn storage_root() -> Result<PathBuf> {
if let Ok(path) = std::env::var("RIVETKIT_STORAGE_PATH") {
return Ok(PathBuf::from(path).join(".rivetkit"));
}
let home = std::env::var("HOME")
.map(PathBuf::from)
.or_else(|_| std::env::current_dir())
.context("locate home directory for engine storage path")?;
Ok(home.join(".rivetkit"))
}
fn ensure_dir(path: &Path) -> Result<()> {
std::fs::create_dir_all(path).with_context(|| format!("create directory `{}`", path.display()))
}
fn open_log_file(path: &Path) -> Result<std::fs::File> {
std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)
.with_context(|| format!("open log file `{}`", path.display()))
}
fn log_timestamp() -> String {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
format!("{}", now.as_secs())
}
async fn wait_for_engine_health(health_url: &str) -> Result<EngineHealthResponse> {
const HEALTH_MAX_WAIT: Duration = Duration::from_secs(60);
const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
const HEALTH_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
const HEALTH_MAX_BACKOFF: Duration = Duration::from_secs(1);
let client = Client::builder()
.build()
.context("build reqwest client for engine health check")?;
let deadline = Instant::now() + HEALTH_MAX_WAIT;
let mut attempt = 0u32;
let mut backoff = HEALTH_INITIAL_BACKOFF;
loop {
attempt += 1;
let last_error = match client
.get(health_url)
.timeout(HEALTH_REQUEST_TIMEOUT)
.send()
.await
{
Ok(response) if response.status().is_success() => {
let health = response
.json::<EngineHealthResponse>()
.await
.context("decode engine health response")?;
return Ok(health);
}
Ok(response) => format!("unexpected status {}", response.status()),
Err(error) => error.to_string(),
};
if Instant::now() >= deadline {
return Err(EngineProcessError::HealthCheckFailed {
attempts: attempt,
reason: last_error,
}
.build());
}
tokio::time::sleep(backoff).await;
backoff = std::cmp::min(backoff * 2, HEALTH_MAX_BACKOFF);
}
}
async fn wait_for_engine_ready(endpoint: &str) -> Result<()> {
const READY_MAX_WAIT: Duration = Duration::from_secs(60);
const READY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
const READY_BACKOFF: Duration = Duration::from_millis(250);
let url = format!(
"{}/actors?namespace=default&name=rivetkit-readiness-probe&key=probe",
endpoint.trim_end_matches('/')
);
let token = admin_token();
let client = Client::builder()
.build()
.context("build reqwest client for engine readiness check")?;
let deadline = Instant::now() + READY_MAX_WAIT;
let mut attempt = 0u32;
loop {
attempt += 1;
let last_error = match client
.get(&url)
.bearer_auth(&token)
.timeout(READY_REQUEST_TIMEOUT)
.send()
.await
{
Ok(response) if response.status().is_success() => return Ok(()),
Ok(response) => {
let status = response.status();
let body = response.text().await.unwrap_or_default();
format!("unexpected status {status}: {body}")
}
Err(error) => error.to_string(),
};
tracing::debug!(attempt, reason = %last_error, "waiting for engine readiness");
if Instant::now() >= deadline {
return Err(EngineProcessError::NotReady {
attempts: attempt,
reason: last_error,
}
.build());
}
tokio::time::sleep(READY_BACKOFF).await;
}
}
async fn terminate_failed_spawn(child: &mut Child) -> Result<()> {
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
if child
.try_wait()
.context("check engine process status")?
.is_some()
{
return Ok(());
}
child
.start_kill()
.context("kill half-started engine process")?;
match tokio::time::timeout(SHUTDOWN_TIMEOUT, child.wait()).await {
Ok(result) => {
let status = result.context("wait for half-started engine to exit")?;
tracing::info!(?status, "half-started engine process exited");
Ok(())
}
Err(_) => {
tracing::warn!("half-started engine process did not exit within timeout");
Ok(())
}
}
}
fn invalid_endpoint(endpoint: &str, reason: &str) -> anyhow::Error {
EngineProcessError::InvalidEndpoint {
endpoint: endpoint.to_owned(),
reason: reason.to_owned(),
}
.build()
}