rama 0.4.0

modular service framework
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

use crate::h1_server::{StreamReadHalf, fixture, init_tracing};
use pin_project_lite::pin_project;
use rama::telemetry::tracing::error;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::sync::mpsc;
use tokio::time::Sleep;

pin_project! {
    #[derive(Debug)]
    pub struct UnbufferedStream {
        #[pin]
        read_half: StreamReadHalf,
        #[pin]
        pending_write: Option<Pin<Box<Sleep>>>,
        write_tx: mpsc::UnboundedSender<Vec<u8>>,
        poll_cnt: usize,
    }
}

impl UnbufferedStream {
    fn new(
        read_rx: mpsc::UnboundedReceiver<Vec<u8>>,
        write_tx: mpsc::UnboundedSender<Vec<u8>>,
    ) -> Self {
        Self {
            read_half: StreamReadHalf::new(read_rx),
            write_tx,
            pending_write: None,
            poll_cnt: 0,
        }
    }

    /// Create a new server stream and client pair.
    /// Returns a server stream (Read+Write) and a client (rx/tx channels).
    fn new_pair() -> (Self, fixture::Client) {
        let (client_tx, server_rx) = mpsc::unbounded_channel();
        let (server_tx, client_rx) = mpsc::unbounded_channel();
        let server = Self::new(server_rx, server_tx);
        let client = fixture::Client {
            rx: client_rx,
            tx: client_tx,
        };
        (server, client)
    }
}

const WRITE_DELAY: Duration = Duration::from_millis(100);

impl AsyncRead for UnbufferedStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        self.as_mut().project().read_half.poll_read(cx, buf)
    }
}

impl AsyncWrite for UnbufferedStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        self.poll_cnt += 1;
        let poll_cnt = self.poll_cnt;
        if let Some(sleep) = self.pending_write.as_mut() {
            let sleep = sleep.as_mut();
            if poll_cnt > 4 {
                return Poll::Ready(Err(io::Error::other("We are being hot polled!")));
            }
            ready!(Future::poll(sleep, cx));
            let mut this = self.as_mut().project();
            this.pending_write.set(None);
            *this.poll_cnt = 0;
        }
        let len = buf.len();
        {
            let mut this = self.as_mut().project();
            let buf = Vec::from(&buf[..buf.len()]);
            // Send data through the channel - this should always be ready for unbounded channels
            let Ok(_) = this.write_tx.send(buf) else {
                error!("UnbufferedStream::poll_write failed - channel closed");
                return Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::BrokenPipe,
                    "Write channel closed",
                )));
            };
            this.pending_write
                .set(Some(Box::pin(tokio::time::sleep(WRITE_DELAY))))
        }
        let Some(sleep) = self.pending_write.as_mut() else {
            panic!("Sleep should have just been set");
        };
        let sleep = sleep.as_mut();
        let Poll::Pending = Future::poll(sleep, cx) else {
            panic!("Sleep always be pending on first poll")
        };
        Poll::Ready(Ok(len))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn body_test() {
    init_tracing();
    let (server, client) = UnbufferedStream::new_pair();
    // The per-chunk timeout is a STALL watchdog, not a tight pacing assertion —
    // the unbuffered/backpressure behaviour is asserted by the `poll_cnt > 4`
    // hot-poll guard and the final byte-count check, not by chunk latency. A
    // thin `WRITE_DELAY * 2` (200ms) margin over the injected 100ms write pacing
    // flaked on slow/loaded CI (notably Windows), where connection setup before
    // the first chunk alone can exceed it. Give a generous multiple so a healthy
    // run never trips while a genuine stall is still caught.
    let config = fixture::TestConfig::with_timeout(WRITE_DELAY * 30);
    fixture::run(server, client, config).await;
}