Skip to main content

imsg_session/
loop_util.rs

1//! Shared accept-loop driver used by MNS session and relay modules.
2
3use std::future::Future;
4
5use futures::StreamExt;
6use tokio::io::{AsyncRead, AsyncWrite};
7use tokio::sync::watch;
8
9/// `drain` returns `false` to halt early, `true` to accept next.
10pub(crate) async fn run_accept_loop<S, T, F, Fut>(
11    mut stream_source: S,
12    mut cancel: watch::Receiver<bool>,
13    mut drain: F,
14) where
15    S: futures::Stream<Item = T> + Unpin,
16    T: AsyncRead + AsyncWrite + Unpin,
17    F: FnMut(T, watch::Receiver<bool>) -> Fut,
18    Fut: Future<Output = bool>,
19{
20    loop {
21        if *cancel.borrow() {
22            return;
23        }
24        let stream = tokio::select! {
25            biased;
26            result = cancel.changed() => {
27                if result.is_err() || *cancel.borrow_and_update() { return; }
28                continue;
29            }
30            maybe = stream_source.next() => match maybe {
31                Some(s) => s,
32                None => return,
33            }
34        };
35        if !drain(stream, cancel.clone()).await {
36            return;
37        }
38    }
39}