add io query system

ntex-io/src/filter.rs

@@ -1,4 +1,4 @@
-use std::{io, rc::Rc, task::Context, task::Poll};
+use std::{any, io, rc::Rc, task::Context, task::Poll};
 
 use ntex_bytes::BytesMut;
 
@@ -26,6 +26,16 @@ impl Filter for DefaultFilter {
 
         Poll::Ready(Ok(()))
     }
+
+    fn query(&self, id: any::TypeId) -> Option<Box<dyn any::Any>> {
+        if let Some(hnd) = self.0.handle.take() {
+            let res = hnd.query(id);
+            self.0.handle.set(Some(hnd));
+            res
+        } else {
+            None
+        }
+    }
 }
 
 impl ReadFilter for DefaultFilter {
@@ -140,6 +150,10 @@ impl Filter for NullFilter {
     fn shutdown(&self, _: &IoRef) -> Poll<Result<(), io::Error>> {
         Poll::Ready(Ok(()))
     }
+
+    fn query(&self, _: any::TypeId) -> Option<Box<dyn any::Any>> {
+        None
+    }
 }
 
 impl ReadFilter for NullFilter {

ntex-io/src/lib.rs

@@ -7,6 +7,7 @@ mod filter;
 mod state;
 mod tasks;
 mod time;
+pub mod types;
 mod utils;
 
 #[cfg(feature = "tokio")]
@@ -14,7 +15,7 @@ mod tokio_impl;
 
 use ntex_bytes::BytesMut;
 use ntex_codec::{Decoder, Encoder};
-use ntex_util::{channel::oneshot::Receiver, future::Either, time::Millis};
+use ntex_util::time::Millis;
 
 pub use self::dispatcher::Dispatcher;
 pub use self::filter::DefaultFilter;
@@ -58,12 +59,7 @@ pub trait WriteFilter {
 pub trait Filter: ReadFilter + WriteFilter + 'static {
     fn shutdown(&self, st: &IoRef) -> Poll<Result<(), io::Error>>;
 
-    fn query(
-        &self,
-        id: TypeId,
-    ) -> Either<Option<Box<dyn Any>>, Receiver<Option<Box<dyn Any>>>> {
-        Either::Left(None)
-    }
+    fn query(&self, id: TypeId) -> Option<Box<dyn Any>>;
 }
 
 pub trait FilterFactory<F: Filter>: Sized {
@@ -76,7 +72,11 @@ pub trait FilterFactory<F: Filter>: Sized {
 }
 
 pub trait IoStream {
-    fn start(self, _: ReadContext, _: WriteContext);
+    fn start(self, _: ReadContext, _: WriteContext) -> Box<dyn Handle>;
+}
+
+pub trait Handle {
+    fn query(&self, id: TypeId) -> Option<Box<dyn Any>>;
 }
 
 /// Framed transport item

ntex-io/src/state.rs

@@ -1,6 +1,6 @@
 use std::cell::{Cell, RefCell};
 use std::task::{Context, Poll};
-use std::{fmt, future::Future, hash, io, mem, ops::Deref, pin::Pin, ptr, rc::Rc};
+use std::{any, fmt, future::Future, hash, io, mem, ops::Deref, pin::Pin, ptr, rc::Rc};
 
 use ntex_bytes::{BytesMut, PoolId, PoolRef};
 use ntex_codec::{Decoder, Encoder};
@@ -9,7 +9,7 @@ use ntex_util::{future::poll_fn, future::Either, task::LocalWaker};
 
 use super::filter::{DefaultFilter, NullFilter};
 use super::tasks::{ReadContext, WriteContext};
-use super::{Filter, FilterFactory, IoStream};
+use super::{types, Filter, FilterFactory, Handle, IoStream};
 
 bitflags::bitflags! {
     pub struct Flags: u16 {
@@ -66,6 +66,7 @@ pub(crate) struct IoStateInner {
     pub(super) read_buf: Cell<Option<BytesMut>>,
     pub(super) write_buf: Cell<Option<BytesMut>>,
     pub(super) filter: Cell<&'static dyn Filter>,
+    pub(super) handle: Cell<Option<Box<dyn Handle>>>,
     on_disconnect: RefCell<Vec<Option<LocalWaker>>>,
 }
 
@@ -198,6 +199,7 @@ impl Io {
             read_buf: Cell::new(None),
             write_buf: Cell::new(None),
             filter: Cell::new(NullFilter::get()),
+            handle: Cell::new(None),
             on_disconnect: RefCell::new(Vec::new()),
         });
 
@@ -211,7 +213,8 @@ impl Io {
         let io_ref = IoRef(inner);
 
         // start io tasks
-        io.start(ReadContext(io_ref.clone()), WriteContext(io_ref.clone()));
+        let hnd = io.start(ReadContext(io_ref.clone()), WriteContext(io_ref.clone()));
+        io_ref.0.handle.set(Some(hnd));
 
         Io(io_ref, FilterItem::Ptr(Box::into_raw(filter)))
     }
@@ -363,8 +366,12 @@ impl IoRef {
 
     #[inline]
     /// Query specific data
-    pub fn query<T: 'static>(&self) -> Option<T> {
-        todo!()
+    pub fn query<T: 'static>(&self) -> types::QueryItem<T> {
+        if let Some(item) = self.0.filter.get().query(any::TypeId::of::<T>()) {
+            types::QueryItem::new(item)
+        } else {
+            types::QueryItem::empty()
+        }
     }
 }
 
@@ -853,18 +860,17 @@ impl<'a> ReadRef<'a> {
         U: Decoder,
     {
         let mut buf = self.0.read_buf.take();
-        let result = if let Some(ref mut buf) = buf {
-            let result = codec.decode(buf);
+        if let Some(ref mut b) = buf {
+            let result = codec.decode(b);
             if result.as_ref().map(|v| v.is_none()).unwrap_or(false) {
                 self.0.remove_flags(Flags::RD_READY);
             }
+            self.0.read_buf.set(buf);
             result
         } else {
             self.0.remove_flags(Flags::RD_READY);
             Ok(None)
-        };
-        self.0.read_buf.set(buf);
-        result
+        }
     }
 
     #[inline]

ntex-io/src/testing.rs

@@ -1,13 +1,13 @@
 use std::cell::{Cell, RefCell};
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll, Waker};
-use std::{cmp, fmt, future::Future, io, mem, pin::Pin, rc::Rc};
+use std::{any, cmp, fmt, future::Future, io, mem, net, pin::Pin, rc::Rc};
 
 use ntex_bytes::{Buf, BufMut, BytesMut};
 use ntex_util::future::poll_fn;
 use ntex_util::time::{sleep, Millis, Sleep};
 
-use crate::{IoStream, ReadContext, WriteContext, WriteReadiness};
+use crate::{types, Handle, IoStream, ReadContext, WriteContext, WriteReadiness};
 
 #[derive(Default)]
 struct AtomicWaker(Arc<Mutex<RefCell<Option<Waker>>>>);
@@ -30,6 +30,7 @@ impl fmt::Debug for AtomicWaker {
 #[derive(Debug)]
 pub struct IoTest {
     tp: Type,
+    peer_addr: Option<net::SocketAddr>,
     state: Arc<Cell<State>>,
     local: Arc<Mutex<RefCell<Channel>>>,
     remote: Arc<Mutex<RefCell<Channel>>>,
@@ -102,12 +103,14 @@ impl IoTest {
         (
             IoTest {
                 tp: Type::Client,
+                peer_addr: None,
                 local: local.clone(),
                 remote: remote.clone(),
                 state: state.clone(),
             },
             IoTest {
                 state,
+                peer_addr: None,
                 tp: Type::Server,
                 local: remote,
                 remote: local,
@@ -128,6 +131,12 @@ impl IoTest {
         self.remote.lock().unwrap().borrow().is_closed()
     }
 
+    /// Set peer addr
+    pub fn set_peer_addr(mut self, addr: net::SocketAddr) -> Self {
+        self.peer_addr = Some(addr);
+        self
+    }
+
     /// Set read to Pending state
     pub fn read_pending(&self) {
         self.remote.lock().unwrap().borrow_mut().read = IoState::Pending;
@@ -317,6 +326,7 @@ impl Clone for IoTest {
             local: self.local.clone(),
             remote: self.remote.clone(),
             state: self.state.clone(),
+            peer_addr: self.peer_addr,
         }
     }
 }
@@ -444,7 +454,7 @@ mod tokio {
 }
 
 impl IoStream for IoTest {
-    fn start(self, read: ReadContext, write: WriteContext) {
+    fn start(self, read: ReadContext, write: WriteContext) -> Box<dyn Handle> {
         let io = Rc::new(self);
 
         ntex_util::spawn(ReadTask {
@@ -452,10 +462,23 @@ impl IoStream for IoTest {
             state: read,
         });
         ntex_util::spawn(WriteTask {
-            io,
+            io: io.clone(),
             state: write,
             st: IoWriteState::Processing(None),
         });
+
+        Box::new(io)
+    }
+}
+
+impl Handle for Rc<IoTest> {
+    fn query(&self, id: any::TypeId) -> Option<Box<dyn any::Any>> {
+        if id == any::TypeId::of::<types::PeerAddr>() {
+            if let Some(addr) = self.peer_addr {
+                return Some(Box::new(addr));
+            }
+        }
+        None
     }
 }
 

ntex-io/src/tokio_impl.rs

@@ -1,31 +1,45 @@
 use std::task::{Context, Poll};
-use std::{cell::RefCell, cmp, future::Future, io, pin::Pin, rc::Rc};
+use std::{any, cell::RefCell, cmp, future::Future, io, pin::Pin, rc::Rc};
 
 use ntex_bytes::{Buf, BufMut};
 use ntex_util::time::{sleep, Sleep};
 use tok_io::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tok_io::net::TcpStream;
 
-use super::{Filter, Io, IoStream, ReadContext, WriteContext, WriteReadiness};
+use super::{
+    types, Filter, Handle, Io, IoStream, ReadContext, WriteContext, WriteReadiness,
+};
 
 impl IoStream for TcpStream {
-    fn start(self, read: ReadContext, write: WriteContext) {
+    fn start(self, read: ReadContext, write: WriteContext) -> Box<dyn Handle> {
         let io = Rc::new(RefCell::new(self));
 
         ntex_util::spawn(ReadTask::new(io.clone(), read));
-        ntex_util::spawn(WriteTask::new(io, write));
+        ntex_util::spawn(WriteTask::new(io.clone(), write));
+        Box::new(io)
     }
 }
 
 #[cfg(unix)]
 impl IoStream for tok_io::net::UnixStream {
-    fn start(self, _read: ReadContext, _write: WriteContext) {
+    fn start(self, _read: ReadContext, _write: WriteContext) -> Box<dyn Handle> {
         let _io = Rc::new(RefCell::new(self));
 
         todo!()
     }
 }
 
+impl Handle for Rc<RefCell<TcpStream>> {
+    fn query(&self, id: any::TypeId) -> Option<Box<dyn any::Any>> {
+        if id == any::TypeId::of::<types::PeerAddr>() {
+            if let Ok(addr) = self.borrow().peer_addr() {
+                return Some(Box::new(addr));
+            }
+        }
+        None
+    }
+}
+
 /// Read io task
 struct ReadTask {
     io: Rc<RefCell<TcpStream>>,

ntex-io/src/types.rs

@@ -0,0 +1,38 @@
+use std::{any, fmt, marker::PhantomData, net::SocketAddr};
+
+pub struct PeerAddr(pub SocketAddr);
+
+impl fmt::Debug for PeerAddr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+pub struct QueryItem<T> {
+    item: Option<Box<dyn any::Any>>,
+    _t: PhantomData<T>,
+}
+
+impl<T: any::Any> QueryItem<T> {
+    pub(crate) fn new(item: Box<dyn any::Any>) -> Self {
+        Self {
+            item: Some(item),
+            _t: PhantomData,
+        }
+    }
+
+    pub(crate) fn empty() -> Self {
+        Self {
+            item: None,
+            _t: PhantomData,
+        }
+    }
+
+    pub fn get(&self) -> Option<&T> {
+        if let Some(ref item) = self.item {
+            item.downcast_ref()
+        } else {
+            None
+        }
+    }
+}