tower/load/
completion.rs

Help
1//! Application-specific request completion semantics.
2
3use futures_core::ready;
4use pin_project_lite::pin_project;
5use 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
33pub 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.
38    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]
45pub struct CompleteOnResponse;
46
47pin_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
60impl<F, C, H> TrackCompletionFuture<F, C, H> {
61    /// Wraps a future, propagating the tracker into its value if successful.
62    pub fn new(completion: C, handle: H, future: F) -> Self {
63        TrackCompletionFuture {
64            future,
65            completion,
66            handle: Some(handle),
67        }
68    }
69}
70
71impl<F, C, H, T, E> Future for TrackCompletionFuture<F, C, H>
72where
73    F: Future<Output = Result<T, E>>,
74    C: TrackCompletion<H, T>,
75{
76    type Output = Result<C::Output, E>;
77
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
88impl<H, V> TrackCompletion<H, V> for CompleteOnResponse {
89    type Output = V;
90
91    fn track_completion(&self, handle: H, value: V) -> V {
92        drop(handle);
93        value
94    }
95}
tower/load/completion.rs

tower/load/
completion.rs
