1use 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;
7use ntex_service::ReadyNotifier;
8
9use crate::{Stream, future::OnDropFn, task::LocalWaker};
10
11const HIGH_WATERMARK: u32 = 32_768;
13const LOW_WATERMARK: u32 = HIGH_WATERMARK / 2;
15
16#[derive(Copy, Clone, Debug, PartialEq, Eq)]
18pub enum Status {
19 Eof,
21 Ready,
23 Dropped,
25}
26
27pub fn channel<E>() -> (Sender<E>, Receiver<E>) {
29 channel_inner(Inner::new(false))
30}
31
32pub fn channel_with_notifier<E>(changed: ReadyNotifier) -> (Sender<E>, Receiver<E>) {
42 channel_inner(Inner {
43 changed: Some(changed),
44 ..Inner::new(false)
45 })
46}
47
48pub fn eof<E>() -> (Sender<E>, Receiver<E>) {
53 channel_inner(Inner::new(true))
54}
55
56fn channel_inner<E>(inner: Inner<E>) -> (Sender<E>, Receiver<E>) {
57 let inner = Rc::new(inner);
58 (
59 Sender {
60 inner: Rc::downgrade(&inner),
61 },
62 Receiver { inner },
63 )
64}
65
66pub fn empty<E>(data: Option<Bytes>) -> Receiver<E> {
68 let rx = Receiver {
69 inner: Rc::new(Inner::new(true)),
70 };
71 if let Some(data) = data {
72 rx.put(data);
73 }
74 rx
75}
76
77#[derive(Debug)]
83pub struct Receiver<E> {
84 inner: Rc<Inner<E>>,
85}
86
87impl<E> Receiver<E> {
88 #[inline]
100 pub fn set_watermarks(&self, high: u32, low: u32) {
101 self.inner.set_watermarks(high, low);
102 }
103
104 #[inline]
108 #[deprecated(since = "4.2.0", note = "Use `Receiver::set_watermarks()` instead")]
109 pub fn max_buffer_size(&self, size: usize) {
110 let size = u32::try_from(size).unwrap_or(u32::MAX);
111 self.inner.set_watermarks(size, size / 2);
112 }
113
114 #[inline]
118 pub fn put(&self, data: Bytes) {
119 self.inner.unread_data(data);
120 }
121
122 #[inline]
123 pub fn is_eof(&self) -> bool {
127 self.inner.flags.get().contains(Flags::EOF)
128 }
129
130 #[inline]
131 pub async fn read(&self) -> Option<Result<Bytes, E>> {
133 poll_fn(|cx| self.poll_read(cx)).await
134 }
135
136 #[inline]
137 pub fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, E>>> {
139 if let Some(data) = self.inner.get_data() {
140 Poll::Ready(Some(Ok(data)))
141 } else if let Some(err) = self.inner.err.take() {
142 self.inner.insert_flag(Flags::EOF);
143 Poll::Ready(Some(Err(err)))
144 } else if self.inner.flags.get().intersects(Flags::EOF | Flags::ERROR) {
145 Poll::Ready(None)
146 } else {
147 self.inner.recv_task.register(cx.waker());
148 Poll::Pending
149 }
150 }
151}
152
153impl<E> Stream for Receiver<E> {
154 type Item = Result<Bytes, E>;
155
156 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
157 self.poll_read(cx)
158 }
159}
160
161impl<E> Drop for Receiver<E> {
162 fn drop(&mut self) {
163 let flags = self.inner.flags.get();
164 self.inner.insert_flag(Flags::RECEIVER_GONE);
165 self.inner.notify(flags, false, true);
166 }
167}
168
169#[derive(Debug)]
174pub struct Sender<E> {
175 inner: Weak<Inner<E>>,
176}
177
178impl<E> Clone for Sender<E> {
179 fn clone(&self) -> Self {
180 Self {
181 inner: self.inner.clone(),
182 }
183 }
184}
185
186impl<E> Drop for Sender<E> {
187 fn drop(&mut self) {
188 if self.inner.weak_count() == 1
189 && let Some(shared) = self.inner.upgrade()
190 {
191 let flags = shared.flags.get();
192 shared.insert_flag(Flags::EOF | Flags::SENDER_GONE);
193 shared.notify(flags, true, true);
195 }
196 }
197}
198
199impl<E> Sender<E> {
200 pub fn is_closed(&self) -> bool {
202 self.inner
203 .upgrade()
204 .is_none_or(|shared| shared.flags.get().contains(Flags::RECEIVER_GONE))
205 }
206
207 pub fn set_error(&self, err: E) {
212 if let Some(shared) = self.inner.upgrade() {
213 shared.set_error(err);
214 }
215 }
216
217 pub fn feed_eof(&self) {
219 if let Some(shared) = self.inner.upgrade() {
220 shared.feed_eof();
221 }
222 }
223
224 pub fn feed_data(&self, data: Bytes) {
229 if let Some(shared) = self.inner.upgrade() {
230 shared.feed_data(data);
231 }
232 }
233
234 pub async fn ready(&self) -> Status {
236 poll_fn(|cx| self.poll_ready(cx)).await
237 }
238
239 pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Status> {
244 if let Some(shared) = self.inner.upgrade() {
245 let status = shared.flags.get().status();
246 if status.is_pending() {
247 shared.send_task.register(cx.waker());
248 }
249 status
250 } else {
251 Poll::Ready(Status::Dropped)
253 }
254 }
255}
256
257bitflags::bitflags! {
258 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
259 struct Flags: u8 {
260 const EOF = 0b0000_0001;
261 const ERROR = 0b0000_0010;
262 const NEED_READ = 0b0000_0100;
263 const SENDER_GONE = 0b0000_1000;
264 const RECEIVER_GONE = 0b0001_0000;
265 }
266}
267
268impl Flags {
269 fn status(self) -> Poll<Status> {
270 if self.intersects(Self::SENDER_GONE | Self::RECEIVER_GONE | Self::ERROR) {
271 Poll::Ready(Status::Dropped)
272 } else if self.contains(Self::EOF) {
273 Poll::Ready(Status::Eof)
274 } else if self.contains(Self::NEED_READ) {
275 Poll::Ready(Status::Ready)
276 } else {
277 Poll::Pending
278 }
279 }
280}
281
282struct Inner<E> {
283 len: Cell<usize>,
284 flags: Cell<Flags>,
285 err: Cell<Option<E>>,
286 items: RefCell<VecDeque<Bytes>>,
287 recv_task: LocalWaker,
288 send_task: LocalWaker,
289 high_watermark: Cell<u32>,
290 low_watermark: Cell<u32>,
291 changed: Option<ReadyNotifier>,
292}
293
294impl<E> Inner<E> {
295 fn new(eof: bool) -> Self {
296 let flags = if eof { Flags::EOF } else { Flags::NEED_READ };
297 Inner {
298 flags: Cell::new(flags),
299 len: Cell::new(0),
300 err: Cell::new(None),
301 items: RefCell::new(VecDeque::new()),
302 recv_task: LocalWaker::new(),
303 send_task: LocalWaker::new(),
304 high_watermark: Cell::new(HIGH_WATERMARK),
305 low_watermark: Cell::new(LOW_WATERMARK),
306 changed: None,
307 }
308 }
309
310 fn insert_flag(&self, f: Flags) {
311 let mut flags = self.flags.get();
312 flags.insert(f);
313 self.flags.set(flags);
314 }
315
316 fn remove_flag(&self, f: Flags) {
317 let mut flags = self.flags.get();
318 flags.remove(f);
319 self.flags.set(flags);
320 }
321
322 fn notify(&self, previous: Flags, recv: bool, send: bool) {
323 let _wake = OnDropFn::new(|| {
324 if recv {
325 self.recv_task.wake();
326 }
327 if send {
328 self.send_task.wake();
329 }
330 });
331 if previous.status() != self.flags.get().status()
332 && let Some(changed) = &self.changed
333 {
334 changed.notify();
335 }
336 }
337
338 fn set_watermarks(&self, high: u32, low: u32) {
339 self.high_watermark.set(high);
340 self.low_watermark.set(low.min(high.saturating_sub(1)));
341
342 let flags = self.flags.get();
343 if flags.intersects(Flags::EOF | Flags::ERROR | Flags::SENDER_GONE) {
344 return;
345 }
346
347 if self.len.get() < high as usize {
348 if !flags.contains(Flags::NEED_READ) {
349 self.insert_flag(Flags::NEED_READ);
350 self.notify(flags, false, true);
351 }
352 } else {
353 self.remove_flag(Flags::NEED_READ);
354 self.notify(flags, false, false);
355 }
356 }
357
358 fn set_error(&self, err: E) {
359 let flags = self.flags.get();
360 if flags.contains(Flags::EOF) {
361 return;
362 }
363 self.err.set(Some(err));
364 self.insert_flag(Flags::ERROR);
365 self.notify(flags, true, true);
366 }
367
368 fn feed_eof(&self) {
369 let flags = self.flags.get();
370 self.insert_flag(Flags::EOF);
371 self.notify(flags, true, true);
372 }
373
374 fn feed_data(&self, data: Bytes) {
375 let flags = self.flags.get();
376 if flags.intersects(Flags::EOF | Flags::ERROR) {
377 return;
378 }
379 let len = self.len.get() + data.len();
380 self.len.set(len);
381 self.items.borrow_mut().push_back(data);
382 if len >= self.high_watermark.get() as usize {
383 self.remove_flag(Flags::NEED_READ);
384 }
385 self.notify(flags, true, false);
386 }
387
388 fn get_data(&self) -> Option<Bytes> {
389 let data = self.items.borrow_mut().pop_front();
390 data.inspect(|data| {
391 let flags = self.flags.get();
392 let len = self.len.get() - data.len();
393
394 self.len.set(len);
396 if len <= self.low_watermark.get() as usize {
397 self.insert_flag(Flags::NEED_READ);
398 self.notify(flags, false, true);
399 }
400 })
401 }
402
403 fn unread_data(&self, data: Bytes) {
404 let flags = self.flags.get();
405 if !data.is_empty() {
406 self.len.set(self.len.get() + data.len());
407 self.items.borrow_mut().push_front(data);
408 if self.len.get() >= self.high_watermark.get() as usize {
409 self.remove_flag(Flags::NEED_READ);
410 }
411 self.notify(flags, false, false);
412 }
413 }
414}
415
416impl<E> fmt::Debug for Inner<E> {
417 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
418 f.debug_struct("Inner")
419 .field("len", &self.len)
420 .field("flags", &self.flags)
421 .field("items", &self.items.borrow())
422 .field("high_watermark", &self.high_watermark)
423 .field("low_watermark", &self.low_watermark)
424 .field("recv_task", &self.recv_task)
425 .field("send_task", &self.send_task)
426 .finish()
427 }
428}
429
430#[cfg(test)]
431mod tests {
432 use ntex_service::{BindCtx, Ctx, Pipeline, ReadyCallback, Service};
433
434 use super::*;
435 use crate::future::lazy;
436
437 struct SenderService(Sender<()>, ReadyNotifier);
438
439 impl Service<(), Bytes> for SenderService {
440 type Res = ();
441 type Error = Status;
442
443 fn bind(&self, ctx: BindCtx<'_>) {
444 ctx.subscribe(&self.1);
445 }
446
447 async fn ready(&self, _: Ctx<'_, Self>) -> Result<(), Status> {
448 match self.0.ready().await {
449 Status::Ready => Ok(()),
450 status => Err(status),
451 }
452 }
453
454 async fn call(&self, data: Bytes, _: Ctx<'_, Self>) -> Result<(), Status> {
455 self.0.feed_data(data);
456 Ok(())
457 }
458 }
459
460 #[ntex::test]
466 async fn cached_sender_readiness_tracks_watermarks() {
467 let changed = ReadyNotifier::new();
468 let (tx, rx) = channel_with_notifier(changed.clone());
469 rx.set_watermarks(4, 2);
470 let srv = Pipeline::new((), SenderService(tx, changed));
471
472 srv.call(Bytes::from_static(b"ab")).await.unwrap();
473 assert!(srv.is_ready());
474 srv.call(Bytes::from_static(b"cd")).await.unwrap();
475 assert!(!srv.is_ready());
476 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
477 assert_eq!(rx.read().await, Some(Ok(Bytes::from_static(b"ab"))));
478 srv.ready().await.unwrap();
479
480 rx.put(Bytes::from_static(b"ef"));
481 assert!(!srv.is_ready());
482 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
483 rx.set_watermarks(8, 4);
484 srv.ready().await.unwrap();
485 rx.set_watermarks(4, 2);
486 assert!(!srv.is_ready());
487 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
488 }
489
490 #[ntex::test]
496 async fn cached_sender_readiness_tracks_terminal_states() {
497 for status in [Status::Eof, Status::Dropped] {
498 let changed = ReadyNotifier::new();
499 let (tx, _rx) = channel_with_notifier(changed.clone());
500 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
501 srv.ready().await.unwrap();
502 if status == Status::Eof {
503 tx.feed_eof();
504 } else {
505 tx.set_error(());
506 }
507 assert!(!srv.is_ready());
508 assert_eq!(srv.ready().await, Err(status));
509 }
510
511 let changed = ReadyNotifier::new();
512 let (tx, rx) = channel_with_notifier(changed.clone());
513 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
514 let on_drop = ReadyCallback::new(move |_| {
515 assert!(tx.is_closed());
516 assert_eq!(
517 tx.poll_ready(&mut Context::from_waker(std::task::Waker::noop())),
518 Poll::Ready(Status::Dropped)
519 );
520 });
521 srv.on_notify(&on_drop);
522 srv.ready().await.unwrap();
523 drop(rx);
524 assert!(!srv.is_ready());
525 assert_eq!(srv.ready().await, Err(Status::Dropped));
526 }
527
528 #[ntex::test]
534 async fn callback_panic_does_not_skip_stream_wakes() {
535 use std::panic::{AssertUnwindSafe, catch_unwind};
536 use std::sync::{Arc, atomic::AtomicUsize, atomic::Ordering};
537 use std::task::{Wake, Waker};
538
539 #[derive(Default)]
540 struct WakeCount(AtomicUsize);
541 impl Wake for WakeCount {
542 fn wake(self: Arc<Self>) {
543 self.0.fetch_add(1, Ordering::Relaxed);
544 }
545 }
546
547 let changed = ReadyNotifier::new();
548 let (tx, rx) = channel_with_notifier(changed.clone());
549 rx.set_watermarks(1, 0);
550 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
551 srv.ready().await.unwrap();
552 let armed = Rc::new(Cell::new(true));
553 let armed2 = armed.clone();
554 let callback = ReadyCallback::new(move |_| {
555 assert!(!armed2.replace(false), "callback panic");
556 });
557 srv.on_notify(&callback);
558 let wakes = Arc::new(WakeCount::default());
559 let waker = Waker::from(wakes.clone());
560 let mut cx = Context::from_waker(&waker);
561
562 assert!(rx.poll_read(&mut cx).is_pending());
563 assert!(catch_unwind(AssertUnwindSafe(|| tx.feed_data(b"x".into()))).is_err());
564 assert!(!srv.is_ready());
565 assert_eq!(wakes.0.load(Ordering::Relaxed), 1);
566 assert!(tx.poll_ready(&mut cx).is_pending());
567
568 armed.set(true);
569 assert!(catch_unwind(AssertUnwindSafe(|| rx.poll_read(&mut cx))).is_err());
570 assert_eq!(wakes.0.load(Ordering::Relaxed), 2);
571 srv.ready().await.unwrap();
572 }
573
574 #[ntex::test]
580 async fn test_sender_drop_wakes_receiver() {
581 let (tx, rx) = channel::<()>();
582 let handle = crate::spawn(async move { rx.read().await.is_none() });
583 crate::time::sleep(crate::time::Millis(10)).await;
584
585 drop(tx);
586 let res = crate::time::timeout(crate::time::Millis(1000), handle).await;
587 assert!(matches!(res, Ok(Ok(true))));
588 }
589
590 #[ntex::test]
595 async fn test_eof() {
596 let (tx, rx) = eof::<()>();
597 rx.set_watermarks(100, 50);
598 assert!(rx.read().await.is_none());
599 assert_eq!(tx.ready().await, Status::Eof);
600 }
601
602 #[ntex::test]
608 async fn test_closed() {
609 let (tx, rx) = channel::<()>();
611 assert!(!tx.is_closed());
612 drop(rx);
613 assert!(tx.is_closed());
614
615 let (tx, rx) = channel::<()>();
617 drop(tx);
618 assert_eq!(rx.read().await, None);
619 }
620
621 #[ntex::test]
626 async fn test_unread_data() {
627 let (_, payload) = channel::<()>();
628
629 payload.put(Bytes::from("data"));
630 assert_eq!(payload.inner.len.get(), 4);
631 assert_eq!(
632 Bytes::from("data"),
633 poll_fn(|cx| payload.poll_read(cx)).await.unwrap().unwrap()
634 );
635 }
636
637 #[ntex::test]
643 async fn buffer_size_updates_sender_readiness() {
644 let (tx, rx) = channel::<()>();
645 rx.set_watermarks(4, 2);
646 tx.feed_data(Bytes::from_static(b"data"));
647
648 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
649 assert!(rx.inner.send_task.is_set());
650
651 rx.set_watermarks(5, 2);
652 assert!(!rx.inner.send_task.is_set());
653 assert_eq!(
654 lazy(|cx| tx.poll_ready(cx)).await,
655 Poll::Ready(Status::Ready)
656 );
657
658 rx.set_watermarks(4, 2);
659 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
660 }
661
662 #[ntex::test]
668 async fn sender_resumes_at_low_watermark() {
669 let (tx, rx) = channel::<()>();
670 rx.set_watermarks(8, 4);
671 for _ in 0..4 {
672 tx.feed_data(Bytes::from_static(b"da"));
673 }
674 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
675
676 assert!(rx.read().await.is_some());
678 assert!(rx.inner.send_task.is_set());
679 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
680
681 assert!(rx.read().await.is_some());
683 assert!(!rx.inner.send_task.is_set());
684 assert_eq!(
685 lazy(|cx| tx.poll_ready(cx)).await,
686 Poll::Ready(Status::Ready)
687 );
688
689 let (tx, rx) = channel::<()>();
691 rx.set_watermarks(4, 10);
692 for _ in 0..3 {
693 tx.feed_data(Bytes::from_static(b"da"));
694 }
695 assert!(rx.read().await.is_some());
696 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
697 assert!(rx.read().await.is_some());
698 assert_eq!(
699 lazy(|cx| tx.poll_ready(cx)).await,
700 Poll::Ready(Status::Ready)
701 );
702 }
703
704 #[ntex::test]
710 async fn custom_low_watermark() {
711 let (tx, rx) = channel::<()>();
712 rx.set_watermarks(8, 2);
713 for _ in 0..4 {
714 tx.feed_data(Bytes::from_static(b"da"));
715 }
716 assert!(rx.read().await.is_some());
717 assert!(rx.read().await.is_some());
718 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
719 assert!(rx.read().await.is_some());
720 assert_eq!(
721 lazy(|cx| tx.poll_ready(cx)).await,
722 Poll::Ready(Status::Ready)
723 );
724 }
725
726 #[ntex::test]
731 #[allow(deprecated)]
732 async fn deprecated_max_buffer_size() {
733 let (_tx, rx) = channel::<()>();
734 rx.max_buffer_size(10);
735 assert_eq!(rx.inner.high_watermark.get(), 10);
736 assert_eq!(rx.inner.low_watermark.get(), 5);
737 }
738
739 #[ntex::test]
744 async fn test_sender_clone() {
745 let (sender, payload) = channel::<()>();
746 assert!(!payload.is_eof());
747 let sender2 = sender.clone();
748 assert!(!payload.is_eof());
749 drop(sender2);
750 assert!(!payload.is_eof());
751 drop(sender);
752 assert!(payload.is_eof());
753 }
754
755 #[ntex::test]
760 async fn test_ready_terminal_states() {
761 let (tx, _rx) = channel::<()>();
762 assert_eq!(tx.ready().await, Status::Ready);
763 tx.set_error(());
764 assert_eq!(tx.ready().await, Status::Dropped);
765
766 let (tx, _rx) = channel::<()>();
767 tx.feed_eof();
768 assert_eq!(tx.ready().await, Status::Eof);
769
770 let (tx, rx) = channel::<()>();
771 drop(rx);
772 assert_eq!(tx.ready().await, Status::Dropped);
773 }
774
775 #[ntex::test]
781 async fn test_empty() {
782 let rx = empty::<()>(None);
783 assert!(rx.is_eof());
784 assert!(rx.read().await.is_none());
785
786 let mut rx = empty::<()>(Some(Bytes::from_static(b"data")));
787 assert!(format!("{rx:?}").contains("high_watermark"));
788 assert_eq!(
789 crate::future::stream_recv(&mut rx).await.unwrap().unwrap(),
790 Bytes::from_static(b"data")
791 );
792 assert!(crate::future::stream_recv(&mut rx).await.is_none());
793 }
794
795 #[ntex::test]
801 async fn test_error() {
802 let (tx, rx) = channel::<&'static str>();
803 tx.feed_data(Bytes::from_static(b"data"));
804 tx.set_error("err");
805
806 assert_eq!(
808 rx.read().await.unwrap().unwrap(),
809 Bytes::from_static(b"data")
810 );
811 assert_eq!(rx.read().await.unwrap().unwrap_err(), "err");
812 assert!(rx.is_eof());
813 assert!(rx.read().await.is_none());
814 }
815
816 #[ntex::test]
823 async fn eof_and_error_are_terminal() {
824 let (tx, rx) = channel::<&'static str>();
825 tx.feed_data(Bytes::from_static(b"data"));
826 tx.feed_eof();
827 tx.feed_data(Bytes::from_static(b"late"));
828 tx.set_error("late");
829 assert_eq!(tx.ready().await, Status::Eof);
830 assert_eq!(
831 rx.read().await.unwrap().unwrap(),
832 Bytes::from_static(b"data")
833 );
834 assert!(rx.read().await.is_none());
835 tx.feed_data(Bytes::from_static(b"later"));
836 assert!(rx.read().await.is_none());
837
838 let (tx, rx) = channel::<&'static str>();
839 tx.set_error("err");
840 tx.feed_data(Bytes::from_static(b"late"));
841 assert_eq!(rx.read().await.unwrap().unwrap_err(), "err");
842 tx.set_error("late");
843 assert!(rx.read().await.is_none());
844 }
845}