rama_http/layer/trace/
body.rs

1use super::{DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, OnBodyChunk, OnEos, OnFailure};
2use crate::body::{Frame, SizeHint, StreamingBody};
3use crate::layer::classify::ClassifyEos;
4use pin_project_lite::pin_project;
5use rama_core::futures::ready;
6use rama_core::telemetry::tracing::Span;
7use std::{
8    fmt,
9    pin::Pin,
10    task::{Context, Poll},
11    time::Instant,
12};
13
14pin_project! {
15    /// Response body for [`Trace`].
16    ///
17    /// [`Trace`]: super::Trace
18    pub struct ResponseBody<B, C, OnBodyChunk = DefaultOnBodyChunk, OnEos = DefaultOnEos, OnFailure = DefaultOnFailure> {
19        #[pin]
20        pub(crate) inner: B,
21        pub(crate) classify_eos: Option<C>,
22        pub(crate) on_eos: Option<(OnEos, Instant)>,
23        pub(crate) on_body_chunk: OnBodyChunk,
24        pub(crate) on_failure: Option<OnFailure>,
25        pub(crate) start: Instant,
26        pub(crate) span: Span,
27    }
28}
29
30impl<B, C, OnBodyChunkT, OnEosT, OnFailureT> StreamingBody
31    for ResponseBody<B, C, OnBodyChunkT, OnEosT, OnFailureT>
32where
33    B: StreamingBody<Error: fmt::Display>,
34    C: ClassifyEos,
35    OnEosT: OnEos,
36    OnBodyChunkT: OnBodyChunk<B::Data>,
37    OnFailureT: OnFailure<C::FailureClass>,
38{
39    type Data = B::Data;
40    type Error = B::Error;
41
42    fn poll_frame(
43        self: Pin<&mut Self>,
44        cx: &mut Context<'_>,
45    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
46        let mut this = self.project();
47        let _guard = this.span.enter();
48        let result = ready!(this.inner.as_mut().poll_frame(cx));
49
50        let latency = this.start.elapsed();
51        *this.start = Instant::now();
52
53        match result {
54            Some(Ok(frame)) => {
55                let frame = match frame.into_data() {
56                    Ok(chunk) => {
57                        this.on_body_chunk.on_body_chunk(&chunk, latency, this.span);
58                        Frame::data(chunk)
59                    }
60                    Err(frame) => frame,
61                };
62
63                let frame = match frame.into_trailers() {
64                    Ok(trailers) => {
65                        if let Some((classify_eos, on_failure)) =
66                            this.classify_eos.take().zip(this.on_failure.take())
67                            && let Err(failure_class) = classify_eos.classify_eos(Some(&trailers))
68                        {
69                            on_failure.on_failure(failure_class, latency, this.span);
70                        }
71                        if let Some((on_eos, stream_start)) = this.on_eos.take() {
72                            on_eos.on_eos(Some(&trailers), stream_start.elapsed(), this.span);
73                        }
74                        Frame::trailers(trailers)
75                    }
76                    Err(frame) => frame,
77                };
78
79                // If the inner body signals end-of-stream after this frame,
80                // fire on_eos now since the consumer may not poll again (e.g.
81                // when Content-Length is exact).
82                if this.inner.is_end_stream() {
83                    if let Some((classify_eos, on_failure)) =
84                        this.classify_eos.take().zip(this.on_failure.take())
85                        && let Err(failure_class) = classify_eos.classify_eos(None)
86                    {
87                        on_failure.on_failure(failure_class, latency, this.span);
88                    }
89                    if let Some((on_eos, stream_start)) = this.on_eos.take() {
90                        on_eos.on_eos(None, stream_start.elapsed(), this.span);
91                    }
92                }
93
94                Poll::Ready(Some(Ok(frame)))
95            }
96            Some(Err(err)) => {
97                if let Some((classify_eos, on_failure)) =
98                    this.classify_eos.take().zip(this.on_failure.take())
99                {
100                    let failure_class = classify_eos.classify_error(&err);
101                    on_failure.on_failure(failure_class, latency, this.span);
102                }
103
104                Poll::Ready(Some(Err(err)))
105            }
106            None => {
107                if let Some((on_eos, stream_start)) = this.on_eos.take() {
108                    on_eos.on_eos(None, stream_start.elapsed(), this.span);
109                }
110
111                Poll::Ready(None)
112            }
113        }
114    }
115
116    fn is_end_stream(&self) -> bool {
117        self.inner.is_end_stream()
118    }
119
120    fn size_hint(&self) -> SizeHint {
121        self.inner.size_hint()
122    }
123}
rama_http/layer/trace/body.rs

rama_http/layer/trace/
body.rs
