sail-rs 0.2.19

//! Per-endpoint workerproxy gRPC channel cache.
//!
//! Channels are opened lazily (no I/O at construction) and cached per
//! endpoint so repeated exec calls reuse the HTTP/2 connection.
//! `invalidate` drops the cached handle WITHOUT closing it: tonic channels are
//! cheap clones over a shared connection, so in-flight RPCs on existing clones
//! keep running while the next call dials a fresh channel. A draining backend
//! surfaces as a retryable `Unavailable` ("draining") status; the exec retry
//! loop treats it as transient, invalidates the cached channel, and redials,
//! repeating until it reaches a non-draining backend or the retry budget runs
//! out.

use std::collections::HashMap;
#[cfg(test)]
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
use std::time::Duration;

use tonic::transport::{Channel, ClientTlsConfig, Endpoint};

use crate::error::SailError;

/// Thread-safe cache of lazily-dialed gRPC channels keyed by endpoint target.
#[derive(Default)]
pub(crate) struct ChannelCache {
    channels: Mutex<HashMap<String, Channel>>,
    /// Counts channel constructions, so the channel-eviction tests can observe
    /// dial vs cache-hit.
    #[cfg(test)]
    generation: AtomicU64,
}

impl ChannelCache {
    /// Create an empty cache.
    pub(crate) fn new() -> ChannelCache {
        ChannelCache::default()
    }

    /// Return the cached channel for `target`, dialing and caching a new lazy
    /// channel on a miss. TLS is used for non-loopback targets.
    pub(crate) fn get(&self, target: &str) -> Result<Channel, SailError> {
        let mut channels = self
            .channels
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(channel) = channels.get(target) {
            return Ok(channel.clone());
        }
        let channel = open_channel(target)?;
        #[cfg(test)]
        self.generation.fetch_add(1, Ordering::Relaxed);
        channels.insert(target.to_string(), channel.clone());
        Ok(channel)
    }

    /// Drop the cached channel for this endpoint so the next call dials
    /// fresh. Used when a workerproxy advertises it is draining.
    pub(crate) fn invalidate(&self, target: &str) {
        self.channels
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .remove(target);
    }

    /// Number of channels constructed so far, for the channel-eviction tests.
    #[cfg(test)]
    fn generation(&self) -> u64 {
        self.generation.load(Ordering::Relaxed)
    }
}

/// Local development workerproxies (and the test suite's fakes) listen on
/// loopback without TLS; deployed endpoints always terminate TLS.
fn is_insecure_target(target: &str) -> bool {
    let normalized = target.trim().to_lowercase();
    ["localhost:", "127.0.0.1:", "[::1]:"]
        .iter()
        .any(|prefix| normalized.starts_with(prefix))
}

fn open_channel(target: &str) -> Result<Channel, SailError> {
    let trimmed = target.trim();
    let invalid = |e: tonic::transport::Error| SailError::Config {
        message: format!("invalid workerproxy endpoint {trimmed:?}: {e}"),
    };
    let endpoint = if is_insecure_target(target) {
        Endpoint::from_shared(format!("http://{trimmed}")).map_err(invalid)?
    } else {
        Endpoint::from_shared(format!("https://{trimmed}"))
            .map_err(invalid)?
            .tls_config(ClientTlsConfig::new().with_native_roots())
            .map_err(|e| SailError::Internal {
                message: format!("TLS config failed for {trimmed:?}: {e}"),
            })?
    };
    // Probe a quiet connection so a half-open workerproxy (restarted or
    // force-deleted mid-exec) is detected in seconds instead of wedging a
    // long-poll WaitSailboxExec or a quiet StreamSailboxExec on a dead socket
    // until the kernel's multi-minute TCP timeout. Without this the exec
    // retry/poll fallback never fires. Ping even while idle, since the SDK only
    // receives on the stream; the 15s interval stays above the workerproxy's
    // 10s keepalive MinTime so the server never GOAWAYs for pinging too often.
    let endpoint = endpoint
        .http2_keep_alive_interval(Duration::from_secs(15))
        .keep_alive_timeout(Duration::from_secs(10))
        .keep_alive_while_idle(true);
    Ok(endpoint.connect_lazy())
}

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

    #[test]
    fn insecure_targets() {
        assert!(is_insecure_target("localhost:50052"));
        assert!(is_insecure_target(" LOCALHOST:50052 "));
        assert!(is_insecure_target("127.0.0.1:1"));
        assert!(is_insecure_target("[::1]:50052"));
        assert!(!is_insecure_target("workerproxy.sailresearch.com:443"));
        assert!(!is_insecure_target("localhost.evil.com:443"));
    }

    // connect_lazy still binds the channel to the ambient tokio runtime,
    // so even the no-I/O path needs a reactor context.
    #[tokio::test]
    async fn cache_reuses_and_invalidates() {
        let cache = ChannelCache::new();
        cache.get("localhost:50052").unwrap();
        cache.get("localhost:50052").unwrap();
        assert_eq!(cache.generation(), 1);
        cache.invalidate("localhost:50052");
        cache.get("localhost:50052").unwrap();
        assert_eq!(cache.generation(), 2);
    }
}