axol 0.2.0

Axol Web Framework
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use crate::Result;

use axol_http::{
    Body,
    body::{BodyComponent, BodyStream},
};
use futures::{Stream, ready};
use std::{
    pin::Pin,
    task::{Context, Poll},
    time::Instant,
};
use tracing::{Level, Span};

pub struct TraceBody {
    /// Entered around each poll. Held here rather than via the unmaintained `tracing-futures`
    /// stream instrumentation, which is the only thing that crate was pulled in for.
    span: Span,
    inner: BodyStream,
    body_start: Instant,
    body_size: usize,
    is_response: bool,
}

impl TraceBody {
    pub fn wrap(span: Span, body: Body, is_response: bool) -> Body {
        let size_hint = match &body {
            Body::Bytes(x) => Some(x.len()),
            Body::Stream {
                size_hint,
                stream: _,
            } => *size_hint,
        };
        let stream = Self {
            span,
            inner: body.into_stream(),
            body_start: Instant::now(),
            body_size: 0,
            is_response,
        };
        Body::Stream {
            size_hint,
            stream: Box::pin(stream),
        }
    }
}

impl Stream for TraceBody {
    type Item = Result<BodyComponent, anyhow::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let span = self.span.clone();
        let _enter = span.enter();
        let next = self.inner.as_mut().poll_next(cx);

        let Some(next) = ready!(next) else {
            let body_elapsed_ms = self.body_start.elapsed().as_secs_f64() * 1000.0;

            if self.is_response {
                span.record("http.response.body.size", self.body_size);
                span.record("http.response.body.elapsed_ms", body_elapsed_ms);
            } else {
                span.record("http.request.body.size", self.body_size);
                span.record("http.request.body.elapsed_ms", body_elapsed_ms);
            }

            return Poll::Ready(None);
        };

        match &next {
            Ok(BodyComponent::Data(data)) => {
                self.body_size += data.len();
            }
            Ok(BodyComponent::Trailers(trailers)) => {
                tracing::event!(Level::DEBUG, ?trailers, "body trailers");
            }
            Err(err) => {
                tracing::event!(Level::ERROR, ?err, "body error");
            }
        }

        Poll::Ready(Some(next))
    }
}