tidaldb/tidal-net/src/client.rs

//! gRPC client connection pool for shipping WAL segments to peers.

use std::collections::HashMap;

use tidaldb::replication::{shard::ShardId, transport::WalSegmentPayload};
use tonic::transport::Channel;

use crate::{
    circuit_breaker::CircuitBreaker,
    config::GrpcTransportConfig,
    error::GrpcTransportError,
    proto::{ShipSegmentRequest, wal_shipping_client::WalShippingClient},
    tls,
};

/// A connection to a single peer shard.
struct PeerConnection {
    client: WalShippingClient<Channel>,
    circuit_breaker: CircuitBreaker,
}

/// Connection pool managing one gRPC channel per peer shard.
pub struct PeerPool {
    peers: HashMap<ShardId, PeerConnection>,
}

impl PeerPool {
    /// Build a pool with lazy connections to all configured peers.
    ///
    /// # Errors
    ///
    /// Returns [`GrpcTransportError`] if the config fails its numeric invariant
    /// check ([`GrpcTransportConfig::validate`]), if TLS is unset but plaintext
    /// was not explicitly opted into (`insecure == false`), if a peer URI is
    /// invalid, or if the client TLS configuration cannot be built.
    pub fn new(config: &GrpcTransportConfig) -> Result<Self, GrpcTransportError> {
        config.validate()?;

        // Refuse to build a PLAINTEXT channel unless plaintext was explicitly
        // opted into. With no TLS config the scheme below falls back to `http`,
        // so silently honoring `insecure == false` would ship WAL bytes in the
        // clear while the operator believes the peer is authenticated. The server
        // (server::start_server) rejects this same asymmetry; the client must
        // match it so a half-configured deployment fails loudly on BOTH ends
        // instead of one end quietly downgrading to plaintext. `TlsConfig` is the
        // right variant: it is classified `is_permanent` (see error.rs), so the
        // shipper quarantines rather than retrying a config that cannot self-heal.
        if config.tls.is_none() && !config.insecure {
            return Err(GrpcTransportError::TlsConfig(
                "TLS not configured and insecure not set; refusing to ship WAL segments over \
                 plaintext"
                    .into(),
            ));
        }

        let mut peers = HashMap::new();

        for (&shard_id, addr) in &config.peers {
            let scheme = if config.tls.is_some() {
                "https"
            } else {
                "http"
            };
            let uri = format!("{scheme}://{addr}");

            // Build the endpoint with hard connect/request timeouts and HTTP/2
            // keep-alive so a stalled peer fails fast (transient error → circuit
            // breaker) instead of hanging the single-threaded WAL shipper. These
            // apply on both the plaintext and TLS paths.
            let mut endpoint = Channel::from_shared(uri)
                .map_err(|e| GrpcTransportError::Internal(format!("invalid URI: {e}")))?
                .connect_timeout(config.connect_timeout)
                .timeout(config.request_timeout)
                .tcp_keepalive(Some(config.keep_alive_interval))
                .http2_keep_alive_interval(config.keep_alive_interval)
                .keep_alive_timeout(config.keep_alive_timeout)
                .keep_alive_while_idle(true);

            // Configure client TLS if available.
            if let Some(ref tls_config) = config.tls {
                let client_tls = tls::client_tls_config(tls_config)?;
                endpoint = endpoint
                    .tls_config(client_tls)
                    .map_err(GrpcTransportError::TonicTransport)?;
            }

            let channel = endpoint.connect_lazy();
            // Raise tonic's default 4 MiB codec limits to the configured max
            // payload on BOTH directions. The client is the ENCODER for the
            // outbound ShipSegmentRequest (and the DECODER for any future
            // server-streaming response). Without this the client encoder trips
            // FIRST on a full-size WAL segment (16 MiB default, up to 64 MiB)
            // before bytes ever leave the process, surfacing as a codec error
            // the shipper would retry forever. Both ends must agree on the same
            // ceiling the application already validates against.
            let client = WalShippingClient::new(channel)
                .max_encoding_message_size(config.max_payload_bytes)
                .max_decoding_message_size(config.max_payload_bytes);
            let circuit_breaker = CircuitBreaker::new(
                config.circuit_breaker_threshold,
                config.circuit_breaker_reset,
            );

            peers.insert(
                shard_id,
                PeerConnection {
                    client,
                    circuit_breaker,
                },
            );
        }

        Ok(Self { peers })
    }

    /// Send a WAL segment to a peer shard.
    ///
    /// # Errors
    ///
    /// Returns [`GrpcTransportError`] if the peer is unknown, the circuit breaker
    /// is open, the follower signals backpressure, or the gRPC call fails.
    // `peer` is a plain `&PeerConnection` borrow into `self.peers` (no Drop of
    // its own) used both before the await (clone the client) and after (record
    // success/failure on its circuit breaker), so it legitimately spans the
    // whole fn — significant_drop_tightening is a false positive for a borrow.
    #[allow(clippy::significant_drop_tightening)]
    pub async fn send_to(
        &self,
        shard: ShardId,
        payload: WalSegmentPayload,
    ) -> Result<(), GrpcTransportError> {
        let peer = self
            .peers
            .get(&shard)
            .ok_or(GrpcTransportError::PeerUnreachable(shard))?;

        // Check circuit breaker.
        peer.circuit_breaker
            .check()
            .map_err(|_| GrpcTransportError::CircuitOpen(shard))?;

        let request: ShipSegmentRequest = payload.into();

        // Clone the client (tonic clients are cheap clones over a shared channel).
        let mut client = peer.client.clone();

        match client.ship_segment(request).await {
            Ok(response) => {
                // The follower signals backpressure with `accepted=false` when
                // its inbound channel is full: the segment was NOT enqueued.
                // This is normal flow control from a healthy-but-busy peer, not
                // a connection failure — so it must NOT count toward opening the
                // circuit breaker (that would convert transient congestion into
                // a self-inflicted `circuit_breaker_reset` replication stall,
                // amplifying follower lag exactly when the leader must keep
                // shipping). We route it through `record_backpressure`, which
                // leaves Closed/Open breaker state untouched (it only resolves an
                // outstanding half-open probe), and return a transient
                // `SegmentRejected` so the shipper does not advance its cursor
                // and retries on the next poll. Only the genuine transport error
                // arm below calls `record_failure`.
                if response.into_inner().accepted {
                    peer.circuit_breaker.record_success();
                    Ok(())
                } else {
                    peer.circuit_breaker.record_backpressure();
                    Err(GrpcTransportError::SegmentRejected(shard))
                }
            }
            Err(status) => {
                peer.circuit_breaker.record_failure();
                Err(GrpcTransportError::Grpc(Box::new(status)))
            }
        }
    }
}

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

    #[test]
    fn rejects_plaintext_when_insecure_unset() {
        // insecure=false with no TLS config must error rather than silently
        // building a plaintext channel — symmetric with the server's rejection.
        // No peers means the early-return guard fires before any endpoint is
        // built, so this needs no tokio reactor.
        let config = GrpcTransportConfig {
            tls: None,
            insecure: false,
            ..GrpcTransportConfig::default()
        };
        // `let-else` instead of `expect_err` so the test does not require
        // `PeerPool: Debug` (it holds tonic channel handles that are not Debug).
        let Err(err) = PeerPool::new(&config) else {
            panic!("plaintext without opt-in must be rejected");
        };
        assert!(
            matches!(err, GrpcTransportError::TlsConfig(_)),
            "expected a permanent TlsConfig error, got {err:?}"
        );
    }

    #[test]
    fn allows_plaintext_when_insecure_set() {
        // The explicit opt-in path must still succeed (no peers => no endpoints).
        let config = GrpcTransportConfig {
            tls: None,
            insecure: true,
            ..GrpcTransportConfig::default()
        };
        assert!(PeerPool::new(&config).is_ok());
    }

    #[test]
    fn rejects_invalid_config() {
        // PeerPool::new must run the config invariant check; a zero capacity is
        // caught here instead of panicking later in the mpsc channel build.
        let config = GrpcTransportConfig {
            channel_capacity: 0,
            ..GrpcTransportConfig::default()
        };
        assert!(PeerPool::new(&config).is_err());
    }
}