rama_http/layer/trace/

mod.rs

//! Middleware that adds high level [tracing] to a [`Service`].
//!
//! # Example
//!
//! Adding tracing to your service can be as simple as:
//!
//! ```rust
//! use rama_http::{Body, Request, Response};
//! use rama_core::service::service_fn;
//! use rama_core::{Layer, Service};
//! use rama_http::layer::trace::TraceLayer;
//! use std::convert::Infallible;
//!
//! async fn handle(request: Request) -> Result<Response, Infallible> {
//!     Ok(Response::new(Body::from("foo")))
//! }
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Setup tracing
//! tracing_subscriber::fmt::init();
//!
//! let mut service = TraceLayer::new_for_http().into_layer(service_fn(handle));
//!
//! let request = Request::new(Body::from("foo"));
//!
//! let response = service
//!     .serve(request)
//!     .await?;
//! # Ok(())
//! # }
//! ```
//!
//! If you run this application with `RUST_LOG=rama=trace cargo run` you should see logs like:
//!
//! ```text
//! Mar 05 20:50:28.523 DEBUG request{method=GET path="/foo"}: rama_http::layer::trace::on_request: started processing request
//! Mar 05 20:50:28.524 DEBUG request{method=GET path="/foo"}: rama_http::layer::trace::on_response: finished processing request latency=1 ms status=200
//! ```
//!
//! # Customization
//!
//! [`Trace`] comes with good defaults but also supports customizing many aspects of the output.
//!
//! The default behaviour supports some customization:
//!
//! ```rust
//! use rama_http::{Body, Request, Response, HeaderMap, StatusCode};
//! use rama_core::service::service_fn;
//! use rama_core::{Service, Layer};
//! use rama_core::telemetry::tracing::Level;
//! use rama_http::layer::trace::{
//!     TraceLayer, DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse,
//! };
//! use rama_utils::latency::LatencyUnit;
//! use std::time::Duration;
//! use std::convert::Infallible;
//!
//! # async fn handle(request: Request) -> Result<Response, Infallible> {
//! #     Ok(Response::new(Body::from("foo")))
//! # }
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # tracing_subscriber::fmt::init();
//! #
//! let service = (
//!     TraceLayer::new_for_http()
//!         .make_span_with(
//!             DefaultMakeSpan::new().with_include_headers(true)
//!         )
//!         .on_request(
//!             DefaultOnRequest::new().with_level(Level::INFO)
//!         )
//!         .on_response(
//!             DefaultOnResponse::new()
//!                 .with_level(Level::INFO)
//!                 .with_latency_unit(LatencyUnit::Micros)
//!         ),
//!         // and so on for `on_eos`, `on_body_chunk`, and `on_failure`
//! ).into_layer(service_fn(handle));
//! # let mut service = service;
//! # let response = service
//! #     .serve(Request::new(Body::from("foo")))
//! #     .await?;
//! # Ok(())
//! # }
//! ```
//!
//! However for maximum control you can provide callbacks:
//!
//! ```rust
//! use rama_http::{Body, Request, Response, HeaderMap, StatusCode};
//! use rama_core::service::service_fn;
//! use rama_core::{Service, Layer};
//! use rama_core::telemetry::tracing::{self, Span};
//! use rama_http::layer::{classify::ServerErrorsFailureClass, trace::TraceLayer};
//! use std::time::Duration;
//! use std::convert::Infallible;
//! use rama_core::bytes::Bytes;
//!
//! # async fn handle(request: Request) -> Result<Response, Infallible> {
//! #     Ok(Response::new(Body::from("foo")))
//! # }
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # tracing_subscriber::fmt::init();
//! #
//! let service = (
//!     TraceLayer::new_for_http()
//!         .make_span_with(|request: &Request| {
//!             tracing::debug_span!("http-request")
//!         })
//!         .on_request(|request: &Request, _span: &Span| {
//!             tracing::debug!(
//!                 http.method = %request.method(),
//!                 url.path = %request.uri().path_or_root().as_ref(),
//!                 "started request",
//!             )
//!         })
//!         .on_response(|response: &Response, latency: Duration, _span: &Span| {
//!             tracing::debug!("response generated in {:?}", latency)
//!         })
//!         .on_body_chunk(|chunk: &Bytes, latency: Duration, _span: &Span| {
//!             tracing::debug!("sending {} bytes", chunk.len())
//!         })
//!         .on_eos(|trailers: Option<&HeaderMap>, stream_duration: Duration, _span: &Span| {
//!             tracing::debug!("stream closed after {:?}", stream_duration)
//!         })
//!         .on_failure(|error: ServerErrorsFailureClass, latency: Duration, _span: &Span| {
//!             tracing::debug!("something went wrong")
//!         })
//! ).into_layer(service_fn(handle));
//! # let mut service = service;
//! # let response = service
//! #     .serve(Request::new(Body::from("foo")))
//! #     .await?;
//! # Ok(())
//! # }
//! ```
//!
//! ## Disabling something
//!
//! Setting the behaviour to `()` will be disable that particular step:
//!
//! ```rust
//! use rama_http::{Body, Request, Response, StatusCode};
//! use rama_core::service::service_fn;
//! use rama_core::{Service, Layer};
//! use rama_core::telemetry::tracing::{self, Span};
//! use rama_http::layer::{classify::ServerErrorsFailureClass, trace::TraceLayer};
//! use std::time::Duration;
//! # use std::convert::Infallible;
//!
//! # async fn handle(request: Request) -> Result<Response, Infallible> {
//! #     Ok(Response::new(Body::from("foo")))
//! # }
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # tracing_subscriber::fmt::init();
//! #
//! let service = (
//!     // This configuration will only emit events on failures
//!     TraceLayer::new_for_http()
//!         .on_request(())
//!         .on_response(())
//!         .on_body_chunk(())
//!         .on_eos(())
//!         .on_failure(|error: ServerErrorsFailureClass, latency: Duration, _span: &Span| {
//!             tracing::debug!("something went wrong")
//!         })
//! ).into_layer(service_fn(handle));
//! # let mut service = service;
//! # let response = service
//! #     .serve(Request::new(Body::from("foo")))
//! #     .await?;
//! # Ok(())
//! # }
//! ```
//!
//! # When the callbacks are called
//!
//! ### `on_request`
//!
//! The `on_request` callback is called when the request arrives at the
//! middleware in [`Service::serve`] just prior to passing the request to the
//! inner service.
//!
//! ### `on_response`
//!
//! The `on_response` callback is called when the inner service's response
//! future completes with `Ok(response)` regardless if the response is
//! classified as a success or a failure.
//!
//! For example if you're using [`ServerErrorsAsFailures`] as your classifier
//! and the inner service responds with `500 Internal Server Error` then the
//! `on_response` callback is still called. `on_failure` would _also_ be called
//! in this case since the response was classified as a failure.
//!
//! ### `on_body_chunk`
//!
//! The `on_body_chunk` callback is called when the response body produces a new
//! chunk, that is when [`crate::StreamingBody::poll_frame`] returns `Poll::Ready(Some(Ok(chunk)))`.
//!
//! `on_body_chunk` is called even if the chunk is empty.
//!
//! ### `on_eos`
//!
//! The `on_eos` callback is called when a streaming response body ends, that is
//! when `http_body::Body::poll_frame` returns `Poll::Ready(None)`.
//!
//! `on_eos` is called even if the trailers produced are `None`.
//!
//! ### `on_failure`
//!
//! The `on_failure` callback is called when:
//!
//! - The inner [`Service`]'s response future resolves to an error.
//! - A response is classified as a failure.
//! - [`crate::StreamingBody::poll_frame`] returns an error.
//! - An end-of-stream is classified as a failure.
//!
//! # Recording fields on the span
//!
//! All callbacks receive a reference to the [tracing] [`Span`], corresponding to this request,
//! produced by the closure passed to [`TraceLayer::make_span_with`]. It can be used to [record
//! field values][record] that weren't known when the span was created.
//!
//! ```rust
//! use rama_http::{Body, Request, Response, HeaderMap, StatusCode};
//! use rama_core::service::service_fn;
//! use rama_core::Layer;
//! use rama_http::layer::trace::TraceLayer;
//! use rama_core::telemetry::tracing::{self, Span};
//! use std::time::Duration;
//! use std::convert::Infallible;
//!
//! # async fn handle(request: Request) -> Result<Response, Infallible> {
//! #     Ok(Response::new(Body::from("foo")))
//! # }
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # tracing_subscriber::fmt::init();
//! #
//! let service = (
//!     TraceLayer::new_for_http()
//!         .make_span_with(|request: &Request| {
//!             tracing::debug_span!(
//!                 "http-request",
//!                 status_code = tracing::field::Empty,
//!             )
//!         })
//!         .on_response(|response: &Response, _latency: Duration, span: &Span| {
//!             span.record("status_code", &tracing::field::display(response.status()));
//!
//!             tracing::debug!("response generated")
//!         }),
//! ).into_layer(service_fn(handle));
//! # Ok(())
//! # }
//! ```
//!
//! # Providing classifiers
//!
//! Tracing requires determining if a response is a success or failure. [`MakeClassifier`] is used
//! to create a classifier for the incoming request. See the docs for [`MakeClassifier`] and
//! [`ClassifyResponse`] for more details on classification.
//!
//! A [`MakeClassifier`] can be provided when creating a [`TraceLayer`]:
//!
//! ```rust
//! use rama_http::{Body, Request, Response};
//! use rama_core::service::service_fn;
//! use rama_core::Layer;
//! use rama_http::layer::{
//!     trace::TraceLayer,
//!     classify::{
//!         MakeClassifier, ClassifyResponse, ClassifiedResponse, NeverClassifyEos,
//!         SharedClassifier,
//!     },
//! };
//! use std::convert::Infallible;
//!
//! # async fn handle(request: Request) -> Result<Response, Infallible> {
//! #     Ok(Response::new(Body::from("foo")))
//! # }
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # tracing_subscriber::fmt::init();
//! #
//! // Our `MakeClassifier` that always crates `MyClassifier` classifiers.
//! #[derive(Copy, Clone)]
//! struct MyMakeClassify;
//!
//! impl MakeClassifier for MyMakeClassify {
//!     type Classifier = MyClassifier;
//!     type FailureClass = &'static str;
//!     type ClassifyEos = NeverClassifyEos<&'static str>;
//!
//!     fn make_classifier<B>(&self, req: &Request<B>) -> Self::Classifier {
//!         MyClassifier
//!     }
//! }
//!
//! // A classifier that classifies failures as `"something went wrong..."`.
//! #[derive(Copy, Clone)]
//! struct MyClassifier;
//!
//! impl ClassifyResponse for MyClassifier {
//!     type FailureClass = &'static str;
//!     type ClassifyEos = NeverClassifyEos<&'static str>;
//!
//!     fn classify_response<B>(
//!         self,
//!         res: &Response<B>
//!     ) -> ClassifiedResponse<Self::FailureClass, Self::ClassifyEos> {
//!         // Classify based on the status code.
//!         if res.status().is_server_error() {
//!             ClassifiedResponse::Ready(Err("something went wrong..."))
//!         } else {
//!             ClassifiedResponse::Ready(Ok(()))
//!         }
//!     }
//!
//!     fn classify_error<E>(self, error: &E) -> Self::FailureClass
//!     where
//!         E: std::fmt::Display,
//!     {
//!         "something went wrong..."
//!     }
//! }
//!
//! let service = (
//!     // Create a trace layer that uses our classifier.
//!     TraceLayer::new(MyMakeClassify),
//! ).into_layer(service_fn(handle));
//!
//! // Since `MyClassifier` is `Clone` we can also use `SharedClassifier`
//! // to avoid having to define a separate `MakeClassifier`.
//! let service = TraceLayer::new(SharedClassifier::new(MyClassifier)).into_layer(service_fn(handle));
//! # Ok(())
//! # }
//! ```
//!
//! [`TraceLayer`] comes with convenience methods for using common classifiers:
//!
//! - [`TraceLayer::new_for_http`] classifies based on the status code. It doesn't consider
//!   streaming responses.
//! - [`TraceLayer::new_for_grpc`] classifies based on the gRPC protocol and supports streaming
//!   responses.
//!
//! [tracing]: https://crates.io/crates/tracing
//! [`Service`]: rama_core::Service
//! [`Service::serve`]: rama_core::Service::serve
//! [`MakeClassifier`]: crate::layer::classify::MakeClassifier
//! [`ClassifyResponse`]: crate::layer::classify::ClassifyResponse
//! [record]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.record
//! [`TraceLayer::make_span_with`]: crate::layer::trace::TraceLayer::make_span_with
//! [`Span`]: rama_core::telemetry::tracing::Span
//! [`ServerErrorsAsFailures`]: crate::layer::classify::ServerErrorsAsFailures

use rama_core::telemetry::tracing::Level;
use std::{fmt, time::Duration};

#[doc(inline)]
pub use self::{
    body::ResponseBody,
    layer::TraceLayer,
    make_span::{DefaultMakeSpan, MakeSpan},
    on_body_chunk::{DefaultOnBodyChunk, OnBodyChunk},
    on_eos::{DefaultOnEos, OnEos},
    on_failure::{DefaultOnFailure, OnFailure},
    on_request::{DefaultOnRequest, OnRequest},
    on_response::{DefaultOnResponse, OnResponse},
    service::Trace,
};

use crate::layer::classify::{GrpcErrorsAsFailures, ServerErrorsAsFailures, SharedClassifier};
use rama_utils::latency::LatencyUnit;

/// MakeClassifier for HTTP requests.
pub type HttpMakeClassifier = SharedClassifier<ServerErrorsAsFailures>;

/// MakeClassifier for gRPC requests.
pub type GrpcMakeClassifier = SharedClassifier<GrpcErrorsAsFailures>;

macro_rules! event_dynamic_lvl {
    ( target: $target:expr, parent: $parent:expr, $lvl:expr, $($tt:tt)* ) => {
        {
            use ::rama_core::telemetry::tracing;
            match $lvl {
                tracing::Level::ERROR => {
                    tracing::event!(target: $target, parent: $parent, tracing::Level::ERROR, $($tt)*);
                }
                tracing::Level::WARN => {
                    tracing::event!(target: $target, parent: $parent, tracing::Level::WARN, $($tt)*);
                }
                tracing::Level::INFO => {
                    tracing::event!(target: $target, parent: $parent, tracing::Level::INFO, $($tt)*);
                }
                tracing::Level::DEBUG => {
                    tracing::event!(target: $target, parent: $parent, tracing::Level::INFO, $($tt)*);
                }
                tracing::Level::TRACE => {
                    tracing::event!(target: $target, parent: $parent, tracing::Level::TRACE, $($tt)*);
                }
            }
        }
    };
    ( target: $target:expr, $lvl:expr, $($tt:tt)* ) => {
        use ::rama_core::telemetry::tracing;
        match $lvl {
            tracing::Level::ERROR => {
                tracing::event!(target: $target, tracing::Level::ERROR, $($tt)*);
            }
            tracing::Level::WARN => {
                tracing::event!(target: $target, tracing::Level::WARN, $($tt)*);
            }
            tracing::Level::INFO => {
                tracing::event!(target: $target, tracing::Level::INFO, $($tt)*);
            }
            tracing::Level::DEBUG => {
                tracing::event!(target: $target, tracing::Level::DEBUG, $($tt)*);
            }
            tracing::Level::TRACE => {
                tracing::event!(target: $target, tracing::Level::TRACE, $($tt)*);
            }
        }
    };
    ( parent: $parent:expr, $lvl:expr, $($tt:tt)* ) => {
        use ::rama_core::telemetry::tracing;
        match $lvl {
            tracing::Level::ERROR => {
                tracing::event!(parent: $parent, tracing::Level::ERROR, $($tt)*);
            }
            tracing::Level::WARN => {
                tracing::event!(parent: $parent, tracing::Level::WARN, $($tt)*);
            }
            tracing::Level::INFO => {
                tracing::event!(parent: $parent, tracing::Level::INFO, $($tt)*);
            }
            tracing::Level::DEBUG => {
                tracing::event!(parent: $parent, tracing::Level::DEBUG, $($tt)*);
            }
            tracing::Level::TRACE => {
                tracing::event!(parent: $parent, tracing::Level::TRACE, $($tt)*);
            }
        }
    };
    ( $lvl:expr, $($tt:tt)* ) => {
        {
            use ::rama_core::telemetry::tracing;
            match $lvl {
                tracing::Level::ERROR => {
                    tracing::event!(tracing::Level::ERROR, $($tt)*);
                }
                tracing::Level::WARN => {
                    tracing::event!(tracing::Level::WARN, $($tt)*);
                }
                tracing::Level::INFO => {
                    tracing::event!(tracing::Level::INFO, $($tt)*);
                }
                tracing::Level::DEBUG => {
                    tracing::event!(tracing::Level::DEBUG, $($tt)*);
                }
                tracing::Level::TRACE => {
                    tracing::event!(tracing::Level::TRACE, $($tt)*);
                }
            }
        }
    };
}

mod body;
mod layer;
mod make_span;
mod on_body_chunk;
mod on_eos;
mod on_failure;
mod on_request;
mod on_response;
mod service;

const DEFAULT_MESSAGE_LEVEL: Level = Level::DEBUG;
const DEFAULT_ERROR_LEVEL: Level = Level::ERROR;

struct Latency {
    unit: LatencyUnit,
    duration: Duration,
}

impl fmt::Display for Latency {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.unit {
            LatencyUnit::Seconds => write!(f, "{} s", self.duration.as_secs_f64()),
            LatencyUnit::Millis => write!(f, "{} ms", self.duration.as_millis()),
            LatencyUnit::Micros => write!(f, "{} μs", self.duration.as_micros()),
            LatencyUnit::Nanos => write!(f, "{} ns", self.duration.as_nanos()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::body::util::BodyExt as _;
    use crate::layer::classify::{
        ClassifiedResponse, ClassifyEos, ClassifyResponse, MakeClassifier, ServerErrorsFailureClass,
    };
    use crate::{Body, HeaderMap, Request, Response};
    use rama_core::bytes::Bytes;
    use rama_core::error::BoxError;
    use rama_core::service::service_fn;
    use rama_core::telemetry::tracing::{self, Span};
    use rama_core::{Layer, Service};
    use std::sync::OnceLock;
    use std::{
        sync::atomic::{AtomicU32, Ordering},
        time::Duration,
    };

    macro_rules! lazy_atomic_u32 {
        ($($name:ident),+) => {
            $(
                #[expect(non_snake_case)]
                fn $name() -> &'static AtomicU32 {
                    static $name: OnceLock<AtomicU32> = OnceLock::new();
                    $name.get_or_init(|| AtomicU32::new(0))
                }
            )+
        };
    }

    #[tokio::test]
    async fn unary_request() {
        lazy_atomic_u32!(
            ON_REQUEST_COUNT,
            ON_RESPONSE_COUNT,
            ON_BODY_CHUNK_COUNT,
            ON_EOS,
            ON_FAILURE
        );

        let trace_layer = TraceLayer::new_for_http()
            .make_span_with(|_req: &Request| {
                tracing::info_span!("test-span", foo = tracing::field::Empty)
            })
            .on_request(|_req: &Request, span: &Span| {
                span.record("foo", 42);
                ON_REQUEST_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_response(|_res: &Response, _latency: Duration, _span: &Span| {
                ON_RESPONSE_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {
                ON_BODY_CHUNK_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::AcqRel);
                },
            )
            .on_failure(
                |_class: ServerErrorsFailureClass, _latency: Duration, _span: &Span| {
                    ON_FAILURE().fetch_add(1, Ordering::AcqRel);
                },
            );

        let svc = trace_layer.into_layer(service_fn(echo));

        let res = svc.serve(Request::new(Body::from("foobar"))).await.unwrap();

        assert_eq!(1, ON_REQUEST_COUNT().load(Ordering::Acquire), "request");
        assert_eq!(1, ON_RESPONSE_COUNT().load(Ordering::Acquire), "request");
        assert_eq!(
            0,
            ON_BODY_CHUNK_COUNT().load(Ordering::Acquire),
            "body chunk"
        );
        assert_eq!(0, ON_EOS().load(Ordering::Acquire), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::Acquire), "failure");

        res.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(
            1,
            ON_BODY_CHUNK_COUNT().load(Ordering::Acquire),
            "body chunk"
        );
        assert_eq!(1, ON_EOS().load(Ordering::Acquire), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::Acquire), "failure");
    }

    #[tokio::test]
    async fn streaming_response() {
        lazy_atomic_u32!(
            ON_REQUEST_COUNT,
            ON_RESPONSE_COUNT,
            ON_BODY_CHUNK_COUNT,
            ON_EOS,
            ON_FAILURE
        );

        let trace_layer = TraceLayer::new_for_http()
            .on_request(|_req: &Request, _span: &Span| {
                ON_REQUEST_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_response(|_res: &Response, _latency: Duration, _span: &Span| {
                ON_RESPONSE_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {
                ON_BODY_CHUNK_COUNT().fetch_add(1, Ordering::AcqRel);
            })
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::AcqRel);
                },
            )
            .on_failure(
                |_class: ServerErrorsFailureClass, _latency: Duration, _span: &Span| {
                    ON_FAILURE().fetch_add(1, Ordering::AcqRel);
                },
            );

        let svc = trace_layer.into_layer(service_fn(streaming_body));

        let res = svc.serve(Request::new(Body::empty())).await.unwrap();

        assert_eq!(1, ON_REQUEST_COUNT().load(Ordering::Acquire), "request");
        assert_eq!(1, ON_RESPONSE_COUNT().load(Ordering::Acquire), "request");
        assert_eq!(
            0,
            ON_BODY_CHUNK_COUNT().load(Ordering::Acquire),
            "body chunk"
        );
        assert_eq!(0, ON_EOS().load(Ordering::Acquire), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::Acquire), "failure");

        res.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(
            3,
            ON_BODY_CHUNK_COUNT().load(Ordering::Acquire),
            "body chunk"
        );
        assert_eq!(1, ON_EOS().load(Ordering::Acquire), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::Acquire), "failure");
    }

    #[tokio::test]
    async fn classify_eos_on_trailers_success() {
        lazy_atomic_u32!(ON_EOS, ON_FAILURE);

        let trace_layer = TraceLayer::new(TestClassify::new(false))
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::SeqCst);
                },
            )
            .on_failure(|_class: &'static str, _latency: Duration, _span: &Span| {
                ON_FAILURE().fetch_add(1, Ordering::SeqCst);
            });

        let svc = trace_layer.into_layer(service_fn(body_with_trailers));

        let res = svc.serve(Request::new(Body::empty())).await.unwrap();

        res.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::SeqCst), "failure");
    }

    #[tokio::test]
    async fn classify_eos_on_trailers_failure() {
        lazy_atomic_u32!(ON_EOS, ON_FAILURE);

        let trace_layer = TraceLayer::new(TestClassify::new(true))
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::SeqCst);
                },
            )
            .on_failure(|_class: &'static str, _latency: Duration, _span: &Span| {
                ON_FAILURE().fetch_add(1, Ordering::SeqCst);
            });

        let svc = trace_layer.into_layer(service_fn(body_with_trailers));

        let res = svc.serve(Request::new(Body::empty())).await.unwrap();

        res.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
        assert_eq!(1, ON_FAILURE().load(Ordering::SeqCst), "failure");
    }

    #[tokio::test]
    async fn classify_eos_on_empty_stream() {
        lazy_atomic_u32!(ON_EOS, ON_FAILURE);

        let trace_layer = TraceLayer::new(TestClassify::new(true))
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::SeqCst);
                },
            )
            .on_failure(|_class: &'static str, _latency: Duration, _span: &Span| {
                ON_FAILURE().fetch_add(1, Ordering::SeqCst);
            });

        let svc = trace_layer.into_layer(service_fn(streaming_body));

        let res = svc.serve(Request::new(Body::empty())).await.unwrap();

        res.into_body().collect().await.unwrap().to_bytes();
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
        assert_eq!(0, ON_FAILURE().load(Ordering::SeqCst), "failure");
    }

    #[tokio::test]
    async fn on_eos_fires_for_content_length_body() {
        // Simulates the scenario where a consumer stops polling after receiving
        // all bytes (as hyper does when Content-Length is exact). We poll only
        // the data frame and never poll to None.

        lazy_atomic_u32!(ON_BODY_CHUNK_COUNT, ON_EOS);

        let trace_layer = TraceLayer::new_for_http()
            .on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {
                ON_BODY_CHUNK_COUNT().fetch_add(1, Ordering::SeqCst);
            })
            .on_eos(
                |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                    ON_EOS().fetch_add(1, Ordering::SeqCst);
                },
            );

        let svc = trace_layer.into_layer(service_fn(echo));

        let res = svc.serve(Request::new(Body::from("hello"))).await.unwrap();

        let mut body = res.into_body();

        // Poll only the data frame (simulating a content-length aware consumer)
        let frame = body.frame().await.unwrap().unwrap();
        assert!(frame.data_ref().is_some());

        // on_eos should have fired immediately after the data frame since
        // is_end_stream() is true for Full bodies after yielding their data.
        assert_eq!(
            1,
            ON_BODY_CHUNK_COUNT().load(Ordering::SeqCst),
            "body chunk"
        );
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
    }

    #[tokio::test]
    async fn on_eos_fires_for_streaming_body_on_none() {
        // Streaming bodies (no content-length) don't report is_end_stream()
        // until polled to None. Verify on_eos still fires via the None path.
        lazy_atomic_u32!(ON_EOS);

        let trace_layer = TraceLayer::new_for_http().on_eos(
            |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                ON_EOS().fetch_add(1, Ordering::SeqCst);
            },
        );

        let svc = trace_layer.into_layer(service_fn(streaming_body));

        let res = svc.serve(Request::new(Body::empty())).await.unwrap();

        res.into_body().collect().await.unwrap();
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
    }

    #[tokio::test]
    async fn on_eos_not_called_twice() {
        // When is_end_stream() fires on_eos after a data frame, a subsequent
        // poll returning None must not fire on_eos again.
        lazy_atomic_u32!(ON_EOS);

        let trace_layer = TraceLayer::new_for_http().on_eos(
            |_trailers: Option<&HeaderMap>, _latency: Duration, _span: &Span| {
                ON_EOS().fetch_add(1, Ordering::SeqCst);
            },
        );

        let svc = trace_layer.into_layer(service_fn(echo));

        let res = svc.serve(Request::new(Body::from("hello"))).await.unwrap();

        // Consume the body fully (polls data frame then None)
        res.into_body().collect().await.unwrap();

        // on_eos should fire exactly once, not twice
        assert_eq!(1, ON_EOS().load(Ordering::SeqCst), "eos");
    }

    async fn echo(req: Request) -> Result<Response, BoxError> {
        Ok(Response::new(req.into_body()))
    }

    async fn streaming_body(_req: Request) -> Result<Response, BoxError> {
        use rama_core::futures::stream::iter;

        let stream = iter(vec![
            Ok::<_, BoxError>(Bytes::from("one")),
            Ok::<_, BoxError>(Bytes::from("two")),
            Ok::<_, BoxError>(Bytes::from("three")),
        ]);

        let body = Body::from_stream(stream);

        Ok(Response::new(body))
    }

    async fn body_with_trailers(_req: Request<Body>) -> Result<Response<Body>, BoxError> {
        let mut trailers = HeaderMap::new();
        trailers.insert("x-test-trailer", "value".parse().unwrap());
        let body = Body::from(Bytes::from("data")).with_trailer_headers(trailers);
        Ok(Response::new(body))
    }

    #[derive(Clone)]
    struct TestClassify {
        reject: bool,
    }

    impl TestClassify {
        fn new(reject: bool) -> Self {
            Self { reject }
        }
    }

    impl MakeClassifier for TestClassify {
        type FailureClass = &'static str;
        type ClassifyEos = TestClassifyEos;
        type Classifier = TestClassifyResponse;

        fn make_classifier<B>(&self, _req: &Request<B>) -> Self::Classifier {
            TestClassifyResponse {
                reject: self.reject,
            }
        }
    }

    #[derive(Clone)]
    struct TestClassifyResponse {
        reject: bool,
    }

    impl ClassifyResponse for TestClassifyResponse {
        type FailureClass = &'static str;
        type ClassifyEos = TestClassifyEos;

        fn classify_response<B>(
            self,
            _res: &Response<B>,
        ) -> ClassifiedResponse<Self::FailureClass, Self::ClassifyEos> {
            ClassifiedResponse::RequiresEos(TestClassifyEos {
                reject: self.reject,
            })
        }

        fn classify_error<E>(self, _error: &E) -> Self::FailureClass
        where
            E: std::fmt::Display,
        {
            "error"
        }
    }

    #[derive(Clone)]
    struct TestClassifyEos {
        reject: bool,
    }

    impl ClassifyEos for TestClassifyEos {
        type FailureClass = &'static str;

        fn classify_eos(self, _trailers: Option<&HeaderMap>) -> Result<(), Self::FailureClass> {
            if self.reject {
                Err("classified as failure")
            } else {
                Ok(())
            }
        }

        fn classify_error<E>(self, _error: &E) -> Self::FailureClass
        where
            E: std::fmt::Display,
        {
            "error"
        }
    }

    /// Regression test for https://github.com/tower-rs/tower-http/issues/655
    ///
    /// Reproduces the reported bug: when a subscriber's filter disables the
    /// request span but still enables events, the events appear without any
    /// span context. This happens because `Span::enter()` on a disabled span
    /// is a no-op, so events relying on ambient context have no parent.
    ///
    /// The fix (using explicit `parent: span`) ensures events always reference
    /// the request span, even when it's disabled. A subscriber that records
    /// disabled spans will still see the correct parent relationship.
    #[test]
    fn events_have_span_context_when_span_is_disabled() {
        use parking_lot::Mutex;
        use std::sync::Arc;
        use tracing::subscriber::with_default;
        use tracing_subscriber::{Layer as _, layer::SubscriberExt, registry::LookupSpan};

        /// A filter that disables spans (by rejecting at the span level)
        /// but allows all events through. This simulates the scenario where
        /// a per-layer EnvFilter disables the request span's callsite.
        struct DisableSpansFilter;

        impl<S: tracing::Subscriber> tracing_subscriber::layer::Filter<S> for DisableSpansFilter {
            fn enabled(
                &self,
                meta: &tracing::Metadata<'_>,
                _cx: &tracing_subscriber::layer::Context<'_, S>,
            ) -> bool {
                // Disable spans, keep events
                !meta.is_span()
            }
        }

        /// Records (event_message, has_any_parent) pairs.
        #[derive(Clone)]
        struct RecordingLayer {
            events: Arc<Mutex<Vec<(String, bool)>>>,
        }

        impl<S> tracing_subscriber::Layer<S> for RecordingLayer
        where
            S: tracing::Subscriber + for<'a> LookupSpan<'a>,
        {
            fn on_event(
                &self,
                event: &tracing::Event<'_>,
                ctx: tracing_subscriber::layer::Context<'_, S>,
            ) {
                let mut msg = String::new();
                event.record(&mut MessageVisitor(&mut msg));

                // Check if the event has ANY parent: explicit or contextual
                let has_parent = event.parent().is_some() || ctx.event_span(event).is_some();

                self.events.lock().push((msg, has_parent));
            }
        }

        struct MessageVisitor<'a>(&'a mut String);
        impl tracing::field::Visit for MessageVisitor<'_> {
            fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
                if field.name() == "message" {
                    *self.0 = format!("{value:?}");
                }
            }
        }

        let events = Arc::new(Mutex::new(Vec::new()));
        let layer = RecordingLayer {
            events: events.clone(),
        };
        let subscriber = tracing_subscriber::registry().with(layer.with_filter(DisableSpansFilter));

        // Use with_default to guarantee cleanup even on panic, avoiding
        // cross-test subscriber pollution.
        with_default(subscriber, || {
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap();

            rt.block_on(async {
                let svc = TraceLayer::new_for_http().into_layer(service_fn(echo));

                let res = svc.serve(Request::new(Body::from("test"))).await.unwrap();

                res.into_body().collect().await.unwrap();
            });
        });

        let events = events.lock();
        let request_events: Vec<_> = events
            .iter()
            .filter(|(msg, _)| {
                msg.contains("started processing request")
                    || msg.contains("finished processing request")
            })
            .collect();

        assert!(
            request_events.len() >= 2,
            "expected on_request and on_response events to fire"
        );

        // The bug: without explicit parent, these events have no span context
        // at all when the request span is disabled. With the fix, they still
        // reference the span (even though it's disabled).
        for (msg, has_parent) in &request_events {
            assert!(
                *has_parent,
                "event {msg:?} has no span context. When the request span is \
                     disabled by a filter, events must still reference it via \
                     explicit parent so subscribers can associate them correctly."
            );
        }
    }
}