krabka_raft/
observed_stream.rs

1//! A stream that passes every byte through and reports each poll's outcome,
2//! the plumbing that the listeners' idle windows share.
3
4use {
5    std::{
6        io::{self, IoSlice},
7        pin::Pin,
8        task::{Context, Poll},
9    },
10    tokio::io::{AsyncRead, AsyncWrite, ReadBuf},
11};
12
13/// What an [`ObservedStream`] reports to.
14///
15/// Each method sees one poll of the wrapped stream and, for a write, flush or
16/// shutdown, returns what the caller gets. The defaults pass every outcome on
17/// unchanged.
18pub trait StreamObserver {
19    /// A read poll moved `filled` bytes into the caller's buffer.
20    fn on_read(&mut self, _filled: usize) {}
21
22    /// A write or vectored write polled `result`.
23    fn on_write(
24        &mut self,
25        _cx: &mut Context<'_>,
26        result: Poll<io::Result<usize>>,
27    ) -> Poll<io::Result<usize>> {
28        result
29    }
30
31    /// A flush or shutdown polled `result`.
32    fn on_flush(
33        &mut self,
34        _cx: &mut Context<'_>,
35        result: Poll<io::Result<()>>,
36    ) -> Poll<io::Result<()>> {
37        result
38    }
39}
40
41/// A stream that reads and writes through `inner` and shows `observer` every
42/// poll.
43pub struct ObservedStream<S, O> {
44    inner: S,
45    observer: O,
46}
47
48impl<S, O> ObservedStream<S, O> {
49    /// Wraps `inner`, reporting to `observer`.
50    pub fn new(inner: S, observer: O) -> Self { Self { inner, observer } }
51
52    /// The wrapped stream, for what a caller reads off the socket itself,
53    /// such as the peer address.
54    pub fn get_ref(&self) -> &S { &self.inner }
55
56    /// What the stream reports to.
57    pub fn observer(&self) -> &O { &self.observer }
58}
59
60impl<S: AsyncRead + Unpin, O: StreamObserver + Unpin> AsyncRead for ObservedStream<S, O> {
61    fn poll_read(
62        self: Pin<&mut Self>,
63        cx: &mut Context<'_>,
64        buf: &mut ReadBuf<'_>,
65    ) -> Poll<io::Result<()>> {
66        let this = self.get_mut();
67        let before = buf.filled().len();
68        let result = Pin::new(&mut this.inner).poll_read(cx, buf);
69        this.observer.on_read(buf.filled().len() - before);
70        result
71    }
72}
73
74impl<S: AsyncWrite + Unpin, O: StreamObserver + Unpin> AsyncWrite for ObservedStream<S, O> {
75    fn poll_write(
76        self: Pin<&mut Self>,
77        cx: &mut Context<'_>,
78        buf: &[u8],
79    ) -> Poll<io::Result<usize>> {
80        let this = self.get_mut();
81        let result = Pin::new(&mut this.inner).poll_write(cx, buf);
82        this.observer.on_write(cx, result)
83    }
84
85    fn poll_write_vectored(
86        self: Pin<&mut Self>,
87        cx: &mut Context<'_>,
88        bufs: &[IoSlice<'_>],
89    ) -> Poll<io::Result<usize>> {
90        let this = self.get_mut();
91        let result = Pin::new(&mut this.inner).poll_write_vectored(cx, bufs);
92        this.observer.on_write(cx, result)
93    }
94
95    fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored() }
96
97    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
98        let this = self.get_mut();
99        let result = Pin::new(&mut this.inner).poll_flush(cx);
100        this.observer.on_flush(cx, result)
101    }
102
103    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
104        let this = self.get_mut();
105        let result = Pin::new(&mut this.inner).poll_shutdown(cx);
106        this.observer.on_flush(cx, result)
107    }
108}
krabka_raft/observed_stream.rs

krabka_raft/
observed_stream.rs
