ntex_io/
ioref.rs

1use std::{any, fmt, hash, io};
2
3use ntex_bytes::{BytesVec, PoolRef};
4use ntex_codec::{Decoder, Encoder};
5use ntex_util::time::Seconds;
6
7use crate::{
8    timer, types, Decoded, Filter, FilterCtx, Flags, IoRef, OnDisconnect, WriteBuf,
9};
10
11impl IoRef {
12    #[inline]
13    #[doc(hidden)]
14    /// Get current state flags
15    pub fn flags(&self) -> Flags {
16        self.0.flags.get()
17    }
18
19    #[inline]
20    /// Get current filter
21    pub(crate) fn filter(&self) -> &dyn Filter {
22        self.0.filter()
23    }
24
25    #[inline]
26    /// Get memory pool
27    pub fn memory_pool(&self) -> PoolRef {
28        self.0.pool.get()
29    }
30
31    #[inline]
32    /// Check if io stream is closed
33    pub fn is_closed(&self) -> bool {
34        self.0
35            .flags
36            .get()
37            .intersects(Flags::IO_STOPPING | Flags::IO_STOPPED)
38    }
39
40    #[inline]
41    /// Check if write back-pressure is enabled
42    pub fn is_wr_backpressure(&self) -> bool {
43        self.0.flags.get().contains(Flags::BUF_W_BACKPRESSURE)
44    }
45
46    #[inline]
47    /// Wake dispatcher task
48    pub fn wake(&self) {
49        self.0.dispatch_task.wake();
50    }
51
52    #[inline]
53    /// Gracefully close connection
54    ///
55    /// Notify dispatcher and initiate io stream shutdown process.
56    pub fn close(&self) {
57        self.0.insert_flags(Flags::DSP_STOP);
58        self.0.init_shutdown();
59    }
60
61    #[inline]
62    /// Force close connection
63    ///
64    /// Dispatcher does not wait for uncompleted responses. Io stream get terminated
65    /// without any graceful period.
66    pub fn force_close(&self) {
67        log::trace!("{}: Force close io stream object", self.tag());
68        self.0.insert_flags(
69            Flags::DSP_STOP
70                | Flags::IO_STOPPED
71                | Flags::IO_STOPPING
72                | Flags::IO_STOPPING_FILTERS,
73        );
74        self.0.read_task.wake();
75        self.0.write_task.wake();
76        self.0.dispatch_task.wake();
77    }
78
79    #[inline]
80    /// Gracefully shutdown io stream
81    pub fn want_shutdown(&self) {
82        if !self
83            .0
84            .flags
85            .get()
86            .intersects(Flags::IO_STOPPED | Flags::IO_STOPPING | Flags::IO_STOPPING_FILTERS)
87        {
88            log::trace!(
89                "{}: Initiate io shutdown {:?}",
90                self.tag(),
91                self.0.flags.get()
92            );
93            self.0.insert_flags(Flags::IO_STOPPING_FILTERS);
94            self.0.read_task.wake();
95        }
96    }
97
98    #[inline]
99    /// Query filter specific data
100    pub fn query<T: 'static>(&self) -> types::QueryItem<T> {
101        if let Some(item) = self.filter().query(any::TypeId::of::<T>()) {
102            types::QueryItem::new(item)
103        } else {
104            types::QueryItem::empty()
105        }
106    }
107
108    #[inline]
109    /// Encode and write item to a buffer and wake up write task
110    pub fn encode<U>(&self, item: U::Item, codec: &U) -> Result<(), <U as Encoder>::Error>
111    where
112        U: Encoder,
113    {
114        if !self.is_closed() {
115            self.with_write_buf(|buf| {
116                // make sure we've got room
117                self.memory_pool().resize_write_buf(buf);
118
119                // encode item and wake write task
120                codec.encode_vec(item, buf)
121            })
122            // .with_write_buf() could return io::Error<Result<(), U::Error>>,
123            // in that case mark io as failed
124            .unwrap_or_else(|err| {
125                log::trace!(
126                    "{}: Got io error while encoding, error: {:?}",
127                    self.tag(),
128                    err
129                );
130                self.0.io_stopped(Some(err));
131                Ok(())
132            })
133        } else {
134            log::trace!("{}: Io is closed/closing, skip frame encoding", self.tag());
135            Ok(())
136        }
137    }
138
139    #[inline]
140    /// Attempts to decode a frame from the read buffer
141    pub fn decode<U>(
142        &self,
143        codec: &U,
144    ) -> Result<Option<<U as Decoder>::Item>, <U as Decoder>::Error>
145    where
146        U: Decoder,
147    {
148        self.0
149            .buffer
150            .with_read_destination(self, |buf| codec.decode_vec(buf))
151    }
152
153    #[inline]
154    /// Attempts to decode a frame from the read buffer
155    pub fn decode_item<U>(
156        &self,
157        codec: &U,
158    ) -> Result<Decoded<<U as Decoder>::Item>, <U as Decoder>::Error>
159    where
160        U: Decoder,
161    {
162        self.0.buffer.with_read_destination(self, |buf| {
163            let len = buf.len();
164            codec.decode_vec(buf).map(|item| Decoded {
165                item,
166                remains: buf.len(),
167                consumed: len - buf.len(),
168            })
169        })
170    }
171
172    #[inline]
173    /// Write bytes to a buffer and wake up write task
174    pub fn write(&self, src: &[u8]) -> io::Result<()> {
175        self.with_write_buf(|buf| buf.extend_from_slice(src))
176    }
177
178    #[inline]
179    /// Get access to write buffer
180    pub fn with_buf<F, R>(&self, f: F) -> io::Result<R>
181    where
182        F: FnOnce(&WriteBuf<'_>) -> R,
183    {
184        let ctx = FilterCtx::new(self, &self.0.buffer);
185        let result = ctx.write_buf(f);
186        self.0.filter().process_write_buf(ctx)?;
187        Ok(result)
188    }
189
190    #[inline]
191    /// Get mut access to source write buffer
192    pub fn with_write_buf<F, R>(&self, f: F) -> io::Result<R>
193    where
194        F: FnOnce(&mut BytesVec) -> R,
195    {
196        if self.0.flags.get().contains(Flags::IO_STOPPED) {
197            Err(self.0.error_or_disconnected())
198        } else {
199            let result = self.0.buffer.with_write_source(self, f);
200            self.0
201                .filter()
202                .process_write_buf(FilterCtx::new(self, &self.0.buffer))?;
203            Ok(result)
204        }
205    }
206
207    #[doc(hidden)]
208    #[inline]
209    /// Get mut access to destination write buffer
210    pub fn with_write_dest_buf<F, R>(&self, f: F) -> R
211    where
212        F: FnOnce(Option<&mut BytesVec>) -> R,
213    {
214        self.0.buffer.with_write_destination(self, f)
215    }
216
217    #[inline]
218    /// Get mut access to source read buffer
219    pub fn with_read_buf<F, R>(&self, f: F) -> R
220    where
221        F: FnOnce(&mut BytesVec) -> R,
222    {
223        self.0.buffer.with_read_destination(self, f)
224    }
225
226    #[inline]
227    /// Wakeup dispatcher
228    pub fn notify_dispatcher(&self) {
229        self.0.dispatch_task.wake();
230        log::trace!("{}: Timer, notify dispatcher", self.tag());
231    }
232
233    #[inline]
234    /// Wakeup dispatcher and send keep-alive error
235    pub fn notify_timeout(&self) {
236        self.0.notify_timeout()
237    }
238
239    #[inline]
240    /// current timer handle
241    pub fn timer_handle(&self) -> timer::TimerHandle {
242        self.0.timeout.get()
243    }
244
245    #[inline]
246    /// Start timer
247    pub fn start_timer(&self, timeout: Seconds) -> timer::TimerHandle {
248        let cur_hnd = self.0.timeout.get();
249
250        if !timeout.is_zero() {
251            if cur_hnd.is_set() {
252                let hnd = timer::update(cur_hnd, timeout, self);
253                if hnd != cur_hnd {
254                    log::trace!("{}: Update timer {:?}", self.tag(), timeout);
255                    self.0.timeout.set(hnd);
256                }
257                hnd
258            } else {
259                log::trace!("{}: Start timer {:?}", self.tag(), timeout);
260                let hnd = timer::register(timeout, self);
261                self.0.timeout.set(hnd);
262                hnd
263            }
264        } else {
265            if cur_hnd.is_set() {
266                self.0.timeout.set(timer::TimerHandle::ZERO);
267                timer::unregister(cur_hnd, self);
268            }
269            timer::TimerHandle::ZERO
270        }
271    }
272
273    #[inline]
274    /// Stop timer
275    pub fn stop_timer(&self) {
276        let hnd = self.0.timeout.get();
277        if hnd.is_set() {
278            log::trace!("{}: Stop timer", self.tag());
279            self.0.timeout.set(timer::TimerHandle::ZERO);
280            timer::unregister(hnd, self)
281        }
282    }
283
284    #[inline]
285    /// Get tag
286    pub fn tag(&self) -> &'static str {
287        self.0.tag.get()
288    }
289
290    #[inline]
291    /// Set tag
292    pub fn set_tag(&self, tag: &'static str) {
293        self.0.tag.set(tag)
294    }
295
296    #[inline]
297    /// Notify when io stream get disconnected
298    pub fn on_disconnect(&self) -> OnDisconnect {
299        OnDisconnect::new(self.0.clone())
300    }
301}
302
303impl Eq for IoRef {}
304
305impl PartialEq for IoRef {
306    #[inline]
307    fn eq(&self, other: &Self) -> bool {
308        self.0.eq(&other.0)
309    }
310}
311
312impl hash::Hash for IoRef {
313    #[inline]
314    fn hash<H: hash::Hasher>(&self, state: &mut H) {
315        self.0.hash(state);
316    }
317}
318
319impl fmt::Debug for IoRef {
320    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321        f.debug_struct("IoRef")
322            .field("state", self.0.as_ref())
323            .finish()
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use std::cell::{Cell, RefCell};
330    use std::{future::poll_fn, future::Future, pin::Pin, rc::Rc, task::Poll};
331
332    use ntex_bytes::Bytes;
333    use ntex_codec::BytesCodec;
334    use ntex_util::future::lazy;
335    use ntex_util::time::{sleep, Millis};
336
337    use super::*;
338    use crate::{testing::IoTest, FilterCtx, FilterReadStatus, Io};
339
340    const BIN: &[u8] = b"GET /test HTTP/1\r\n\r\n";
341    const TEXT: &str = "GET /test HTTP/1\r\n\r\n";
342
343    #[ntex::test]
344    async fn utils() {
345        let (client, server) = IoTest::create();
346        client.remote_buffer_cap(1024);
347        client.write(TEXT);
348
349        let state = Io::new(server);
350        assert_eq!(state.get_ref(), state.get_ref());
351
352        let msg = state.recv(&BytesCodec).await.unwrap().unwrap();
353        assert_eq!(msg, Bytes::from_static(BIN));
354        assert_eq!(state.get_ref(), state.as_ref().clone());
355        assert!(format!("{:?}", state).find("Io {").is_some());
356        assert!(format!("{:?}", state.get_ref()).find("IoRef {").is_some());
357
358        let res = poll_fn(|cx| Poll::Ready(state.poll_recv(&BytesCodec, cx))).await;
359        assert!(res.is_pending());
360        client.write(TEXT);
361        sleep(Millis(50)).await;
362        let res = poll_fn(|cx| Poll::Ready(state.poll_recv(&BytesCodec, cx))).await;
363        if let Poll::Ready(msg) = res {
364            assert_eq!(msg.unwrap(), Bytes::from_static(BIN));
365        }
366
367        client.read_error(io::Error::other("err"));
368        let msg = state.recv(&BytesCodec).await;
369        assert!(msg.is_err());
370        assert!(state.flags().contains(Flags::IO_STOPPED));
371
372        let (client, server) = IoTest::create();
373        client.remote_buffer_cap(1024);
374        let state = Io::new(server);
375
376        client.read_error(io::Error::other("err"));
377        let res = poll_fn(|cx| Poll::Ready(state.poll_recv(&BytesCodec, cx))).await;
378        if let Poll::Ready(msg) = res {
379            assert!(msg.is_err());
380            assert!(state.flags().contains(Flags::IO_STOPPED));
381            assert!(state.flags().contains(Flags::DSP_STOP));
382        }
383
384        let (client, server) = IoTest::create();
385        client.remote_buffer_cap(1024);
386        let state = Io::new(server);
387        state.write(b"test").unwrap();
388        let buf = client.read().await.unwrap();
389        assert_eq!(buf, Bytes::from_static(b"test"));
390
391        client.write(b"test");
392        state.read_ready().await.unwrap();
393        let buf = state.decode(&BytesCodec).unwrap().unwrap();
394        assert_eq!(buf, Bytes::from_static(b"test"));
395
396        client.write_error(io::Error::other("err"));
397        let res = state.send(Bytes::from_static(b"test"), &BytesCodec).await;
398        assert!(res.is_err());
399        assert!(state.flags().contains(Flags::IO_STOPPED));
400
401        let (client, server) = IoTest::create();
402        client.remote_buffer_cap(1024);
403        let state = Io::new(server);
404        state.force_close();
405        assert!(state.flags().contains(Flags::DSP_STOP));
406        assert!(state.flags().contains(Flags::IO_STOPPING));
407    }
408
409    #[ntex::test]
410    async fn read_readiness() {
411        let (client, server) = IoTest::create();
412        client.remote_buffer_cap(1024);
413
414        let io = Io::new(server);
415        assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
416
417        client.write(TEXT);
418        assert_eq!(io.read_ready().await.unwrap(), Some(()));
419        assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
420
421        let item = io.with_read_buf(|buffer| buffer.split());
422        assert_eq!(item, Bytes::from_static(BIN));
423
424        client.write(TEXT);
425        sleep(Millis(50)).await;
426        assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_ready());
427        assert!(lazy(|cx| io.poll_read_ready(cx)).await.is_pending());
428    }
429
430    #[ntex::test]
431    #[allow(clippy::unit_cmp)]
432    async fn on_disconnect() {
433        let (client, server) = IoTest::create();
434        let state = Io::new(server);
435        let mut waiter = state.on_disconnect();
436        assert_eq!(
437            lazy(|cx| Pin::new(&mut waiter).poll(cx)).await,
438            Poll::Pending
439        );
440        let mut waiter2 = waiter.clone();
441        assert_eq!(
442            lazy(|cx| Pin::new(&mut waiter2).poll(cx)).await,
443            Poll::Pending
444        );
445        client.close().await;
446        assert_eq!(waiter.await, ());
447        assert_eq!(waiter2.await, ());
448
449        let mut waiter = state.on_disconnect();
450        assert_eq!(
451            lazy(|cx| Pin::new(&mut waiter).poll(cx)).await,
452            Poll::Ready(())
453        );
454
455        let (client, server) = IoTest::create();
456        let state = Io::new(server);
457        let mut waiter = state.on_disconnect();
458        assert_eq!(
459            lazy(|cx| Pin::new(&mut waiter).poll(cx)).await,
460            Poll::Pending
461        );
462        client.read_error(io::Error::other("err"));
463        assert_eq!(waiter.await, ());
464    }
465
466    #[ntex::test]
467    async fn write_to_closed_io() {
468        let (client, server) = IoTest::create();
469        let state = Io::new(server);
470        client.close().await;
471
472        assert!(state.is_closed());
473        assert!(state.write(TEXT.as_bytes()).is_err());
474        assert!(state
475            .with_write_buf(|buf| buf.extend_from_slice(BIN))
476            .is_err());
477    }
478
479    #[derive(Debug)]
480    struct Counter<F> {
481        layer: F,
482        idx: usize,
483        in_bytes: Rc<Cell<usize>>,
484        out_bytes: Rc<Cell<usize>>,
485        read_order: Rc<RefCell<Vec<usize>>>,
486        write_order: Rc<RefCell<Vec<usize>>>,
487    }
488
489    impl<F: Filter> Filter for Counter<F> {
490        fn process_read_buf(
491            &self,
492            ctx: FilterCtx<'_>,
493            nbytes: usize,
494        ) -> io::Result<FilterReadStatus> {
495            self.read_order.borrow_mut().push(self.idx);
496            self.in_bytes.set(self.in_bytes.get() + nbytes);
497            self.layer.process_read_buf(ctx, nbytes)
498        }
499
500        fn process_write_buf(&self, ctx: FilterCtx<'_>) -> io::Result<()> {
501            self.write_order.borrow_mut().push(self.idx);
502            self.out_bytes.set(
503                self.out_bytes.get()
504                    + ctx.write_buf(|buf| {
505                        buf.with_src(|b| b.as_ref().map(|b| b.len()).unwrap_or_default())
506                    }),
507            );
508            self.layer.process_write_buf(ctx)
509        }
510
511        crate::forward_ready!(layer);
512        crate::forward_query!(layer);
513        crate::forward_shutdown!(layer);
514    }
515
516    #[ntex::test]
517    async fn filter() {
518        let in_bytes = Rc::new(Cell::new(0));
519        let out_bytes = Rc::new(Cell::new(0));
520        let read_order = Rc::new(RefCell::new(Vec::new()));
521        let write_order = Rc::new(RefCell::new(Vec::new()));
522
523        let (client, server) = IoTest::create();
524        let io = Io::new(server).map_filter(|layer| Counter {
525            layer,
526            idx: 1,
527            in_bytes: in_bytes.clone(),
528            out_bytes: out_bytes.clone(),
529            read_order: read_order.clone(),
530            write_order: write_order.clone(),
531        });
532
533        client.remote_buffer_cap(1024);
534        client.write(TEXT);
535        let msg = io.recv(&BytesCodec).await.unwrap().unwrap();
536        assert_eq!(msg, Bytes::from_static(BIN));
537
538        io.send(Bytes::from_static(b"test"), &BytesCodec)
539            .await
540            .unwrap();
541        let buf = client.read().await.unwrap();
542        assert_eq!(buf, Bytes::from_static(b"test"));
543
544        assert_eq!(in_bytes.get(), BIN.len());
545        assert_eq!(out_bytes.get(), 4);
546    }
547
548    #[ntex::test]
549    async fn boxed_filter() {
550        let in_bytes = Rc::new(Cell::new(0));
551        let out_bytes = Rc::new(Cell::new(0));
552        let read_order = Rc::new(RefCell::new(Vec::new()));
553        let write_order = Rc::new(RefCell::new(Vec::new()));
554
555        let (client, server) = IoTest::create();
556        let state = Io::new(server)
557            .map_filter(|layer| Counter {
558                layer,
559                idx: 2,
560                in_bytes: in_bytes.clone(),
561                out_bytes: out_bytes.clone(),
562                read_order: read_order.clone(),
563                write_order: write_order.clone(),
564            })
565            .map_filter(|layer| Counter {
566                layer,
567                idx: 1,
568                in_bytes: in_bytes.clone(),
569                out_bytes: out_bytes.clone(),
570                read_order: read_order.clone(),
571                write_order: write_order.clone(),
572            });
573        let state = state.seal();
574
575        client.remote_buffer_cap(1024);
576        client.write(TEXT);
577        let msg = state.recv(&BytesCodec).await.unwrap().unwrap();
578        assert_eq!(msg, Bytes::from_static(BIN));
579
580        state
581            .send(Bytes::from_static(b"test"), &BytesCodec)
582            .await
583            .unwrap();
584        let buf = client.read().await.unwrap();
585        assert_eq!(buf, Bytes::from_static(b"test"));
586
587        assert_eq!(in_bytes.get(), BIN.len() * 2);
588        assert_eq!(out_bytes.get(), 8);
589        assert_eq!(
590            state.with_write_dest_buf(|b| b.map(|b| b.len()).unwrap_or(0)),
591            0
592        );
593
594        // refs
595        assert_eq!(Rc::strong_count(&in_bytes), 3);
596        drop(state);
597        assert_eq!(Rc::strong_count(&in_bytes), 1);
598        assert_eq!(*read_order.borrow(), &[1, 2][..]);
599        assert_eq!(*write_order.borrow(), &[1, 2][..]);
600    }
601}
ntex_io/ioref.rs

ntex_io/
ioref.rs
