1 //! Application-specific request completion semantics.
2 
3 use futures_core::ready;
4 use pin_project_lite::pin_project;
5 use std::{
6     future::Future,
7     pin::Pin,
8     task::{Context, Poll},
9 };
10 
11 /// Attaches `H`-typed completion tracker to `V` typed values.
12 ///
13 /// Handles (of type `H`) are intended to be RAII guards that primarily implement [`Drop`] and update
14 /// load metric state as they are dropped. This trait allows implementors to "forward" the handle
15 /// to later parts of the request-handling pipeline, so that the handle is only dropped when the
16 /// request has truly completed.
17 ///
18 /// This utility allows load metrics to have a protocol-agnostic means to track streams past their
19 /// initial response future. For example, if `V` represents an HTTP response type, an
20 /// implementation could add `H`-typed handles to each response's extensions to detect when all the
21 /// response's extensions have been dropped.
22 ///
23 /// A base `impl<H, V> TrackCompletion<H, V> for CompleteOnResponse` is provided to drop the handle
24 /// once the response future is resolved. This is appropriate when a response is discrete and
25 /// cannot comprise multiple messages.
26 ///
27 /// In many cases, the `Output` type is simply `V`. However, [`TrackCompletion`] may alter the type
28 /// in order to instrument it appropriately. For example, an HTTP [`TrackCompletion`] may modify
29 /// the body type: so a [`TrackCompletion`] that takes values of type
30 /// [`http::Response<A>`][response] may output values of type [`http::Response<B>`][response].
31 ///
32 /// [response]: https://docs.rs/http/latest/http/response/struct.Response.html
33 pub trait TrackCompletion<H, V>: Clone {
34     /// The instrumented value type.
35     type Output;
36 
37     /// Attaches a `H`-typed handle to a `V`-typed value.
track_completion(&self, handle: H, value: V) -> Self::Output38     fn track_completion(&self, handle: H, value: V) -> Self::Output;
39 }
40 
41 /// A [`TrackCompletion`] implementation that considers the request completed when the response
42 /// future is resolved.
43 #[derive(Clone, Copy, Debug, Default)]
44 #[non_exhaustive]
45 pub struct CompleteOnResponse;
46 
47 pin_project! {
48     /// Attaches a `C`-typed completion tracker to the result of an `F`-typed [`Future`].
49     #[derive(Debug)]
50     pub struct TrackCompletionFuture<F, C, H> {
51         #[pin]
52         future: F,
53         handle: Option<H>,
54         completion: C,
55     }
56 }
57 
58 // ===== impl InstrumentFuture =====
59 
60 impl<F, C, H> TrackCompletionFuture<F, C, H> {
61     /// Wraps a future, propagating the tracker into its value if successful.
new(completion: C, handle: H, future: F) -> Self62     pub fn new(completion: C, handle: H, future: F) -> Self {
63         TrackCompletionFuture {
64             future,
65             completion,
66             handle: Some(handle),
67         }
68     }
69 }
70 
71 impl<F, C, H, T, E> Future for TrackCompletionFuture<F, C, H>
72 where
73     F: Future<Output = Result<T, E>>,
74     C: TrackCompletion<H, T>,
75 {
76     type Output = Result<C::Output, E>;
77 
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>78     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
79         let this = self.project();
80         let rsp = ready!(this.future.poll(cx))?;
81         let h = this.handle.take().expect("handle");
82         Poll::Ready(Ok(this.completion.track_completion(h, rsp)))
83     }
84 }
85 
86 // ===== CompleteOnResponse =====
87 
88 impl<H, V> TrackCompletion<H, V> for CompleteOnResponse {
89     type Output = V;
90 
track_completion(&self, handle: H, value: V) -> V91     fn track_completion(&self, handle: H, value: V) -> V {
92         drop(handle);
93         value
94     }
95 }
96