Re-register task on every future poll

actix-ioframe/CHANGES.md

@@ -1,5 +1,11 @@
 # Changes
 
+
+## [0.1.1] - 2019-10-14
+
+* Re-register task on every dispatcher poll.
+
+
 ## [0.1.0] - 2019-09-25
 
 * Initial release

actix-ioframe/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "actix-ioframe"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Nikolay Kim <fafhrd91@gmail.com>"]
 description = "Actix framed service"
 keywords = ["network", "framework", "async", "futures"]

actix-ioframe/src/cell.rs

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

actix-ioframe/src/dispatcher.rs

@@ -154,7 +154,6 @@ where
                     };
 
                     let mut cell = self.inner.clone();
-                    unsafe { cell.get_mut().task.register() };
                     tokio_current_thread::spawn(
                         self.service
                             .call(Item::new(self.state.clone(), self.sink.clone(), item))
@@ -275,6 +274,8 @@ where
     type Error = ServiceError<S::Error, U>;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        unsafe { self.inner.get_ref().task.register() };
+
         match mem::replace(&mut self.dispatch_state, FramedState::Processing) {
             FramedState::Processing => {
                 if self.poll_read() || self.poll_write() {