1use std::cell::{Cell, UnsafeCell};
2use std::future::{Future, poll_fn};
3use std::task::{Context, Poll};
4use std::{fmt, hash, io, marker, mem, ops, pin::Pin, ptr, rc::Rc};
5
6use ntex_bytes::{BytePageSize, BytesMut};
7use ntex_codec::{Decoder, Encoder};
8use ntex_service::cfg::{Cfg, SharedCfg};
9use ntex_util::{future::Either, task::LocalWaker, time::Sleep};
10
11use crate::buf::Stack;
12use crate::cfg::IoConfig;
13use crate::ctx::IoContext;
14use crate::filter::{Base, Filter, Layer};
15use crate::filterptr::FilterPtr;
16use crate::flags::Flags;
17use crate::ops::{Id, IoManager, TimerHandle};
18use crate::seal::{IoBoxed, Sealed};
19use crate::utils::Extensions;
20use crate::{Decoded, FilterLayer, Handle, IoStatusUpdate, IoStream, RecvError};
21
22pub struct Io<F = Base>(UnsafeCell<IoRef>, marker::PhantomData<F>);
24
25#[derive(Clone)]
26pub struct IoRef(pub(super) Rc<IoState>);
27
28pub(crate) struct IoState {
29 filter: FilterPtr,
30 pub(super) id: Cell<Id>,
31 pub(super) cfg: Cfg<IoConfig>,
32 pub(super) flags: Flags,
33 pub(super) error: Cell<Option<io::Error>>,
34 pub(super) read_task: LocalWaker,
35 pub(super) write_task: LocalWaker,
36 dispatch_task: LocalWaker,
37 pub(super) buffer: Stack,
38 pub(super) handle: Cell<Option<Box<dyn Handle>>>,
39 pub(super) timeout: Cell<TimerHandle>,
40 pub(super) shutdown_timeout: Cell<Option<Sleep>>,
41 pub(super) extensions: Extensions,
42}
43
44impl IoState {
45 pub(super) fn id(&self) -> Id {
46 self.id.get()
47 }
48
49 pub(super) fn tag(&self) -> &'static str {
50 self.cfg.tag()
51 }
52
53 pub(super) fn filter(&self) -> &dyn Filter {
54 self.filter.get()
55 }
56
57 pub(super) fn notify_timeout(&self) {
58 if self.flags.check_dispatcher_timeout_unset() {
59 self.wake_dispatch_task();
60 log::trace!("{}: Timer, notify dispatcher", self.cfg.tag());
61 }
62 }
63
64 pub(super) fn notify_disconnect(&self) {
65 self.extensions.notify_disconnect();
66 }
67
68 pub(super) fn error(&self) -> Option<io::Error> {
70 if let Some(err) = self.error.take() {
71 self.error
72 .set(Some(io::Error::new(err.kind(), format!("{err}"))));
73 Some(err)
74 } else {
75 None
76 }
77 }
78
79 pub(super) fn error_or_disconnected(&self) -> io::Error {
81 self.error()
82 .unwrap_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "Disconnected"))
83 }
84
85 pub(super) fn filters_stopped(&self) {
86 self.wake_read_task();
87 self.wake_write_task();
88 self.wake_dispatch_task();
89 self.flags.set_filters_stopped();
90 }
91
92 pub(super) fn terminate_connection(&self, err: Option<io::Error>) {
93 if !self.flags.is_terminated() {
94 log::trace!("{}: Terminate io with error {:?}", self.cfg.tag(), err);
95 if err.is_some() {
96 self.error.set(err);
97 }
98 self.flags.set_terminate();
99 self.wake_read_task();
100 self.wake_write_task();
101 self.wake_dispatch_task();
102 self.notify_disconnect();
103 self.handle.take();
104 }
105 }
106
107 pub(super) fn start_shutdown(&self) {
109 if !self.flags.is_stopping_any() {
110 log::trace!("{}: Initiate io shutdown {:?}", self.cfg.tag(), self.flags);
111 self.flags.set_filter_stopping();
112 self.wake_read_task();
113 self.wake_write_task();
114 }
115 }
116
117 pub(super) fn get_read_buf(&self) -> BytesMut {
118 self.cfg.read_buf().get()
119 }
120
121 pub(super) fn is_rd_backpressure_needed(&self, size: usize) -> bool {
122 size >= self.cfg.read_buf().high
123 }
124
125 pub(super) fn is_wr_backpressure_needed(&self, size: usize) -> bool {
126 size >= self.cfg.write_buf().high
127 }
128
129 pub(super) fn should_disable_wr_backpressure(&self, size: usize) -> bool {
130 size <= self.cfg.write_buf().half
131 }
132
133 pub(super) fn wake_read_task(&self) {
134 self.read_task.wake();
135 }
136
137 pub(super) fn wake_write_task(&self) {
138 #[cfg(feature = "trace")]
139 log::trace!("{}: Wake write task, flags:{:?}", self.tag(), self.flags);
140 self.write_task.wake();
141 }
142
143 pub(super) fn wake_dispatch_task(&self) {
144 self.dispatch_task.wake();
145 }
146}
147
148impl Eq for IoState {}
149
150impl PartialEq for IoState {
151 #[inline]
152 fn eq(&self, other: &Self) -> bool {
153 ptr::eq(self, other)
154 }
155}
156
157impl hash::Hash for IoState {
158 #[inline]
159 fn hash<H: hash::Hasher>(&self, state: &mut H) {
160 (ptr::from_ref(self) as usize).hash(state);
161 }
162}
163
164impl fmt::Debug for IoState {
165 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166 let err = self.error.take();
167 let res = f
168 .debug_struct("IoState")
169 .field("id", &self.id)
170 .field("flags", &self.flags)
171 .field("filter", &self.filter.is_set())
172 .field("timeout", &self.timeout)
173 .field("error", &err)
174 .field("buffer", &self.buffer)
175 .field("cfg", &self.cfg)
176 .finish();
177 self.error.set(err);
178 res
179 }
180}
181
182impl Io {
183 pub fn new<I: IoStream, T: Into<SharedCfg>>(io: I, cfg: T) -> Self {
185 let cfg = cfg.into().get::<IoConfig>();
186 let size = cfg.write_page_size();
187 let flags = Flags::new(cfg.write_buf_threshold() > 0);
188
189 let inner = Rc::new(IoState {
190 cfg,
191 flags,
192 id: Cell::new(Id::default()),
193 filter: FilterPtr::null(),
194 error: Cell::new(None),
195 dispatch_task: LocalWaker::new(),
196 read_task: LocalWaker::new(),
197 write_task: LocalWaker::new(),
198 buffer: Stack::new(size),
199 handle: Cell::new(None),
200 timeout: Cell::new(TimerHandle::default()),
201 shutdown_timeout: Cell::new(None),
202 extensions: Extensions::default(),
203 });
204 inner.filter.set(Base::new(IoRef(inner.clone())));
205
206 let ioref = IoRef(inner);
207 ioref.0.id.set(IoManager::register(&ioref));
208
209 let hnd = io.start(IoContext::new(ioref.clone()));
211 ioref.0.handle.set(Some(hnd));
212
213 Io(UnsafeCell::new(ioref), marker::PhantomData)
214 }
215}
216
217impl<I: IoStream> From<I> for Io {
218 #[inline]
219 fn from(io: I) -> Io {
220 Io::new(io, SharedCfg::default())
221 }
222}
223
224impl IoRef {
225 fn create_empty() -> IoRef {
226 IoRef(Rc::new(IoState {
227 id: Cell::new(Id::default()),
228 cfg: SharedCfg::default().get::<IoConfig>(),
229 filter: FilterPtr::null(),
230 flags: Flags::new_stopped(),
231 error: Cell::new(None),
232 dispatch_task: LocalWaker::new(),
233 read_task: LocalWaker::new(),
234 write_task: LocalWaker::new(),
235 buffer: Stack::new(BytePageSize::Size16),
236 handle: Cell::new(None),
237 timeout: Cell::new(TimerHandle::default()),
238 shutdown_timeout: Cell::new(None),
239 extensions: Extensions::default(),
240 }))
241 }
242}
243
244impl<F> Io<F> {
245 #[inline]
246 pub fn get_ref(&self) -> IoRef {
248 self.io_ref().clone()
249 }
250
251 #[inline]
252 #[must_use]
253 pub fn take(&self) -> Self {
257 Self(UnsafeCell::new(self.take_io_ref()), marker::PhantomData)
258 }
259
260 fn take_io_ref(&self) -> IoRef {
261 unsafe { mem::replace(&mut *self.0.get(), IoRef::create_empty()) }
262 }
263
264 fn st(&self) -> &IoState {
265 unsafe { &(*self.0.get()).0 }
266 }
267
268 fn io_ref(&self) -> &IoRef {
269 unsafe { &*self.0.get() }
270 }
271
272 #[inline]
273 pub fn set_config<T: Into<SharedCfg>>(&self, cfg: T) {
275 unsafe {
276 let cfg = cfg.into().get::<IoConfig>();
277 self.st().buffer.set_page_size(cfg.write_page_size());
278 self.st().cfg.replace(cfg);
279 }
280 }
281}
282
283impl<F: FilterLayer, T: Filter> Io<Layer<F, T>> {
284 #[inline]
285 pub fn filter(&self) -> &F {
287 &self.st().filter.filter::<Layer<F, T>>().0
288 }
289}
290
291impl<F: Filter> Io<F> {
292 #[inline]
293 pub fn seal(self) -> Io<Sealed> {
295 let state = self.take_io_ref();
296 state.0.filter.seal::<F>();
297
298 Io(UnsafeCell::new(state), marker::PhantomData)
299 }
300
301 #[inline]
302 pub fn boxed(self) -> IoBoxed {
304 self.seal().into()
305 }
306
307 #[inline]
308 pub fn add_filter<U>(self, nf: U) -> Io<Layer<U, F>>
310 where
311 U: FilterLayer,
312 {
313 self.with_callbacks(|cb| cb.before_processing(&self));
314
315 if let Err(e) = self.st().buffer.process_write_buf_no_cb(&self) {
318 self.st().terminate_connection(Some(e));
319 }
320
321 let state = self.take_io_ref();
322
323 unsafe { &mut *(Rc::as_ptr(&state.0).cast_mut()) }
328 .buffer
329 .add_layer(state.0.cfg.write_page_size());
330
331 state.0.filter.add_filter::<F, U>(nf);
333
334 let io = Io(UnsafeCell::new(state), marker::PhantomData);
335
336 if let Err(e) = io.st().buffer.process_read_buf_no_cb(&io, 0) {
338 io.st().terminate_connection(Some(e));
339 }
340 io.with_callbacks(|cb| cb.after_processing(&io));
341
342 io
343 }
344
345 pub fn map_filter<U, R>(self, f: U) -> Io<R>
347 where
348 U: FnOnce(F) -> R,
349 R: Filter,
350 {
351 self.with_callbacks(|cb| cb.before_processing(&self));
352
353 if let Err(e) = self.st().buffer.process_write_buf(&self) {
356 self.st().terminate_connection(Some(e));
357 }
358
359 let state = self.take_io_ref();
360 state.0.filter.map_filter::<F, U, R>(f);
361
362 let io = Io(UnsafeCell::new(state), marker::PhantomData);
363 io.with_callbacks(|cb| cb.after_processing(&io));
364 io
365 }
366}
367
368impl<F> Io<F> {
369 pub async fn recv<U>(
371 &self,
372 codec: &U,
373 ) -> Result<Option<U::Item>, Either<U::Error, io::Error>>
374 where
375 U: Decoder,
376 {
377 loop {
378 return match poll_fn(|cx| self.poll_recv(codec, cx)).await {
379 Ok(item) => Ok(Some(item)),
380 Err(RecvError::KeepAlive) => Err(Either::Right(io::Error::new(
381 io::ErrorKind::TimedOut,
382 "Timeout",
383 ))),
384 Err(RecvError::WriteBackpressure) => {
385 poll_fn(|cx| self.poll_flush(cx, false))
386 .await
387 .map_err(Either::Right)?;
388 continue;
389 }
390 Err(RecvError::Decoder(err)) => Err(Either::Left(err)),
391 Err(RecvError::PeerGone(Some(err))) => Err(Either::Right(err)),
392 Err(RecvError::PeerGone(None)) => Ok(None),
393 };
394 }
395 }
396
397 pub async fn read(&self, dst: &mut [u8]) -> io::Result<()> {
403 loop {
404 let completed = self.with_read_buf(|buf| {
405 if buf.len() >= dst.len() {
406 let _ = io::Read::read(buf, dst).expect("Cannot fail");
407 true
408 } else {
409 false
410 }
411 });
412 if completed {
413 return Ok(());
414 }
415 if self.read_ready().await?.is_none() {
417 return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Disconnected"));
418 }
419 }
420 }
421
422 #[inline]
423 pub async fn read_ready(&self) -> io::Result<Option<()>> {
425 poll_fn(|cx| self.poll_read_ready(cx)).await
426 }
427
428 #[inline]
429 pub async fn read_notify(&self) -> io::Result<Option<()>> {
431 poll_fn(|cx| self.poll_read_notify(cx)).await
432 }
433
434 #[inline]
435 pub fn pause(&self) {
437 let st = self.st();
438 if !st.flags.is_read_paused() {
439 st.wake_read_task();
440 st.flags.set_read_paused();
441 }
442 }
443
444 #[inline]
445 pub async fn send<U>(
447 &self,
448 item: U::Item,
449 codec: &U,
450 ) -> Result<(), Either<U::Error, io::Error>>
451 where
452 U: Encoder,
453 {
454 self.encode(item, codec).map_err(Either::Left)?;
455
456 poll_fn(|cx| self.poll_flush(cx, true))
457 .await
458 .map_err(Either::Right)?;
459
460 Ok(())
461 }
462
463 #[inline]
464 pub async fn flush(&self, full: bool) -> io::Result<()> {
468 poll_fn(|cx| self.poll_flush(cx, full)).await
469 }
470
471 #[inline]
472 pub async fn shutdown(&self) -> io::Result<()> {
474 poll_fn(|cx| self.poll_shutdown(cx)).await
475 }
476
477 #[inline]
478 pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<Option<()>>> {
493 let st = self.st();
494
495 if st.flags.is_closed() {
496 Poll::Ready(Ok(None))
497 } else {
498 let ready = st.flags.is_read_ready();
499
500 if st.flags.is_read_paused_or_backpressure() {
503 st.flags.unset_all_read_flags();
504 st.flags.unset_read_paused();
505 st.wake_read_task();
506 if ready {
507 Poll::Ready(Ok(Some(())))
508 } else {
509 st.dispatch_task.register(cx.waker());
510 Poll::Pending
511 }
512 } else if ready {
513 Poll::Ready(Ok(Some(())))
514 } else {
515 if st.flags.is_read_paused() {
516 st.wake_read_task();
517 st.flags.unset_read_paused();
518 }
519 st.dispatch_task.register(cx.waker());
520 Poll::Pending
521 }
522 }
523 }
524
525 #[inline]
526 pub fn poll_read_notify(&self, cx: &mut Context<'_>) -> Poll<io::Result<Option<()>>> {
528 let st = self.st();
529 if st.flags.is_stopping() {
530 Poll::Ready(Ok(None))
531 } else if st.flags.check_read_notifed() {
532 Poll::Ready(Ok(Some(())))
533 } else {
534 st.flags.set_read_notify();
535 if self.poll_read_ready(cx).is_ready() {
536 st.dispatch_task.register(cx.waker());
537 }
538 st.dispatch_task.register(cx.waker());
539 Poll::Pending
540 }
541 }
542
543 #[inline]
544 pub fn poll_recv<U>(
549 &self,
550 codec: &U,
551 cx: &mut Context<'_>,
552 ) -> Poll<Result<U::Item, RecvError<U>>>
553 where
554 U: Decoder,
555 {
556 let decoded = self.poll_recv_decode(codec, cx)?;
557
558 if let Some(item) = decoded.item {
559 Poll::Ready(Ok(item))
560 } else {
561 Poll::Pending
562 }
563 }
564
565 #[inline]
566 pub fn poll_recv_decode<U>(
571 &self,
572 codec: &U,
573 cx: &mut Context<'_>,
574 ) -> Result<Decoded<U::Item>, RecvError<U>>
575 where
576 U: Decoder,
577 {
578 let st = self.st();
579 st.flags.unset_read_ready();
580
581 let decoded = self
582 .decode_item(codec)
583 .map_err(|err| RecvError::Decoder(err))?;
584
585 if decoded.item.is_some() {
586 Ok(decoded)
587 } else if st.flags.is_stopping() {
588 Err(RecvError::PeerGone(st.error()))
589 } else if st.flags.check_dispatcher_timeout() {
590 Err(RecvError::KeepAlive)
591 } else if st.flags.is_wr_backpressure() {
592 Err(RecvError::WriteBackpressure)
593 } else {
594 match self.poll_read_ready(cx) {
595 Poll::Pending | Poll::Ready(Ok(Some(()))) => {
596 #[cfg(feature = "trace")]
597 if decoded.remains != 0 {
598 log::trace!("{}: Not enough data to decode next frame", self.tag());
599 }
600 Ok(decoded)
601 }
602 Poll::Ready(Err(e)) => Err(RecvError::PeerGone(Some(e))),
603 Poll::Ready(Ok(None)) => Err(RecvError::PeerGone(None)),
604 }
605 }
606 }
607
608 #[inline]
609 pub fn poll_flush(&self, cx: &mut Context<'_>, full: bool) -> Poll<io::Result<()>> {
614 let st = self.st();
615
616 st.buffer.process_write_buf_force(self)?;
618 self.consolidate_write_state(false);
619
620 let len = st.buffer.write_buf_size();
621 if len > 0 {
622 if st.flags.is_closed() {
623 return Poll::Ready(Err(st.error_or_disconnected()));
624 } else if full {
625 st.flags.set_wants_write_flush();
626 st.dispatch_task.register(cx.waker());
627 return Poll::Pending;
628 } else if st.is_wr_backpressure_needed(len) {
629 st.flags.set_wr_backpressure();
630 st.dispatch_task.register(cx.waker());
631 return Poll::Pending;
632 }
633 }
634 if st.flags.is_closed() && !st.flags.is_write_flush() {
635 Poll::Ready(Err(st.error_or_disconnected()))
636 } else {
637 st.flags.unset_wr_backpressure_and_flush();
638 Poll::Ready(Ok(()))
639 }
640 }
641
642 #[inline]
643 pub fn poll_shutdown(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
645 let st = self.st();
646
647 if st.flags.is_stopping() {
648 if let Some(err) = st.error() {
649 Poll::Ready(Err(err))
650 } else {
651 Poll::Ready(Ok(()))
652 }
653 } else {
654 if !st.flags.is_stopping_filters() {
655 st.start_shutdown();
656 }
657 st.flags.unset_all_read_flags();
658 st.flags.unset_read_paused();
659
660 st.wake_read_task();
661 st.dispatch_task.register(cx.waker());
662 Poll::Pending
663 }
664 }
665
666 #[inline]
667 pub fn poll_read_pause(&self, cx: &mut Context<'_>) -> Poll<IoStatusUpdate> {
671 self.pause();
672 self.poll_status_update(cx)
673 }
674
675 #[inline]
676 pub fn poll_status_update(&self, cx: &mut Context<'_>) -> Poll<IoStatusUpdate> {
678 let st = self.st();
679 st.dispatch_task.register(cx.waker());
680 if st.flags.is_closed() {
681 Poll::Ready(IoStatusUpdate::PeerGone(st.error()))
682 } else if st.flags.check_dispatcher_timeout() {
683 Poll::Ready(IoStatusUpdate::KeepAlive)
684 } else if st.flags.is_wr_backpressure() {
685 if st.should_disable_wr_backpressure(st.buffer.write_buf_size()) {
687 st.flags.unset_wr_backpressure();
688 }
689 Poll::Ready(IoStatusUpdate::WriteBackpressure)
690 } else {
691 Poll::Pending
692 }
693 }
694
695 #[inline]
696 pub fn poll_dispatch(&self, cx: &mut Context<'_>) {
698 self.st().dispatch_task.register(cx.waker());
699 }
700}
701
702impl<F> AsRef<IoRef> for Io<F> {
703 #[inline]
704 fn as_ref(&self) -> &IoRef {
705 self.io_ref()
706 }
707}
708
709impl<F> Eq for Io<F> {}
710
711impl<F> PartialEq for Io<F> {
712 #[inline]
713 fn eq(&self, other: &Self) -> bool {
714 self.io_ref().eq(other.io_ref())
715 }
716}
717
718impl<F> hash::Hash for Io<F> {
719 #[inline]
720 fn hash<H: hash::Hasher>(&self, state: &mut H) {
721 self.io_ref().hash(state);
722 }
723}
724
725impl<F> fmt::Debug for Io<F> {
726 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
727 f.debug_struct("Io").field("state", self.st()).finish()
728 }
729}
730
731impl<F> ops::Deref for Io<F> {
732 type Target = IoRef;
733
734 #[inline]
735 fn deref(&self) -> &Self::Target {
736 self.io_ref()
737 }
738}
739
740impl<F> Drop for Io<F> {
741 fn drop(&mut self) {
742 let st = self.st();
743 self.stop_timer();
744
745 if st.filter.is_set() {
746 if !st.flags.is_terminated() {
749 log::trace!("{}: Io is dropped, terminate connection", st.tag());
750 }
751
752 st.terminate_connection(None);
753 st.filter.drop_filter::<F>();
754 }
755
756 IoManager::unregister(self.io_ref());
757 }
758}
759
760#[derive(Debug)]
761#[must_use = "OnDisconnect do nothing unless polled"]
763pub struct OnDisconnect {
764 token: usize,
765 inner: Rc<IoState>,
766}
767
768impl OnDisconnect {
769 pub(super) fn new(inner: Rc<IoState>) -> Self {
770 Self::new_inner(inner.flags.is_stopping(), inner)
771 }
772
773 fn new_inner(disconnected: bool, inner: Rc<IoState>) -> Self {
774 let token = if disconnected {
775 usize::MAX
776 } else {
777 inner.extensions.register_disconnect()
778 };
779 Self { token, inner }
780 }
781
782 #[inline]
783 pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
785 if self.token == usize::MAX || self.inner.flags.is_stopping() {
786 Poll::Ready(())
787 } else {
788 self.inner
789 .extensions
790 .poll_disconnect(self.token, cx.waker())
791 }
792 }
793}
794
795impl Clone for OnDisconnect {
796 fn clone(&self) -> Self {
797 if self.token == usize::MAX {
798 OnDisconnect::new_inner(true, self.inner.clone())
799 } else {
800 OnDisconnect::new_inner(false, self.inner.clone())
801 }
802 }
803}
804
805impl Future for OnDisconnect {
806 type Output = ();
807
808 #[inline]
809 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
810 self.poll_ready(cx)
811 }
812}
813
814#[cfg(test)]
815mod tests {
816 use ntex_bytes::{BufMut, BytePages, Bytes, BytesMut};
817 use ntex_codec::BytesCodec;
818 use ntex_util::{future::lazy, time::Millis, time::sleep};
819
820 use super::*;
821 use crate::{
822 FilterBuf, IoContext, IoTaskStatus, Readiness, ops::Iops, testing::IoTest,
823 };
824
825 const BIN: &[u8] = b"GET /test HTTP/1\r\n\r\n";
826 const TEXT: &str = "GET /test HTTP/1\r\n\r\n";
827 const BIN2: &[u8] = b"12345678901234561234567890123456";
828
829 #[ntex::test]
830 async fn test_basics() {
831 let (client, server) = IoTest::create();
832 client.remote_buffer_cap(1024);
833
834 let server = Io::from(server);
835 assert!(server.eq(&server));
836 assert!(server.io_ref().eq(server.io_ref()));
837 }
838
839 #[ntex::test]
840 async fn test_recv() {
841 let (client, server) = IoTest::create();
842 client.remote_buffer_cap(1024);
843
844 let server = Io::new(server, SharedCfg::new("SRV"));
845
846 server.st().notify_timeout();
847 let err = server.recv(&BytesCodec).await.err().unwrap();
848 assert!(format!("{err:?}").contains("Timeout"));
849
850 client.write(TEXT);
851 server.st().flags.set_wr_backpressure();
852 let item = server.recv(&BytesCodec).await.ok().unwrap().unwrap();
853 assert_eq!(item, TEXT);
854 }
855
856 #[ntex::test]
857 async fn test_read() {
858 let (client, server) = IoTest::create();
859 client.remote_buffer_cap(1024);
860
861 let server = Io::new(server, SharedCfg::new("SRV"));
862
863 client.write(b"1234");
864 let mut buf: [u8; 4] = [0, 0, 0, 0];
865 server.read(&mut buf).await.unwrap();
866 assert_eq!(&buf, b"1234");
867
868 let fut = ntex_rt::spawn(async move {
870 let mut buf: [u8; 4] = [0, 0, 0, 0];
871 server.read(&mut buf).await
872 });
873 client.close().await;
874 let err = fut.await.unwrap().err().unwrap();
875 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
876 }
877
878 #[ntex::test]
879 async fn test_send() {
880 let (client, server) = IoTest::create();
881 client.remote_buffer_cap(1024);
882
883 let server = Io::from(server);
884 assert!(server.eq(&server));
885
886 server
887 .send(Bytes::from_static(BIN), &BytesCodec)
888 .await
889 .ok()
890 .unwrap();
891 let item = client.read_any();
892 assert_eq!(item, TEXT);
893 }
894
895 #[ntex::test]
896 async fn read() {
897 let io = Io::new(
898 IoTest::create().0,
899 SharedCfg::new("SRV").add(IoConfig::default().set_read_buf(8, 4, 16)),
900 );
901 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
902 assert!(io.st().dispatch_task.is_set());
903
904 let ctx = IoContext::new(io.get_ref());
905
906 assert_eq!(
908 lazy(|cx| ctx.poll_read_ready(cx)).await,
909 Poll::Ready(Readiness::Ready)
910 );
911 assert!(io.st().read_task.is_set());
912 assert!(!io.st().flags.is_read_ready());
913 assert!(!io.st().flags.is_rd_backpressure());
914
915 ctx.update_read_status(BytesMut::copy_from_slice(b"1234567890"), Ok(10));
917
918 assert!(!io.st().dispatch_task.is_set());
920 assert!(io.st().flags.is_read_paused());
922 assert!(io.st().flags.is_read_ready());
924 assert!(io.st().flags.is_rd_backpressure());
926 assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
928
929 assert_eq!(io.with_read_buf(|buf| buf.split_to(1)), b"1");
931 assert!(io.st().flags.is_read_ready());
933 assert!(io.st().flags.is_rd_backpressure());
935
936 assert!(io.st().read_task.is_set());
938
939 assert_eq!(io.with_read_buf(|buf| buf.split_to(1)), b"2");
941 assert!(io.st().flags.is_rd_backpressure());
943
944 assert_eq!(io.with_read_buf(|buf| buf.split_to(4)), b"3456");
946 assert!(!io.st().flags.is_read_paused());
948 assert!(!io.st().flags.is_read_ready());
950 assert!(!io.st().flags.is_rd_backpressure());
952 assert!(!io.st().read_task.is_set());
954 assert_eq!(
955 lazy(|cx| ctx.poll_read_ready(cx)).await,
956 Poll::Ready(Readiness::Ready)
957 );
958
959 lazy(|cx| io.poll_dispatch(cx)).await;
961
962 ctx.update_read_status(BytesMut::copy_from_slice(b"1234"), Ok(4));
964
965 assert!(!io.st().dispatch_task.is_set());
967 assert!(io.st().flags.is_read_paused());
969 assert!(io.st().flags.is_read_ready());
971 assert!(io.st().flags.is_rd_backpressure());
973 assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
975
976 assert_eq!(io.with_read_buf(|buf| buf.split_to(4)), b"7890");
978 assert!(!io.st().flags.is_rd_backpressure());
980
981 lazy(|cx| io.poll_dispatch(cx)).await;
983
984 ctx.update_read_status(BytesMut::copy_from_slice(b"567"), Ok(3));
986
987 assert!(!io.st().flags.is_read_paused());
989 assert!(io.st().flags.is_read_ready());
991 assert!(!io.st().flags.is_rd_backpressure());
993 assert_eq!(
995 lazy(|cx| ctx.poll_read_ready(cx)).await,
996 Poll::Ready(Readiness::Ready)
997 );
998
999 assert_eq!(io.with_read_buf(BytesMut::take), b"1234567");
1001 assert!(!io.st().flags.is_read_paused());
1003 assert!(!io.st().flags.is_read_ready());
1005 assert!(io.st().read_task.is_set());
1007
1008 io.terminate();
1010 assert!(!io.st().read_task.is_set());
1012 assert_eq!(
1014 lazy(|cx| ctx.poll_read_ready(cx)).await,
1015 Poll::Ready(Readiness::Terminate)
1016 );
1017 }
1018
1019 #[ntex::test]
1020 async fn read_notify() {
1021 let io = Io::new(
1022 IoTest::create().0,
1023 SharedCfg::new("SRV").add(IoConfig::default().set_read_buf(8, 4, 16)),
1024 );
1025 assert!(!io.st().flags.is_read_notify());
1026 assert!(lazy(|cx| io.poll_read_notify(cx)).await.is_pending());
1027 assert!(io.st().dispatch_task.is_set());
1028 assert!(io.st().flags.is_read_notify());
1029
1030 let ctx = IoContext::new(io.get_ref());
1031
1032 ctx.update_read_status(BytesMut::copy_from_slice(b"1"), Ok(1));
1034
1035 assert!(!io.st().dispatch_task.is_set());
1036 assert!(io.st().flags.is_read_ready());
1038 assert!(io.st().flags.is_read_notify());
1039 assert!(io.st().flags.is_read_notified());
1041 let res = lazy(|cx| io.poll_read_notify(cx)).await;
1042 assert!(matches!(res, Poll::Ready(Ok(Some(())))));
1043
1044 assert!(!io.st().dispatch_task.is_set());
1046 assert!(io.st().flags.is_read_ready());
1048
1049 assert!(lazy(|cx| io.poll_read_notify(cx)).await.is_pending());
1051 assert!(io.st().dispatch_task.is_set());
1052 assert!(io.st().flags.is_read_notify());
1053 assert!(io.st().flags.is_read_ready());
1054 assert_eq!(
1056 lazy(|cx| ctx.poll_read_ready(cx)).await,
1057 Poll::Ready(Readiness::Ready)
1058 );
1059
1060 ctx.update_read_status(BytesMut::copy_from_slice(b"2345678"), Ok(7));
1062 assert!(io.st().flags.is_rd_backpressure());
1064
1065 assert!(io.st().flags.is_read_ready());
1067 assert!(io.st().flags.is_read_notify());
1068 assert!(io.st().flags.is_read_notified());
1070 let res = lazy(|cx| io.poll_read_notify(cx)).await;
1071 assert!(matches!(res, Poll::Ready(Ok(Some(())))));
1072 assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
1074 assert!(io.st().read_task.is_set());
1076
1077 assert!(lazy(|cx| io.poll_read_notify(cx)).await.is_pending());
1079 assert!(!io.st().flags.is_rd_backpressure());
1081 assert!(!io.st().flags.is_read_ready());
1082 assert!(!io.st().flags.is_read_paused());
1083 assert!(!io.st().read_task.is_set());
1085 assert_eq!(
1087 lazy(|cx| ctx.poll_read_ready(cx)).await,
1088 Poll::Ready(Readiness::Ready)
1089 );
1090
1091 ctx.update_read_status(BytesMut::copy_from_slice(b"1"), Ok(1));
1093 assert!(!io.st().dispatch_task.is_set());
1094 assert!(io.st().flags.is_read_ready());
1096 assert!(io.st().flags.is_read_notify());
1097 assert!(io.st().flags.is_read_paused());
1098 assert!(io.st().flags.is_rd_backpressure());
1099 assert!(io.st().flags.is_read_notified());
1101 assert!(matches!(
1102 lazy(|cx| io.poll_read_notify(cx)).await,
1103 Poll::Ready(Ok(Some(())))
1104 ));
1105
1106 io.terminate();
1108 let res = lazy(|cx| io.poll_read_notify(cx)).await;
1109 assert!(matches!(res, Poll::Ready(Ok(None))), "{res:?}");
1110 }
1111
1112 #[ntex::test]
1113 async fn read_readiness() {
1114 let (client, server) = IoTest::create();
1115 client.remote_buffer_cap(1024);
1116
1117 let io = Io::from(server);
1118 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
1119
1120 client.write(TEXT);
1121 assert_eq!(io.read_ready().await.unwrap(), Some(()));
1122 assert!(matches!(
1123 lazy(|cx| io.poll_read_ready(cx)).await,
1124 Poll::Ready(Ok(Some(())))
1125 ));
1126
1127 let item = io.with_read_buf(BytesMut::take);
1128 assert_eq!(item, Bytes::from_static(BIN));
1129
1130 client.write(TEXT);
1131 sleep(Millis(50)).await;
1132 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_ready());
1133 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_ready());
1134 }
1135
1136 #[ntex::test]
1137 async fn read_backpressure() {
1138 let (client, server) = IoTest::create();
1139
1140 let io = Io::new(
1141 server,
1142 SharedCfg::new("SRV").add(IoConfig::default().set_read_buf(64, 32, 12)),
1143 );
1144 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
1145
1146 client.write(BIN2);
1147 client.write(BIN2);
1148 sleep(Millis(50)).await;
1149 assert!(io.flags().is_read_ready());
1150 assert!(io.flags().is_rd_backpressure());
1151 let _item = io.recv(&BytesCodec).await.ok().unwrap().unwrap();
1152 assert!(!io.flags().is_read_ready());
1153 assert!(!io.flags().is_rd_backpressure());
1154
1155 client.write(BIN2);
1156 client.write(BIN2);
1157 sleep(Millis(50)).await;
1158 assert!(io.flags().is_read_ready());
1159 assert!(io.flags().is_rd_backpressure());
1160 assert_eq!(io.read_ready().await.unwrap(), Some(()));
1161 }
1162
1163 #[ntex::test]
1164 async fn write() {
1165 let io = Io::new(
1166 IoTest::create().0,
1167 SharedCfg::new("SRV").add(IoConfig::default().set_write_buf(8, 4, 16)),
1168 );
1169 assert!(lazy(|cx| io.poll_status_update(cx)).await.is_pending());
1170 assert!(io.st().dispatch_task.is_set());
1171 assert!(io.st().flags.is_direct_wr_enabled());
1172
1173 let ctx = IoContext::new(io.get_ref());
1174
1175 assert_eq!(lazy(|cx| ctx.poll_write_ready(cx)).await, Poll::Pending);
1177 assert!(io.st().write_task.is_set());
1178 assert!(io.st().flags.is_write_paused());
1179 assert!(!io.st().flags.is_wr_backpressure());
1180
1181 io.with_write_buf(|buf| buf.put_slice(b"1234")).unwrap();
1183 assert_eq!(lazy(|cx| ctx.poll_write_ready(cx)).await, Poll::Pending);
1184 assert!(io.st().flags.is_write_paused());
1186 assert!(io.st().flags.is_wr_send_scheduled());
1188 assert!(!io.st().flags.is_wr_backpressure());
1190 assert!(io.st().dispatch_task.is_set());
1192
1193 io.with_write_buf(|buf| buf.put_slice(b"5678")).unwrap();
1195 assert!(io.st().flags.is_wr_backpressure());
1197 assert!(!io.st().dispatch_task.is_set());
1199 assert!(io.st().write_task.is_set());
1201 assert!(matches!(
1203 lazy(|cx| io.poll_status_update(cx)).await,
1204 Poll::Ready(IoStatusUpdate::WriteBackpressure)
1205 ));
1206 assert!(lazy(|cx| io.poll_flush(cx, false)).await.is_pending());
1208 assert!(!io.st().flags.is_write_flush());
1210
1211 Iops::run();
1213 assert!(!io.st().flags.is_wr_send_scheduled());
1215 assert!(!io.st().flags.is_write_paused());
1217 assert!(!io.st().write_task.is_set());
1219 assert_eq!(
1221 lazy(|cx| ctx.poll_write_ready(cx)).await,
1222 Poll::Ready(Readiness::Ready)
1223 );
1224
1225 assert_eq!(ctx.with_write_buf(|buf| buf.split_to(4).freeze()), b"1234");
1227 assert_eq!(ctx.update_write_status(Ok(true)), IoTaskStatus::Io);
1229 assert_eq!(
1231 lazy(|cx| ctx.poll_write_ready(cx)).await,
1232 Poll::Ready(Readiness::Ready)
1233 );
1234 assert!(!io.st().flags.is_write_paused());
1236 assert!(io.st().flags.is_wr_backpressure());
1238 assert!(matches!(
1240 lazy(|cx| io.poll_status_update(cx)).await,
1241 Poll::Ready(IoStatusUpdate::WriteBackpressure)
1242 ));
1243 assert!(!io.st().flags.is_wr_backpressure());
1245 assert!(lazy(|cx| io.poll_status_update(cx)).await.is_pending());
1246 assert!(matches!(
1248 lazy(|cx| io.poll_flush(cx, false)).await,
1249 Poll::Ready(Ok(()))
1250 ));
1251
1252 io.with_write_buf(|buf| buf.put_slice(b"1234")).unwrap();
1254 assert!(lazy(|cx| io.poll_flush(cx, true)).await.is_pending());
1255 assert!(io.st().flags.is_write_flush());
1257 assert!(io.st().flags.is_wr_backpressure());
1259
1260 Iops::run();
1262 assert_eq!(ctx.with_write_buf(BytePages::freeze), b"56781234");
1263 assert!(!io.st().flags.is_wr_send_scheduled());
1265 assert_eq!(ctx.update_write_status(Ok(true)), IoTaskStatus::Pause);
1267 assert!(io.st().flags.is_write_paused());
1269 assert!(io.st().flags.is_write_flush());
1271 assert!(io.st().flags.is_wr_backpressure());
1273 assert!(!io.st().dispatch_task.is_set());
1275
1276 assert!(matches!(
1278 lazy(|cx| io.poll_flush(cx, false)).await,
1279 Poll::Ready(Ok(()))
1280 ));
1281 assert!(!io.st().flags.is_write_flush());
1283 assert!(!io.st().flags.is_wr_backpressure());
1285
1286 io.terminate();
1288 assert!(!io.st().write_task.is_set());
1290 assert_eq!(
1292 lazy(|cx| ctx.poll_write_ready(cx)).await,
1293 Poll::Ready(Readiness::Terminate)
1294 );
1295 let Poll::Ready(Err(err)) = lazy(|cx| io.poll_flush(cx, false)).await else {
1297 panic!()
1298 };
1299 assert_eq!(err.kind(), io::ErrorKind::NotConnected);
1300 assert!(matches!(
1302 lazy(|cx| io.poll_status_update(cx)).await,
1303 Poll::Ready(IoStatusUpdate::PeerGone(None))
1304 ));
1305 }
1306
1307 #[ntex::test]
1308 async fn write_backpressure() {
1309 let (client, server) = IoTest::create();
1310 client.remote_buffer_cap(0);
1311
1312 let io = Io::new(
1313 server,
1314 SharedCfg::new("SRV").add(IoConfig::default().set_write_buf(16, 8, 12)),
1315 );
1316 assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
1317 assert!(io.flags().is_write_paused());
1318 assert!(!io.flags().is_wr_backpressure());
1319 assert!(!io.is_wr_backpressure());
1320
1321 io.encode_slice(BIN2).unwrap();
1322 assert!(Iops::is_registered(&io));
1323 assert!(io.flags().is_wr_backpressure());
1324
1325 client.remote_buffer_cap(1024);
1326 let item = client.read().await.unwrap();
1327 assert_eq!(item, BIN2);
1328 assert!(io.flags().is_wr_backpressure());
1329 assert!(matches!(
1330 lazy(|cx| io.poll_status_update(cx)).await,
1331 Poll::Ready(IoStatusUpdate::WriteBackpressure)
1332 ));
1333 assert!(!io.flags().is_wr_backpressure());
1334 assert!(matches!(
1335 lazy(|cx| io.poll_flush(cx, false)).await,
1336 Poll::Ready(Ok(()))
1337 ));
1338 assert!(!io.flags().is_wr_backpressure());
1339 }
1340
1341 #[ntex::test]
1342 async fn shutdown_flushes_write_buf_with_read_backpressure() {
1343 let (client, server) = IoTest::create();
1346 client.remote_buffer_cap(0);
1348
1349 let io = Io::new(
1350 server,
1351 SharedCfg::new("SRV").add(
1352 IoConfig::default()
1353 .set_read_buf(8, 4, 16)
1354 .set_disconnect_timeout(ntex_util::time::Seconds(2)),
1355 ),
1356 );
1357
1358 io.encode_slice(b"response-tail").unwrap();
1360 sleep(Millis(50)).await;
1361 assert_eq!(io.st().buffer.write_buf_size(), 13);
1362
1363 client.write("0123456789");
1366 sleep(Millis(50)).await;
1367 assert!(io.flags().is_read_paused());
1368 assert!(io.flags().is_rd_backpressure());
1369 assert_eq!(io.st().buffer.write_buf_size(), 13);
1370
1371 io.close();
1373 sleep(Millis(50)).await;
1374
1375 client.remote_buffer_cap(1024);
1377
1378 let data = ntex_util::time::timeout(Millis(2000), client.read())
1383 .await
1384 .expect("write buffer was dropped during shutdown")
1385 .unwrap();
1386 assert_eq!(&data[..], b"response-tail");
1387
1388 ntex_util::time::timeout(Millis(4000), io.on_disconnect())
1391 .await
1392 .expect("io stream did not disconnect after flush");
1393 }
1394
1395 #[ntex::test]
1396 async fn shutdown() {
1397 #[derive(Debug)]
1399 struct F;
1400
1401 impl FilterLayer for F {
1402 fn process_read_buf(&self, _: &FilterBuf<'_>) -> io::Result<()> {
1403 Ok(())
1404 }
1405 fn process_write_buf(&self, _: &FilterBuf<'_>) -> io::Result<()> {
1406 Ok(())
1407 }
1408 }
1409
1410 let io = Io::new(
1411 IoTest::create().0,
1412 SharedCfg::new("SRV").add(IoConfig::default().set_write_buf(8, 4, 16)),
1413 );
1414 let st = io.st();
1415 assert!(lazy(|cx| io.poll_status_update(cx)).await.is_pending());
1416 assert!(st.dispatch_task.is_set());
1417 assert!(!st.flags.is_closed());
1418 assert!(!st.flags.is_stopping_filters());
1419
1420 let ctx = IoContext::new(io.get_ref());
1421
1422 io.close();
1424 assert!(!st.flags.is_closed());
1425 assert!(st.flags.is_stopping_filters());
1426 let err = io.with_write_buf(|_| 1).unwrap_err();
1428 assert_eq!(err.kind(), io::ErrorKind::Other);
1429
1430 let io = io.add_filter(F);
1431 let layer = Layer::new(F, Base::new(io.get_ref()));
1432
1433 let st = io.st();
1434 st.buffer.with_write_src(|p| p.put_slice(b"123"));
1435 assert_eq!(st.buffer.write_buf_size(), 3);
1436 let res = st.buffer.with_filter(io.as_ref(), |f| layer.shutdown(f));
1437 assert!(matches!(res, Ok(Poll::Ready(()))));
1438 assert_eq!(st.buffer.write_buf_size(), 0);
1439
1440 ctx.stop(None);
1442 assert!(st.flags.is_closed());
1443 assert!(st.flags.is_terminated());
1444 assert!(st.flags.is_stopping_filters());
1445
1446 let err = io.with_write_buf(|_| 1).unwrap_err();
1447 assert_eq!(err.kind(), io::ErrorKind::NotConnected);
1448 }
1449}