nordisk/skade

skade-flight/src/server.rs

//! Bind and serve. TLS rides tonic's `tls-ring` behind the `tls` feature.

use std::net::SocketAddr;
use std::sync::Arc;

use arrow_flight::flight_service_server::FlightServiceServer;
use tonic::transport::Server;

use crate::auth::{Auth, interceptor};
use crate::engine::Engine;
use crate::service::SkadeFlightService;

#[derive(Debug, Clone, Default)]
pub struct ServeOptions {
    /// PEM cert + key for TLS; optional client CA for mTLS.
    #[cfg(feature = "tls")]
    pub tls: Option<TlsOptions>,
}

#[cfg(feature = "tls")]
#[derive(Debug, Clone)]
pub struct TlsOptions {
    pub cert_pem: Vec<u8>,
    pub key_pem: Vec<u8>,
    pub client_ca_pem: Option<Vec<u8>>,
}

fn builder(opts: ServeOptions) -> Result<Server, tonic::transport::Error> {
    #[allow(unused_mut)]
    let mut b = Server::builder();
    // HTTP/2 max frame size. h2's default is 16 KiB: a 100 MB scan is ~6 500
    // DATA frames, each a state-machine turn. MEASURED 2026-09-09, 2M rows × 6
    // cols on loopback, both peers at 1 MiB: 23.5M → 39.8M rows/s, raw wire
    // 1.2 → 2.0 GB/s. Window size and batch size made no difference — this did.
    // The frame size a peer may SEND is what the RECEIVER advertises, so this
    // setting governs what clients send us (DoPut, handshake) and the client's
    // own setting governs our DoGet stream; gRPC C++ (the ODBC driver) and the
    // Go ADBC driver both advertise more than h2's default.
    // SKADE_FLIGHT_H2_FRAME_KB overrides (clamped to h2's legal 16 KiB..16 MiB-1).
    let frame = std::env::var("SKADE_FLIGHT_H2_FRAME_KB")
        .ok()
        .and_then(|v| v.parse::<u32>().ok())
        .map(|kb| kb << 10)
        .unwrap_or(1 << 20)
        .clamp(16 * 1024, (1 << 24) - 1);
    b = b.max_frame_size(frame);
    #[cfg(feature = "tls")]
    if let Some(t) = opts.tls {
        use tonic::transport::{Certificate, Identity, ServerTlsConfig};
        let mut cfg = ServerTlsConfig::new().identity(Identity::from_pem(t.cert_pem, t.key_pem));
        if let Some(ca) = t.client_ca_pem {
            cfg = cfg.client_ca_root(Certificate::from_pem(ca));
        }
        b = b.tls_config(cfg)?;
    }
    #[cfg(not(feature = "tls"))]
    let _ = opts;
    Ok(b)
}

/// Serve until `shutdown` resolves (or forever with `futures::future::pending()`).
pub async fn serve(
    engine: Arc<Engine>,
    auth: Arc<Auth>,
    addr: SocketAddr,
    opts: ServeOptions,
    shutdown: impl std::future::Future<Output = ()>,
) -> Result<(), tonic::transport::Error> {
    let svc = SkadeFlightService::new(engine, auth.clone());
    let svc = FlightServiceServer::with_interceptor(svc, interceptor(auth));
    builder(opts)?.add_service(svc).serve_with_shutdown(addr, shutdown).await
}

/// Bind an ephemeral or fixed listener first (so the caller knows the port),
/// then serve on it. What the tests, the bench and a `--listen 127.0.0.1:0` use.
pub async fn serve_on(
    engine: Arc<Engine>,
    auth: Arc<Auth>,
    listener: tokio::net::TcpListener,
    opts: ServeOptions,
    shutdown: impl std::future::Future<Output = ()>,
) -> Result<(), tonic::transport::Error> {
    let svc = SkadeFlightService::new(engine, auth.clone());
    let svc = FlightServiceServer::with_interceptor(svc, interceptor(auth));
    // tonic's own incoming, so accepted sockets get TCP_NODELAY like `serve`'s
    // do. MEASURED 2026-09-09: a hand-rolled accept loop handed tonic raw
    // sockets and every small RPC paid Nagle + delayed-ACK — 42 ms per
    // GetFlightInfo on loopback, 83 ms per GetSqlInfo (two round trips).
    let incoming = tonic::transport::server::TcpIncoming::from(listener).with_nodelay(Some(true));
    builder(opts)?.add_service(svc).serve_with_incoming_shutdown(incoming, shutdown).await
}