fix last tracing test

tower
Conrad Ludgate 2022-11-15 17:07:36 +00:00
parent 571b9abc49
commit df4990d62e
No known key found for this signature in database
GPG Key ID: 197E3CACA1C980B5
4 changed files with 32 additions and 25 deletions

View File

@ -62,16 +62,17 @@
//! use tracing::Span; //! use tracing::Span;
//! use std::time::{Duration, Instant}; //! use std::time::{Duration, Instant};
//! //!
//! pub struct TimeTrace; //! pub struct TimeTrace(Instant);
//! //!
//! impl ReqwestOtelSpanBackend for TimeTrace { //! impl ReqwestOtelSpanBackend for TimeTrace {
//! fn on_request_start(req: &Request, extension: &mut Extensions) -> Span { //! fn on_request_start(req: &Request, _extension: &mut Extensions) -> (Self, Span) {
//! extension.insert(Instant::now()); //! let now = Self(Instant::now());
//! reqwest_otel_span!(name="example-request", req, time_elapsed = tracing::field::Empty) //! let span = reqwest_otel_span!(name="example-request", req, time_elapsed = tracing::field::Empty);
//! (now, span)
//! } //! }
//! //!
//! fn on_request_end(span: &Span, outcome: &Result<Response, Error>, extension: &mut Extensions) { //! fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>) {
//! let time_elapsed = extension.get::<Instant>().unwrap().elapsed().as_millis() as i64; //! let time_elapsed = self.0.elapsed().as_millis() as i64;
//! default_on_request_end(span, outcome); //! default_on_request_end(span, outcome);
//! span.record("time_elapsed", &time_elapsed); //! span.record("time_elapsed", &time_elapsed);
//! } //! }

View File

@ -47,14 +47,14 @@ where
fn layer(&self, inner: Svc) -> Self::Service { fn layer(&self, inner: Svc) -> Self::Service {
TracingMiddlewareService { TracingMiddlewareService {
service: inner, service: inner,
layer: *self, _layer: *self,
} }
} }
} }
/// Middleware Service for tracing requests using the current Opentelemetry Context. /// Middleware Service for tracing requests using the current Opentelemetry Context.
pub struct TracingMiddlewareService<S: ReqwestOtelSpanBackend, Svc> { pub struct TracingMiddlewareService<S: ReqwestOtelSpanBackend, Svc> {
layer: TracingMiddleware<S>, _layer: TracingMiddleware<S>,
service: Svc, service: Svc,
} }
@ -77,7 +77,7 @@ where
request, request,
mut extensions, mut extensions,
} = req; } = req;
let request_span = ReqwestOtelSpan::on_request_start(&request, &mut extensions); let (backend, span) = ReqwestOtelSpan::on_request_start(&request, &mut extensions);
// Adds tracing headers to the given request to propagate the OpenTelemetry context to downstream revivers of the request. // Adds tracing headers to the given request to propagate the OpenTelemetry context to downstream revivers of the request.
// Spans added by downstream consumers will be part of the same trace. // Spans added by downstream consumers will be part of the same trace.
#[cfg(any( #[cfg(any(
@ -96,8 +96,8 @@ where
}); });
TracingMiddlewareFuture { TracingMiddlewareFuture {
layer: self.layer, span,
span: request_span, backend: Some(backend),
future, future,
} }
} }
@ -105,8 +105,8 @@ where
pin_project!( pin_project!(
pub struct TracingMiddlewareFuture<S: ReqwestOtelSpanBackend, F> { pub struct TracingMiddlewareFuture<S: ReqwestOtelSpanBackend, F> {
layer: TracingMiddleware<S>,
span: Span, span: Span,
backend: Option<S>,
#[pin] #[pin]
future: F, future: F,
} }
@ -123,7 +123,10 @@ impl<S: ReqwestOtelSpanBackend, F: Future<Output = Result<Response, Error>>> Fut
let _guard = this.span.enter(); let _guard = this.span.enter();
ready!(this.future.poll(cx)) ready!(this.future.poll(cx))
}; };
S::on_request_end(this.span, &outcome); this.backend
.take()
.expect("poll should not be called after completion")
.on_request_end(this.span, &outcome);
Poll::Ready(outcome) Poll::Ready(outcome)
} }
} }

View File

@ -39,12 +39,12 @@ pub const HTTP_USER_AGENT: &str = "http.user_agent";
/// Check out [`reqwest_otel_span`] documentation for examples. /// Check out [`reqwest_otel_span`] documentation for examples.
/// ///
/// [`TracingMiddleware`]: crate::middleware::TracingMiddleware. /// [`TracingMiddleware`]: crate::middleware::TracingMiddleware.
pub trait ReqwestOtelSpanBackend { pub trait ReqwestOtelSpanBackend: Sized {
/// Initalized a new span before the request is executed. /// Initalized a new span before the request is executed.
fn on_request_start(req: &Request, extension: &mut Extensions) -> Span; fn on_request_start(req: &Request, extension: &mut Extensions) -> (Self, Span);
/// Runs after the request call has executed. /// Runs after the request call has executed.
fn on_request_end(span: &Span, outcome: &Result<Response, Error>); fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>);
} }
/// Populates default success/failure fields for a given [`reqwest_otel_span!`] span. /// Populates default success/failure fields for a given [`reqwest_otel_span!`] span.
@ -95,15 +95,15 @@ pub fn default_on_request_failure(span: &Span, e: &Error) {
pub struct DefaultSpanBackend; pub struct DefaultSpanBackend;
impl ReqwestOtelSpanBackend for DefaultSpanBackend { impl ReqwestOtelSpanBackend for DefaultSpanBackend {
fn on_request_start(req: &Request, ext: &mut Extensions) -> Span { fn on_request_start(req: &Request, ext: &mut Extensions) -> (DefaultSpanBackend, Span) {
let name = ext let name = ext
.get::<OtelName>() .get::<OtelName>()
.map(|on| on.0.as_ref()) .map(|on| on.0.as_ref())
.unwrap_or("reqwest-http-client"); .unwrap_or("reqwest-http-client");
reqwest_otel_span!(name = name, req) (Self, reqwest_otel_span!(name = name, req))
} }
fn on_request_end(span: &Span, outcome: &Result<Response, Error>) { fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>) {
default_on_request_end(span, outcome) default_on_request_end(span, outcome)
} }
} }
@ -119,16 +119,19 @@ fn get_header_value(key: &str, headers: &HeaderMap) -> String {
pub struct SpanBackendWithUrl; pub struct SpanBackendWithUrl;
impl ReqwestOtelSpanBackend for SpanBackendWithUrl { impl ReqwestOtelSpanBackend for SpanBackendWithUrl {
fn on_request_start(req: &Request, ext: &mut Extensions) -> Span { fn on_request_start(req: &Request, ext: &mut Extensions) -> (Self, Span) {
let name = ext let name = ext
.get::<OtelName>() .get::<OtelName>()
.map(|on| on.0.as_ref()) .map(|on| on.0.as_ref())
.unwrap_or("reqwest-http-client"); .unwrap_or("reqwest-http-client");
reqwest_otel_span!(name = name, req, http.url = %remove_credentials(req.url())) (
Self,
reqwest_otel_span!(name = name, req, http.url = %remove_credentials(req.url())),
)
} }
fn on_request_end(span: &Span, outcome: &Result<Response, Error>) { fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>) {
default_on_request_end(span, outcome) default_on_request_end(span, outcome)
} }
} }

View File

@ -41,11 +41,11 @@
/// pub struct CustomReqwestOtelSpanBackend; /// pub struct CustomReqwestOtelSpanBackend;
/// ///
/// impl ReqwestOtelSpanBackend for CustomReqwestOtelSpanBackend { /// impl ReqwestOtelSpanBackend for CustomReqwestOtelSpanBackend {
/// fn on_request_start(req: &Request, _extension: &mut Extensions) -> Span { /// fn on_request_start(req: &Request, _extension: &mut Extensions) -> (Self, Span) {
/// reqwest_otel_span!(name = "reqwest-http-request", req) /// (Self, reqwest_otel_span!(name = "reqwest-http-request", req))
/// } /// }
/// ///
/// fn on_request_end(span: &Span, outcome: &Result<Response, Error>) { /// fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>) {
/// default_on_request_end(span, outcome) /// default_on_request_end(span, outcome)
/// } /// }
/// } /// }