gunnar-grpc 1.5.0

//! **The handlers run where they cannot starve the connections** (R5.1-G,
//! 2026-10-07).
//!
//! A control-plane handler is an `async fn` that makes synchronous calls into
//! the store — a `StorageInfo` disk walk, a txlog fsync, a snapshot — and holds
//! the worker it runs on until each one returns. tonic runs a request's
//! handler on the runtime that serves its connection, so enough slow calls at
//! once held every worker of that runtime, nothing polled the HTTP/2
//! connection, a client's keepalive PING went unanswered, and the client
//! declared the connection dead — failing every call in flight on it, the slow
//! ones and the quick ones behind them alike. Measured on test2: the console's
//! `KeepAliveTimedOut` stalls, and a paid `Entitlement.Set` that came back to
//! monetize as an `Internal` stream reset.
//!
//! [`OffloadLayer`] moves each request's whole service call — the auth gate
//! (an interceptor, so inside the call) and the handler — onto the runtime
//! [`ServiceContext::handlers`](crate::ServiceContext::handlers) names. The
//! connection keeps its own runtime, whose workers only ever do I/O, so a PING
//! is answered however long the calls take; the calls themselves queue for the
//! handler runtime's workers exactly as they queued for the shared ones.
//!
//! Dropping the response future (the client reset the stream, the connection
//! closed) aborts the spawned call, as dropping the in-place future used to
//! cancel it: the call is cancelled at its next `.await`, never orphaned.

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use tonic::body::Body;
use tonic::codegen::http;
use tonic::codegen::Service;
use tracing::Instrument;

/// The layer. `None` runs every call in place, on the connection's runtime —
/// a bare context, a test router; `gunnar-server` always hands in a runtime.
#[derive(Clone, Debug, Default)]
pub struct OffloadLayer {
    handlers: Option<tokio::runtime::Handle>,
}

impl OffloadLayer {
    /// Run every call on `handlers` (`None`: in place).
    pub fn new(handlers: Option<tokio::runtime::Handle>) -> OffloadLayer {
        OffloadLayer { handlers }
    }
}

impl<S> tower_layer::Layer<S> for OffloadLayer {
    type Service = Offload<S>;

    fn layer(&self, inner: S) -> Offload<S> {
        Offload {
            inner,
            handlers: self.handlers.clone(),
        }
    }
}

/// One gRPC call, run on the handler runtime.
#[derive(Clone, Debug)]
pub struct Offload<S> {
    inner: S,
    handlers: Option<tokio::runtime::Handle>,
}

/// Aborts the spawned call when the response future is dropped, so a call the
/// client gave up on is cancelled exactly as it was before it moved runtimes.
struct AbortOnDrop<T>(tokio::task::JoinHandle<T>);

impl<T> Drop for AbortOnDrop<T> {
    fn drop(&mut self) {
        self.0.abort();
    }
}

impl<S, ResBody> Service<http::Request<Body>> for Offload<S>
where
    S: Service<http::Request<Body>, Response = http::Response<ResBody>> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Send + 'static,
    ResBody: Default + Send + 'static,
{
    type Response = http::Response<ResBody>;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, S::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: http::Request<Body>) -> Self::Future {
        // The tower clone dance (as in `trace::RequestSpan`): `self.inner` is
        // the copy `poll_ready` readied, so it is the one that serves this
        // request; a clone stays behind for the next.
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);
        let Some(handlers) = self.handlers.as_ref() else {
            return Box::pin(inner.call(request));
        };
        // The request's span (`grpc.request`, entered by the layer outside
        // this one) goes with the call, so the audit record the gate writes on
        // the handler runtime still carries `req`.
        let span = tracing::Span::current();
        // `inner.call` itself runs on the handler runtime too, not just the
        // future it returns: the gate is an interceptor, and an interceptor
        // runs when the call is made.
        let task = AbortOnDrop(handlers.spawn(async move { inner.call(request).await }.instrument(span)));
        Box::pin(async move {
            let mut task = task;
            match (&mut task.0).await {
                Ok(answer) => answer,
                Err(e) if e.is_panic() => {
                    tracing::error!(
                        target: crate::AUDIT_TARGET,
                        event = "grpc.handler.panicked",
                        "a control-plane handler panicked; the call is answered INTERNAL"
                    );
                    Ok(tonic::Status::internal("the control-plane handler for this call panicked")
                        .into_http())
                }
                // Cancelled: the handler runtime is shutting down with the
                // server. Retryable, and says why.
                Err(_) => Ok(tonic::Status::unavailable(
                    "the control plane is shutting down; this call did not run",
                )
                .into_http()),
            }
        })
    }
}

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

    /// A service that answers with the name of the thread it ran on.
    #[derive(Clone)]
    struct WhereAmI;

    impl Service<http::Request<Body>> for WhereAmI {
        type Response = http::Response<Body>;
        type Error = std::convert::Infallible;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: http::Request<Body>) -> Self::Future {
            let at = std::thread::current().name().unwrap_or("").to_owned();
            Box::pin(async move {
                let mut r = http::Response::new(Body::default());
                r.headers_mut()
                    .insert("x-thread", http::HeaderValue::from_str(&at).unwrap());
                Ok(r)
            })
        }
    }

    fn rt(name: &str) -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(1)
            .thread_name(name)
            .enable_all()
            .build()
            .unwrap()
    }

    fn thread_of(answer: &http::Response<Body>) -> String {
        answer.headers()["x-thread"].to_str().unwrap().to_owned()
    }

    #[test]
    fn the_call_itself_runs_on_the_handler_runtime() {
        let conn = rt("conn");
        let handlers = rt("handler");
        let mut svc = tower_layer::Layer::layer(&OffloadLayer::new(Some(handlers.handle().clone())), WhereAmI);
        let answer = conn.block_on(async move { svc.call(http::Request::new(Body::default())).await }).unwrap();
        assert_eq!(thread_of(&answer), "handler", "the call ran on the connection's runtime");
    }

    #[test]
    fn without_a_handler_runtime_the_call_runs_in_place() {
        let conn = rt("conn");
        let mut svc = tower_layer::Layer::layer(&OffloadLayer::new(None), WhereAmI);
        let answer = conn.block_on(async move { svc.call(http::Request::new(Body::default())).await }).unwrap();
        assert_ne!(thread_of(&answer), "handler");
    }

    /// A service whose call never finishes; counts the calls that started and
    /// the ones still alive.
    #[derive(Clone)]
    struct Forever {
        started: std::sync::Arc<std::sync::atomic::AtomicUsize>,
        alive: std::sync::Arc<std::sync::atomic::AtomicUsize>,
    }

    struct Alive(std::sync::Arc<std::sync::atomic::AtomicUsize>);

    impl Drop for Alive {
        fn drop(&mut self) {
            self.0.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
        }
    }

    impl Service<http::Request<Body>> for Forever {
        type Response = http::Response<Body>;
        type Error = std::convert::Infallible;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: http::Request<Body>) -> Self::Future {
            let alive = std::sync::Arc::clone(&self.alive);
            let started = std::sync::Arc::clone(&self.started);
            Box::pin(async move {
                started.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                alive.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                let _alive = Alive(alive);
                std::future::pending::<()>().await;
                unreachable!()
            })
        }
    }

    /// The client reset the stream: the call it no longer waits for is
    /// cancelled on the handler runtime, not left running there.
    #[test]
    fn a_call_whose_answer_nobody_waits_for_is_cancelled() {
        let conn = rt("conn");
        let handlers = rt("handler");
        let alive = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let started = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let mut svc = tower_layer::Layer::layer(
            &OffloadLayer::new(Some(handlers.handle().clone())),
            Forever { started: started.clone(), alive: alive.clone() },
        );
        conn.block_on(async {
            let call = svc.call(http::Request::new(Body::default()));
            // Let it start on the handler runtime, then give up on it.
            let _ = tokio::time::timeout(std::time::Duration::from_millis(200), call).await;
        });
        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
        while alive.load(std::sync::atomic::Ordering::SeqCst) != 0 && std::time::Instant::now() < deadline {
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
        assert_eq!(started.load(std::sync::atomic::Ordering::SeqCst), 1, "the call never started, so this proved nothing");
        assert_eq!(alive.load(std::sync::atomic::Ordering::SeqCst), 0, "the abandoned call is still running");
    }
}