Skip to main content

ntex_io/
io.rs

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
22/// Interface object to the underlying I/O stream
23pub 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    /// Get the current I/O error.
69    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    /// Returns the current I/O error, or creates a `NotConnected` error.
80    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    /// Gracefully shuts down the read and write I/O tasks.
108    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    /// Creates a new `Io` instance.
184    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        // start io tasks
210        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    /// Get an instance of `IoRef`.
247    pub fn get_ref(&self) -> IoRef {
248        self.io_ref().clone()
249    }
250
251    #[inline]
252    #[must_use]
253    /// Takes the current I/O object.
254    ///
255    /// After this call, the I/O object is no longer valid for use.
256    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    /// Updates the shared I/O configuration.
274    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    /// Returns a reference to a filter.
286    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    /// Convert the current I/O stream into a sealed version.
294    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    /// Convert the current I/O stream into a boxed version.
303    pub fn boxed(self) -> IoBoxed {
304        self.seal().into()
305    }
306
307    #[inline]
308    /// Adds a new processing layer to the current filter chain.
309    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        // Write buffer processing may be delayed,
316        // call the filter chain to process pending writes
317        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        // Add the buffers layer.
324        //
325        // Safety: no references into the buffer storage are retained.
326        // All APIs first remove the buffer from storage before processing it.
327        unsafe { &mut *(Rc::as_ptr(&state.0).cast_mut()) }
328            .buffer
329            .add_layer(state.0.cfg.write_page_size());
330
331        // Replace current filter
332        state.0.filter.add_filter::<F, U>(nf);
333
334        let io = Io(UnsafeCell::new(state), marker::PhantomData);
335
336        // push read data into new filter
337        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    /// Wraps the current layer with a wrapper.
346    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        // Write buffer processing may be delayed,
354        // call the filter chain to process pending writes
355        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    /// Reads from the incoming I/O stream and decodes a codec item.
370    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    /// Reads bytes from this I/O stream into the specified buffer.
398    ///
399    /// If there is not enough data available, waits for incoming data.
400    /// Returns an error of kind [`io::ErrorKind::UnexpectedEof`] if the stream
401    /// is disconnected before `dst` is completely filled.
402    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            // `read_ready` resolves with `None` once the io is closed/stopped.
416            if self.read_ready().await?.is_none() {
417                return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Disconnected"));
418            }
419        }
420    }
421
422    #[inline]
423    /// Waits until the I/O stream is ready for reading.
424    pub async fn read_ready(&self) -> io::Result<Option<()>> {
425        poll_fn(|cx| self.poll_read_ready(cx)).await
426    }
427
428    #[inline]
429    /// Waits until the I/O stream receives new data.
430    pub async fn read_notify(&self) -> io::Result<Option<()>> {
431        poll_fn(|cx| self.poll_read_notify(cx)).await
432    }
433
434    #[inline]
435    /// Pauses the read task.
436    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    /// Encodes an item and sends it to the peer, fully flushing the write buffer.
446    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    /// Wakes the write task and requests a flush of buffered data.
465    ///
466    /// This is the async version of `.poll_flush()` method.
467    pub async fn flush(&self, full: bool) -> io::Result<()> {
468        poll_fn(|cx| self.poll_flush(cx, full)).await
469    }
470
471    #[inline]
472    /// Gracefully shuts down the I/O stream.
473    pub async fn shutdown(&self) -> io::Result<()> {
474        poll_fn(|cx| self.poll_shutdown(cx)).await
475    }
476
477    #[inline]
478    /// Polls for read readiness.
479    ///
480    /// If the I/O stream is not currently ready for reading,
481    /// this method will store a clone of the `Waker` from the provided `Context`.
482    /// When the I/O stream becomes ready for reading, `wake()` will be called on the waker.
483    ///
484    /// # Returns
485    ///
486    /// The function returns:
487    ///
488    /// - `Poll::Pending` if the I/O stream is not ready for reading.
489    /// - `Poll::Ready(Ok(Some(())))` if the I/O stream is ready for reading.
490    /// - `Poll::Ready(Ok(None))` if the I/O stream is disconnected.
491    /// - `Poll::Ready(Err(e))` if an error is encountered.
492    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 the dispatcher requests more data but no read occurs,
501            // restart the read task.
502            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    /// Polls the I/O stream for availability of incoming data.
527    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    /// Decode codec item from incoming bytes stream.
545    ///
546    /// Wake read task and request to read more data if data is not enough for decoding.
547    /// If error get returned this method does not register waker for later wake up action.
548    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    /// Decode codec item from incoming bytes stream.
567    ///
568    /// Wake read task and request to read more data if data is not enough for decoding.
569    /// If error get returned this method does not register waker for later wake up action.
570    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    /// Wakes the write task and instructs it to flush data.
610    ///
611    /// If `full` is true, wakes the dispatcher when all data has been flushed;
612    /// otherwise, it wakes when the write buffer size falls below the low-watermark size.
613    pub fn poll_flush(&self, cx: &mut Context<'_>, full: bool) -> Poll<io::Result<()>> {
614        let st = self.st();
615
616        // flush filter state
617        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    /// Gracefully shuts down the I/O stream.
644    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    /// Pauses the read task.
668    ///
669    /// Returns status updates.
670    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    /// Polls for available status updates.
677    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            // write backpressure is enabled and write buf smaller than half
686            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    /// Registers a dispatch task.
697    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            // filter is unsafe and must be dropped explicitly,
747            // and won't be dropped without special attention
748            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/// The `OnDisconnect` future resolves when the I/O stream is disconnected.
762#[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    /// Checks if the I/O stream is disconnected.
784    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        // disconnect during read
869        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        // Ready
907        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        // == Enable backpressure
916        ctx.update_read_status(BytesMut::copy_from_slice(b"1234567890"), Ok(10));
917
918        // dispatcher is woken
919        assert!(!io.st().dispatch_task.is_set());
920        // read task is paused
921        assert!(io.st().flags.is_read_paused());
922        // read buffer is ready
923        assert!(io.st().flags.is_read_ready());
924        // read backpressure is enabled
925        assert!(io.st().flags.is_rd_backpressure());
926        // read task paused
927        assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
928
929        // read one byte
930        assert_eq!(io.with_read_buf(|buf| buf.split_to(1)), b"1");
931        // read buffer is ready
932        assert!(io.st().flags.is_read_ready());
933        // read backpressure is enabled
934        assert!(io.st().flags.is_rd_backpressure());
935
936        // read task is set
937        assert!(io.st().read_task.is_set());
938
939        // read one more byte
940        assert_eq!(io.with_read_buf(|buf| buf.split_to(1)), b"2");
941        // read backpressure is enabled
942        assert!(io.st().flags.is_rd_backpressure());
943
944        // read 4 bytes. buf size is 4, less that half of high watermark
945        assert_eq!(io.with_read_buf(|buf| buf.split_to(4)), b"3456");
946        // read task is not paused anymore
947        assert!(!io.st().flags.is_read_paused());
948        // read buffer is not ready
949        assert!(!io.st().flags.is_read_ready());
950        // read backpressure is disabled
951        assert!(!io.st().flags.is_rd_backpressure());
952        // read task is woken
953        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        // register dispatcher task
960        lazy(|cx| io.poll_dispatch(cx)).await;
961
962        // == Enable backpressure, 4 bytes in buffer + 4 more
963        ctx.update_read_status(BytesMut::copy_from_slice(b"1234"), Ok(4));
964
965        // dispatcher is woken
966        assert!(!io.st().dispatch_task.is_set());
967        // read task is paused
968        assert!(io.st().flags.is_read_paused());
969        // read buffer is ready
970        assert!(io.st().flags.is_read_ready());
971        // read backpressure is enabled
972        assert!(io.st().flags.is_rd_backpressure());
973        // read task paused
974        assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
975
976        // read 4 bytes. buf size is 4, less that half of high watermark
977        assert_eq!(io.with_read_buf(|buf| buf.split_to(4)), b"7890");
978        // read backpressure is disabled
979        assert!(!io.st().flags.is_rd_backpressure());
980
981        // register dispatcher task
982        lazy(|cx| io.poll_dispatch(cx)).await;
983
984        // == No backpressure, 4 bytes in buffer + 3 more
985        ctx.update_read_status(BytesMut::copy_from_slice(b"567"), Ok(3));
986
987        // read task is paused
988        assert!(!io.st().flags.is_read_paused());
989        // read buffer is ready
990        assert!(io.st().flags.is_read_ready());
991        // read backpressure is enabled
992        assert!(!io.st().flags.is_rd_backpressure());
993        // read task ready
994        assert_eq!(
995            lazy(|cx| ctx.poll_read_ready(cx)).await,
996            Poll::Ready(Readiness::Ready)
997        );
998
999        // read 4 bytes. buf size is 4, less that half of high watermark
1000        assert_eq!(io.with_read_buf(BytesMut::take), b"1234567");
1001        // read task is paused
1002        assert!(!io.st().flags.is_read_paused());
1003        // read buffer is ready
1004        assert!(!io.st().flags.is_read_ready());
1005        // read task is not woken
1006        assert!(io.st().read_task.is_set());
1007
1008        // == Terminate
1009        io.terminate();
1010        // read task is woken
1011        assert!(!io.st().read_task.is_set());
1012        // read task ready
1013        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        // incoming bytes
1033        ctx.update_read_status(BytesMut::copy_from_slice(b"1"), Ok(1));
1034
1035        assert!(!io.st().dispatch_task.is_set());
1036        // rd buffer is ready
1037        assert!(io.st().flags.is_read_ready());
1038        assert!(io.st().flags.is_read_notify());
1039        // dispatcher is notified
1040        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        // disapcher is not set
1045        assert!(!io.st().dispatch_task.is_set());
1046        // rd buffer is ready
1047        assert!(io.st().flags.is_read_ready());
1048
1049        // == start notification process again
1050        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        // read task ready
1055        assert_eq!(
1056            lazy(|cx| ctx.poll_read_ready(cx)).await,
1057            Poll::Ready(Readiness::Ready)
1058        );
1059
1060        // == enable packpressure
1061        ctx.update_read_status(BytesMut::copy_from_slice(b"2345678"), Ok(7));
1062        // read backpressure is enabled
1063        assert!(io.st().flags.is_rd_backpressure());
1064
1065        // rd buffer is ready
1066        assert!(io.st().flags.is_read_ready());
1067        assert!(io.st().flags.is_read_notify());
1068        // dispatcher is notified
1069        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        // read task paused
1073        assert_eq!(lazy(|cx| ctx.poll_read_ready(cx)).await, Poll::Pending);
1074        // read task is set
1075        assert!(io.st().read_task.is_set());
1076
1077        // == start notification process again
1078        assert!(lazy(|cx| io.poll_read_notify(cx)).await.is_pending());
1079        // read flags active
1080        assert!(!io.st().flags.is_rd_backpressure());
1081        assert!(!io.st().flags.is_read_ready());
1082        assert!(!io.st().flags.is_read_paused());
1083        // read task is woken
1084        assert!(!io.st().read_task.is_set());
1085        // read task ready
1086        assert_eq!(
1087            lazy(|cx| ctx.poll_read_ready(cx)).await,
1088            Poll::Ready(Readiness::Ready)
1089        );
1090
1091        // incoming bytes
1092        ctx.update_read_status(BytesMut::copy_from_slice(b"1"), Ok(1));
1093        assert!(!io.st().dispatch_task.is_set());
1094        // rd buffer is ready
1095        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        // dispatcher is notified
1100        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        // == Terminate
1107        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        // == No write work
1176        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        // write
1182        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        // write task is paused
1185        assert!(io.st().flags.is_write_paused());
1186        // send-buf op is scheduled
1187        assert!(io.st().flags.is_wr_send_scheduled());
1188        // back-pressure is not enabled
1189        assert!(!io.st().flags.is_wr_backpressure());
1190        // dispatch is not woken up
1191        assert!(io.st().dispatch_task.is_set());
1192
1193        // == enable wr backpressure
1194        io.with_write_buf(|buf| buf.put_slice(b"5678")).unwrap();
1195        // back-pressure is enabled
1196        assert!(io.st().flags.is_wr_backpressure());
1197        // dispatch is woken up
1198        assert!(!io.st().dispatch_task.is_set());
1199        // write task is set
1200        assert!(io.st().write_task.is_set());
1201        // dispatcher gets WriteBackpressure
1202        assert!(matches!(
1203            lazy(|cx| io.poll_status_update(cx)).await,
1204            Poll::Ready(IoStatusUpdate::WriteBackpressure)
1205        ));
1206        // flush write buffer
1207        assert!(lazy(|cx| io.poll_flush(cx, false)).await.is_pending());
1208        // full flush is not enabled
1209        assert!(!io.st().flags.is_write_flush());
1210
1211        // run send-buf ops
1212        Iops::run();
1213        // send-buf op is not scheduled
1214        assert!(!io.st().flags.is_wr_send_scheduled());
1215        // write task is not paused
1216        assert!(!io.st().flags.is_write_paused());
1217        // write task has been woken up
1218        assert!(!io.st().write_task.is_set());
1219        // write task can proceed
1220        assert_eq!(
1221            lazy(|cx| ctx.poll_write_ready(cx)).await,
1222            Poll::Ready(Readiness::Ready)
1223        );
1224
1225        // wrote 4 bytes to io
1226        assert_eq!(ctx.with_write_buf(|buf| buf.split_to(4).freeze()), b"1234");
1227        // continue to write
1228        assert_eq!(ctx.update_write_status(Ok(true)), IoTaskStatus::Io);
1229        // write task can proceed
1230        assert_eq!(
1231            lazy(|cx| ctx.poll_write_ready(cx)).await,
1232            Poll::Ready(Readiness::Ready)
1233        );
1234        // write task is not paused
1235        assert!(!io.st().flags.is_write_paused());
1236        // back-pressure is enabled
1237        assert!(io.st().flags.is_wr_backpressure());
1238        // dispatcher gets WriteBackpressure, buf wr-backpressure flags is removed
1239        assert!(matches!(
1240            lazy(|cx| io.poll_status_update(cx)).await,
1241            Poll::Ready(IoStatusUpdate::WriteBackpressure)
1242        ));
1243        // back-pressure is disabled
1244        assert!(!io.st().flags.is_wr_backpressure());
1245        assert!(lazy(|cx| io.poll_status_update(cx)).await.is_pending());
1246        // write buffer is flushed
1247        assert!(matches!(
1248            lazy(|cx| io.poll_flush(cx, false)).await,
1249            Poll::Ready(Ok(()))
1250        ));
1251
1252        // full flush write buffer
1253        io.with_write_buf(|buf| buf.put_slice(b"1234")).unwrap();
1254        assert!(lazy(|cx| io.poll_flush(cx, true)).await.is_pending());
1255        // full flush is enabled
1256        assert!(io.st().flags.is_write_flush());
1257        // back-pressure is enabled
1258        assert!(io.st().flags.is_wr_backpressure());
1259
1260        // wrote all data
1261        Iops::run();
1262        assert_eq!(ctx.with_write_buf(BytePages::freeze), b"56781234");
1263        // write task is not paused, so send-buf op is not scheduled
1264        assert!(!io.st().flags.is_wr_send_scheduled());
1265        // update status, no more work
1266        assert_eq!(ctx.update_write_status(Ok(true)), IoTaskStatus::Pause);
1267        // write task is paused
1268        assert!(io.st().flags.is_write_paused());
1269        // flush is still enabled
1270        assert!(io.st().flags.is_write_flush());
1271        // back-pressure is still enabled
1272        assert!(io.st().flags.is_wr_backpressure());
1273        // dispatch is woken up
1274        assert!(!io.st().dispatch_task.is_set());
1275
1276        // write buffer is flushed
1277        assert!(matches!(
1278            lazy(|cx| io.poll_flush(cx, false)).await,
1279            Poll::Ready(Ok(()))
1280        ));
1281        // full flush is disabled
1282        assert!(!io.st().flags.is_write_flush());
1283        // back-pressure is disabled
1284        assert!(!io.st().flags.is_wr_backpressure());
1285
1286        // == Terminate
1287        io.terminate();
1288        // read task is woken
1289        assert!(!io.st().write_task.is_set());
1290        // write task ready
1291        assert_eq!(
1292            lazy(|cx| ctx.poll_write_ready(cx)).await,
1293            Poll::Ready(Readiness::Terminate)
1294        );
1295        // flush returns error
1296        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        // statis returns error
1301        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        // Graceful shutdown must flush the pending write buffer even if
1344        // the peer keeps sending data (read backpressure is enabled).
1345        let (client, server) = IoTest::create();
1346        // remote side does not accept any data yet, write task stalls
1347        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        // queue response data; remote is stalled so it stays in the write buffer
1359        io.encode_slice(b"response-tail").unwrap();
1360        sleep(Millis(50)).await;
1361        assert_eq!(io.st().buffer.write_buf_size(), 13);
1362
1363        // peer keeps sending, crossing the read high watermark,
1364        // read task gets paused with back-pressure enabled
1365        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        // start graceful shutdown while the write buffer is not empty
1372        io.close();
1373        sleep(Millis(50)).await;
1374
1375        // peer starts draining the connection
1376        client.remote_buffer_cap(1024);
1377
1378        // all previously queued data must be written before io stream is closed.
1379        // Without the fix graceful shutdown completes immediately (read is paused
1380        // with back-pressure), dropping the buffered write data, so the read here
1381        // returns nothing instead of the queued response tail.
1382        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        // the connection still closes gracefully afterwards (within the
1389        // disconnect timeout) instead of hanging
1390        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        // layer drops all unprocessed data after filter shutdown
1398        #[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        // == init shutdown
1423        io.close();
1424        assert!(!st.flags.is_closed());
1425        assert!(st.flags.is_stopping_filters());
1426        // encoding is not allowed in shutting down stage
1427        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        // == terminate
1441        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}