sail/

worker.rs

//! Per-sailbox worker-proxy client: the gRPC operations that terminate at a
//! sailbox's own worker proxy (exec wait/cancel, listeners, files), sharing one
//! lazily-dialed, drain-aware channel cache and the API-key credential.
//!
//! Transient-failure retry and drain-aware channel eviction live here; the
//! interrupt UX (Ctrl-C cascades, AbortSignal) stays in the language wrappers.

use std::sync::Arc;
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use tokio::sync::Mutex as AsyncMutex;
use tokio_stream::wrappers::ReceiverStream;
use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue};
use tonic::transport::Channel;
use tonic::{Code, Request, Status};

use crate::channels::ChannelCache;
use crate::error::SailError;
use crate::pb::workerproxy::v1 as pb;
use pb::worker_proxy_service_client::WorkerProxyServiceClient;

/// How many file chunks to buffer between the caller and the wire, for
/// backpressure on both read and write.
const FILE_READ_CHANNEL_CAP: usize = 4;

/// Chunk size for streaming a file write: each `FileWriter::write_chunk` call
/// becomes one gRPC message, so this keeps a chunk under the transport's message
/// limit. Bindings stream their source in pieces of this size.
pub const FILE_WRITE_CHUNK_BYTES: usize = 1 << 20;

/// Initial backoff in seconds before the first transient-RPC retry; doubled on
/// each subsequent attempt.
pub const EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS: f64 = 0.2;
/// Ceiling in seconds for the exponential backoff between transient-RPC retries.
pub const EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS: f64 = 2.0;
/// Per-attempt deadline for a unary exec RPC. Caps a single attempt so a
/// stalled (not dead) connection times out and the retry loop runs again
/// within the overall budget, instead of one await consuming it all.
pub const EXEC_RPC_ATTEMPT_TIMEOUT_SECONDS: f64 = 10.0;

/// gRPC status-message fragments that mark a mid-flight connection drop (a peer
/// rolling/lameducking during a deploy) rather than a permanent failure, so the
/// RPC is worth retrying. Shared by the exec and imagebuilder retry paths.
pub(crate) const TRANSIENT_TRANSPORT_FRAGMENTS: &[&str] = &[
    "endpoint closing",
    "error reading server preface",
    "connection reset",
    "socket closed",
    "transport is closing",
];

/// Whether a gRPC status message names a transient transport drop (see
/// [`TRANSIENT_TRANSPORT_FRAGMENTS`]). Case-insensitive.
pub(crate) fn is_transient_transport_message(message: &str) -> bool {
    let details = message.to_lowercase();
    TRANSIENT_TRANSPORT_FRAGMENTS
        .iter()
        .any(|fragment| details.contains(fragment))
}

/// Authoritative buffered result from polling `WaitSailboxExec`. Callers
/// inspect `status` (e.g. a terminal failure) and shape the public result.
#[derive(Debug, Clone)]
pub struct WaitOutcome {
    /// Terminal exec status as the `SailboxExecStatus` proto enum value.
    pub status: i32,
    /// Buffered stdout from the persisted row.
    pub stdout: String,
    /// Buffered stderr from the persisted row.
    pub stderr: String,
    /// The command's exit code.
    pub return_code: i32,
    /// Whether the command was killed for exceeding its timeout.
    pub timed_out: bool,
    /// Whether stdout overflowed the server output ring and lost its oldest
    /// bytes.
    pub stdout_truncated: bool,
    /// Whether stderr overflowed the server output ring and lost its oldest
    /// bytes.
    pub stderr_truncated: bool,
}

/// A sailbox ingress listener. Parsed from the sailbox-API listener JSON, whose
/// optional address fields are absent until the route is active.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(default)]
pub struct Listener {
    /// The in-guest port traffic is forwarded to.
    pub guest_port: u32,
    /// Publicly reachable URL for the listener.
    pub public_url: String,
    /// Wire protocol exposed, e.g. `tcp` or `http`.
    pub protocol: String,
    /// Route status as the proto enum name, e.g. `LISTENER_ROUTE_STATUS_ACTIVE`.
    pub route_status: String,
    /// Public hostname the listener is reachable at.
    pub public_host: String,
    /// Public port the listener is reachable at.
    pub public_port: u32,
}

/// Optional settings for [`Client::upload_file`](crate::Client::upload_file).
/// `Default` writes the file without creating missing parent directories and
/// leaves its mode to the guest's default.
#[derive(Debug, Clone, Default)]
pub struct UploadOptions {
    /// Create missing parent directories before writing.
    pub create_parents: bool,
    /// Unix mode bits for the written file; `None` leaves the guest default.
    pub mode: Option<u32>,
}

/// A streaming reader over a `ReadSailboxFile` response. A background task
/// pumps chunks into a bounded channel (so a slow consumer applies
/// backpressure rather than buffering the whole file); [`FileReader::next`]
/// yields the next chunk, `None` at end of file.
pub struct FileReader {
    rx: AsyncMutex<mpsc::Receiver<Result<Vec<u8>, SailError>>>,
}

impl FileReader {
    /// Yield the next file chunk, or `None` at end of file. A stream error
    /// surfaces as `Some(Err(..))`.
    pub async fn next(&self) -> Option<Result<Vec<u8>, SailError>> {
        self.rx.lock().await.recv().await
    }
}

/// Client for the worker-proxy gRPC operations of a single sailbox, holding the
/// shared lazily-dialed channel cache and the bearer credential applied to every
/// request.
#[doc(hidden)]
pub struct WorkerProxy {
    channels: ChannelCache,
    authorization: AsciiMetadataValue,
}

/// A `retry_timeout <= 0` deadline is "now" so the very first transient
/// failure short-circuits, preserving the no-retry-by-default contract
/// for programmatic cancel callers. A non-finite or overflowing timeout
/// (e.g. `float("inf")` from Python) means "retry forever".
pub(crate) fn retry_deadline(retry_timeout: f64) -> Instant {
    let now = Instant::now();
    if retry_timeout <= 0.0 {
        return now;
    }
    Duration::try_from_secs_f64(retry_timeout)
        .ok()
        .and_then(|timeout| now.checked_add(timeout))
        .unwrap_or(now + RETRY_FOREVER_FALLBACK)
}

/// A ~100-year stand-in for "retry forever" when the requested deadline is so
/// large it overflows `Instant`.
const RETRY_FOREVER_FALLBACK: Duration = Duration::from_hours(876_000);

/// Bound a single unary-RPC attempt: at most [`EXEC_RPC_ATTEMPT_TIMEOUT_SECONDS`]
/// and at most half the remaining budget. Halving leaves headroom for at least
/// one retry. Otherwise, on a budget smaller than the ceiling, a single attempt
/// hanging on a half-open connection would consume the whole budget and the loop
/// would give up without ever redialing.
pub(crate) fn rpc_attempt_timeout(deadline: Instant) -> Duration {
    (deadline.saturating_duration_since(Instant::now()) / 2)
        .min(Duration::from_secs_f64(EXEC_RPC_ATTEMPT_TIMEOUT_SECONDS))
        .max(Duration::from_millis(1))
}

pub(crate) fn is_workerproxy_draining(status: &Status) -> bool {
    status.code() == Code::Unavailable && status.message().to_lowercase().contains("draining")
}

/// Whether a retryable failure should also drop the cached channel before the
/// next attempt. The connection itself is suspect when the target is draining,
/// on a server-enforced deadline (`DeadlineExceeded`), or on any client-side
/// transport failure (a half-open socket, keepalive timeout, failed connect, or
/// a fired per-attempt `set_timeout`, all of which tonic surfaces as a status
/// carrying a transport `source`). Reusing such a channel would burn every retry
/// on the dead connection, so dial fresh instead. A server-sent transient (no
/// source, not draining, not a deadline) keeps the connection.
pub(crate) fn should_invalidate_channel(status: &Status) -> bool {
    use std::error::Error;
    is_workerproxy_draining(status)
        || status.code() == Code::DeadlineExceeded
        || status.source().is_some()
}

pub(crate) fn should_retry_transient_exec_rpc(status: &Status, deadline: Instant) -> bool {
    use std::error::Error;
    if Instant::now() >= deadline {
        return false;
    }
    match status.code() {
        Code::Unavailable | Code::DeadlineExceeded => true,
        // A fired per-attempt `set_timeout` surfaces as `Cancelled` carrying the
        // timeout as a source, as does a client-side transport cancel; both are
        // transient and retryable. A server-sent `CANCELLED` (no source) is a
        // deliberate cancellation and is left alone.
        Code::Cancelled => status.source().is_some(),
        Code::Unknown => is_transient_transport_message(status.message()),
        _ => false,
    }
}

/// Sleep before the next retry attempt, bounded by the max per-retry delay
/// and the remaining budget; returns the doubled delay for the next round.
pub(crate) async fn sleep_before_retry(delay: f64, deadline: Instant) -> f64 {
    let mut sleep_for = delay.min(EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
    let remaining = deadline
        .saturating_duration_since(Instant::now())
        .as_secs_f64();
    if remaining <= 0.0 {
        return delay;
    }
    sleep_for = sleep_for.min(remaining);
    tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
    (delay * 2.0).min(EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
}

impl WorkerProxy {
    /// Build a worker proxy that authenticates with `api_key`. Fails if the key
    /// cannot form a valid gRPC `authorization` metadata value.
    pub fn new(api_key: &str) -> Result<WorkerProxy, SailError> {
        let authorization = format!("Bearer {api_key}")
            .parse()
            .map_err(|_| SailError::Config {
                message: "SAIL_API_KEY contains characters invalid in a gRPC metadata value"
                    .to_string(),
            })?;
        Ok(WorkerProxy {
            channels: ChannelCache::new(),
            authorization,
        })
    }

    pub(crate) fn channels(&self) -> &ChannelCache {
        &self.channels
    }

    pub(crate) fn client_for(
        &self,
        endpoint: &str,
    ) -> Result<WorkerProxyServiceClient<Channel>, SailError> {
        let channel = self.channels.get(endpoint)?;
        Ok(WorkerProxyServiceClient::new(channel))
    }

    pub(crate) fn request_for<T>(
        &self,
        message: T,
        extra_metadata: &[(String, String)],
        timeout: Option<Duration>,
    ) -> Result<Request<T>, SailError> {
        let mut request = Request::new(message);
        request
            .metadata_mut()
            .insert("authorization", self.authorization.clone());
        for (key, value) in extra_metadata {
            let key: AsciiMetadataKey = key.parse().map_err(|_| SailError::Config {
                message: format!("invalid gRPC metadata key {key:?}"),
            })?;
            let value: AsciiMetadataValue = value.parse().map_err(|_| SailError::Config {
                message: format!("invalid gRPC metadata value for key {key:?}"),
            })?;
            request.metadata_mut().insert(key, value);
        }
        if let Some(timeout) = timeout {
            request.set_timeout(timeout);
        }
        Ok(request)
    }

    /// Wait for an exec to finish. The retry deadline starts at the FIRST
    /// transient error, not at the call: a wait legitimately blocks for as
    /// long as the guest command runs, so only consecutive failure time is
    /// budgeted.
    pub async fn wait_exec(
        &self,
        endpoint: &str,
        sailbox_id: &str,
        exec_request_id: &str,
        retry_timeout: f64,
    ) -> Result<WaitOutcome, SailError> {
        let mut deadline: Option<Instant> = None;
        let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
        loop {
            let message = pb::WaitSailboxExecRequest {
                sailbox_id: sailbox_id.to_string(),
                exec_request_id: exec_request_id.to_string(),
            };
            let request = self.request_for(message, &[], None)?;
            match self.client_for(endpoint)?.wait_sailbox_exec(request).await {
                Ok(resp) => {
                    let resp = resp.into_inner();
                    return Ok(WaitOutcome {
                        status: resp.status,
                        stdout: resp.stdout,
                        stderr: resp.stderr,
                        return_code: resp.return_code,
                        timed_out: resp.timed_out,
                        stdout_truncated: resp.stdout_truncated,
                        stderr_truncated: resp.stderr_truncated,
                    });
                }
                Err(status) => {
                    let deadline = *deadline.get_or_insert_with(|| retry_deadline(retry_timeout));
                    if !should_retry_transient_exec_rpc(&status, deadline) {
                        return Err(SailError::from_exec_status(&status));
                    }
                    tracing::warn!(code = ?status.code(), endpoint, "retrying transient worker-proxy RPC");
                    if should_invalidate_channel(&status) {
                        self.channels.invalidate(endpoint);
                    }
                    delay = sleep_before_retry(delay, deadline).await;
                }
            }
        }
    }

    /// Signal the guest command: SIGINT by default, SIGKILL if force.
    /// `retry_timeout > 0` budgets retries across the saild
    /// registration-gap window AND bounds each attempt with a per-call
    /// gRPC deadline so a stuck connection cannot hang past the budget.
    pub async fn cancel_exec(
        &self,
        endpoint: &str,
        sailbox_id: &str,
        exec_request_id: &str,
        force: bool,
        retry_timeout: f64,
    ) -> Result<(), SailError> {
        let deadline = retry_deadline(retry_timeout);
        let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
        loop {
            let per_attempt_timeout = if retry_timeout > 0.0 {
                // Cap each attempt so a half-open connection times out and the
                // loop redials, instead of one attempt consuming the whole
                // budget. retry_timeout == 0 keeps a single uncapped attempt
                // (programmatic no-retry cancel).
                Some(rpc_attempt_timeout(deadline))
            } else {
                None
            };
            let message = pb::CancelSailboxExecRequest {
                sailbox_id: sailbox_id.to_string(),
                exec_request_id: exec_request_id.to_string(),
                force,
            };
            let request = self.request_for(message, &[], per_attempt_timeout)?;
            match self
                .client_for(endpoint)?
                .cancel_sailbox_exec(request)
                .await
            {
                Ok(_) => return Ok(()),
                Err(status) => {
                    if !should_retry_transient_exec_rpc(&status, deadline) {
                        return Err(SailError::from_exec_status(&status));
                    }
                    tracing::warn!(code = ?status.code(), endpoint, "retrying transient worker-proxy RPC");
                    if should_invalidate_channel(&status) {
                        self.channels.invalidate(endpoint);
                    }
                    delay = sleep_before_retry(delay, deadline).await;
                }
            }
        }
    }

    /// Open a streaming read of a guest file. The returned [`FileReader`] yields
    /// chunks as they arrive; errors surface from its `next`.
    ///
    /// # Runtime
    ///
    /// Spawns the background pump on the calling task's tokio runtime, so call it
    /// from within one (every binding does, via the shared runtime; an async host
    /// from its own). The dialed channel co-locates on that runtime.
    pub fn read_file(self: &Arc<Self>, endpoint: &str, sailbox_id: &str, path: &str) -> FileReader {
        let (tx, rx) = mpsc::channel(FILE_READ_CHANNEL_CAP);
        let worker = Arc::clone(self);
        let endpoint = endpoint.to_string();
        let message = pb::ReadSailboxFileRequest {
            sailbox_id: sailbox_id.to_string(),
            path: path.to_string(),
        };
        tokio::spawn(async move {
            let request = match worker.request_for(message, &[], None) {
                Ok(request) => request,
                Err(err) => {
                    let _ = tx.send(Err(err)).await;
                    return;
                }
            };
            let mut client = match worker.client_for(&endpoint) {
                Ok(client) => client,
                Err(err) => {
                    let _ = tx.send(Err(err)).await;
                    return;
                }
            };
            let mut stream = match client.read_sailbox_file(request).await {
                Ok(resp) => resp.into_inner(),
                Err(status) => {
                    let _ = tx.send(Err(SailError::from_file_rpc_status(&status))).await;
                    return;
                }
            };
            loop {
                match stream.message().await {
                    Ok(Some(resp)) => {
                        if !resp.data.is_empty() && tx.send(Ok(resp.data)).await.is_err() {
                            return; // the reader was dropped
                        }
                    }
                    Ok(None) => return,
                    Err(status) => {
                        let _ = tx.send(Err(SailError::from_file_rpc_status(&status))).await;
                        return;
                    }
                }
            }
        });
        FileReader {
            rx: AsyncMutex::new(rx),
        }
    }

    /// Open a streaming write to a guest file. The caller feeds chunks via
    /// [`FileWriter::write_chunk`] and ends with [`FileWriter::finish`], so a
    /// large source is never buffered whole. The first chunk carries the
    /// path/flags; the rest carry data only.
    ///
    /// # Runtime
    ///
    /// Spawns the streaming RPC on the calling task's tokio runtime, so call it
    /// from within one. The dialed channel co-locates on that runtime.
    pub fn write_file(
        self: &Arc<Self>,
        endpoint: &str,
        sailbox_id: &str,
        path: &str,
        create_parents: bool,
        mode: Option<u32>,
    ) -> FileWriter {
        let (tx, rx) = mpsc::channel(FILE_READ_CHANNEL_CAP);
        let worker = Arc::clone(self);
        let endpoint = endpoint.to_string();
        // Build the client and run the RPC inside the spawned task so the
        // channel is dialed in the runtime's reactor context.
        let task = tokio::spawn(async move {
            let request = worker.request_for(ReceiverStream::new(rx), &[], None)?;
            worker
                .client_for(&endpoint)?
                .write_sailbox_file(request)
                .await
                .map(|_| ())
                .map_err(|status| SailError::from_file_rpc_status(&status))
        });
        FileWriter {
            tx: Some(tx),
            task: Some(task),
            first: true,
            sailbox_id: sailbox_id.to_string(),
            path: path.to_string(),
            create_parents,
            mode,
        }
    }
}

/// A streaming write to a guest file. Chunks feed a bounded channel that backs
/// the client-streaming RPC, so a slow uplink applies backpressure rather than
/// buffering the source. The RPC result surfaces from [`FileWriter::finish`].
pub struct FileWriter {
    tx: Option<mpsc::Sender<pb::WriteSailboxFileRequest>>,
    task: Option<tokio::task::JoinHandle<Result<(), SailError>>>,
    first: bool,
    sailbox_id: String,
    path: String,
    create_parents: bool,
    mode: Option<u32>,
}

impl FileWriter {
    fn build(&mut self, data: Vec<u8>) -> pb::WriteSailboxFileRequest {
        let header = self.first;
        self.first = false;
        pb::WriteSailboxFileRequest {
            sailbox_id: if header {
                self.sailbox_id.clone()
            } else {
                String::new()
            },
            path: if header {
                self.path.clone()
            } else {
                String::new()
            },
            data,
            create_parents: header && self.create_parents,
            mode: if header { self.mode } else { None },
        }
    }

    /// Stop the stream and return the RPC's result.
    async fn join(&mut self) -> Result<(), SailError> {
        self.tx = None; // dropping the sender ends the client stream
        match self.task.take() {
            Some(task) => task.await.unwrap_or_else(|join_err| {
                Err(SailError::Internal {
                    message: format!("file write task failed: {join_err}"),
                })
            }),
            None => Ok(()),
        }
    }

    /// Send the next chunk of file data. If the RPC has already ended, returns
    /// its result instead.
    pub async fn write_chunk(&mut self, data: Vec<u8>) -> Result<(), SailError> {
        let request = self.build(data);
        match &self.tx {
            // A send error means the RPC task already ended; surface its result.
            Some(tx) if tx.send(request).await.is_ok() => Ok(()),
            _ => self.join().await,
        }
    }

    /// Finish the write and return the RPC's result, creating an empty file when
    /// no chunks were sent.
    pub async fn finish(&mut self) -> Result<(), SailError> {
        if self.first {
            // No chunks were written: send a header-only message so an empty
            // file is still created.
            let request = self.build(Vec::new());
            if let Some(tx) = &self.tx {
                let _ = tx.send(request).await;
            }
        }
        self.join().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn transient_codes_retry_within_deadline() {
        let deadline = Instant::now() + Duration::from_secs(5);
        assert!(should_retry_transient_exec_rpc(
            &Status::unavailable("x"),
            deadline
        ));
        assert!(should_retry_transient_exec_rpc(
            &Status::deadline_exceeded("x"),
            deadline
        ));
        assert!(should_retry_transient_exec_rpc(
            &Status::unknown("HTTP/2 connection reset by remote"),
            deadline
        ));
        assert!(!should_retry_transient_exec_rpc(
            &Status::unknown("guest exploded"),
            deadline
        ));
        assert!(!should_retry_transient_exec_rpc(
            &Status::not_found("x"),
            deadline
        ));
    }

    #[test]
    fn fired_attempt_timeout_retries_but_server_cancel_does_not() {
        let deadline = Instant::now() + Duration::from_secs(5);
        // A fired per-attempt set_timeout reaches us as Cancelled carrying the
        // timeout as a source: retry it.
        let timed_out = Status::from_error(Box::new(tonic::TimeoutExpired(())));
        assert_eq!(timed_out.code(), Code::Cancelled);
        assert!(should_retry_transient_exec_rpc(&timed_out, deadline));
        // A server-sent CANCELLED has no source: do not retry.
        assert!(!should_retry_transient_exec_rpc(
            &Status::cancelled("client went away"),
            deadline
        ));
    }

    #[test]
    fn expired_deadline_never_retries() {
        let deadline = Instant::now();
        assert!(!should_retry_transient_exec_rpc(
            &Status::unavailable("x"),
            deadline
        ));
    }

    #[test]
    fn draining_detection_is_case_insensitive_and_code_scoped() {
        assert!(is_workerproxy_draining(&Status::unavailable(
            "workerproxy DRAINING for deploy"
        )));
        assert!(!is_workerproxy_draining(&Status::internal("draining")));
        assert!(!is_workerproxy_draining(&Status::unavailable("lameduck")));
    }

    #[test]
    fn rpc_attempt_timeout_caps_and_leaves_retry_headroom() {
        let now = Instant::now();
        // Far deadline: capped at the per-attempt ceiling.
        let far = rpc_attempt_timeout(now + Duration::from_mins(1));
        assert!(far <= Duration::from_secs_f64(EXEC_RPC_ATTEMPT_TIMEOUT_SECONDS));
        assert!(far > Duration::from_secs(5));
        // Budget below the ceiling: a single attempt takes less than the whole
        // budget (~half), so a retry can still land if it hangs.
        let small = rpc_attempt_timeout(now + Duration::from_secs(5));
        assert!(small < Duration::from_secs(5));
        assert!(small <= Duration::from_secs(3));
        // Past deadline: never zero, so the attempt still fires and fails fast.
        assert_eq!(
            rpc_attempt_timeout(now.checked_sub(Duration::from_secs(1)).unwrap()),
            Duration::from_millis(1)
        );
    }

    #[test]
    fn invalidate_on_draining_or_transport_failure_only() {
        // Draining: drop the channel even though it is a clean server status.
        assert!(should_invalidate_channel(&Status::unavailable(
            "workerproxy is draining"
        )));
        // A client-side transport failure carries a source: the connection is
        // suspect, so dial fresh.
        let io = std::io::Error::new(std::io::ErrorKind::ConnectionReset, "socket closed");
        assert!(should_invalidate_channel(&Status::from_error(Box::new(io))));
        // A plain server-sent transient keeps the connection.
        assert!(!should_invalidate_channel(&Status::unavailable(
            "try again"
        )));
    }
}