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) {
209 if let Some(shared) = self.inner.upgrade() {
210 shared.set_error(err);
211 }
212 }
213
214 pub fn feed_eof(&self) {
216 if let Some(shared) = self.inner.upgrade() {
217 shared.feed_eof();
218 }
219 }
220
221 pub fn feed_data(&self, data: Bytes) {
225 if let Some(shared) = self.inner.upgrade() {
226 shared.feed_data(data);
227 }
228 }
229
230 pub async fn ready(&self) -> Status {
232 poll_fn(|cx| self.poll_ready(cx)).await
233 }
234
235 pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Status> {
240 if let Some(shared) = self.inner.upgrade() {
241 let status = shared.flags.get().status();
242 if status.is_pending() {
243 shared.send_task.register(cx.waker());
244 }
245 status
246 } else {
247 Poll::Ready(Status::Dropped)
249 }
250 }
251}
252
253bitflags::bitflags! {
254 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
255 struct Flags: u8 {
256 const EOF = 0b0000_0001;
257 const ERROR = 0b0000_0010;
258 const NEED_READ = 0b0000_0100;
259 const SENDER_GONE = 0b0000_1000;
260 const RECEIVER_GONE = 0b0001_0000;
261 }
262}
263
264impl Flags {
265 fn status(self) -> Poll<Status> {
266 if self.intersects(Self::SENDER_GONE | Self::RECEIVER_GONE | Self::ERROR) {
267 Poll::Ready(Status::Dropped)
268 } else if self.contains(Self::EOF) {
269 Poll::Ready(Status::Eof)
270 } else if self.contains(Self::NEED_READ) {
271 Poll::Ready(Status::Ready)
272 } else {
273 Poll::Pending
274 }
275 }
276}
277
278struct Inner<E> {
279 len: Cell<usize>,
280 flags: Cell<Flags>,
281 err: Cell<Option<E>>,
282 items: RefCell<VecDeque<Bytes>>,
283 recv_task: LocalWaker,
284 send_task: LocalWaker,
285 high_watermark: Cell<u32>,
286 low_watermark: Cell<u32>,
287 changed: Option<ReadyNotifier>,
288}
289
290impl<E> Inner<E> {
291 fn new(eof: bool) -> Self {
292 let flags = if eof { Flags::EOF } else { Flags::NEED_READ };
293 Inner {
294 flags: Cell::new(flags),
295 len: Cell::new(0),
296 err: Cell::new(None),
297 items: RefCell::new(VecDeque::new()),
298 recv_task: LocalWaker::new(),
299 send_task: LocalWaker::new(),
300 high_watermark: Cell::new(HIGH_WATERMARK),
301 low_watermark: Cell::new(LOW_WATERMARK),
302 changed: None,
303 }
304 }
305
306 fn insert_flag(&self, f: Flags) {
307 let mut flags = self.flags.get();
308 flags.insert(f);
309 self.flags.set(flags);
310 }
311
312 fn remove_flag(&self, f: Flags) {
313 let mut flags = self.flags.get();
314 flags.remove(f);
315 self.flags.set(flags);
316 }
317
318 fn notify(&self, previous: Flags, recv: bool, send: bool) {
319 let _wake = OnDropFn::new(|| {
320 if recv {
321 self.recv_task.wake();
322 }
323 if send {
324 self.send_task.wake();
325 }
326 });
327 if previous.status() != self.flags.get().status()
328 && let Some(changed) = &self.changed
329 {
330 changed.notify();
331 }
332 }
333
334 fn set_watermarks(&self, high: u32, low: u32) {
335 self.high_watermark.set(high);
336 self.low_watermark.set(low.min(high.saturating_sub(1)));
337
338 let flags = self.flags.get();
339 if flags.intersects(Flags::EOF | Flags::ERROR | Flags::SENDER_GONE) {
340 return;
341 }
342
343 if self.len.get() < high as usize {
344 if !flags.contains(Flags::NEED_READ) {
345 self.insert_flag(Flags::NEED_READ);
346 self.notify(flags, false, true);
347 }
348 } else {
349 self.remove_flag(Flags::NEED_READ);
350 self.notify(flags, false, false);
351 }
352 }
353
354 fn set_error(&self, err: E) {
355 let flags = self.flags.get();
356 self.err.set(Some(err));
357 self.insert_flag(Flags::ERROR);
358 self.notify(flags, true, true);
359 }
360
361 fn feed_eof(&self) {
362 let flags = self.flags.get();
363 self.insert_flag(Flags::EOF);
364 self.notify(flags, true, true);
365 }
366
367 fn feed_data(&self, data: Bytes) {
368 let flags = self.flags.get();
369 let len = self.len.get() + data.len();
370 self.len.set(len);
371 self.items.borrow_mut().push_back(data);
372 if len >= self.high_watermark.get() as usize {
373 self.remove_flag(Flags::NEED_READ);
374 }
375 self.notify(flags, true, false);
376 }
377
378 fn get_data(&self) -> Option<Bytes> {
379 let data = self.items.borrow_mut().pop_front();
380 data.inspect(|data| {
381 let flags = self.flags.get();
382 let len = self.len.get() - data.len();
383
384 self.len.set(len);
386 if len <= self.low_watermark.get() as usize {
387 self.insert_flag(Flags::NEED_READ);
388 self.notify(flags, false, true);
389 }
390 })
391 }
392
393 fn unread_data(&self, data: Bytes) {
394 let flags = self.flags.get();
395 if !data.is_empty() {
396 self.len.set(self.len.get() + data.len());
397 self.items.borrow_mut().push_front(data);
398 if self.len.get() >= self.high_watermark.get() as usize {
399 self.remove_flag(Flags::NEED_READ);
400 }
401 self.notify(flags, false, false);
402 }
403 }
404}
405
406impl<E> fmt::Debug for Inner<E> {
407 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
408 f.debug_struct("Inner")
409 .field("len", &self.len)
410 .field("flags", &self.flags)
411 .field("items", &self.items.borrow())
412 .field("high_watermark", &self.high_watermark)
413 .field("low_watermark", &self.low_watermark)
414 .field("recv_task", &self.recv_task)
415 .field("send_task", &self.send_task)
416 .finish()
417 }
418}
419
420#[cfg(test)]
421mod tests {
422 use ntex_service::{BindCtx, Ctx, Pipeline, ReadyCallback, Service};
423
424 use super::*;
425 use crate::future::lazy;
426
427 struct SenderService(Sender<()>, ReadyNotifier);
428
429 impl Service<(), Bytes> for SenderService {
430 type Res = ();
431 type Error = Status;
432
433 fn bind(&self, ctx: BindCtx<'_>) {
434 ctx.subscribe(&self.1);
435 }
436
437 async fn ready(&self, _: Ctx<'_, Self>) -> Result<(), Status> {
438 match self.0.ready().await {
439 Status::Ready => Ok(()),
440 status => Err(status),
441 }
442 }
443
444 async fn call(&self, data: Bytes, _: Ctx<'_, Self>) -> Result<(), Status> {
445 self.0.feed_data(data);
446 Ok(())
447 }
448 }
449
450 #[ntex::test]
456 async fn cached_sender_readiness_tracks_watermarks() {
457 let changed = ReadyNotifier::new();
458 let (tx, rx) = channel_with_notifier(changed.clone());
459 rx.set_watermarks(4, 2);
460 let srv = Pipeline::new((), SenderService(tx, changed));
461
462 srv.call(Bytes::from_static(b"ab")).await.unwrap();
463 assert!(srv.is_ready());
464 srv.call(Bytes::from_static(b"cd")).await.unwrap();
465 assert!(!srv.is_ready());
466 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
467 assert_eq!(rx.read().await, Some(Ok(Bytes::from_static(b"ab"))));
468 srv.ready().await.unwrap();
469
470 rx.put(Bytes::from_static(b"ef"));
471 assert!(!srv.is_ready());
472 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
473 rx.set_watermarks(8, 4);
474 srv.ready().await.unwrap();
475 rx.set_watermarks(4, 2);
476 assert!(!srv.is_ready());
477 assert!(lazy(|cx| srv.poll_ready(cx)).await.is_pending());
478 }
479
480 #[ntex::test]
486 async fn cached_sender_readiness_tracks_terminal_states() {
487 for status in [Status::Eof, Status::Dropped] {
488 let changed = ReadyNotifier::new();
489 let (tx, _rx) = channel_with_notifier(changed.clone());
490 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
491 srv.ready().await.unwrap();
492 if status == Status::Eof {
493 tx.feed_eof();
494 } else {
495 tx.set_error(());
496 }
497 assert!(!srv.is_ready());
498 assert_eq!(srv.ready().await, Err(status));
499 }
500
501 let changed = ReadyNotifier::new();
502 let (tx, rx) = channel_with_notifier(changed.clone());
503 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
504 let on_drop = ReadyCallback::new(move |_| {
505 assert!(tx.is_closed());
506 assert_eq!(
507 tx.poll_ready(&mut Context::from_waker(std::task::Waker::noop())),
508 Poll::Ready(Status::Dropped)
509 );
510 });
511 srv.on_notify(&on_drop);
512 srv.ready().await.unwrap();
513 drop(rx);
514 assert!(!srv.is_ready());
515 assert_eq!(srv.ready().await, Err(Status::Dropped));
516 }
517
518 #[ntex::test]
524 async fn callback_panic_does_not_skip_stream_wakes() {
525 use std::panic::{AssertUnwindSafe, catch_unwind};
526 use std::sync::{Arc, atomic::AtomicUsize, atomic::Ordering};
527 use std::task::{Wake, Waker};
528
529 #[derive(Default)]
530 struct WakeCount(AtomicUsize);
531 impl Wake for WakeCount {
532 fn wake(self: Arc<Self>) {
533 self.0.fetch_add(1, Ordering::Relaxed);
534 }
535 }
536
537 let changed = ReadyNotifier::new();
538 let (tx, rx) = channel_with_notifier(changed.clone());
539 rx.set_watermarks(1, 0);
540 let srv = Pipeline::new((), SenderService(tx.clone(), changed));
541 srv.ready().await.unwrap();
542 let armed = Rc::new(Cell::new(true));
543 let armed2 = armed.clone();
544 let callback = ReadyCallback::new(move |_| {
545 assert!(!armed2.replace(false), "callback panic");
546 });
547 srv.on_notify(&callback);
548 let wakes = Arc::new(WakeCount::default());
549 let waker = Waker::from(wakes.clone());
550 let mut cx = Context::from_waker(&waker);
551
552 assert!(rx.poll_read(&mut cx).is_pending());
553 assert!(catch_unwind(AssertUnwindSafe(|| tx.feed_data(b"x".into()))).is_err());
554 assert!(!srv.is_ready());
555 assert_eq!(wakes.0.load(Ordering::Relaxed), 1);
556 assert!(tx.poll_ready(&mut cx).is_pending());
557
558 armed.set(true);
559 assert!(catch_unwind(AssertUnwindSafe(|| rx.poll_read(&mut cx))).is_err());
560 assert_eq!(wakes.0.load(Ordering::Relaxed), 2);
561 srv.ready().await.unwrap();
562 }
563
564 #[ntex::test]
570 async fn test_sender_drop_wakes_receiver() {
571 let (tx, rx) = channel::<()>();
572 let handle = crate::spawn(async move { rx.read().await.is_none() });
573 crate::time::sleep(crate::time::Millis(10)).await;
574
575 drop(tx);
576 let res = crate::time::timeout(crate::time::Millis(1000), handle).await;
577 assert!(matches!(res, Ok(Ok(true))));
578 }
579
580 #[ntex::test]
585 async fn test_eof() {
586 let (tx, rx) = eof::<()>();
587 rx.set_watermarks(100, 50);
588 assert!(rx.read().await.is_none());
589 assert_eq!(tx.ready().await, Status::Eof);
590 }
591
592 #[ntex::test]
598 async fn test_closed() {
599 let (tx, rx) = channel::<()>();
601 assert!(!tx.is_closed());
602 drop(rx);
603 assert!(tx.is_closed());
604
605 let (tx, rx) = channel::<()>();
607 drop(tx);
608 assert_eq!(rx.read().await, None);
609 }
610
611 #[ntex::test]
616 async fn test_unread_data() {
617 let (_, payload) = channel::<()>();
618
619 payload.put(Bytes::from("data"));
620 assert_eq!(payload.inner.len.get(), 4);
621 assert_eq!(
622 Bytes::from("data"),
623 poll_fn(|cx| payload.poll_read(cx)).await.unwrap().unwrap()
624 );
625 }
626
627 #[ntex::test]
633 async fn buffer_size_updates_sender_readiness() {
634 let (tx, rx) = channel::<()>();
635 rx.set_watermarks(4, 2);
636 tx.feed_data(Bytes::from_static(b"data"));
637
638 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
639 assert!(rx.inner.send_task.is_set());
640
641 rx.set_watermarks(5, 2);
642 assert!(!rx.inner.send_task.is_set());
643 assert_eq!(
644 lazy(|cx| tx.poll_ready(cx)).await,
645 Poll::Ready(Status::Ready)
646 );
647
648 rx.set_watermarks(4, 2);
649 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
650 }
651
652 #[ntex::test]
658 async fn sender_resumes_at_low_watermark() {
659 let (tx, rx) = channel::<()>();
660 rx.set_watermarks(8, 4);
661 for _ in 0..4 {
662 tx.feed_data(Bytes::from_static(b"da"));
663 }
664 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
665
666 assert!(rx.read().await.is_some());
668 assert!(rx.inner.send_task.is_set());
669 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
670
671 assert!(rx.read().await.is_some());
673 assert!(!rx.inner.send_task.is_set());
674 assert_eq!(
675 lazy(|cx| tx.poll_ready(cx)).await,
676 Poll::Ready(Status::Ready)
677 );
678
679 let (tx, rx) = channel::<()>();
681 rx.set_watermarks(4, 10);
682 for _ in 0..3 {
683 tx.feed_data(Bytes::from_static(b"da"));
684 }
685 assert!(rx.read().await.is_some());
686 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
687 assert!(rx.read().await.is_some());
688 assert_eq!(
689 lazy(|cx| tx.poll_ready(cx)).await,
690 Poll::Ready(Status::Ready)
691 );
692 }
693
694 #[ntex::test]
700 async fn custom_low_watermark() {
701 let (tx, rx) = channel::<()>();
702 rx.set_watermarks(8, 2);
703 for _ in 0..4 {
704 tx.feed_data(Bytes::from_static(b"da"));
705 }
706 assert!(rx.read().await.is_some());
707 assert!(rx.read().await.is_some());
708 assert!(lazy(|cx| tx.poll_ready(cx)).await.is_pending());
709 assert!(rx.read().await.is_some());
710 assert_eq!(
711 lazy(|cx| tx.poll_ready(cx)).await,
712 Poll::Ready(Status::Ready)
713 );
714 }
715
716 #[ntex::test]
721 #[allow(deprecated)]
722 async fn deprecated_max_buffer_size() {
723 let (_tx, rx) = channel::<()>();
724 rx.max_buffer_size(10);
725 assert_eq!(rx.inner.high_watermark.get(), 10);
726 assert_eq!(rx.inner.low_watermark.get(), 5);
727 }
728
729 #[ntex::test]
734 async fn test_sender_clone() {
735 let (sender, payload) = channel::<()>();
736 assert!(!payload.is_eof());
737 let sender2 = sender.clone();
738 assert!(!payload.is_eof());
739 drop(sender2);
740 assert!(!payload.is_eof());
741 drop(sender);
742 assert!(payload.is_eof());
743 }
744
745 #[ntex::test]
750 async fn test_ready_terminal_states() {
751 let (tx, _rx) = channel::<()>();
752 assert_eq!(tx.ready().await, Status::Ready);
753 tx.set_error(());
754 assert_eq!(tx.ready().await, Status::Dropped);
755
756 let (tx, _rx) = channel::<()>();
757 tx.feed_eof();
758 assert_eq!(tx.ready().await, Status::Eof);
759
760 let (tx, rx) = channel::<()>();
761 drop(rx);
762 assert_eq!(tx.ready().await, Status::Dropped);
763 }
764
765 #[ntex::test]
771 async fn test_empty() {
772 let rx = empty::<()>(None);
773 assert!(rx.is_eof());
774 assert!(rx.read().await.is_none());
775
776 let mut rx = empty::<()>(Some(Bytes::from_static(b"data")));
777 assert!(format!("{rx:?}").contains("high_watermark"));
778 assert_eq!(
779 crate::future::stream_recv(&mut rx).await.unwrap().unwrap(),
780 Bytes::from_static(b"data")
781 );
782 assert!(crate::future::stream_recv(&mut rx).await.is_none());
783 }
784
785 #[ntex::test]
791 async fn test_error() {
792 let (tx, rx) = channel::<&'static str>();
793 tx.feed_data(Bytes::from_static(b"data"));
794 tx.set_error("err");
795
796 assert_eq!(
798 rx.read().await.unwrap().unwrap(),
799 Bytes::from_static(b"data")
800 );
801 assert_eq!(rx.read().await.unwrap().unwrap_err(), "err");
802 assert!(rx.is_eof());
803 assert!(rx.read().await.is_none());
804 }
805}