talksik/tonic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(tonic): pass trace_fn the request rather than just the headers (#634)

Browse Source
This commit is contained in:
David Pedersen
2021-05-12 18:24:09 +02:00
committed by GitHub
parent 2bf14e1d7d
commit 7862a2259d
1 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tonic/src/transport/server/mod.rs
+20 -12
View File
@@ -33,7 +33,7 @@ use futures_util::{
future::{self, Either as FutureEither, MapErr}, future::{self, Either as FutureEither, MapErr},
TryFutureExt, TryFutureExt,
}; };
use http::{HeaderMap, Request, Response}; use http::{Request, Response};
use hyper::{server::accept, Body}; use hyper::{server::accept, Body};
use std::{ use std::{
fmt, fmt,
@@ -48,7 +48,7 @@ use tower::{limit::concurrency::ConcurrencyLimitLayer, util::Either, Service, Se
use tracing_futures::{Instrument, Instrumented}; use tracing_futures::{Instrument, Instrumented};
type BoxService = tower::util::BoxService<Request<Body>, Response<BoxBody>, crate::Error>; type BoxService = tower::util::BoxService<Request<Body>, Response<BoxBody>, crate::Error>;
type TraceInterceptor = Arc<dyn Fn(&HeaderMap) -> tracing::Span + Send + Sync + 'static>; type TraceInterceptor = Arc<dyn Fn(&http::Request<()>) -> tracing::Span + Send + Sync + 'static>;
const DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS: u64 = 20; const DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS: u64 = 20;
@@ -290,10 +290,10 @@ impl Server {
} }
} }
/// Intercept inbound headers and add a [`tracing::Span`] to each response future. /// Intercept inbound requests and add a [`tracing::Span`] to each response future.
pub fn trace_fn<F>(self, f: F) -> Self pub fn trace_fn<F>(self, f: F) -> Self
where where
F: Fn(&HeaderMap) -> tracing::Span + Send + Sync + 'static, F: Fn(&http::Request<()>) -> tracing::Span + Send + Sync + 'static,
{ {
Server { Server {
trace_interceptor: Some(Arc::new(f)), trace_interceptor: Some(Arc::new(f)),
@@ -361,7 +361,7 @@ impl Server {
IE: Into<crate::Error>, IE: Into<crate::Error>,
F: Future<Output = ()>, F: Future<Output = ()>,
{ {
let span = self.trace_interceptor.clone(); let trace_interceptor = self.trace_interceptor.clone();
let concurrency_limit = self.concurrency_limit; let concurrency_limit = self.concurrency_limit;
let init_connection_window_size = self.init_connection_window_size; let init_connection_window_size = self.init_connection_window_size;
let init_stream_window_size = self.init_stream_window_size; let init_stream_window_size = self.init_stream_window_size;
@@ -381,7 +381,7 @@ impl Server {
inner: svc, inner: svc,
concurrency_limit, concurrency_limit,
timeout, timeout,
span, trace_interceptor,
}; };
let server = hyper::Server::builder(incoming) let server = hyper::Server::builder(incoming)
@@ -582,7 +582,7 @@ impl fmt::Debug for Server {
struct Svc<S> { struct Svc<S> {
inner: S, inner: S,
span: Option<TraceInterceptor>, trace_interceptor: Option<TraceInterceptor>,
conn_info: ConnectionInfo, conn_info: ConnectionInfo,
} }
@@ -602,8 +602,16 @@ where
} }
fn call(&mut self, mut req: Request<Body>) -> Self::Future { fn call(&mut self, mut req: Request<Body>) -> Self::Future {
let span = if let Some(trace_interceptor) = &self.span { let span = if let Some(trace_interceptor) = &self.trace_interceptor {
trace_interceptor(req.headers()) let (parts, body) = req.into_parts();
let bodyless_request = Request::from_parts(parts, ());
let span = trace_interceptor(&bodyless_request);
let (parts, _) = bodyless_request.into_parts();
req = Request::from_parts(parts, body);
span
} else { } else {
tracing::Span::none() tracing::Span::none()
}; };
@@ -624,7 +632,7 @@ struct MakeSvc<S> {
concurrency_limit: Option<usize>, concurrency_limit: Option<usize>,
timeout: Option<Duration>, timeout: Option<Duration>,
inner: S, inner: S,
span: Option<TraceInterceptor>, trace_interceptor: Option<TraceInterceptor>,
} }
impl<S> Service<&ServerIo> for MakeSvc<S> impl<S> Service<&ServerIo> for MakeSvc<S>
@@ -650,7 +658,7 @@ where
let svc = self.inner.clone(); let svc = self.inner.clone();
let concurrency_limit = self.concurrency_limit; let concurrency_limit = self.concurrency_limit;
let timeout = self.timeout; let timeout = self.timeout;
let span = self.span.clone(); let trace_interceptor = self.trace_interceptor.clone();
Box::pin(async move { Box::pin(async move {
let svc = ServiceBuilder::new() let svc = ServiceBuilder::new()
@@ -661,7 +669,7 @@ where
let svc = BoxService::new(Svc { let svc = BoxService::new(Svc {
inner: svc, inner: svc,
span, trace_interceptor,
conn_info, conn_info,
}); });