refactor connect stage

actix-ioframe/src/cell.rs

@@ -29,10 +29,6 @@ impl<T> Cell<T> {
         }
     }
 
-    pub fn get_ref(&self) -> &T {
-        unsafe { &*self.inner.as_ref().get() }
-    }
-
     pub fn get_mut(&mut self) -> &mut T {
         unsafe { &mut *self.inner.as_ref().get() }
     }

actix-ioframe/src/connect.rs

@@ -1,140 +1,45 @@
+use std::marker::PhantomData;
+
 use actix_codec::{AsyncRead, AsyncWrite, Decoder, Encoder, Framed};
 use futures::unsync::mpsc;
 
-use crate::cell::Cell;
 use crate::dispatcher::FramedMessage;
 use crate::sink::Sink;
 
 pub struct Connect<Io, St = (), Codec = ()> {
-    io: IoItem<Io, Codec>,
+    io: Io,
-    state: St,
+    _t: PhantomData<(St, Codec)>,
 }
 
-enum IoItem<Io, Codec> {
+impl<Io> Connect<Io>
-    Io(Io),
-    Framed(Framed<Io, Codec>),
-}
-
-impl<Io, C> IoItem<Io, C>
 where
     Io: AsyncRead + AsyncWrite,
 {
-    fn into_codec<Codec>(self, codec: Codec) -> IoItem<Io, Codec>
-    where
-        Codec: Encoder + Decoder,
-    {
-        match self {
-            IoItem::Io(io) => IoItem::Framed(Framed::new(io, codec)),
-            IoItem::Framed(framed) => IoItem::Framed(framed.into_framed(codec)),
-        }
-    }
-
-    fn as_framed(&mut self) -> &mut Framed<Io, C>
-    where
-        C: Encoder + Decoder,
-    {
-        match self {
-            IoItem::Io(_) => panic!("Codec is not set"),
-            IoItem::Framed(ref mut framed) => framed,
-        }
-    }
-
-    fn into_framed(self) -> Framed<Io, C>
-    where
-        C: Encoder + Decoder,
-    {
-        match self {
-            IoItem::Io(_) => panic!("Codec is not set"),
-            IoItem::Framed(framed) => framed,
-        }
-    }
-}
-
-impl<Io> Connect<Io> {
     pub(crate) fn new(io: Io) -> Self {
         Self {
-            io: IoItem::Io(io),
+            io,
-            state: (),
+            _t: PhantomData,
         }
     }
-}
 
-impl<Io, S, C> Connect<Io, S, C>
+    pub fn codec<Codec>(self, codec: Codec) -> ConnectResult<Io, (), Codec>
-where
-    Io: AsyncRead + AsyncWrite,
-{
-    pub fn codec<Codec>(self, codec: Codec) -> Connect<Io, S, Codec>
     where
         Codec: Encoder + Decoder,
     {
-        Connect {
-            io: self.io.into_codec(codec),
-            state: self.state,
-        }
-    }
-
-    pub fn state<St>(self, state: St) -> Connect<Io, St, C> {
-        Connect { state, io: self.io }
-    }
-}
-
-impl<Io, S, C> Connect<Io, S, C>
-where
-    C: Encoder + Decoder,
-    Io: AsyncRead + AsyncWrite,
-{
-    pub fn into_result(self) -> ConnectResult<Io, S, C> {
         let (tx, rx) = mpsc::unbounded();
         let sink = Sink::new(tx);
 
         ConnectResult {
-            state: Cell::new(self.state),
+            state: (),
-            framed: self.io.into_framed(),
+            framed: Framed::new(self.io, codec),
             rx,
             sink,
         }
     }
 }
 
-impl<Io, St, Codec> futures::Stream for Connect<Io, St, Codec>
-where
-    Codec: Encoder + Decoder,
-    Io: AsyncRead + AsyncWrite,
-{
-    type Item = <Codec as Decoder>::Item;
-    type Error = <Codec as Decoder>::Error;
-
-    fn poll(&mut self) -> futures::Poll<Option<Self::Item>, Self::Error> {
-        self.io.as_framed().poll()
-    }
-}
-
-impl<Io, St, Codec> futures::Sink for Connect<Io, St, Codec>
-where
-    Codec: Encoder + Decoder,
-    Io: AsyncRead + AsyncWrite,
-{
-    type SinkItem = <Codec as Encoder>::Item;
-    type SinkError = <Codec as Encoder>::Error;
-
-    fn start_send(
-        &mut self,
-        item: Self::SinkItem,
-    ) -> futures::StartSend<Self::SinkItem, Self::SinkError> {
-        self.io.as_framed().start_send(item)
-    }
-
-    fn poll_complete(&mut self) -> futures::Poll<(), Self::SinkError> {
-        self.io.as_framed().poll_complete()
-    }
-
-    fn close(&mut self) -> futures::Poll<(), Self::SinkError> {
-        self.io.as_framed().close()
-    }
-}
-
 pub struct ConnectResult<Io, St, Codec: Encoder + Decoder> {
-    pub(crate) state: Cell<St>,
+    pub(crate) state: St,
     pub(crate) framed: Framed<Io, Codec>,
     pub(crate) rx: mpsc::UnboundedReceiver<FramedMessage<<Codec as Encoder>::Item>>,
     pub(crate) sink: Sink<<Codec as Encoder>::Item>,
@@ -145,6 +50,16 @@ impl<Io, St, Codec: Encoder + Decoder> ConnectResult<Io, St, Codec> {
     pub fn sink(&self) -> &Sink<<Codec as Encoder>::Item> {
         &self.sink
     }
+
+    #[inline]
+    pub fn state<S>(self, state: S) -> ConnectResult<Io, S, Codec> {
+        ConnectResult {
+            state,
+            framed: self.framed,
+            rx: self.rx,
+            sink: self.sink,
+        }
+    }
 }
 
 impl<Io, St, Codec> futures::Stream for ConnectResult<Io, St, Codec>

actix-ioframe/src/dispatcher.rs

@@ -13,6 +13,7 @@ use crate::cell::Cell;
 use crate::error::ServiceError;
 use crate::item::Item;
 use crate::sink::Sink;
+use crate::state::State;
 
 type Request<S, U> = Item<S, U>;
 type Response<U> = <U as Encoder>::Item;
@@ -37,8 +38,8 @@ where
 {
     service: S,
     sink: Sink<<U as Encoder>::Item>,
-    state: Cell<St>,
+    state: State<St>,
-    dispatch_state: State<S, U>,
+    dispatch_state: FramedState<S, U>,
     framed: Framed<T, U>,
     rx: Option<mpsc::UnboundedReceiver<FramedMessage<<U as Encoder>::Item>>>,
     inner: Cell<FramedDispatcherInner<<U as Encoder>::Item, S::Error>>,
@@ -56,7 +57,7 @@ where
 {
     pub(crate) fn new<F: IntoService<S>>(
         framed: Framed<T, U>,
-        state: Cell<St>,
+        state: State<St>,
         service: F,
         rx: mpsc::UnboundedReceiver<FramedMessage<<U as Encoder>::Item>>,
         sink: Sink<<U as Encoder>::Item>,
@@ -67,7 +68,7 @@ where
             sink,
             rx: Some(rx),
             service: service.into_service(),
-            dispatch_state: State::Processing,
+            dispatch_state: FramedState::Processing,
             inner: Cell::new(FramedDispatcherInner {
                 buf: VecDeque::new(),
                 task: AtomicTask::new(),
@@ -76,7 +77,7 @@ where
     }
 }
 
-enum State<S: Service, U: Encoder + Decoder> {
+enum FramedState<S: Service, U: Encoder + Decoder> {
     Processing,
     Error(ServiceError<S::Error, U>),
     FramedError(ServiceError<S::Error, U>),
@@ -84,22 +85,22 @@ enum State<S: Service, U: Encoder + Decoder> {
     Stopping,
 }
 
-impl<S: Service, U: Encoder + Decoder> State<S, U> {
+impl<S: Service, U: Encoder + Decoder> FramedState<S, U> {
     fn stop(&mut self, tx: Option<oneshot::Sender<()>>) {
         match self {
-            State::FlushAndStop(ref mut vec) => {
+            FramedState::FlushAndStop(ref mut vec) => {
                 if let Some(tx) = tx {
                     vec.push(tx)
                 }
             }
-            State::Processing => {
+            FramedState::Processing => {
-                *self = State::FlushAndStop(if let Some(tx) = tx {
+                *self = FramedState::FlushAndStop(if let Some(tx) = tx {
                     vec![tx]
                 } else {
                     Vec::new()
                 })
             }
-            State::Error(_) | State::FramedError(_) | State::Stopping => {
+            FramedState::Error(_) | FramedState::FramedError(_) | FramedState::Stopping => {
                 if let Some(tx) = tx {
                     let _ = tx.send(());
                 }
@@ -131,12 +132,12 @@ where
                         Ok(Async::Ready(Some(el))) => el,
                         Err(err) => {
                             self.dispatch_state =
-                                State::FramedError(ServiceError::Decoder(err));
+                                FramedState::FramedError(ServiceError::Decoder(err));
                             return true;
                         }
                         Ok(Async::NotReady) => return false,
                         Ok(Async::Ready(None)) => {
-                            self.dispatch_state = State::Stopping;
+                            self.dispatch_state = FramedState::Stopping;
                             return true;
                         }
                     };
@@ -161,7 +162,7 @@ where
                 }
                 Ok(Async::NotReady) => return false,
                 Err(err) => {
-                    self.dispatch_state = State::Error(ServiceError::Service(err));
+                    self.dispatch_state = FramedState::Error(ServiceError::Service(err));
                     return true;
                 }
             }
@@ -180,13 +181,14 @@ where
                         Ok(msg) => {
                             if let Err(err) = self.framed.force_send(msg) {
                                 self.dispatch_state =
-                                    State::FramedError(ServiceError::Encoder(err));
+                                    FramedState::FramedError(ServiceError::Encoder(err));
                                 return true;
                             }
                             buf_empty = inner.buf.is_empty();
                         }
                         Err(err) => {
-                            self.dispatch_state = State::Error(ServiceError::Service(err));
+                            self.dispatch_state =
+                                FramedState::Error(ServiceError::Service(err));
                             return true;
                         }
                     }
@@ -197,7 +199,7 @@ where
                         Ok(Async::Ready(Some(FramedMessage::Message(msg)))) => {
                             if let Err(err) = self.framed.force_send(msg) {
                                 self.dispatch_state =
-                                    State::FramedError(ServiceError::Encoder(err));
+                                    FramedState::FramedError(ServiceError::Encoder(err));
                                 return true;
                             }
                         }
@@ -231,7 +233,8 @@ where
                     Ok(Async::NotReady) => break,
                     Err(err) => {
                         debug!("Error sending data: {:?}", err);
-                        self.dispatch_state = State::FramedError(ServiceError::Encoder(err));
+                        self.dispatch_state =
+                            FramedState::FramedError(ServiceError::Encoder(err));
                         return true;
                     }
                     Ok(Async::Ready(_)) => (),
@@ -259,32 +262,32 @@ where
     type Error = ServiceError<S::Error, U>;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match mem::replace(&mut self.dispatch_state, State::Processing) {
+        match mem::replace(&mut self.dispatch_state, FramedState::Processing) {
-            State::Processing => {
+            FramedState::Processing => {
                 if self.poll_read() || self.poll_write() {
                     self.poll()
                 } else {
                     Ok(Async::NotReady)
                 }
             }
-            State::Error(err) => {
+            FramedState::Error(err) => {
                 if self.framed.is_write_buf_empty()
                     || (self.poll_write() || self.framed.is_write_buf_empty())
                 {
                     Err(err)
                 } else {
-                    self.dispatch_state = State::Error(err);
+                    self.dispatch_state = FramedState::Error(err);
                     Ok(Async::NotReady)
                 }
             }
-            State::FlushAndStop(mut vec) => {
+            FramedState::FlushAndStop(mut vec) => {
                 if !self.framed.is_write_buf_empty() {
                     match self.framed.poll_complete() {
                         Err(err) => {
                             debug!("Error sending data: {:?}", err);
                         }
                         Ok(Async::NotReady) => {
-                            self.dispatch_state = State::FlushAndStop(vec);
+                            self.dispatch_state = FramedState::FlushAndStop(vec);
                             return Ok(Async::NotReady);
                         }
                         Ok(Async::Ready(_)) => (),
@@ -295,8 +298,8 @@ where
                 }
                 Ok(Async::Ready(()))
             }
-            State::FramedError(err) => Err(err),
+            FramedState::FramedError(err) => Err(err),
-            State::Stopping => Ok(Async::Ready(())),
+            FramedState::Stopping => Ok(Async::Ready(())),
         }
     }
 }

actix-ioframe/src/item.rs

@@ -1,13 +1,14 @@
+use std::cell::{Ref, RefMut};
 use std::fmt;
 use std::ops::{Deref, DerefMut};
 
 use actix_codec::{Decoder, Encoder};
 
-use crate::cell::Cell;
 use crate::sink::Sink;
+use crate::state::State;
 
-pub struct Item<S, Codec: Encoder + Decoder> {
+pub struct Item<St, Codec: Encoder + Decoder> {
-    state: Cell<S>,
+    state: State<St>,
     sink: Sink<<Codec as Encoder>::Item>,
     item: <Codec as Decoder>::Item,
 }
@@ -17,7 +18,7 @@ where
     Codec: Encoder + Decoder,
 {
     pub(crate) fn new(
-        state: Cell<St>,
+        state: State<St>,
         sink: Sink<<Codec as Encoder>::Item>,
         item: <Codec as Decoder>::Item,
     ) -> Self {
@@ -25,12 +26,12 @@ where
     }
 
     #[inline]
-    pub fn state(&self) -> &St {
+    pub fn state(&self) -> Ref<St> {
         self.state.get_ref()
     }
 
     #[inline]
-    pub fn state_mut(&mut self) -> &mut St {
+    pub fn state_mut(&mut self) -> RefMut<St> {
         self.state.get_mut()
     }
 
@@ -43,6 +44,17 @@ where
     pub fn into_inner(self) -> <Codec as Decoder>::Item {
         self.item
     }
+
+    #[inline]
+    pub fn into_parts(
+        self,
+    ) -> (
+        State<St>,
+        Sink<<Codec as Encoder>::Item>,
+        <Codec as Decoder>::Item,
+    ) {
+        (self.state, self.sink, self.item)
+    }
 }
 
 impl<St, Codec> Deref for Item<St, Codec>

actix-ioframe/src/lib.rs

@@ -5,9 +5,11 @@ mod error;
 mod item;
 mod service;
 mod sink;
+mod state;
 
 pub use self::connect::{Connect, ConnectResult};
 pub use self::error::ServiceError;
 pub use self::item::Item;
 pub use self::service::{Builder, NewServiceBuilder, ServiceBuilder};
 pub use self::sink::Sink;
+pub use self::state::State;

actix-ioframe/src/service.rs

@@ -9,6 +9,7 @@ use crate::connect::{Connect, ConnectResult};
 use crate::dispatcher::FramedDispatcher;
 use crate::error::ServiceError;
 use crate::item::Item;
+use crate::state::State;
 
 type RequestItem<S, U> = Item<S, U>;
 type ResponseItem<U> = Option<<U as Encoder>::Item>;
@@ -295,7 +296,7 @@ where
                 match fut.poll()? {
                     Async::Ready(res) => {
                         self.inner = FramedServiceImplResponseInner::Handler(
-                            handler.new_service(res.state.get_ref()),
+                            handler.new_service(&res.state),
                             Some(res),
                         );
                         self.poll()
@@ -309,7 +310,11 @@ where
                         let res = res.take().unwrap();
                         self.inner =
                             FramedServiceImplResponseInner::Dispatcher(FramedDispatcher::new(
-                                res.framed, res.state, handler, res.rx, res.sink,
+                                res.framed,
+                                State::new(res.state),
+                                handler,
+                                res.rx,
+                                res.sink,
                             ));
                         self.poll()
                     }

actix-ioframe/src/state.rs

@@ -0,0 +1,30 @@
+use std::cell::{Ref, RefCell, RefMut};
+use std::rc::Rc;
+
+/// Connection state
+///
+/// Connection state is an arbitrary data attached to the each incoming message.
+#[derive(Debug)]
+pub struct State<T>(Rc<RefCell<T>>);
+
+impl<T> State<T> {
+    pub(crate) fn new(st: T) -> Self {
+        State(Rc::new(RefCell::new(st)))
+    }
+
+    #[inline]
+    pub fn get_ref(&self) -> Ref<T> {
+        self.0.borrow()
+    }
+
+    #[inline]
+    pub fn get_mut(&mut self) -> RefMut<T> {
+        self.0.borrow_mut()
+    }
+}
+
+impl<T> Clone for State<T> {
+    fn clone(&self) -> Self {
+        State(self.0.clone())
+    }
+}