tonic-side-effect 0.4.0

Tower service that can monitor if a Tonic request frame was produced prior to error.
Documentation 
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use hyper::body::{Body, Frame, SizeHint};
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use tonic::body::Body as TonicBody;
use tonic::transport::Channel;
use tower_service::Service;

/// Resettable handle for indicating if a frame has been produced.
#[derive(Clone, Debug, Default)]
pub struct FrameSignal(Arc<AtomicBool>);

impl FrameSignal {
    fn signal(&self) {
        self.0.store(true, Ordering::Release)
    }

    pub fn new() -> Self {
        Self(Arc::new(AtomicBool::new(false)))
    }

    pub fn is_signalled(&self) -> bool {
        self.0.load(Ordering::Acquire)
    }

    pub fn reset(&self) {
        self.0.store(false, Ordering::Release)
    }
}

pin_project! {
    struct RequestFrameMonitorBody<B> {
        #[pin]
        inner: B,
        frame_signal: FrameSignal,
    }
}

impl<B> Body for RequestFrameMonitorBody<B>
where
    B: Body,
{
    type Data = B::Data;
    type Error = B::Error;

    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let this = self.project();
        match this.inner.poll_frame(cx) {
            Poll::Ready(Some(res)) => match res {
                Ok(frame) => {
                    this.frame_signal.signal();
                    Poll::Ready(Some(Ok(frame)))
                }
                Err(status) => Poll::Ready(Some(Err(status))),
            },
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }
}

/// Service for monitoring if an HTTP request frame was ever emitted.
#[derive(Clone, Debug)]
pub struct RequestFrameMonitor<S = Channel>
where
    S: Clone,
{
    /// Wrapped channel to monitor.
    inner: S,

    /// Signal indicating if request frame has been produced.
    frame_signal: FrameSignal,
}

impl<S: Clone> RequestFrameMonitor<S> {
    pub fn new(inner: S, frame_signal: FrameSignal) -> Self {
        Self {
            inner,
            frame_signal: frame_signal.clone(),
        }
    }
}

impl<S> Service<http::Request<TonicBody>> for RequestFrameMonitor<S>
where
    S: Service<http::Request<TonicBody>> + Clone,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: http::Request<TonicBody>) -> Self::Future {
        let (head, body) = req.into_parts();
        let body = TonicBody::new(RequestFrameMonitorBody {
            inner: body,
            frame_signal: self.frame_signal.clone(),
        });
        // See <https://docs.rs/tower/latest/tower/trait.Service.html#be-careful-when-cloning-inner-services>
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);
        inner.call(http::Request::from_parts(head, body))
    }
}