GRPC TraceLayer example with custom methods

starsoccer commented 1 year ago

Following the example here, I am able to create a very bare bones example of logging with grpc, but Id like to be able to define custom methods for on_request, and on_response but cant seem to figure out the correct way to do so.

Ive tried tweaking the ServiceBuilder in the above example to look as follows:

    let channel = ServiceBuilder::new()
        // Decompress response bodies
        .layer(DecompressionLayer::new())
        // Set a `User-Agent` header
        .layer(SetRequestHeaderLayer::overriding(
            header::USER_AGENT,
            HeaderValue::from_static("tonic-key-value-store"),
        ))
        // Log all requests and responses
        .layer(
           /// SECTION I CHANGED BELOW
            TraceLayer::new_for_grpc().make_span_with(DefaultMakeSpan::new().include_headers(true))
            .on_request(|request: &Request<Body>, _span: &Span| {
                tracing::debug!("started")
            }),
        )
        // Build our final `Service`
        .service(channel);

This provides an error saying:

error[E0631]: type mismatch in closure arguments
   --> src\main.rs:311:8
    |
303 |             .on_request(|request: &Request<Body>, _span: &Span| {
    |                         --------------------------------------- found signature defined here
...
311 |     Ok(KeyValueStoreClient::new(channel))
    |        ^^^^^^^^^^^^^^^^^^^^^^^^ expected due to this
    |
    = note: expected closure signature `for<'r, 's> fn(&'r hyper::Request<http_body::combinators::box_body::BoxBody<bytes::Bytes, Status>>, &'s Span) -> _`
               found closure signature `for<'r, 's> fn(&'r tonic::Request<hyper::Body>, &'s Span) -> _`
    = note: required for `[closure@src\main.rs:303:25: 303:64]` to implement `OnRequest<http_body::combinators::box_body::BoxBody<bytes::Bytes, Status>>`
    = note: required for `tower_http::trace::Trace<Channel, SharedClassifier<GrpcErrorsAsFailures>, DefaultMakeSpan, [closure@src\main.rs:303:25: 303:64]>` to implement `Service<hyper::Request<http_body::combinators::box_body::BoxBody<bytes::Bytes, Status>>>`
    = note: required for `Decompression<SetRequestHeader<tower_http::trace::Trace<Channel, SharedClassifier<GrpcErrorsAsFailures>, DefaultMakeSpan, [closure@src\main.rs:303:25: 303:64]>, HeaderValue>>` to implement `GrpcService<http_body::combinators::box_body::BoxBody<bytes::Bytes, Status>>`

erebe commented 1 year ago

Hello,

Your function need to be generic over the RequestBody. Or in your case be specific to http_body::combinators::box_body::BoxBody<bytes::Bytes, Status>>

You can use free function to introduce the generic.

    fn on_request<B>(self, request: &Request<B>,  _span: &Span) {
         tracing::debug!("started")
    }

    TraceLayer::new_for_grpc()
    .make_span_with(DefaultMakeSpan::new().include_headers(true))
    .on_request(on_request)

or if you want to use a lambda

    TraceLayer::new_for_grpc()
    .make_span_with(DefaultMakeSpan::new().include_headers(true))
    .on_request(|request: &Request<_>, _span: &Span| {
                    tracing::debug!("started")
                }
   )

starsoccer commented 1 year ago

Thanks the above works, but how would I go about actually logging the request body or response body? For requests the body seems to just be Body(Streaming) and the response is just UnsyncBoxBody

tower-rs / tower-http

GRPC TraceLayer example with custom methods #366