use std::{
collections::HashMap,
env,
os::unix::net::UnixListener,
sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
time::Duration,
};
use serde_json::json;
use tokio::{
io::AsyncWriteExt,
net::{UnixListener as TokioUnixListener, UnixStream as TokioUnixStream},
sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError},
};
use crate::{
configuration::load_bundle_group_memberships,
relay::{
BundleCatalog, append_startup_failure, serve_connection, shutdown_bundle_runtime,
startup_bundle, wait_for_async_delivery_shutdown,
},
runtime::{
bootstrap::{
RelayRuntimeLock, acquire_relay_runtime_lock, bind_relay_listener,
relay_runtime_lock_is_held,
},
error::RuntimeError,
inscriptions::{configure_process_inscriptions, emit_inscription, relay_inscriptions_path},
paths::{
BundleRuntimePaths, RelayRuntimePaths, RuntimeRootOverrides, RuntimeRoots,
ensure_bundle_runtime_directory, ensure_relay_runtime_directory,
},
signals::{install_shutdown_signal_handlers, shutdown_requested},
starter::ensure_starter_configuration_layout,
},
};
use crate::commands::{
RelayHostArguments, RelayHostStartupBundle, RelayHostStartupSummary, RuntimeArguments, shared,
};
use super::summary::{
build_startup_summary, failed_startup_bundle, hosted_startup_bundle, render_startup_summary,
skipped_startup_bundle, startup_summary_payload,
};
#[derive(Clone, Debug)]
enum RelayHostStartupMode {
Autostart,
ProcessOnly,
}
#[derive(Debug)]
struct HostedBundle {
paths: BundleRuntimePaths,
}
enum RelayHostPreparation {
NoHostedBundles,
Serve(Box<RelayHostServePlan>),
}
#[derive(Debug)]
struct RelayHostServePlan {
summary: RelayHostStartupSummary,
hosted_bundles: Vec<HostedBundle>,
relay_paths: RelayRuntimePaths,
listener: UnixListener,
runtime_lock: RelayRuntimeLock,
}
#[derive(Debug)]
struct RelayConnectionMetrics {
active_connections: AtomicUsize,
rejected_connections: AtomicUsize,
}
impl RelayConnectionMetrics {
fn new() -> Self {
Self {
active_connections: AtomicUsize::new(0),
rejected_connections: AtomicUsize::new(0),
}
}
}
const RELAY_MAX_CONNECTIONS: usize = 512;
const RELAY_PRE_HELLO_IDLE_TIMEOUT_MS: u64 = 2_000;
const RELAY_SHUTDOWN_POLL_INTERVAL_MS: u64 = 100;
pub(super) async fn run_relay_host(arguments: RelayHostArguments) -> Result<(), RuntimeError> {
let _signal_handlers = install_shutdown_signal_handlers()?;
let (roots, preparation) = tokio::task::spawn_blocking(move || {
let roots = resolve_runtime_roots(arguments.runtime)?;
let preparation = prepare_relay_host(&roots, arguments.no_autostart)?;
Ok::<_, RuntimeError>((roots, preparation))
})
.await
.map_err(|source| supervisor_join_error("relay host startup", source))??;
match preparation {
RelayHostPreparation::NoHostedBundles => Ok(()),
RelayHostPreparation::Serve(plan) => {
let RelayHostServePlan {
summary,
hosted_bundles,
relay_paths,
listener,
runtime_lock,
} = *plan;
serve_relay_host(
roots,
summary,
hosted_bundles,
relay_paths,
listener,
runtime_lock,
)
.await
}
}
}
fn resolve_runtime_roots(runtime: RuntimeArguments) -> Result<RuntimeRoots, RuntimeError> {
let current_directory = env::current_dir()
.map_err(|source| RuntimeError::io("resolve current working directory", source))?;
let overrides = RuntimeRootOverrides {
configuration_root: runtime.configuration_root,
state_root: runtime.state_root,
inscriptions_root: runtime.inscriptions_root,
repository_root: runtime.repository_root.or(Some(current_directory)),
};
let roots = RuntimeRoots::resolve(&overrides)?;
ensure_starter_configuration_layout(&roots.configuration_root)?;
Ok(roots)
}
fn prepare_relay_host(
roots: &RuntimeRoots,
no_autostart: bool,
) -> Result<RelayHostPreparation, RuntimeError> {
let memberships = load_bundle_group_memberships(&roots.configuration_root)
.map_err(shared::map_bundle_load_error)?;
if let Some(first_bundle) = memberships.first()
&& let Err(source) = configure_process_inscriptions(&relay_inscriptions_path(
&roots.inscriptions_root,
first_bundle.bundle_name.as_str(),
))
{
return Err(source);
}
let relay_paths = RelayRuntimePaths::resolve(&roots.state_root);
ensure_relay_runtime_directory(&relay_paths)?;
if relay_runtime_lock_is_held(&relay_paths)? {
let outcomes = memberships
.into_iter()
.map(|membership| {
skipped_startup_bundle(
membership.bundle_name.as_str(),
"lock_held",
"relay runtime lock is already held".to_string(),
)
})
.collect();
let summary = build_startup_summary(
if no_autostart {
"process_only"
} else {
"autostart"
},
outcomes,
);
emit_inscription("relay.startup.summary", &startup_summary_payload(&summary));
render_startup_summary(&summary);
return Ok(RelayHostPreparation::NoHostedBundles);
}
let runtime_lock = acquire_relay_runtime_lock(&relay_paths)?;
let mut outcomes = Vec::with_capacity(memberships.len());
let mut hosted_bundles = Vec::<HostedBundle>::with_capacity(memberships.len());
for membership in memberships {
let startup_mode = if no_autostart || !membership.autostart {
RelayHostStartupMode::ProcessOnly
} else {
RelayHostStartupMode::Autostart
};
let (mut outcome, hosted_bundle) =
host_selected_bundle(roots, membership.bundle_name.as_str(), startup_mode);
if no_autostart {
outcome = skipped_startup_bundle(
membership.bundle_name.as_str(),
"process_only",
"relay started without bundle autostart".to_string(),
);
}
outcomes.push(outcome);
if let Some(hosted_bundle) = hosted_bundle {
hosted_bundles.push(hosted_bundle);
}
}
let summary = build_startup_summary(
if no_autostart {
"process_only"
} else {
"autostart"
},
outcomes,
);
if hosted_bundles.is_empty() {
if no_autostart {
emit_inscription("relay.startup.summary", &startup_summary_payload(&summary));
render_startup_summary(&summary);
return Ok(RelayHostPreparation::NoHostedBundles);
}
if summary.failed_bundle_count > 0 {
return Err(RuntimeError::validation(
"runtime_startup_failed",
format!(
"failed to start relay for {} bundle(s)",
summary.failed_bundle_count
),
));
}
return Ok(RelayHostPreparation::NoHostedBundles);
}
let listener = bind_relay_listener(&relay_paths)?;
Ok(RelayHostPreparation::Serve(Box::new(RelayHostServePlan {
summary,
hosted_bundles,
relay_paths,
listener,
runtime_lock,
})))
}
async fn serve_relay_host(
roots: RuntimeRoots,
summary: RelayHostStartupSummary,
hosted_bundles: Vec<HostedBundle>,
relay_paths: RelayRuntimePaths,
listener: UnixListener,
runtime_lock: RelayRuntimeLock,
) -> Result<(), RuntimeError> {
emit_inscription("relay.startup.summary", &startup_summary_payload(&summary));
render_startup_summary(&summary);
let stop_requested = Arc::new(AtomicBool::new(false));
let max_connections = relay_max_connections();
let connection_permits = Arc::new(Semaphore::new(max_connections));
let bundle_paths_for_cleanup: Vec<BundleRuntimePaths> = hosted_bundles
.iter()
.map(|hosted| hosted.paths.clone())
.collect();
let catalog: BundleCatalog = Arc::new(
hosted_bundles
.into_iter()
.map(|hosted| (hosted.paths.bundle_name.clone(), hosted.paths))
.collect::<HashMap<_, _>>(),
);
remove_relay_ready_sentinel(&relay_paths);
let accept_handle = {
let configuration_root = roots.configuration_root.clone();
let stop_requested = Arc::clone(&stop_requested);
let connection_permits = Arc::clone(&connection_permits);
let catalog = Arc::clone(&catalog);
tokio::spawn(run_relay_accept_loop(
configuration_root,
listener,
catalog,
stop_requested,
connection_permits,
max_connections,
))
};
write_relay_ready_sentinel(&relay_paths)?;
let accept_outcome = supervise_accept_loop(&stop_requested, accept_handle).await;
if shutdown_requested() {
emit_inscription("relay.shutdown.signal", &json!({"signal": "termination"}));
}
stop_requested.store(true, Ordering::SeqCst);
let cleanup_relay_paths = relay_paths.clone();
tokio::task::spawn_blocking(move || {
cleanup_relay_host(cleanup_relay_paths, bundle_paths_for_cleanup)
})
.await
.map_err(|source| supervisor_join_error("relay host cleanup", source))??;
drop(runtime_lock);
accept_outcome
}
async fn supervise_accept_loop(
stop_requested: &Arc<AtomicBool>,
accept_handle: tokio::task::JoinHandle<Result<(), RuntimeError>>,
) -> Result<(), RuntimeError> {
let mut accept_handle = accept_handle;
loop {
if shutdown_requested() || stop_requested.load(Ordering::SeqCst) {
return accept_loop_join_outcome(accept_handle.await);
}
tokio::select! {
biased;
joined = &mut accept_handle => {
return accept_loop_join_outcome(joined);
}
() = tokio::time::sleep(Duration::from_millis(RELAY_SHUTDOWN_POLL_INTERVAL_MS)) => {}
}
}
}
fn accept_loop_join_outcome(
result: Result<Result<(), RuntimeError>, tokio::task::JoinError>,
) -> Result<(), RuntimeError> {
match result {
Ok(Ok(())) => Ok(()),
Ok(Err(error)) => Err(error),
Err(join_error) => Err(RuntimeError::validation(
"internal_unexpected_failure",
format!("relay accept loop task panicked: {join_error}"),
)),
}
}
fn supervisor_join_error(context: &str, source: tokio::task::JoinError) -> RuntimeError {
RuntimeError::validation(
"internal_unexpected_failure",
format!("{context} task failed to join: {source}"),
)
}
fn cleanup_relay_host(
relay_paths: RelayRuntimePaths,
bundle_paths: Vec<BundleRuntimePaths>,
) -> Result<(), RuntimeError> {
let async_workers_remaining = if shutdown_requested() {
wait_for_async_delivery_shutdown(Duration::from_millis(1_500))
} else {
0
};
remove_relay_ready_sentinel(&relay_paths);
shared::remove_relay_socket_file(&relay_paths.relay_socket)?;
for paths in bundle_paths {
let shutdown =
shutdown_bundle_runtime(&paths.tmux_socket).map_err(shared::map_reconcile_error)?;
emit_inscription(
"relay.shutdown.complete",
&json!({
"bundle_name": paths.bundle_name,
"pruned_count": shutdown.pruned_sessions.len(),
"killed_tmux_server": shutdown.killed_tmux_server,
"pruned_sessions": shutdown.pruned_sessions,
"async_workers_remaining": async_workers_remaining,
}),
);
}
Ok(())
}
fn write_relay_ready_sentinel(paths: &RelayRuntimePaths) -> Result<(), RuntimeError> {
let sentinel = &paths.relay_ready_sentinel;
let temporary = sentinel.with_extension("ready.tmp");
std::fs::write(&temporary, b"").map_err(|source| {
RuntimeError::io(
format!("write relay ready sentinel {}", temporary.display()),
source,
)
})?;
std::fs::rename(&temporary, sentinel).map_err(|source| {
RuntimeError::io(
format!("publish relay ready sentinel {}", sentinel.display()),
source,
)
})
}
fn remove_relay_ready_sentinel(paths: &RelayRuntimePaths) {
let _ = std::fs::remove_file(&paths.relay_ready_sentinel);
}
async fn run_relay_accept_loop(
configuration_root: std::path::PathBuf,
listener: UnixListener,
bundle_catalog: BundleCatalog,
stop_requested: Arc<AtomicBool>,
connection_permits: Arc<Semaphore>,
max_connections: usize,
) -> Result<(), RuntimeError> {
listener
.set_nonblocking(true)
.map_err(|source| RuntimeError::io("set relay socket non-blocking".to_string(), source))?;
let listener = TokioUnixListener::from_std(listener).map_err(|source| {
RuntimeError::io(
"register relay socket with async runtime".to_string(),
source,
)
})?;
let metrics = Arc::new(RelayConnectionMetrics::new());
let accept_outcome = loop {
if shutdown_requested() || stop_requested.load(Ordering::SeqCst) {
break Ok(());
}
tokio::select! {
biased;
accepted = listener.accept() => {
match accepted {
Ok((stream, _address)) => {
if shutdown_requested() || stop_requested.load(Ordering::SeqCst) {
break Ok(());
}
match Arc::clone(&connection_permits).try_acquire_owned() {
Ok(permit) => {
spawn_connection_worker(
permit,
stream,
configuration_root.clone(),
Arc::clone(&bundle_catalog),
Arc::clone(&metrics),
);
}
Err(TryAcquireError::NoPermits) => {
reject_overloaded_connection(stream, &metrics).await;
}
Err(TryAcquireError::Closed) => break Ok(()),
}
emit_connection_metrics(max_connections, &metrics);
}
Err(source) => {
break Err(RuntimeError::io(
"accept relay socket connection".to_string(),
source,
));
}
}
}
() = tokio::time::sleep(Duration::from_millis(RELAY_SHUTDOWN_POLL_INTERVAL_MS)) => {}
}
};
if accept_outcome.is_ok() && (shutdown_requested() || stop_requested.load(Ordering::SeqCst)) {
emit_inscription(
"relay.shutdown.connection_workers_detached",
&json!({
"max_connections": max_connections,
"active_connections": metrics.active_connections.load(Ordering::SeqCst),
"rejected_connections": metrics.rejected_connections.load(Ordering::SeqCst),
}),
);
}
accept_outcome
}
fn spawn_connection_worker(
permit: OwnedSemaphorePermit,
stream: TokioUnixStream,
configuration_root: std::path::PathBuf,
bundle_catalog: BundleCatalog,
metrics: Arc<RelayConnectionMetrics>,
) {
metrics.active_connections.fetch_add(1, Ordering::SeqCst);
let pre_hello_idle_timeout = relay_pre_hello_idle_timeout();
tokio::spawn(async move {
let result = serve_connection(
stream,
&configuration_root,
&bundle_catalog,
pre_hello_idle_timeout,
)
.await;
metrics.active_connections.fetch_sub(1, Ordering::SeqCst);
drop(permit);
if let Err(source) = result {
emit_inscription(
"relay.request_failed",
&json!({
"error": source.to_string(),
}),
);
}
});
}
fn relay_max_connections() -> usize {
parse_env_positive_usize("AGENTMUX_RELAY_MAX_CONNECTIONS").unwrap_or(RELAY_MAX_CONNECTIONS)
}
fn relay_pre_hello_idle_timeout() -> Duration {
let timeout_ms = parse_env_positive_usize("AGENTMUX_RELAY_PRE_HELLO_IDLE_TIMEOUT_MS")
.and_then(|value| u64::try_from(value).ok())
.unwrap_or(RELAY_PRE_HELLO_IDLE_TIMEOUT_MS);
Duration::from_millis(timeout_ms)
}
fn parse_env_positive_usize(name: &str) -> Option<usize> {
env::var(name)
.ok()
.and_then(|value| value.parse::<usize>().ok())
.filter(|value| *value > 0)
}
async fn reject_overloaded_connection(
mut stream: TokioUnixStream,
metrics: &RelayConnectionMetrics,
) {
metrics.rejected_connections.fetch_add(1, Ordering::SeqCst);
emit_inscription(
"relay.connection.rejected",
&json!({
"reason_code": "runtime_connection_limit_reached",
"reason": "relay connection limit reached",
}),
);
let response = crate::relay::RelayResponse::Error {
error: crate::relay::RelayError {
code: "runtime_connection_limit_reached".to_string(),
message: "relay connection limit reached".to_string(),
details: None,
},
};
let frame = json!({
"frame": "response",
"response": response,
});
if let Ok(mut encoded) = serde_json::to_vec(&frame) {
encoded.push(b'\n');
let _ = stream.write_all(&encoded).await;
let _ = stream.flush().await;
}
let _ = stream.shutdown().await;
}
fn emit_connection_metrics(max_connections: usize, metrics: &RelayConnectionMetrics) {
emit_inscription(
"relay.connection_pool.metrics",
&json!({
"max_connections": max_connections,
"active_connections": metrics.active_connections.load(Ordering::SeqCst),
"rejected_connections": metrics.rejected_connections.load(Ordering::SeqCst),
}),
);
}
fn host_selected_bundle(
roots: &RuntimeRoots,
bundle_name: &str,
startup_mode: RelayHostStartupMode,
) -> (RelayHostStartupBundle, Option<HostedBundle>) {
let paths = match BundleRuntimePaths::resolve(&roots.state_root, bundle_name) {
Ok(paths) => paths,
Err(source) => return (failed_startup_bundle(bundle_name, source), None),
};
emit_inscription(
"relay.startup",
&json!({
"bundle_name": paths.bundle_name,
"tmux_socket": paths.tmux_socket,
"configuration_root": roots.configuration_root,
"state_root": roots.state_root,
"inscriptions_root": roots.inscriptions_root,
}),
);
if let Err(source) = ensure_bundle_runtime_directory(&paths) {
return (failed_startup_bundle(bundle_name, source), None);
}
let mut startup_report = None;
if let RelayHostStartupMode::Autostart = startup_mode {
let report = match startup_bundle(
&roots.configuration_root,
&paths.bundle_name,
&paths.runtime_directory,
)
.map_err(shared::map_reconcile_error)
{
Ok(report) => report,
Err(source) => return (failed_startup_bundle(bundle_name, source), None),
};
for failure in &report.failed_startups {
let persisted = match append_startup_failure(&paths.runtime_directory, failure.clone())
{
Ok(value) => value,
Err(cause) => {
return (
RelayHostStartupBundle {
bundle_name: bundle_name.to_string(),
outcome: "failed".to_string(),
reason_code: Some("runtime_startup_failed".to_string()),
reason: Some(format!(
"failed to persist startup failure history: {cause}"
)),
},
None,
);
}
};
emit_inscription(
"relay.session_start_failed",
&json!({
"bundle_name": persisted.bundle_name,
"session_id": persisted.session_id,
"transport": persisted.transport,
"code": persisted.code,
"reason": persisted.reason,
"timestamp": persisted.timestamp,
"sequence": persisted.sequence,
"details": persisted.details,
}),
);
}
startup_report = Some(report);
}
let startup_bundle = match startup_mode {
RelayHostStartupMode::Autostart => {
let report = startup_report.expect("autostart report must be present");
if report.ready_session_count > 0 {
hosted_startup_bundle(bundle_name)
} else {
RelayHostStartupBundle {
bundle_name: bundle_name.to_string(),
outcome: "failed".to_string(),
reason_code: Some("runtime_startup_failed".to_string()),
reason: Some("zero configured sessions reached ready state".to_string()),
}
}
}
RelayHostStartupMode::ProcessOnly => skipped_startup_bundle(
bundle_name,
"process_only",
"relay started without bundle autostart".to_string(),
),
};
(startup_bundle, Some(HostedBundle { paths }))
}