picomux 0.2.0

A protocol to multiplex multiple AsyncRead/Write streams over a single one, inspired by smux and http2
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use std::sync::{
    Arc, OnceLock,
    atomic::{AtomicBool, AtomicU64, Ordering},
};

use crossbeam_queue::SegQueue;
use geph5_rt::{Task, spawn};
use tokio::io::{AsyncWrite, AsyncWriteExt};

use crate::frame::Frame;

/// A writer for outgoing data.
#[derive(Clone)]
pub struct Outgoing {
    inner: Arc<Inner>,
    err: Arc<OnceLock<anyhow::Error>>,
    _task: Arc<Task<()>>,
}

impl Outgoing {
    /// Creates a new outgoing writer.
    pub fn new(write: impl AsyncWrite + Send + Unpin + 'static) -> Self {
        let inner = Arc::new(Inner::default());
        let err_cell = Arc::new(OnceLock::new());
        // DIAG watchdog: detect an outgoing_loop parked while frames are queued.
        {
            let weak = Arc::downgrade(&inner);
            spawn(async move {
                let mut last_written = 0u64;
                let mut stuck_ticks = 0u32;
                loop {
                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
                    let Some(inner) = weak.upgrade() else { return };
                    let q = inner.queue.len();
                    let w = inner.written.load(Ordering::Relaxed);
                    let iw = inner.in_write.load(Ordering::Relaxed);
                    if q > 0 && w == last_written && !iw {
                        stuck_ticks += 1;
                        tracing::warn!(
                            queue_len = q,
                            written = w,
                            stuck_secs = stuck_ticks * 2,
                            "OUTGOING-STUCK: frames queued but loop not writing"
                        );
                    } else {
                        stuck_ticks = 0;
                    }
                    last_written = w;
                }
            })
            .detach();
        }
        Self {
            inner: inner.clone(),
            err: err_cell.clone(),
            _task: Arc::new(spawn(async move {
                if let Err(err) = outgoing_loop(write, inner).await {
                    err_cell.set(err).unwrap();
                }
            })),
        }
    }

    /// Send a frame to the outgoing writer, returning once it's fully written.
    pub async fn send(&self, outgoing: Frame) -> anyhow::Result<()> {
        self.enqueue(outgoing);
        // wait until the queue is empty
        self.inner
            .shrink_signal
            .wait_until(|| {
                if let Some(err) = self.err.get() {
                    return Some(Err(anyhow::anyhow!("{:?}", err)));
                }
                if self.inner.queue.len() < 10 {
                    Some(anyhow::Ok(()))
                } else {
                    None
                }
            })
            .await?;
        Ok(())
    }

    /// Number of frames currently queued but not yet written to the wire.
    pub fn queue_len(&self) -> usize {
        self.inner.queue.len()
    }

    /// Total number of frames the outgoing loop has fully written to the wire.
    pub fn written(&self) -> u64 {
        self.inner.written.load(Ordering::Relaxed)
    }

    /// Whether the outgoing loop is currently blocked inside `write_all`.
    pub fn in_write(&self) -> bool {
        self.inner.in_write.load(Ordering::Relaxed)
    }

    /// Infallibly, non-blockingly enqueues a frame to be sent to the outgoing writer.
    pub fn enqueue(&self, outgoing: Frame) {
        tracing::trace!(
            command = outgoing.header.command,
            stream_id = outgoing.header.stream_id,
            body_len = outgoing.header.body_len,
            "sending outgoing frame"
        );
        self.inner.queue.push(outgoing);
        self.inner.grow_signal.notify_one();
    }
}

#[derive(Default)]
struct Inner {
    queue: SegQueue<Frame>,
    grow_signal: async_event::Event,
    shrink_signal: async_event::Event,
    written: AtomicU64,
    in_write: AtomicBool,
    dead: AtomicBool,
}

async fn outgoing_loop(
    mut write: impl AsyncWrite + Send + Unpin + 'static,
    inner: Arc<Inner>,
) -> anyhow::Result<()> {
    scopeguard::defer!(inner.shrink_signal.notify_all());
    loop {
        let next = inner.grow_signal.wait_until(|| inner.queue.pop()).await;
        inner.shrink_signal.notify_all();
        inner.in_write.store(true, Ordering::Relaxed);
        write.write_all(&next.bytes()).await?;
        inner.in_write.store(false, Ordering::Relaxed);
        inner.written.fetch_add(1, Ordering::Relaxed);
    }
}