rama_http/layer/traffic_writer/
request.rs

1use super::{
2    PerMessageFileWriter, WriterMode, capture_body_channel, ensure_traffic_writer_id,
3    write_headers_body_flags,
4};
5use crate::io::write_http_request_streaming;
6use crate::{Body, Request, StreamingBody, body::util::BodyExt as _};
7use rama_core::bytes::Bytes;
8use rama_core::error::{BoxError, ErrorContext as _};
9use rama_core::extensions::{Extension, ExtensionsRef};
10use rama_core::rt::Executor;
11use rama_core::telemetry::tracing::{self, Instrument};
12use rama_core::{Layer, Service};
13use rama_http_types::CaptureBody;
14use std::{fmt::Debug, io, path::PathBuf, sync::Arc};
15use tokio::io::{AsyncWrite, AsyncWriteExt, stderr, stdout};
16use tokio::sync::mpsc::{Sender, UnboundedSender, channel, unbounded_channel};
17
18/// Write a single request entry (request followed by a `\r\n` separator) to the writer.
19async fn write_request_entry<W>(writer: &mut W, req: Request, write_headers: bool, write_body: bool)
20where
21    W: AsyncWrite + Unpin + Send + Sync + 'static,
22{
23    if let Err(err) = write_http_request_streaming(writer, req, write_headers, write_body).await {
24        tracing::error!("failed to write http request to writer: {err:?}")
25    }
26    if let Err(err) = writer.write_all(b"\r\n").await {
27        tracing::error!("failed to write separator to writer: {err:?}")
28    }
29}
30
31macro_rules! drive_request_writer {
32    ($writer:ident, $rx:ident, $write_headers:ident, $write_body:ident) => {{
33        while let Some(req) = $rx.recv().await {
34            write_request_entry(&mut $writer, req, $write_headers, $write_body).await;
35        }
36        if let Err(err) = $writer.flush().await {
37            tracing::error!("failed to flush request writer: {err:?}")
38        }
39    }};
40}
41
42/// A trait for writing http requests.
43pub trait RequestWriter: Send + Sync + 'static {
44    /// Write the HTTP request while its body is streaming.
45    ///
46    /// Implementations must consume or deliberately drop the body while this
47    /// future runs. Retaining it for later applies backpressure once its
48    /// single-frame observational capture queue is full.
49    fn write_request(&self, req: Request) -> impl Future<Output = ()> + Send + '_;
50}
51
52/// Marker struct to indicate that the request should not be printed.
53#[derive(Debug, Clone, Default, Extension)]
54#[extension(tags(http))]
55#[non_exhaustive]
56pub struct DoNotWriteRequest;
57
58impl DoNotWriteRequest {
59    /// Create a new [`DoNotWriteRequest`] marker.
60    #[must_use]
61    pub const fn new() -> Self {
62        Self
63    }
64}
65
66#[derive(Clone)]
67/// Middleware to print Http request in std format.
68///
69/// Request bodies are observed while they stream to the inner service. The
70/// writer receives an owned body stream and processes frames concurrently.
71///
72/// See the [module docs](super) for more details.
73pub struct RequestWriterService<S, W> {
74    inner: S,
75    writer: Arc<W>,
76    executor: Executor,
77}
78
79impl<S, W> RequestWriterService<S, W> {
80    /// Create a new [`RequestWriterService`] with a custom [`RequestWriter`].
81    ///
82    /// Use [`Self::new_with_executor`] when streaming writer tasks must
83    /// participate in graceful shutdown.
84    pub fn new(inner: S, writer: W) -> Self {
85        Self::new_with_executor(inner, writer, Executor::new())
86    }
87
88    /// Create a new [`RequestWriterService`] with a custom [`RequestWriter`]
89    /// and executor for streaming writer tasks.
90    pub fn new_with_executor(inner: S, writer: W, executor: Executor) -> Self {
91        Self {
92            inner,
93            writer: Arc::new(writer),
94            executor,
95        }
96    }
97}
98
99impl<S: Debug, W> Debug for RequestWriterService<S, W> {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        f.debug_struct("RequestWriterService")
102            .field("inner", &self.inner)
103            .field("writer", &format_args!("{}", std::any::type_name::<W>()))
104            .finish()
105    }
106}
107
108impl<S> RequestWriterService<S, UnboundedSender<Request>> {
109    /// Create a new [`RequestWriterService`] that prints requests to an [`AsyncWrite`]r
110    /// over an unbounded channel
111    pub fn writer_unbounded<W>(
112        inner: S,
113        executor: &Executor,
114        mut writer: W,
115        mode: Option<WriterMode>,
116    ) -> Self
117    where
118        W: AsyncWrite + Unpin + Send + Sync + 'static,
119    {
120        let (tx, mut rx) = unbounded_channel();
121        let (write_headers, write_body) = write_headers_body_flags(mode);
122
123        let span =
124            tracing::trace_root_span!("TrafficWriter::request::unbounded", otel.kind = "consumer");
125
126        executor.spawn_task(
127            async move {
128                drive_request_writer!(writer, rx, write_headers, write_body);
129            }
130            .instrument(span),
131        );
132        Self {
133            writer: Arc::new(tx),
134            inner,
135            executor: executor.clone(),
136        }
137    }
138
139    /// Create a new [`RequestWriterService`] that prints requests to stdout
140    /// over an unbounded channel.
141    #[must_use]
142    pub fn stdout_unbounded(inner: S, executor: &Executor, mode: Option<WriterMode>) -> Self {
143        Self::writer_unbounded(inner, executor, stdout(), mode)
144    }
145
146    /// Create a new [`RequestWriterService`] that prints requests to stderr
147    /// over an unbounded channel.
148    #[must_use]
149    pub fn stderr_unbounded(inner: S, executor: &Executor, mode: Option<WriterMode>) -> Self {
150        Self::writer_unbounded(inner, executor, stderr(), mode)
151    }
152}
153
154impl<S> RequestWriterService<S, PerMessageFileWriter> {
155    /// Create a service that streams every request into a unique file below
156    /// `directory` using the portable filename `prefix`.
157    pub async fn file_per_request(
158        inner: S,
159        executor: &Executor,
160        directory: impl Into<PathBuf>,
161        prefix: impl AsRef<str>,
162        mode: Option<WriterMode>,
163    ) -> io::Result<Self> {
164        let writer = PerMessageFileWriter::try_new(directory, prefix)
165            .await?
166            .with_request_mode(mode)
167            .with_response_mode(None);
168        Ok(Self::new_with_executor(inner, writer, executor.clone()))
169    }
170}
171
172impl<S> RequestWriterService<S, Sender<Request>> {
173    /// Create a new [`RequestWriterService`] that prints requests to an [`AsyncWrite`]r
174    /// over a bounded channel with a fixed buffer size.
175    pub fn writer<W>(
176        inner: S,
177        executor: &Executor,
178        mut writer: W,
179        buffer_size: usize,
180        mode: Option<WriterMode>,
181    ) -> Self
182    where
183        W: AsyncWrite + Unpin + Send + Sync + 'static,
184    {
185        let (tx, mut rx) = channel(buffer_size);
186        let (write_headers, write_body) = write_headers_body_flags(mode);
187
188        let span =
189            tracing::trace_root_span!("TrafficWriter::request::bounded", otel.kind = "consumer");
190
191        executor.spawn_task(
192            async move {
193                drive_request_writer!(writer, rx, write_headers, write_body);
194            }
195            .instrument(span),
196        );
197        Self {
198            writer: Arc::new(tx),
199            inner,
200            executor: executor.clone(),
201        }
202    }
203
204    /// Create a new [`RequestWriterService`] that prints requests to stdout
205    /// over a bounded channel with a fixed buffer size.
206    #[must_use]
207    pub fn stdout(
208        inner: S,
209        executor: &Executor,
210        buffer_size: usize,
211        mode: Option<WriterMode>,
212    ) -> Self {
213        Self::writer(inner, executor, stdout(), buffer_size, mode)
214    }
215
216    /// Create a new [`RequestWriterService`] that prints requests to stderr
217    /// over a bounded channel with a fixed buffer size.
218    #[must_use]
219    pub fn stderr(
220        inner: S,
221        executor: &Executor,
222        buffer_size: usize,
223        mode: Option<WriterMode>,
224    ) -> Self {
225        Self::writer(inner, executor, stderr(), buffer_size, mode)
226    }
227}
228
229impl<S, W, ReqBody> Service<Request<ReqBody>> for RequestWriterService<S, W>
230where
231    S: Service<Request, Error: Into<BoxError>>,
232    ReqBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
233    W: RequestWriter,
234{
235    type Error = BoxError;
236    type Output = S::Output;
237
238    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
239        if req.extensions().get_ref::<DoNotWriteRequest>().is_some() {
240            self.inner.serve(req.map(Body::new)).await.into_box_error()
241        } else {
242            ensure_traffic_writer_id(req.extensions());
243            let (parts, body) = req.into_parts();
244            let (capture, captured_body) = capture_body_channel();
245            let captured_parts = parts.clone();
246            let captured_request = Request::from_parts(captured_parts, captured_body);
247            let request = Request::from_parts(
248                parts,
249                Body::new(CaptureBody::new(body.map_err(Into::into), capture)),
250            );
251            let writer = Arc::clone(&self.writer);
252            self.executor.spawn_task(async move {
253                writer.write_request(captured_request).await;
254            });
255            self.inner.serve(request).await.into_box_error()
256        }
257    }
258}
259
260impl RequestWriter for Sender<Request> {
261    async fn write_request(&self, req: Request) {
262        if let Err(err) = self.send(req).await {
263            tracing::error!("failed to send request to channel: {err:?}")
264        }
265    }
266}
267
268impl RequestWriter for UnboundedSender<Request> {
269    async fn write_request(&self, req: Request) {
270        if let Err(err) = self.send(req) {
271            tracing::error!("failed to send request to unbounded channel: {err:?}")
272        }
273    }
274}
275
276impl<F, Fut> RequestWriter for F
277where
278    F: Fn(Request) -> Fut + Send + Sync + 'static,
279    Fut: Future<Output = ()> + Send + 'static,
280{
281    async fn write_request(&self, req: Request) {
282        self(req).await
283    }
284}
285
286#[derive(Clone)]
287/// Middleware to print Http request in std format.
288///
289/// Request bodies are observed while they stream to the inner service. The
290/// writer receives an owned body stream and processes frames concurrently.
291///
292/// See the [module docs](super) for more details.
293pub struct RequestWriterLayer<W> {
294    writer: W,
295    executor: Executor,
296}
297
298impl<W> RequestWriterLayer<W> {
299    /// Create a new [`RequestWriterLayer`] with a custom [`RequestWriter`].
300    ///
301    /// Use [`Self::new_with_executor`] when streaming writer tasks must
302    /// participate in graceful shutdown.
303    pub const fn new(writer: W) -> Self {
304        Self::new_with_executor(writer, Executor::new())
305    }
306
307    /// Create a new [`RequestWriterLayer`] with a custom [`RequestWriter`] and
308    /// executor for streaming writer tasks.
309    pub const fn new_with_executor(writer: W, executor: Executor) -> Self {
310        Self { writer, executor }
311    }
312}
313
314impl<W> Debug for RequestWriterLayer<W> {
315    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316        f.debug_struct("RequestWriterLayer")
317            .field("writer", &format_args!("{}", std::any::type_name::<W>()))
318            .finish()
319    }
320}
321
322impl RequestWriterLayer<UnboundedSender<Request>> {
323    /// Create a new [`RequestWriterLayer`] that prints requests to an [`AsyncWrite`]r
324    /// over an unbounded channel
325    pub fn writer_unbounded<W>(executor: &Executor, mut writer: W, mode: Option<WriterMode>) -> Self
326    where
327        W: AsyncWrite + Unpin + Send + Sync + 'static,
328    {
329        let (tx, mut rx) = unbounded_channel();
330        let (write_headers, write_body) = write_headers_body_flags(mode);
331
332        let span =
333            tracing::trace_root_span!("TrafficWriter::request::unbounded", otel.kind = "consumer");
334
335        executor.spawn_task(
336            async move {
337                drive_request_writer!(writer, rx, write_headers, write_body);
338            }
339            .instrument(span),
340        );
341        Self {
342            writer: tx,
343            executor: executor.clone(),
344        }
345    }
346
347    /// Create a new [`RequestWriterService`] that prints requests to stdout
348    /// over an unbounded channel.
349    #[must_use]
350    pub fn stdout_unbounded(executor: &Executor, mode: Option<WriterMode>) -> Self {
351        Self::writer_unbounded(executor, stdout(), mode)
352    }
353
354    /// Create a new [`RequestWriterService`] that prints requests to stderr
355    /// over an unbounded channel.
356    #[must_use]
357    pub fn stderr_unbounded(executor: &Executor, mode: Option<WriterMode>) -> Self {
358        Self::writer_unbounded(executor, stderr(), mode)
359    }
360}
361
362impl RequestWriterLayer<PerMessageFileWriter> {
363    /// Create a layer that streams every request into a unique file below
364    /// `directory` using the portable filename `prefix`.
365    pub async fn file_per_request(
366        executor: &Executor,
367        directory: impl Into<PathBuf>,
368        prefix: impl AsRef<str>,
369        mode: Option<WriterMode>,
370    ) -> io::Result<Self> {
371        let writer = PerMessageFileWriter::try_new(directory, prefix)
372            .await?
373            .with_request_mode(mode)
374            .with_response_mode(None);
375        Ok(Self::new_with_executor(writer, executor.clone()))
376    }
377}
378
379impl RequestWriterLayer<Sender<Request>> {
380    /// Create a new [`RequestWriterLayer`] that prints requests to an [`AsyncWrite`]r
381    /// over a bounded channel with a fixed buffer size.
382    pub fn writer<W>(
383        executor: &Executor,
384        mut writer: W,
385        buffer_size: usize,
386        mode: Option<WriterMode>,
387    ) -> Self
388    where
389        W: AsyncWrite + Unpin + Send + Sync + 'static,
390    {
391        let (tx, mut rx) = channel(buffer_size);
392        let (write_headers, write_body) = write_headers_body_flags(mode);
393
394        let span =
395            tracing::trace_root_span!("TrafficWriter::request::bounded", otel.kind = "consumer");
396
397        executor.spawn_task(
398            async move {
399                drive_request_writer!(writer, rx, write_headers, write_body);
400            }
401            .instrument(span),
402        );
403        Self {
404            writer: tx,
405            executor: executor.clone(),
406        }
407    }
408
409    /// Create a new [`RequestWriterService`] that prints requests to stdout
410    /// over a bounded channel with a fixed buffer size.
411    #[must_use]
412    pub fn stdout(executor: &Executor, buffer_size: usize, mode: Option<WriterMode>) -> Self {
413        Self::writer(executor, stdout(), buffer_size, mode)
414    }
415
416    /// Create a new [`RequestWriterService`] that prints requests to stderr
417    /// over a bounded channel with a fixed buffer size.
418    #[must_use]
419    pub fn stderr(executor: &Executor, buffer_size: usize, mode: Option<WriterMode>) -> Self {
420        Self::writer(executor, stderr(), buffer_size, mode)
421    }
422}
423
424impl<S, W: Clone> Layer<S> for RequestWriterLayer<W> {
425    type Service = RequestWriterService<S, W>;
426
427    fn layer(&self, inner: S) -> Self::Service {
428        RequestWriterService {
429            inner,
430            writer: Arc::new(self.writer.clone()),
431            executor: self.executor.clone(),
432        }
433    }
434
435    fn into_layer(self, inner: S) -> Self::Service {
436        RequestWriterService {
437            inner,
438            writer: Arc::new(self.writer),
439            executor: self.executor,
440        }
441    }
442}
443
444#[cfg(test)]
445mod tests {
446    use std::{
447        convert::Infallible,
448        sync::{
449            Arc,
450            atomic::{AtomicUsize, Ordering},
451        },
452        time::Duration,
453    };
454
455    use rama_core::{
456        Layer as _, Service as _, bytes::Bytes, extensions::ExtensionsRef as _, rt::Executor,
457        service::service_fn,
458    };
459    use tokio::io::AsyncReadExt as _;
460
461    use super::{RequestWriterLayer, RequestWriterService};
462    use crate::layer::traffic_writer::{TrafficWriterId, WriterMode};
463    use crate::{Body, Request, Response, body::util::BodyExt as _};
464
465    #[tokio::test]
466    async fn shared_request_writer_flushes_when_channel_closes() {
467        let executor = Executor::new();
468        let (output, mut captured) = tokio::io::duplex(128);
469        let layer = RequestWriterLayer::writer_unbounded(
470            &executor,
471            tokio::io::BufWriter::new(output),
472            Some(WriterMode::Headers),
473        );
474        let sender = layer.writer.clone();
475        drop(layer);
476
477        sender.send(Request::new(Body::empty())).unwrap();
478        drop(sender);
479
480        let mut bytes = Vec::new();
481        tokio::time::timeout(Duration::from_secs(1), captured.read_to_end(&mut bytes))
482            .await
483            .expect("request writer should flush and close")
484            .unwrap();
485        assert_eq!(bytes, b"GET / HTTP/1.1\r\n\r\n");
486    }
487
488    #[tokio::test]
489    async fn file_constructors_only_enable_request_capture() {
490        let temp = tempfile::tempdir().unwrap();
491        let executor = Executor::new();
492        let layer = RequestWriterLayer::file_per_request(
493            &executor,
494            temp.path(),
495            "requests",
496            Some(WriterMode::Headers),
497        )
498        .await
499        .unwrap();
500        assert_eq!(layer.writer.request_mode(), Some(WriterMode::Headers));
501        assert_eq!(layer.writer.response_mode(), None);
502
503        let service = RequestWriterService::file_per_request(
504            (),
505            &executor,
506            temp.path(),
507            "requests",
508            Some(WriterMode::Body),
509        )
510        .await
511        .unwrap();
512        assert_eq!(service.writer.request_mode(), Some(WriterMode::Body));
513        assert_eq!(service.writer.response_mode(), None);
514    }
515
516    #[tokio::test]
517    async fn request_writer_shares_private_id_with_inner_request() {
518        let (writer, mut captured_requests) = tokio::sync::mpsc::unbounded_channel();
519        let (seen_id, mut seen_ids) = tokio::sync::mpsc::unbounded_channel();
520        let inner = service_fn(move |request: Request| {
521            let seen_id = seen_id.clone();
522            async move {
523                seen_id
524                    .send(*request.extensions().get_ref::<TrafficWriterId>().unwrap())
525                    .unwrap();
526                Ok::<_, Infallible>(Response::new(Body::empty()))
527            }
528        });
529        let service = RequestWriterLayer::new(writer).into_layer(inner);
530
531        service.serve(Request::new(Body::empty())).await.unwrap();
532        let captured = tokio::time::timeout(Duration::from_secs(1), captured_requests.recv())
533            .await
534            .unwrap()
535            .unwrap();
536        let request_id = tokio::time::timeout(Duration::from_secs(1), seen_ids.recv())
537            .await
538            .unwrap()
539            .unwrap();
540        assert_eq!(
541            captured.extensions().get_ref::<TrafficWriterId>(),
542            Some(&request_id)
543        );
544    }
545
546    #[tokio::test]
547    async fn writer_observes_request_without_collecting_before_inner_service() {
548        let polls = Arc::new(AtomicUsize::new(0));
549        let body_polls = Arc::clone(&polls);
550        let body = Body::from_stream(rama_core::futures::stream::once(async move {
551            body_polls.fetch_add(1, Ordering::Relaxed);
552            Ok::<_, Infallible>(Bytes::from_static(b"streamed request"))
553        }));
554        let (writer, mut written) = tokio::sync::mpsc::unbounded_channel();
555        let inner_polls = Arc::clone(&polls);
556        let service =
557            RequestWriterLayer::new(writer).into_layer(service_fn(move |request: Request| {
558                let polls = Arc::clone(&inner_polls);
559                async move {
560                    assert_eq!(polls.load(Ordering::Relaxed), 0);
561                    assert_eq!(
562                        request.into_body().collect().await.unwrap().to_bytes(),
563                        Bytes::from_static(b"streamed request")
564                    );
565                    Ok::<_, Infallible>(Response::new(Body::empty()))
566                }
567            }));
568
569        let serve = service.serve(Request::new(body));
570        let observe = async {
571            let request = written.recv().await.unwrap();
572            request.into_body().collect().await.unwrap().to_bytes()
573        };
574        let (result, captured) = tokio::time::timeout(std::time::Duration::from_secs(1), async {
575            tokio::join!(serve, observe)
576        })
577        .await
578        .expect("request capture should finish promptly");
579        result.unwrap();
580        assert_eq!(polls.load(Ordering::Relaxed), 1);
581        assert_eq!(captured, Bytes::from_static(b"streamed request"));
582    }
583
584    #[tokio::test]
585    async fn retained_request_body_does_not_delay_the_response() {
586        let (held_body, mut held_bodies) = tokio::sync::mpsc::unbounded_channel();
587        let inner = service_fn(move |request: Request| {
588            let held_body = held_body.clone();
589            async move {
590                held_body.send(request.into_body()).unwrap();
591                Ok::<_, Infallible>(Response::new(Body::empty()))
592            }
593        });
594        let (writer_done, mut writer_completions) = tokio::sync::mpsc::unbounded_channel();
595        let writer = move |request: Request| {
596            let writer_done = writer_done.clone();
597            async move {
598                request.into_body().collect().await.unwrap();
599                writer_done.send(()).unwrap();
600            }
601        };
602        let service = RequestWriterLayer::new(writer).into_layer(inner);
603
604        tokio::time::timeout(
605            std::time::Duration::from_secs(1),
606            service.serve(Request::new(Body::from("retained"))),
607        )
608        .await
609        .expect("the response must not wait for the request writer")
610        .unwrap();
611
612        let body = held_bodies.recv().await.unwrap();
613        assert_eq!(
614            body.collect().await.unwrap().to_bytes(),
615            Bytes::from_static(b"retained")
616        );
617        tokio::time::timeout(std::time::Duration::from_secs(1), writer_completions.recv())
618            .await
619            .expect("the request writer should finish after the body is consumed")
620            .unwrap();
621    }
622
623    #[tokio::test]
624    async fn built_in_writer_emits_the_streaming_request_entry() {
625        let inner = service_fn(|request: Request| async move {
626            assert_eq!(
627                request.into_body().collect().await.unwrap().to_bytes(),
628                Bytes::from_static(b"payload")
629            );
630            Ok::<_, Infallible>(Response::new(Body::empty()))
631        });
632        let executor = Executor::new();
633        let (writer, mut output) = tokio::io::duplex(256);
634        let service =
635            RequestWriterLayer::writer_unbounded(&executor, writer, Some(WriterMode::All))
636                .into_layer(inner);
637
638        service
639            .serve(
640                Request::builder()
641                    .method("POST")
642                    .uri("/upload")
643                    .header("x-test", "yes")
644                    .body(Body::from("payload"))
645                    .unwrap(),
646            )
647            .await
648            .unwrap();
649        drop(service);
650
651        let mut written = Vec::new();
652        tokio::time::timeout(Duration::from_secs(1), output.read_to_end(&mut written))
653            .await
654            .expect("the request writer should shut down promptly")
655            .unwrap();
656        let written = String::from_utf8(written).unwrap();
657        assert!(written.starts_with("POST /upload HTTP/1.1\r\n"));
658        assert!(written.contains("x-test: yes\r\n"));
659        assert!(written.ends_with("\r\npayload\r\n"));
660    }
661}
rama_http/layer/traffic_writer/request.rs

rama_http/layer/traffic_writer/
request.rs
