rustrust-tonic

applying different layers to different services in one tonic server


I have two versions of a layer, EndpointTimer_1 and EndpointTimer_2 which track slightly different metrics. I want EndpointTimer_1 to be applied to Server_1 and EndpointTimer_2 to applied to Server_2.

What's the correct way to use a common ServerBuilder object to facilitate this? Currently I have something like

use tonic::transport::Server;

let mut server_builder = Server::builder()
    .layer(EndpointTimerLayer::new())
    .layer(ConcurrencyGaugeLayer::new());

let provider_1 = provider_1_impl::new(&CONFIG)?;
let provider_2 = provider_2_impl::new(&CONFIG)?;
let router = server_builder
    .add_service(Server_1::new(provider_1))
    .add_service(Server_2::new(provider_2));
router.serve(addr).await?;

and I want something like

use tonic::transport::Server;

let mut server_builder = Server::builder()
    .layer(ConcurrencyGauge::new());

let provider_1 = provider_1_impl::new(&CONFIG)?;
let provider_2 = provider_2_impl::new(&CONFIG)?;
let router = server_builder
    .add_service(Server_1::new(provider_1).layer(EndpointTimerLayer_1::new())
    .add_service(Server_2::new(provider_2).layer(EndpointTimerLayer_2::new());
router.serve(addr).await?;

but I can't figure out the right way to make it work.


Solution

  • Took a bunch of experimenting and reading through tonic source code, and I still have some things I don't understand, but the functional answer is to first, implement the NamedService trait for each layer:

    use tonic::transport::server::NamedService;
    
    impl<T> NamedService for EndpointTimer_1<T>
    where
        T: NamedService,
    {
        const NAME: &'static str = T::NAME;
    }
    

    Note, that the error type for the layer's service must be Infallible, e.g.

    use std::convert::Infallible;
    
    impl<T, RequestBody, ResponseBody> Service<tonic::codegen::http::Request<RequestBody>>
        for EndpointTimer_1<T>
    where
        T: Service<
                tonic::codegen::http::Request<RequestBody>,
                Response = tonic::codegen::http::Response<ResponseBody>,
                Error = Infallible,
            > + Send,
    

    instead of the example-provided type Error = Box<dyn Error + Send + Sync>;. AFAICT this works because the protobuf-generated server is Infallible as well? Unsure, would appreciate an explanation of this by anyone in the know.

    Finally, services can be registered by wrapping them in layers individually:

    let mut server_builder = Server::builder()
        .layer(ConcurrencyGauge::new());
    
    let provider_1 = provider_1_impl::new(&CONFIG)?;
    let server_1 = Server_1::new(provider_1);
    
    let provider_2 = provider_2_impl::new(&CONFIG)?;
    let server_2 = Server_2::new(provider_2);
    
    
    let router = server_builder
        .add_service(EndpointTimer_1::new(Server_1))
        .add_service(EndpointTimer_2::new(Server_2));
    router.serve(addr).await?;
    

    Note that this is using the Timer itself, which takes a Service as an argument to new(), and not the Layer around it.