Replace blocking! macro by generic function (#1305)

1 files changed, 27 insertions, 0 deletions

diff --git a/src/tokio_util.rs b/src/tokio_util.rs
index 31a67574e..6d295453f 100644
--- a/src/tokio_util.rs
+++ b/src/tokio_util.rs
@@ -73,3 +73,30 @@ impl Future for Accept {
     }
   }
 }
+
+/// `futures::future::poll_fn` only support `F: FnMut()->Poll<T, E>`
+/// However, we require that `F: FnOnce()->Poll<T, E>`.
+/// Therefore, we created our version of `poll_fn`.
+pub fn poll_fn<T, E, F>(f: F) -> PollFn<F>
+where
+  F: FnOnce() -> Poll<T, E>,
+{
+  PollFn { inner: Some(f) }
+}
+
+pub struct PollFn<F> {
+  inner: Option<F>,
+}
+
+impl<T, E, F> Future for PollFn<F>
+where
+  F: FnOnce() -> Poll<T, E>,
+{
+  type Item = T;
+  type Error = E;
+
+  fn poll(&mut self) -> Poll<T, E> {
+    let f = self.inner.take().expect("Inner fn has been taken.");
+    f()
+  }
+}