reqwest-tracing: add method for handling cancelled requests

reqwest-tracing/CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- New method on `ReqwestOtelSpanBackend` for when the request does not finish and is cancelled instead. This happens when the request future is dropped instead of polling it to completion.
+
 ## [0.5.2] - 2024-07-15
 
 ### Added

reqwest-tracing/src/middleware.rs

@@ -1,7 +1,9 @@
+use std::marker::PhantomData;
+
 use http::Extensions;
 use reqwest::{Request, Response};
 use reqwest_middleware::{Middleware, Next, Result};
-use tracing::Instrument;
+use tracing::{Instrument, Span};
 
 use crate::{DefaultSpanBackend, ReqwestOtelSpanBackend};
 
@@ -42,6 +44,33 @@ where
         extensions: &mut Extensions,
         next: Next<'_>,
     ) -> Result<Response> {
+        struct CancelGuard<'s, ReqwestOtelSpan: ReqwestOtelSpanBackend> {
+            armed: bool,
+            span: &'s Span,
+            _phantom: PhantomData<ReqwestOtelSpan>,
+        }
+
+        impl<'s, ReqwestOtelSpan: ReqwestOtelSpanBackend> CancelGuard<'s, ReqwestOtelSpan> {
+            fn new(span: &'s Span) -> Self {
+                Self {
+                    armed: true,
+                    span,
+                    _phantom: PhantomData,
+                }
+            }
+            fn disarm(mut self) {
+                self.armed = false;
+            }
+        }
+
+        impl<'s, ReqwestOtelSpan: ReqwestOtelSpanBackend> Drop for CancelGuard<'s, ReqwestOtelSpan> {
+            fn drop(&mut self) {
+                if self.armed {
+                    ReqwestOtelSpan::on_request_cancelled(self.span);
+                }
+            }
+        }
+
         let request_span = ReqwestOtelSpan::on_request_start(&req, extensions);
 
         let outcome_future = async {
@@ -59,8 +88,13 @@ where
                 req
             };
 
+            let guard = CancelGuard::<'_, ReqwestOtelSpan>::new(&request_span);
+
             // Run the request
             let outcome = next.run(req, extensions).await;
+
+            guard.disarm();
+
             ReqwestOtelSpan::on_request_end(&request_span, &outcome, extensions);
             outcome
         };

reqwest-tracing/src/reqwest_otel_span_builder.rs

@@ -67,6 +67,11 @@ pub trait ReqwestOtelSpanBackend {
 
     /// Runs after the request call has executed.
     fn on_request_end(span: &Span, outcome: &Result<Response>, extension: &mut Extensions);
+
+    /// Runs if the request was cancelled before finishing.
+    fn on_request_cancelled(span: &Span) {
+        default_on_request_cancelled(span);
+    }
 }
 
 /// Populates default success/failure fields for a given [`reqwest_otel_span!`] span.
@@ -78,6 +83,11 @@ pub fn default_on_request_end(span: &Span, outcome: &Result<Response>) {
     }
 }
 
+#[inline]
+pub fn default_on_request_cancelled(span: &Span) {
+    span.record(OTEL_STATUS_CODE, "ERROR");
+}
+
 /// Populates default success fields for a given [`reqwest_otel_span!`] span.
 #[inline]
 pub fn default_on_request_success(span: &Span, response: &Response) {
@@ -148,6 +158,10 @@ impl ReqwestOtelSpanBackend for DefaultSpanBackend {
     fn on_request_end(span: &Span, outcome: &Result<Response>, _: &mut Extensions) {
         default_on_request_end(span, outcome)
     }
+
+    fn on_request_cancelled(span: &Span) {
+        default_on_request_cancelled(span);
+    }
 }
 
 /// Similar to [`DefaultSpanBackend`] but also adds the `url.full` attribute to request spans.
@@ -164,6 +178,10 @@ impl ReqwestOtelSpanBackend for SpanBackendWithUrl {
     fn on_request_end(span: &Span, outcome: &Result<Response>, _: &mut Extensions) {
         default_on_request_end(span, outcome)
     }
+
+    fn on_request_cancelled(span: &Span) {
+        default_on_request_cancelled(span);
+    }
 }
 
 /// HTTP Mapping <https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#status>