fix last tracing test

This commit is contained in:
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 std::time::{Duration, Instant};
//!
//! pub struct TimeTrace;
//! pub struct TimeTrace(Instant);
//!
//! impl ReqwestOtelSpanBackend for TimeTrace {
//! fn on_request_start(req: &Request, extension: &mut Extensions) -> Span {
//! extension.insert(Instant::now());
//! reqwest_otel_span!(name="example-request", req, time_elapsed = tracing::field::Empty)
//! fn on_request_start(req: &Request, _extension: &mut Extensions) -> (Self, Span) {
//! let now = Self(Instant::now());
//! 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) {
//! let time_elapsed = extension.get::<Instant>().unwrap().elapsed().as_millis() as i64;
//! fn on_request_end(self, span: &Span, outcome: &Result<Response, Error>) {
//! let time_elapsed = self.0.elapsed().as_millis() as i64;
//! default_on_request_end(span, outcome);
//! span.record("time_elapsed", &time_elapsed);
//! }

View file

@ -47,14 +47,14 @@ where
fn layer(&self, inner: Svc) -> Self::Service {
TracingMiddlewareService {
service: inner,
layer: *self,
_layer: *self,
}
}
}
/// Middleware Service for tracing requests using the current Opentelemetry Context.
pub struct TracingMiddlewareService<S: ReqwestOtelSpanBackend, Svc> {
layer: TracingMiddleware<S>,
_layer: TracingMiddleware<S>,
service: Svc,
}
@ -77,7 +77,7 @@ where
request,
mut extensions,
} = 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.
// Spans added by downstream consumers will be part of the same trace.
#[cfg(any(
@ -96,8 +96,8 @@ where
});
TracingMiddlewareFuture {
layer: self.layer,
span: request_span,
span,
backend: Some(backend),
future,
}
}
@ -105,8 +105,8 @@ where
pin_project!(
pub struct TracingMiddlewareFuture<S: ReqwestOtelSpanBackend, F> {
layer: TracingMiddleware<S>,
span: Span,
backend: Option<S>,
#[pin]
future: F,
}
@ -123,7 +123,10 @@ impl<S: ReqwestOtelSpanBackend, F: Future<Output = Result<Response, Error>>> Fut
let _guard = this.span.enter();
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)
}
}

View file

@ -39,12 +39,12 @@ pub const HTTP_USER_AGENT: &str = "http.user_agent";
/// Check out [`reqwest_otel_span`] documentation for examples.
///
/// [`TracingMiddleware`]: crate::middleware::TracingMiddleware.
pub trait ReqwestOtelSpanBackend {
pub trait ReqwestOtelSpanBackend: Sized {
/// 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.
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.
@ -95,15 +95,15 @@ pub fn default_on_request_failure(span: &Span, e: &Error) {
pub struct 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
.get::<OtelName>()
.map(|on| on.0.as_ref())
.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)
}
}
@ -119,16 +119,19 @@ fn get_header_value(key: &str, headers: &HeaderMap) -> String {
pub struct 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
.get::<OtelName>()
.map(|on| on.0.as_ref())
.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)
}
}

View file

@ -41,11 +41,11 @@
/// pub struct CustomReqwestOtelSpanBackend;
///
/// impl ReqwestOtelSpanBackend for CustomReqwestOtelSpanBackend {
/// fn on_request_start(req: &Request, _extension: &mut Extensions) -> Span {
/// reqwest_otel_span!(name = "reqwest-http-request", req)
/// fn on_request_start(req: &Request, _extension: &mut Extensions) -> (Self, Span) {
/// (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)
/// }
/// }