aws_smithy_http_server/instrumentation/
service.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! A [`Service`] and it's associated [`Future`] providing sensitivity aware logging.
7
8use std::{
9    future::Future,
10    pin::Pin,
11    task::{Context, Poll},
12};
13
14use futures_util::{ready, TryFuture};
15use http::{HeaderMap, Request, Response, StatusCode, Uri};
16use tower::Service;
17use tracing::{debug, debug_span, instrument::Instrumented, Instrument};
18
19use crate::shape_id::ShapeId;
20
21use super::{MakeDebug, MakeDisplay, MakeIdentity};
22
23pin_project_lite::pin_project! {
24    /// A [`Future`] responsible for logging the response status code and headers.
25    struct InnerFuture<Fut, ResponseMakeFmt> {
26        #[pin]
27        inner: Fut,
28        make: ResponseMakeFmt
29    }
30}
31
32impl<Fut, ResponseMakeFmt, T> Future for InnerFuture<Fut, ResponseMakeFmt>
33where
34    Fut: TryFuture<Ok = Response<T>>,
35    Fut: Future<Output = Result<Fut::Ok, Fut::Error>>,
36
37    for<'a> ResponseMakeFmt: MakeDebug<&'a HeaderMap>,
38    for<'a> ResponseMakeFmt: MakeDisplay<StatusCode>,
39{
40    type Output = Fut::Output;
41
42    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
43        let this = self.project();
44        let response = ready!(this.inner.poll(cx))?;
45
46        {
47            let headers = this.make.make_debug(response.headers());
48            let status_code = this.make.make_display(response.status());
49            debug!(?headers, %status_code, "response");
50        }
51
52        Poll::Ready(Ok(response))
53    }
54}
55
56// This is to provide type erasure.
57pin_project_lite::pin_project! {
58    /// An instrumented [`Future`] responsible for logging the response status code and headers.
59    pub struct InstrumentedFuture<Fut, ResponseMakeFmt> {
60        #[pin]
61        inner: Instrumented<InnerFuture<Fut, ResponseMakeFmt>>
62    }
63}
64
65impl<Fut, ResponseMakeFmt, T> Future for InstrumentedFuture<Fut, ResponseMakeFmt>
66where
67    Fut: TryFuture<Ok = Response<T>>,
68    Fut: Future<Output = Result<Fut::Ok, Fut::Error>>,
69
70    for<'a> ResponseMakeFmt: MakeDebug<&'a HeaderMap>,
71    for<'a> ResponseMakeFmt: MakeDisplay<StatusCode>,
72{
73    type Output = Fut::Output;
74
75    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
76        self.project().inner.poll(cx)
77    }
78}
79
80/// A middleware [`Service`] responsible for:
81///   - Opening a [`tracing::debug_span`] for the lifetime of the request, which includes the operation name, the
82///     [`Uri`], and the request headers.
83///   - A [`tracing::debug`] during response, which includes the response status code and headers.
84///
85/// The [`Display`](std::fmt::Display) and [`Debug`] of the request and response components can be modified using
86/// [`request_fmt`](InstrumentOperation::request_fmt) and [`response_fmt`](InstrumentOperation::response_fmt).
87///
88/// # Example
89///
90/// ```
91/// # use aws_smithy_http_server::instrumentation::{sensitivity::{*, uri::*, headers::*}, *};
92/// # use aws_smithy_http_server::shape_id::ShapeId;
93/// # use tower::{Service, service_fn};
94/// # use http::{Request, Response};
95/// # async fn f(request: Request<()>) -> Result<Response<()>, ()> { Ok(Response::new(())) }
96/// # let mut svc = service_fn(f);
97/// # const ID: ShapeId = ShapeId::new("namespace#foo-operation", "namespace", "foo-operation");
98/// let request_fmt = RequestFmt::new()
99///     .label(|index| index == 1, None)
100///     .query(|_| QueryMarker { key: false, value: true });
101/// let response_fmt = ResponseFmt::new().status_code();
102/// let mut svc = InstrumentOperation::new(svc, ID)
103///     .request_fmt(request_fmt)
104///     .response_fmt(response_fmt);
105/// # svc.call(Request::new(()));
106/// ```
107#[derive(Debug, Clone)]
108pub struct InstrumentOperation<S, RequestMakeFmt = MakeIdentity, ResponseMakeFmt = MakeIdentity> {
109    inner: S,
110    operation_id: ShapeId,
111    make_request: RequestMakeFmt,
112    make_response: ResponseMakeFmt,
113}
114
115impl<S> InstrumentOperation<S> {
116    /// Constructs a new [`InstrumentOperation`] with no data redacted.
117    pub fn new(inner: S, operation_id: ShapeId) -> Self {
118        Self {
119            inner,
120            operation_id,
121            make_request: MakeIdentity,
122            make_response: MakeIdentity,
123        }
124    }
125}
126
127impl<S, RequestMakeFmt, ResponseMakeFmt> InstrumentOperation<S, RequestMakeFmt, ResponseMakeFmt> {
128    /// Configures the request format.
129    ///
130    /// The argument is typically [`RequestFmt`](super::sensitivity::RequestFmt).
131    pub fn request_fmt<R>(self, make_request: R) -> InstrumentOperation<S, R, ResponseMakeFmt> {
132        InstrumentOperation {
133            inner: self.inner,
134            operation_id: self.operation_id,
135            make_request,
136            make_response: self.make_response,
137        }
138    }
139
140    /// Configures the response format.
141    ///
142    /// The argument is typically [`ResponseFmt`](super::sensitivity::ResponseFmt).
143    pub fn response_fmt<R>(self, make_response: R) -> InstrumentOperation<S, RequestMakeFmt, R> {
144        InstrumentOperation {
145            inner: self.inner,
146            operation_id: self.operation_id,
147            make_request: self.make_request,
148            make_response,
149        }
150    }
151}
152
153impl<S, U, V, RequestMakeFmt, ResponseMakeFmt> Service<Request<U>>
154    for InstrumentOperation<S, RequestMakeFmt, ResponseMakeFmt>
155where
156    S: Service<Request<U>, Response = Response<V>>,
157
158    for<'a> RequestMakeFmt: MakeDebug<&'a HeaderMap>,
159    for<'a> RequestMakeFmt: MakeDisplay<&'a Uri>,
160
161    ResponseMakeFmt: Clone,
162    for<'a> ResponseMakeFmt: MakeDebug<&'a HeaderMap>,
163    for<'a> ResponseMakeFmt: MakeDisplay<StatusCode>,
164{
165    type Response = S::Response;
166    type Error = S::Error;
167    type Future = InstrumentedFuture<S::Future, ResponseMakeFmt>;
168
169    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
170        self.inner.poll_ready(cx)
171    }
172
173    fn call(&mut self, request: Request<U>) -> Self::Future {
174        let span = {
175            let headers = self.make_request.make_debug(request.headers());
176            let uri = self.make_request.make_display(request.uri());
177            debug_span!("request", operation = %self.operation_id.absolute(), method = %request.method(), %uri, ?headers)
178        };
179
180        InstrumentedFuture {
181            inner: InnerFuture {
182                inner: self.inner.call(request),
183                make: self.make_response.clone(),
184            }
185            .instrument(span),
186        }
187    }
188}
aws_smithy_http_server/instrumentation/service.rs

aws_smithy_http_server/instrumentation/
service.rs
