ntex_util/channel/
bstream.rs

1//! A buffered stream of byte chunks.
2use std::cell::{Cell, RefCell};
3use std::task::{Context, Poll};
4use std::{collections::VecDeque, fmt, future::poll_fn, pin::Pin, rc::Rc, rc::Weak};
5
6use ntex_bytes::Bytes;
7
8use crate::{Stream, task::LocalWaker};
9
10/// max buffer size 32k
11const MAX_BUFFER_SIZE: usize = 32_768;
12
13/// Indicates the current status of a byte stream.
14#[derive(Copy, Clone, Debug, PartialEq, Eq)]
15pub enum Status {
16    /// End of stream reached.
17    Eof,
18    /// Stream is ready to accept more bytes.
19    Ready,
20    /// The receiver side has been dropped.
21    Dropped,
22}
23
24/// Creates a byte stream and returns its sender and receiver.
25pub fn channel<E>() -> (Sender<E>, Receiver<E>) {
26    let inner = Rc::new(Inner::new(false));
27
28    (
29        Sender {
30            inner: Rc::downgrade(&inner),
31        },
32        Receiver { inner },
33    )
34}
35
36/// Creates a byte stream that starts at EOF.
37///
38/// Data may still be added through the returned sender. The receiver yields
39/// that buffered data first and then completes.
40pub fn eof<E>() -> (Sender<E>, Receiver<E>) {
41    let inner = Rc::new(Inner::new(true));
42
43    (
44        Sender {
45            inner: Rc::downgrade(&inner),
46        },
47        Receiver { inner },
48    )
49}
50
51/// Creates a receiver that contains optional data and is already at EOF.
52pub fn empty<E>(data: Option<Bytes>) -> Receiver<E> {
53    let rx = Receiver {
54        inner: Rc::new(Inner::new(true)),
55    };
56    if let Some(data) = data {
57        rx.put(data);
58    }
59    rx
60}
61
62/// A buffered stream of byte chunks.
63///
64/// The receiver yields chunks in insertion order. Its configured buffer size
65/// is a cooperative backpressure threshold for the sender, not a hard memory
66/// limit.
67#[derive(Debug)]
68pub struct Receiver<E> {
69    inner: Rc<Inner<E>>,
70}
71
72impl<E> Receiver<E> {
73    /// Sets the sender backpressure threshold.
74    ///
75    /// Once buffered data reaches this size, [`Sender::poll_ready`] stops
76    /// reporting [`Status::Ready`] until the receiver consumes enough data.
77    /// Sending does not enforce the threshold, so producers must cooperate by
78    /// waiting for readiness. Changing the threshold immediately updates and,
79    /// when needed, wakes sender readiness. The default is 32 KiB.
80    #[inline]
81    pub fn max_buffer_size(&self, size: usize) {
82        self.inner.set_max_buffer_size(size);
83    }
84
85    /// Puts previously read data back at the front of the stream.
86    ///
87    /// This may grow the buffer past its readiness threshold.
88    #[inline]
89    pub fn put(&self, data: Bytes) {
90        self.inner.unread_data(data);
91    }
92
93    #[inline]
94    /// Returns `true` once EOF has been marked.
95    ///
96    /// Buffered chunks may still be available after this returns `true`.
97    pub fn is_eof(&self) -> bool {
98        self.inner.flags.get().contains(Flags::EOF)
99    }
100
101    #[inline]
102    /// Waits for and returns the next chunk, stream error, or EOF.
103    pub async fn read(&self) -> Option<Result<Bytes, E>> {
104        poll_fn(|cx| self.poll_read(cx)).await
105    }
106
107    #[inline]
108    /// Polls for the next chunk, stream error, or EOF.
109    pub fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, E>>> {
110        if let Some(data) = self.inner.get_data() {
111            Poll::Ready(Some(Ok(data)))
112        } else if let Some(err) = self.inner.err.take() {
113            self.inner.insert_flag(Flags::EOF);
114            Poll::Ready(Some(Err(err)))
115        } else if self.inner.flags.get().intersects(Flags::EOF | Flags::ERROR) {
116            Poll::Ready(None)
117        } else {
118            self.inner.recv_task.register(cx.waker());
119            Poll::Pending
120        }
121    }
122}
123
124impl<E> Stream for Receiver<E> {
125    type Item = Result<Bytes, E>;
126
127    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
128        self.poll_read(cx)
129    }
130}
131
132impl<E> Drop for Receiver<E> {
133    fn drop(&mut self) {
134        self.inner.send_task.wake();
135    }
136}
137
138/// Sender side of the byte stream.
139///
140/// Clones share one readiness registration. If several clones poll readiness,
141/// the most recently registered task is the one that will be woken.
142#[derive(Debug)]
143pub struct Sender<E> {
144    inner: Weak<Inner<E>>,
145}
146
147impl<E> Clone for Sender<E> {
148    fn clone(&self) -> Self {
149        Self {
150            inner: self.inner.clone(),
151        }
152    }
153}
154
155impl<E> Drop for Sender<E> {
156    fn drop(&mut self) {
157        if self.inner.weak_count() == 1
158            && let Some(shared) = self.inner.upgrade()
159        {
160            shared.insert_flag(Flags::EOF | Flags::SENDER_GONE);
161        }
162    }
163}
164
165impl<E> Sender<E> {
166    /// Returns `true` if the receiver has been dropped.
167    pub fn is_closed(&self) -> bool {
168        self.inner.strong_count() == 0
169    }
170
171    /// Stores a terminal stream error and wakes both sides.
172    pub fn set_error(&self, err: E) {
173        if let Some(shared) = self.inner.upgrade() {
174            shared.set_error(err);
175        }
176    }
177
178    /// Marks the stream as EOF and wakes both sides.
179    pub fn feed_eof(&self) {
180        if let Some(shared) = self.inner.upgrade() {
181            shared.feed_eof();
182        }
183    }
184
185    /// Adds a chunk to the stream.
186    ///
187    /// This method does not enforce the configured backpressure threshold.
188    pub fn feed_data(&self, data: Bytes) {
189        if let Some(shared) = self.inner.upgrade() {
190            shared.feed_data(data);
191        }
192    }
193
194    /// Waits until the stream needs more data or reaches a terminal state.
195    pub async fn ready(&self) -> Status {
196        poll_fn(|cx| self.poll_ready(cx)).await
197    }
198
199    /// Polls until the stream needs more data or reaches a terminal state.
200    pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Status> {
201        if let Some(shared) = self.inner.upgrade() {
202            let flags = shared.flags.get();
203            if flags.contains(Flags::NEED_READ) {
204                Poll::Ready(Status::Ready)
205            } else if flags.contains(Flags::SENDER_GONE | Flags::ERROR) {
206                Poll::Ready(Status::Dropped)
207            } else if flags.intersects(Flags::EOF) {
208                Poll::Ready(Status::Eof)
209            } else {
210                shared.send_task.register(cx.waker());
211                Poll::Pending
212            }
213        } else {
214            // receiver is gone
215            Poll::Ready(Status::Dropped)
216        }
217    }
218}
219
220bitflags::bitflags! {
221    #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
222    struct Flags: u8 {
223        const EOF         = 0b0000_0001;
224        const ERROR       = 0b0000_0010;
225        const NEED_READ   = 0b0000_0100;
226        const SENDER_GONE = 0b0000_1000;
227    }
228}
229
230struct Inner<E> {
231    len: Cell<usize>,
232    flags: Cell<Flags>,
233    err: Cell<Option<E>>,
234    items: RefCell<VecDeque<Bytes>>,
235    max_buffer_size: Cell<usize>,
236    recv_task: LocalWaker,
237    send_task: LocalWaker,
238}
239
240impl<E> Inner<E> {
241    fn new(eof: bool) -> Self {
242        let flags = if eof { Flags::EOF } else { Flags::NEED_READ };
243        Inner {
244            flags: Cell::new(flags),
245            len: Cell::new(0),
246            err: Cell::new(None),
247            items: RefCell::new(VecDeque::new()),
248            recv_task: LocalWaker::new(),
249            send_task: LocalWaker::new(),
250            max_buffer_size: Cell::new(MAX_BUFFER_SIZE),
251        }
252    }
253
254    fn insert_flag(&self, f: Flags) {
255        let mut flags = self.flags.get();
256        flags.insert(f);
257        self.flags.set(flags);
258    }
259
260    fn remove_flag(&self, f: Flags) {
261        let mut flags = self.flags.get();
262        flags.remove(f);
263        self.flags.set(flags);
264    }
265
266    fn set_max_buffer_size(&self, size: usize) {
267        self.max_buffer_size.set(size);
268
269        let flags = self.flags.get();
270        if flags.intersects(Flags::EOF | Flags::ERROR | Flags::SENDER_GONE) {
271            return;
272        }
273
274        if self.len.get() < size {
275            if !flags.contains(Flags::NEED_READ) {
276                self.insert_flag(Flags::NEED_READ);
277                self.send_task.wake();
278            }
279        } else {
280            self.remove_flag(Flags::NEED_READ);
281        }
282    }
283
284    fn set_error(&self, err: E) {
285        self.err.set(Some(err));
286        self.insert_flag(Flags::ERROR);
287        self.recv_task.wake();
288        self.send_task.wake();
289    }
290
291    fn feed_eof(&self) {
292        self.insert_flag(Flags::EOF);
293        self.recv_task.wake();
294        self.send_task.wake();
295    }
296
297    fn feed_data(&self, data: Bytes) {
298        let len = self.len.get() + data.len();
299        self.len.set(len);
300        self.items.borrow_mut().push_back(data);
301        self.recv_task.wake();
302
303        if len >= self.max_buffer_size.get() {
304            self.remove_flag(Flags::NEED_READ);
305        }
306    }
307
308    fn get_data(&self) -> Option<Bytes> {
309        self.items.borrow_mut().pop_front().inspect(|data| {
310            let len = self.len.get() - data.len();
311
312            // check size of stream buffer,
313            // if stream has more space wake up sender
314            self.len.set(len);
315            if len < self.max_buffer_size.get() {
316                self.insert_flag(Flags::NEED_READ);
317                self.send_task.wake();
318            }
319        })
320    }
321
322    fn unread_data(&self, data: Bytes) {
323        if !data.is_empty() {
324            self.len.set(self.len.get() + data.len());
325            self.items.borrow_mut().push_front(data);
326        }
327    }
328}
329
330impl<E> fmt::Debug for Inner<E> {
331    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
332        f.debug_struct("Inner")
333            .field("len", &self.len)
334            .field("flags", &self.flags)
335            .field("items", &self.items.borrow())
336            .field("max_buffer_size", &self.max_buffer_size)
337            .field("recv_task", &self.recv_task)
338            .field("send_task", &self.send_task)
339            .finish()
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346    use crate::future::lazy;
347
348    #[ntex::test]
349    async fn test_eof() {
350        let (tx, rx) = eof::<()>();
351        rx.max_buffer_size(100);
352        assert!(rx.read().await.is_none());
353        assert_eq!(tx.ready().await, Status::Eof);
354    }
355
356    #[ntex::test]
357    async fn test_closed() {
358        // drop receiver
359        let (tx, rx) = channel::<()>();
360        assert!(!tx.is_closed());
361        drop(rx);
362        assert!(tx.is_closed());
363
364        // drop sender
365        let (tx, rx) = channel::<()>();
366        drop(tx);
367        assert_eq!(rx.read().await, None);
368    }
369
370    #[ntex::test]
371    async fn test_unread_data() {
372        let (_, payload) = channel::<()>();
373
374        payload.put(Bytes::from("data"));
375        assert_eq!(payload.inner.len.get(), 4);
376        assert_eq!(
377            Bytes::from("data"),
378            poll_fn(|cx| payload.poll_read(cx)).await.unwrap().unwrap()
379        );
380    }
381
382    #[ntex::test]
383    async fn buffer_size_updates_sender_readiness() {
384        let (tx, rx) = channel::<()>();
385        rx.max_buffer_size(4);
386        tx.feed_data(Bytes::from_static(b"data"));
387
388        assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
389        assert!(rx.inner.send_task.is_set());
390
391        rx.max_buffer_size(5);
392        assert!(!rx.inner.send_task.is_set());
393        assert_eq!(
394            lazy(|cx| tx.poll_ready(cx)).await,
395            Poll::Ready(Status::Ready)
396        );
397
398        rx.max_buffer_size(4);
399        assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
400    }
401
402    #[ntex::test]
403    async fn test_sender_clone() {
404        let (sender, payload) = channel::<()>();
405        assert!(!payload.is_eof());
406        let sender2 = sender.clone();
407        assert!(!payload.is_eof());
408        drop(sender2);
409        assert!(!payload.is_eof());
410        drop(sender);
411        assert!(payload.is_eof());
412    }
413}
ntex_util/channel/bstream.rs

ntex_util/channel/
bstream.rs
