use std::collections::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::Instant;
use tokio::time::sleep;
use tokio::time::timeout_at;
use super::ConnectionStatus;
use super::ExecServerClient;
use super::ExecServerError;
use super::Inner;
use super::OrderedSessionEvents;
use super::RecoveryPolicy;
use super::SessionState;
use super::disconnected_message;
use super::fail_all_in_flight_work;
use super::handle_server_notification;
use super::is_transport_closed_error;
use crate::client_transport::ExecServerReconnectStrategy;
use crate::process::ExecProcessEvent;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::rpc::RpcClient;
use crate::rpc::RpcClientEvent;
use crate::rpc::SESSION_ALREADY_ATTACHED_ERROR_CODE;
#[cfg(test)]
const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_millis(500);
#[cfg(not(test))]
const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_secs(25);
const SESSION_RECOVERY_RETRY_INTERVAL: Duration = Duration::from_millis(100);
const REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL: Duration = Duration::from_millis(500);
const REGISTRY_RECOVERY_MAX_RETRY_INTERVAL: Duration = Duration::from_secs(5);
impl SessionState {
fn last_published_seq(&self) -> u64 {
self.ordered_events
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.last_published_seq
}
fn recover_events(&self, response: ReadResponse) -> Result<bool, ExecServerError> {
let ReadResponse {
chunks,
next_seq,
exited,
exit_code,
closed,
failure,
sandbox_denied,
} = response;
if let Some(message) = failure {
return Err(ExecServerError::Protocol(format!(
"process failed while recovering: {message}"
)));
}
let target_seq = next_seq.saturating_sub(1);
let published_closed = {
let mut ordered_events = self
.ordered_events
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if ordered_events.failure.is_some()
|| ordered_events.closed_published
|| target_seq <= ordered_events.last_published_seq
{
return Ok(false);
}
let pending_exit = ordered_events.pending.range_mut(..=target_seq).find_map(
|(_, event)| match event {
ExecProcessEvent::Exited {
sandbox_denied: pending_sandbox_denied,
..
} => Some(pending_sandbox_denied),
_ => None,
},
);
let exit_pending = pending_exit.is_some();
if let Some(pending_sandbox_denied) = pending_exit {
*pending_sandbox_denied =
Some(pending_sandbox_denied.unwrap_or(false) || sandbox_denied);
}
let mut exit_known = ordered_events.exit_published || exit_pending;
if closed
&& (matches!(
ordered_events.pending.get(&target_seq),
Some(event) if !matches!(event, ExecProcessEvent::Closed { .. })
) || chunks.iter().any(|chunk| chunk.seq == target_seq))
{
return Err(ExecServerError::Protocol(format!(
"process close sequence {target_seq} conflicts with recovered output"
)));
}
let mut published_closed = false;
for chunk in chunks {
if chunk.seq > target_seq {
return Err(ExecServerError::Protocol(format!(
"recovered process output sequence {} exceeds target sequence {target_seq}",
chunk.seq
)));
}
let next_seq = ordered_events.last_published_seq.saturating_add(1);
if exited && !exit_known && chunk.seq > next_seq {
let exit_code = exit_code.ok_or_else(|| {
ExecServerError::Protocol(
"recovering exited process did not include its exit code".to_string(),
)
})?;
ordered_events
.insert_pending(ExecProcessEvent::Exited {
seq: next_seq,
exit_code,
sandbox_denied: Some(sandbox_denied),
})
.map_err(ExecServerError::Protocol)?;
published_closed |= self.publish_ready(&mut ordered_events);
exit_known = true;
}
if chunk.seq > ordered_events.last_published_seq {
ordered_events
.insert_pending(ExecProcessEvent::Output(chunk))
.map_err(ExecServerError::Protocol)?;
published_closed |= self.publish_ready(&mut ordered_events);
}
}
if closed
&& !ordered_events.closed_published
&& !matches!(
ordered_events.pending.get(&target_seq),
Some(ExecProcessEvent::Closed { .. })
)
{
ordered_events
.insert_pending(ExecProcessEvent::Closed { seq: target_seq })
.map_err(ExecServerError::Protocol)?;
}
let event_count = target_seq.saturating_sub(ordered_events.last_published_seq);
let first_unpublished_seq = ordered_events.last_published_seq.saturating_add(1);
let retained_count = if first_unpublished_seq <= target_seq {
ordered_events
.pending
.range(first_unpublished_seq..=target_seq)
.count() as u64
} else {
0
};
let missing_count = event_count.saturating_sub(retained_count);
if exited && !exit_known {
if missing_count != 1 {
return Err(recovery_gap_error(target_seq));
}
let seq = first_missing_seq(&ordered_events, target_seq);
let exit_code = exit_code.ok_or_else(|| {
ExecServerError::Protocol(
"recovering exited process did not include its exit code".to_string(),
)
})?;
ordered_events
.insert_pending(ExecProcessEvent::Exited {
seq,
exit_code,
sandbox_denied: Some(sandbox_denied),
})
.map_err(ExecServerError::Protocol)?;
} else if missing_count != 0 {
return Err(recovery_gap_error(target_seq));
}
published_closed |= self.publish_ready(&mut ordered_events);
published_closed
};
self.note_change(target_seq);
Ok(published_closed)
}
}
fn first_missing_seq(events: &OrderedSessionEvents, target_seq: u64) -> u64 {
let mut expected = events.last_published_seq.saturating_add(1);
for seq in events
.pending
.range(expected..=target_seq)
.map(|(seq, _)| *seq)
{
if seq != expected {
break;
}
expected = expected.saturating_add(1);
}
expected
}
fn recovery_gap_error(target_seq: u64) -> ExecServerError {
ExecServerError::Protocol(format!(
"process events are no longer retained while recovering through sequence {target_seq}"
))
}
impl Inner {
pub(super) async fn rpc_client(self: &Arc<Self>) -> Result<Arc<RpcClient>, ExecServerError> {
let mut connection_changed = self.connection_changed.subscribe();
loop {
if let Some(message) = self.failure_message() {
return Err(ExecServerError::Disconnected(message));
}
let rpc_client = {
let connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match &connection.status {
ConnectionStatus::Connected(rpc_client) => Some(Arc::clone(rpc_client)),
ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None,
}
};
let Some(rpc_client) = rpc_client else {
let _ = connection_changed.changed().await;
continue;
};
if !rpc_client.is_disconnected() {
return Ok(rpc_client);
}
let _ = connection_changed.changed().await;
}
}
pub(super) fn begin_process_start(&self, expected: &Arc<RpcClient>) -> bool {
let mut connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let ConnectionStatus::Connected(current) = &connection.status else {
return false;
};
if !Arc::ptr_eq(current, expected) || expected.is_disconnected() {
return false;
}
connection.active_process_starts += 1;
true
}
pub(super) fn finish_process_start(&self) {
{
let mut connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if connection.active_process_starts == 0 {
tracing::error!("finished an exec-server process start that was not active");
return;
}
connection.active_process_starts -= 1;
}
self.notify_connection_changed();
}
pub(super) fn is_failed(&self) -> bool {
self.failure_message().is_some()
}
pub(super) fn failure_message(&self) -> Option<String> {
let connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match &connection.status {
ConnectionStatus::Failed(message) => Some(message.clone()),
ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None,
}
}
fn request_recovery(
self: &Arc<Self>,
failed_rpc_client: Arc<RpcClient>,
disconnect_message: String,
) {
let should_recover = {
let mut connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match &connection.status {
ConnectionStatus::Connected(current)
if Arc::ptr_eq(current, &failed_rpc_client) =>
{
connection.set_status(ConnectionStatus::Recovering);
true
}
ConnectionStatus::Connected(_)
| ConnectionStatus::Recovering
| ConnectionStatus::Failed(_) => false,
}
};
if !should_recover {
return;
}
self.notify_connection_changed();
let inner = Arc::clone(self);
tokio::spawn(async move {
inner.recover(disconnect_message).await;
});
}
async fn recover(self: &Arc<Self>, disconnect_message: String) {
let deadline = Instant::now() + SESSION_RECOVERY_TIMEOUT;
self.fail_all_http_body_streams(disconnect_message.clone())
.await;
if timeout_at(deadline, self.wait_for_process_starts())
.await
.is_err()
{
let message = format!(
"{disconnect_message}; failed to resume exec-server session: recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}"
);
self.fail(message).await;
return;
}
if self.reconnect_strategy.is_none() {
self.fail(disconnect_message).await;
return;
}
let Some(session_id) = self.session_id.get().cloned() else {
let message = format!(
"{disconnect_message}; failed to resume exec-server session: missing session id"
);
self.fail(message).await;
return;
};
let uses_registry_backoff = matches!(
self.reconnect_strategy.as_ref(),
Some(ExecServerReconnectStrategy::NoiseRendezvous { .. })
);
let mut registry_retry_attempt = 0;
let last_error = loop {
match timeout_at(deadline, self.resume_once(&session_id)).await {
Ok(Ok(candidate)) => {
if !candidate.is_disconnected() && self.install_recovered_client(candidate) {
return;
}
}
Ok(Err(error)) if !is_retryable_recovery_error(&error) => {
break error.to_string();
}
Ok(Err(_)) => {}
Err(_) => {
break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
}
}
let retry_delay = if uses_registry_backoff {
let delay = registry_recovery_retry_delay(&session_id, registry_retry_attempt);
registry_retry_attempt = registry_retry_attempt.saturating_add(1);
delay
} else {
SESSION_RECOVERY_RETRY_INTERVAL
};
let now = Instant::now();
if now >= deadline {
break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
}
sleep(retry_delay.min(deadline - now)).await;
};
let message =
format!("{disconnect_message}; failed to resume exec-server session: {last_error}");
self.fail(message).await;
}
async fn wait_for_process_starts(&self) {
let mut connection_changed = self.connection_changed.subscribe();
loop {
let starts_are_done = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_process_starts
== 0;
if starts_are_done {
return;
}
let _ = connection_changed.changed().await;
}
}
fn install_recovered_client(&self, rpc_client: Arc<RpcClient>) -> bool {
let installed = {
let mut connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if !matches!(connection.status, ConnectionStatus::Recovering)
|| rpc_client.is_disconnected()
{
false
} else {
connection.set_status(ConnectionStatus::Connected(rpc_client));
true
}
};
if installed {
self.notify_connection_changed();
}
installed
}
fn notify_connection_changed(&self) {
self.connection_changed.send_replace(());
}
async fn resume_once(
self: &Arc<Self>,
session_id: &str,
) -> Result<Arc<RpcClient>, ExecServerError> {
let reconnect_strategy = self
.reconnect_strategy
.as_ref()
.ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?;
let (connection, options) = reconnect_strategy.resume(session_id).await?;
let (rpc_client, events_rx) = RpcClient::new(connection);
let rpc_client = Arc::new(rpc_client);
let client = ExecServerClient {
inner: Arc::clone(self),
recovery_policy: RecoveryPolicy::Wait,
};
client.spawn_rpc_reader(&rpc_client, events_rx);
client.initialize_rpc(&rpc_client, options).await?;
self.recover_processes(&rpc_client).await?;
Ok(rpc_client)
}
async fn recover_processes(
self: &Arc<Self>,
rpc_client: &RpcClient,
) -> Result<(), ExecServerError> {
let sessions = self.sessions.load_full();
for (process_id, session) in sessions.iter() {
if !session.recoverable.load(Ordering::Acquire) {
continue;
}
let response = rpc_client
.call::<_, ReadResponse>(
EXEC_READ_METHOD,
&ReadParams {
process_id: process_id.clone(),
after_seq: Some(session.last_published_seq()),
max_bytes: None,
wait_ms: Some(0),
},
)
.await
.map_err(ExecServerError::from);
let recovered = match response {
Ok(response) => session.recover_events(response),
Err(error) if is_transport_closed_error(&error) => return Err(error),
Err(error) => Err(error),
};
match recovered {
Ok(true) => self.remove_session_if(process_id, session),
Ok(false) => {}
Err(error) => {
let terminated: Result<TerminateResponse, ExecServerError> = rpc_client
.call_for_cleanup(
EXEC_TERMINATE_METHOD,
&TerminateParams {
process_id: process_id.clone(),
},
)
.await
.map_err(ExecServerError::from);
if let Err(terminate_error) = terminated
&& is_transport_closed_error(&terminate_error)
{
return Err(terminate_error);
}
self.remove_session_if(process_id, session);
session.set_failure(format!("failed to recover process {process_id}: {error}"));
}
}
}
Ok(())
}
async fn fail(self: &Arc<Self>, message: String) {
let (message, newly_failed) = {
let mut connection = self
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match &connection.status {
ConnectionStatus::Failed(existing) => (existing.clone(), false),
ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => {
connection.set_status(ConnectionStatus::Failed(message.clone()));
(message, true)
}
}
};
if newly_failed {
self.notify_connection_changed();
fail_all_in_flight_work(self, message.clone()).await;
}
}
}
impl ExecServerClient {
pub(super) fn spawn_rpc_reader(
&self,
rpc_client: &Arc<RpcClient>,
mut events_rx: mpsc::Receiver<RpcClientEvent>,
) {
let inner = Arc::downgrade(&self.inner);
let rpc_client = Arc::downgrade(rpc_client);
tokio::spawn(async move {
while let Some(event) = events_rx.recv().await {
let (Some(inner), Some(rpc_client)) = (inner.upgrade(), rpc_client.upgrade())
else {
return;
};
match event {
RpcClientEvent::Notification(notification) => {
if let Err(error) = handle_server_notification(&inner, notification).await {
rpc_client.close_transport().await;
inner.request_recovery(
rpc_client,
format!("exec-server notification handling failed: {error}"),
);
return;
}
}
RpcClientEvent::Disconnected { reason } => {
inner.request_recovery(rpc_client, disconnected_message(reason.as_deref()));
return;
}
}
}
});
}
}
fn is_retryable_recovery_error(error: &ExecServerError) -> bool {
is_transport_closed_error(error)
|| matches!(
error,
ExecServerError::WebSocketConnectTimeout { .. }
| ExecServerError::WebSocketConnect { .. }
| ExecServerError::InitializeTimedOut { .. }
)
|| is_retryable_registry_error(error)
|| matches!(
error,
ExecServerError::Server { code, .. }
if *code == SESSION_ALREADY_ATTACHED_ERROR_CODE
)
}
fn is_retryable_registry_error(error: &ExecServerError) -> bool {
matches!(
error,
ExecServerError::EnvironmentRegistryRequest(error)
if error.is_connect() || error.is_timeout()
) || matches!(
error,
ExecServerError::EnvironmentRegistryHttp { status, code, .. }
if status.is_server_error()
|| *status == reqwest::StatusCode::REQUEST_TIMEOUT
|| *status == reqwest::StatusCode::TOO_MANY_REQUESTS
|| (*status == reqwest::StatusCode::CONFLICT
&& code.as_deref() == Some("environment_offline"))
)
}
fn registry_recovery_retry_delay(session_id: &str, attempt: u32) -> Duration {
let multiplier = 1_u32.checked_shl(attempt.min(4)).unwrap_or(u32::MAX);
let base_delay = REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL
.saturating_mul(multiplier)
.min(REGISTRY_RECOVERY_MAX_RETRY_INTERVAL);
let base_millis = base_delay.as_millis() as u64;
let mut hasher = DefaultHasher::new();
session_id.hash(&mut hasher);
attempt.hash(&mut hasher);
Duration::from_millis(base_millis + hasher.finish() % (base_millis / 2 + 1))
}
#[cfg(test)]
#[path = "client_recovery_tests.rs"]
mod tests;