From c0e73c727561879e6ace130b356db92b2af1a572 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Wed, 11 Oct 2017 19:20:05 -0700
Subject: [PATCH] provide wait method

---
 src/context.rs | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 0f726312..2941d71c 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -4,7 +4,7 @@ use std::collections::VecDeque;
 use futures::{Async, Stream, Poll};
 
 use bytes::Bytes;
-use actix::{Actor, ActorState, ActorContext, AsyncActorContext};
+use actix::{Actor, ActorState, ActorContext, AsyncContext};
 use actix::fut::ActorFuture;
 use actix::dev::{AsyncContextApi, ActorAddressCell, ActorItemsCell, SpawnHandle};
 
@@ -20,6 +20,7 @@ pub struct HttpContext<A> where A: Actor<Context=HttpContext<A>> + Route,
     items: ActorItemsCell<A>,
     address: ActorAddressCell<A>,
     stream: VecDeque<Frame>,
+    wait: Option<Box<ActorFuture<Item=(), Error=(), Actor=A>>>,
     app_state: Rc<<A as Route>::State>,
 }
 
@@ -47,7 +48,7 @@ impl<A> ActorContext<A> for HttpContext<A> where A: Actor<Context=Self> + Route
     }
 }
 
-impl<A> AsyncActorContext<A> for HttpContext<A> where A: Actor<Context=Self> + Route
+impl<A> AsyncContext<A> for HttpContext<A> where A: Actor<Context=Self> + Route
 {
     fn spawn<F>(&mut self, fut: F) -> SpawnHandle
         where F: ActorFuture<Item=(), Error=(), Actor=A> + 'static
@@ -55,6 +56,12 @@ impl<A> AsyncActorContext<A> for HttpContext<A> where A: Actor<Context=Self> + R
         self.items.spawn(fut)
     }
 
+    fn wait<F>(&mut self, fut: F)
+        where F: ActorFuture<Item=(), Error=(), Actor=A> + 'static
+    {
+        self.wait = Some(Box::new(fut));
+    }
+
     fn cancel_future(&mut self, handle: SpawnHandle) -> bool {
         self.items.cancel_future(handle)
     }
@@ -77,6 +84,7 @@ impl<A> HttpContext<A> where A: Actor<Context=Self> + Route {
             items: ActorItemsCell::default(),
             address: ActorAddressCell::default(),
             stream: VecDeque::new(),
+            wait: None,
             app_state: state,
         }
     }
@@ -139,6 +147,19 @@ impl<A> Stream for HttpContext<A> where A: Actor<Context=Self> + Route
             _ => ()
         }
 
+        // check wait future
+        if self.wait.is_some() && self.act.is_some() {
+            if let Some(ref mut act) = self.act {
+                if let Some(ref mut fut) = self.wait {
+                    match fut.poll(act, ctx) {
+                        Ok(Async::NotReady) => return Ok(Async::NotReady),
+                        _ => (),
+                    }
+                }
+            }
+            self.wait = None;
+        }
+
         let mut prep_stop = false;
         loop {
             let mut not_ready = true;