tower_load/
instrument.rs

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

tower_load/
instrument.rs
