krabka_raft/handshake/
stream.rs

1//! The stream a controller-listener handshake reads and writes, and the clock
2//! that measures its silence for `connections.max.idle.ms`.
3
4use {
5    crate::{ObservedStream, StreamObserver},
6    std::{
7        io,
8        pin::pin,
9        sync::{
10            Arc, Mutex, PoisonError,
11            atomic::{AtomicBool, Ordering},
12        },
13        task::{Context, Poll},
14        time::Duration,
15    },
16    tokio::{
17        net::TcpStream,
18        time::{self, Instant},
19    },
20};
21
22/// An accepted controller-listener connection, handed to
23/// [`RaftListenerHandshake::upgrade`](super::RaftListenerHandshake::upgrade).
24///
25/// Every read or write that moves bytes restarts the listener's idle window,
26/// as Kafka's `Selector` refreshes a channel's idle clock on every ready key.
27/// The window closes a handshake that goes silent, not one that is slow. Once
28/// the listener takes the upgraded stream for its request loop, the clock is
29/// disarmed and the stream only passes bytes through.
30pub type HandshakeStream<S = TcpStream> = ObservedStream<S, HandshakeClock>;
31
32/// The idle clock a [`HandshakeStream`] restarts on every read or write that
33/// moves bytes, or none when no window watches the stream.
34pub struct HandshakeClock(Option<Arc<ActivityClock>>);
35
36impl HandshakeClock {
37    fn progressed(&self) {
38        if let Some(clock) = &self.0 {
39            clock.progressed();
40        }
41    }
42}
43
44impl StreamObserver for HandshakeClock {
45    fn on_read(&mut self, filled: usize) {
46        if filled > 0 {
47            self.progressed();
48        }
49    }
50
51    fn on_write(
52        &mut self,
53        _cx: &mut Context<'_>,
54        result: Poll<io::Result<usize>>,
55    ) -> Poll<io::Result<usize>> {
56        if matches!(result, Poll::Ready(Ok(written)) if written > 0) {
57            self.progressed();
58        }
59        result
60    }
61}
62
63impl<S> HandshakeStream<S> {
64    /// Wraps `inner` and returns the watch that expires once `window` passes
65    /// with no byte read or written.
66    pub(crate) fn watched(inner: S, window: Duration) -> (Self, IdleWatch) {
67        let clock = Arc::new(ActivityClock {
68            armed: AtomicBool::new(true),
69            last_progress: Mutex::new(Instant::now()),
70        });
71        let watch = IdleWatch {
72            clock: Arc::clone(&clock),
73            window,
74        };
75        (Self::new(inner, HandshakeClock(Some(clock))), watch)
76    }
77}
78
79/// A stream no idle window watches, for a listener without
80/// `connections.max.idle.ms` and for tests that drive a handshake directly.
81impl<S> From<S> for HandshakeStream<S> {
82    fn from(inner: S) -> Self { Self::new(inner, HandshakeClock(None)) }
83}
84
85/// When a handshake last moved a byte, and whether anyone still asks.
86struct ActivityClock {
87    armed: AtomicBool,
88    last_progress: Mutex<Instant>,
89}
90
91impl ActivityClock {
92    fn progressed(&self) {
93        if self.armed.load(Ordering::Relaxed) {
94            *self
95                .last_progress
96                .lock()
97                .unwrap_or_else(PoisonError::into_inner) = Instant::now();
98        }
99    }
100
101    fn last_progress(&self) -> Instant {
102        *self
103            .last_progress
104            .lock()
105            .unwrap_or_else(PoisonError::into_inner)
106    }
107}
108
109/// The listener's side of a [`HandshakeStream`]: it expires once the window
110/// passes in silence, and dropping it disarms the stream's clock.
111pub struct IdleWatch {
112    clock: Arc<ActivityClock>,
113    window: Duration,
114}
115
116impl IdleWatch {
117    /// Resolves once `window` passes with no byte read or written. The timer
118    /// is re-armed from the last progress each time it fires early, so the
119    /// stream records progress without touching a timer.
120    pub(crate) async fn expired(&self) {
121        let mut timer = pin!(time::sleep_until(self.clock.last_progress() + self.window));
122        loop {
123            timer.as_mut().await;
124            let deadline = self.clock.last_progress() + self.window;
125            if deadline <= Instant::now() {
126                return;
127            }
128            timer.as_mut().reset(deadline);
129        }
130    }
131
132    /// Restarts the window without a byte moving, as Kafka's `Selector.send`
133    /// refreshes a channel's idle clock when a response is queued.
134    pub(crate) fn refresh(&self) { self.clock.progressed(); }
135}
136
137impl Drop for IdleWatch {
138    fn drop(&mut self) { self.clock.armed.store(false, Ordering::Relaxed); }
139}
krabka_raft/handshake/stream.rs

krabka_raft/handshake/
stream.rs
