From 2e8c2c77337ac3f7a1a0fd2d3c8636a163c8607a Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Mon, 14 Oct 2019 17:55:52 +0600
Subject: [PATCH] Re-register task on every future poll

---
 actix-ioframe/CHANGES.md        | 6 ++++++
 actix-ioframe/Cargo.toml        | 2 +-
 actix-ioframe/src/cell.rs       | 4 ++++
 actix-ioframe/src/dispatcher.rs | 3 ++-
 actix-utils/CHANGES.md          | 5 +++++
 actix-utils/Cargo.toml          | 2 +-
 actix-utils/src/framed.rs       | 3 ++-
 7 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/actix-ioframe/CHANGES.md b/actix-ioframe/CHANGES.md
index 90905f5f..a7aa244f 100644
--- a/actix-ioframe/CHANGES.md
+++ b/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
diff --git a/actix-ioframe/Cargo.toml b/actix-ioframe/Cargo.toml
index 53639e36..97e32152 100644
--- a/actix-ioframe/Cargo.toml
+++ b/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"]
diff --git a/actix-ioframe/src/cell.rs b/actix-ioframe/src/cell.rs
index e517aca1..3e82aee4 100644
--- a/actix-ioframe/src/cell.rs
+++ b/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()
     }
diff --git a/actix-ioframe/src/dispatcher.rs b/actix-ioframe/src/dispatcher.rs
index 57713b10..05a7d05f 100644
--- a/actix-ioframe/src/dispatcher.rs
+++ b/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() {
diff --git a/actix-utils/CHANGES.md b/actix-utils/CHANGES.md
index d2db8811..4e4ac8d7 100644
--- a/actix-utils/CHANGES.md
+++ b/actix-utils/CHANGES.md
@@ -1,5 +1,10 @@
 # Changes
 
+## [0.4.7] - 2019-10-14
+
+* Re-register task on every framed transport poll.
+
+
 ## [0.4.6] - 2019-10-08
 
 * Refactor `Counter` type. register current task in available method.
diff --git a/actix-utils/Cargo.toml b/actix-utils/Cargo.toml
index b687d030..ee3c1dba 100644
--- a/actix-utils/Cargo.toml
+++ b/actix-utils/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "actix-utils"
-version = "0.4.6"
+version = "0.4.7"
 authors = ["Nikolay Kim <fafhrd91@gmail.com>"]
 description = "Actix utils - various actix net related services"
 keywords = ["network", "framework", "async", "futures"]
diff --git a/actix-utils/src/framed.rs b/actix-utils/src/framed.rs
index 7be2c750..2fa2d8e9 100644
--- a/actix-utils/src/framed.rs
+++ b/actix-utils/src/framed.rs
@@ -129,7 +129,6 @@ where
                     };
 
                     let mut cell = self.inner.clone();
-                    cell.get_mut().task.register();
                     tokio_current_thread::spawn(self.service.call(item).then(move |item| {
                         let inner = cell.get_mut();
                         inner.buf.push_back(item);
@@ -293,6 +292,8 @@ where
     type Error = FramedTransportError<S::Error, U>;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        self.inner.get_ref().task.register();
+
         match mem::replace(&mut self.state, TransportState::Processing) {
             TransportState::Processing => {
                 if self.poll_read() || self.poll_write() {