Fix disconnection handling

CHANGES.md

@@ -11,6 +11,8 @@
 
 * Re-use `BinaryBody` for `Frame::Payload`
 
+* Fix disconnection handling.
+
 
 ## 0.1.0 (2017-10-23)
 

Cargo.toml

@@ -49,7 +49,7 @@ tokio-proto = "0.1"
 # h2 = { git = 'https://github.com/carllerche/h2', optional = true }
 
 [dependencies.actix]
-#version = "0.3"
+#version = ">=0.3.1"
 #path = "../actix"
 git = "https://github.com/actix/actix.git"
 default-features = false

src/body.rs

@@ -49,7 +49,7 @@ impl Body {
     }
 
     /// Create body from slice (copy)
-    pub fn from_slice<'a>(s: &'a [u8]) -> Body {
+    pub fn from_slice(s: &[u8]) -> Body {
         Body::Binary(BinaryBody::Bytes(Bytes::from(s)))
     }
 }
@@ -61,19 +61,23 @@ impl<T> From<T> for Body where T: Into<BinaryBody>{
 }
 
 impl BinaryBody {
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
     pub fn len(&self) -> usize {
-        match self {
-            &BinaryBody::Bytes(ref bytes) => bytes.len(),
-            &BinaryBody::Slice(slice) => slice.len(),
-            &BinaryBody::SharedBytes(ref bytes) => bytes.len(),
-            &BinaryBody::ArcSharedBytes(ref bytes) => bytes.len(),
-            &BinaryBody::SharedString(ref s) => s.len(),
-            &BinaryBody::ArcSharedString(ref s) => s.len(),
+        match *self {
+            BinaryBody::Bytes(ref bytes) => bytes.len(),
+            BinaryBody::Slice(slice) => slice.len(),
+            BinaryBody::SharedBytes(ref bytes) => bytes.len(),
+            BinaryBody::ArcSharedBytes(ref bytes) => bytes.len(),
+            BinaryBody::SharedString(ref s) => s.len(),
+            BinaryBody::ArcSharedString(ref s) => s.len(),
         }
     }
 
     /// Create binary body from slice
-    pub fn from_slice<'a>(s: &'a [u8]) -> BinaryBody {
+    pub fn from_slice(s: &[u8]) -> BinaryBody {
         BinaryBody::Bytes(Bytes::from(s))
     }
 }
@@ -164,13 +168,13 @@ impl<'a> From<&'a Arc<String>> for BinaryBody {
 
 impl AsRef<[u8]> for BinaryBody {
     fn as_ref(&self) -> &[u8] {
-        match self {
-            &BinaryBody::Bytes(ref bytes) => bytes.as_ref(),
-            &BinaryBody::Slice(slice) => slice,
-            &BinaryBody::SharedBytes(ref bytes) => bytes.as_ref(),
-            &BinaryBody::ArcSharedBytes(ref bytes) => bytes.as_ref(),
-            &BinaryBody::SharedString(ref s) => s.as_bytes(),
-            &BinaryBody::ArcSharedString(ref s) => s.as_bytes(),
+        match *self {
+            BinaryBody::Bytes(ref bytes) => bytes.as_ref(),
+            BinaryBody::Slice(slice) => slice,
+            BinaryBody::SharedBytes(ref bytes) => bytes.as_ref(),
+            BinaryBody::ArcSharedBytes(ref bytes) => bytes.as_ref(),
+            BinaryBody::SharedString(ref s) => s.as_bytes(),
+            BinaryBody::ArcSharedString(ref s) => s.as_bytes(),
         }
     }
 }

src/context.rs

@@ -10,6 +10,7 @@ use actix::fut::ActorFuture;
 use actix::dev::{AsyncContextApi, ActorAddressCell, ActorItemsCell, ActorWaitCell, SpawnHandle,
                  Envelope, ToEnvelope, RemoteEnvelope};
 
+use task::IoContext;
 use body::BinaryBody;
 use route::{Route, Frame};
 use httpresponse::HttpResponse;
@@ -26,10 +27,20 @@ pub struct HttpContext<A> where A: Actor<Context=HttpContext<A>> + Route,
     stream: VecDeque<Frame>,
     wait: ActorWaitCell<A>,
     app_state: Rc<<A as Route>::State>,
+    disconnected: bool,
 }
 
+impl<A> IoContext for HttpContext<A> where A: Actor<Context=Self> + Route {
 
-impl<A> ActorContext<A> for HttpContext<A> where A: Actor<Context=Self> + Route
+    fn disconnected(&mut self) {
+        self.disconnected = true;
+        if self.state == ActorState::Running {
+            self.state = ActorState::Stopping;
+        }
+    }
+}
+
+impl<A> ActorContext for HttpContext<A> where A: Actor<Context=Self> + Route
 {
     /// Stop actor execution
     fn stop(&mut self) {
@@ -95,6 +106,7 @@ impl<A> HttpContext<A> where A: Actor<Context=Self> + Route {
             wait: ActorWaitCell::default(),
             stream: VecDeque::new(),
             app_state: state,
+            disconnected: false,
         }
     }
 
@@ -124,6 +136,11 @@ impl<A> HttpContext<A> where A: Actor<Context=Self> + Route {
     pub fn write_eof(&mut self) {
         self.stream.push_back(Frame::Payload(None))
     }
+
+    /// Check if connection still open
+    pub fn connected(&self) -> bool {
+        !self.disconnected
+    }
 }
 
 impl<A> HttpContext<A> where A: Actor<Context=Self> + Route {
@@ -157,7 +174,6 @@ impl<A> Stream for HttpContext<A> where A: Actor<Context=Self> + Route
         if self.act.is_none() {
             return Ok(Async::NotReady)
         }
-
         let act: &mut A = unsafe {
             std::mem::transmute(self.act.as_mut().unwrap() as &mut A)
         };

src/resource.rs

@@ -158,7 +158,7 @@ impl<A> Reply<A> where A: Actor + Route
             },
             ReplyItem::Actor(act) => {
                 ctx.set_actor(act);
-                Task::with_stream(ctx)
+                Task::with_context(ctx)
             }
         }
     }

src/server.rs

@@ -171,11 +171,11 @@ pub struct HttpChannel<T: 'static, A: 'static, H: 'static> {
     keepalive_timer: Option<Timeout>,
 }
 
-/*impl<T: 'static, A: 'static> Drop for HttpChannel<T, A> {
+impl<T: 'static, A: 'static, H: 'static> Drop for HttpChannel<T, A, H> {
     fn drop(&mut self) {
         println!("Drop http channel");
     }
-}*/
+}
 
 impl<T, A, H> Actor for HttpChannel<T, A, H>
     where T: AsyncRead + AsyncWrite + 'static,
@@ -205,6 +205,8 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
         }
 
         loop {
+            let mut not_ready = true;
+
             // check in-flight messages
             let mut idx = 0;
             while idx < self.items.len() {
@@ -218,6 +220,7 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
                     match self.items[idx].task.poll_io(&mut self.stream, req)
                     {
                         Ok(Async::Ready(ready)) => {
+                            not_ready = false;
                             let mut item = self.items.pop_front().unwrap();
 
                             // overide keep-alive state
@@ -247,8 +250,10 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
                 } else if !self.items[idx].finished && !self.items[idx].error {
                     match self.items[idx].task.poll() {
                         Ok(Async::NotReady) => (),
-                        Ok(Async::Ready(_)) =>
-                            self.items[idx].finished = true,
+                        Ok(Async::Ready(_)) => {
+                            not_ready = false;
+                            self.items[idx].finished = true;
+                        },
                         Err(_) =>
                             self.items[idx].error = true,
                     }
@@ -267,8 +272,10 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
                 if !self.inactive[idx].finished && !self.inactive[idx].error {
                     match self.inactive[idx].task.poll() {
                         Ok(Async::NotReady) => (),
-                        Ok(Async::Ready(_)) =>
-                            self.inactive[idx].finished = true,
+                        Ok(Async::Ready(_)) => {
+                            not_ready = false;
+                            self.inactive[idx].finished = true
+                        }
                         Err(_) =>
                             self.inactive[idx].error = true,
                     }
@@ -280,6 +287,8 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
             if !self.error && self.items.len() < MAX_PIPELINED_MESSAGES {
                 match self.reader.parse(&mut self.stream) {
                     Ok(Async::Ready((mut req, payload))) => {
+                        not_ready = false;
+
                         // stop keepalive timer
                         self.keepalive_timer.take();
 
@@ -300,6 +309,12 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
                                    finished: false});
                     }
                     Err(err) => {
+                        // notify all tasks
+                        not_ready = false;
+                        for entry in &mut self.items {
+                            entry.task.disconnected()
+                        }
+
                         // kill keepalive
                         self.keepalive = false;
                         self.keepalive_timer.take();
@@ -344,6 +359,10 @@ impl<T, A, H> Future for HttpChannel<T, A, H>
             if self.items.is_empty() && self.inactive.is_empty() && self.error {
                 return Ok(Async::Ready(()))
             }
+
+            if not_ready {
+                return Ok(Async::NotReady)
+            }
         }
     }
 }

src/task.rs

@@ -1,4 +1,4 @@
-use std::{cmp, io};
+use std::{mem, cmp, io};
 use std::rc::Rc;
 use std::fmt::Write;
 use std::collections::VecDeque;
@@ -47,16 +47,27 @@ impl TaskIOState {
     }
 }
 
+enum TaskStream {
+    None,
+    Stream(Box<FrameStream>),
+    Context(Box<IoContext<Item=Frame, Error=io::Error>>),
+}
+
+pub(crate) trait IoContext: Stream<Item=Frame, Error=io::Error> + 'static {
+    fn disconnected(&mut self);
+}
+
 pub struct Task {
     state: TaskRunningState,
     iostate: TaskIOState,
     frames: VecDeque<Frame>,
-    stream: Option<Box<FrameStream>>,
+    stream: TaskStream,
     encoder: Encoder,
     buffer: BytesMut,
     upgrade: bool,
     keepalive: bool,
     prepared: Option<HttpResponse>,
+    disconnected: bool,
     middlewares: Option<Rc<Vec<Box<Middleware>>>>,
 }
 
@@ -71,12 +82,13 @@ impl Task {
             state: TaskRunningState::Running,
             iostate: TaskIOState::Done,
             frames: frames,
-            stream: None,
+            stream: TaskStream::None,
             encoder: Encoder::length(0),
             buffer: BytesMut::new(),
             upgrade: false,
             keepalive: false,
             prepared: None,
+            disconnected: false,
             middlewares: None,
         }
     }
@@ -88,12 +100,30 @@ impl Task {
             state: TaskRunningState::Running,
             iostate: TaskIOState::ReadingMessage,
             frames: VecDeque::new(),
-            stream: Some(Box::new(stream)),
+            stream: TaskStream::Stream(Box::new(stream)),
             encoder: Encoder::length(0),
             buffer: BytesMut::new(),
             upgrade: false,
             keepalive: false,
             prepared: None,
+            disconnected: false,
+            middlewares: None,
+        }
+    }
+
+    pub(crate) fn with_context<C: IoContext>(ctx: C) -> Self
+    {
+        Task {
+            state: TaskRunningState::Running,
+            iostate: TaskIOState::ReadingMessage,
+            frames: VecDeque::new(),
+            stream: TaskStream::Context(Box::new(ctx)),
+            encoder: Encoder::length(0),
+            buffer: BytesMut::new(),
+            upgrade: false,
+            keepalive: false,
+            prepared: None,
+            disconnected: false,
             middlewares: None,
         }
     }
@@ -106,6 +136,15 @@ impl Task {
         self.middlewares = Some(middlewares);
     }
 
+    pub(crate) fn disconnected(&mut self) {
+        let len = self.buffer.len();
+        self.buffer.split_to(len);
+        self.disconnected = true;
+        if let TaskStream::Context(ref mut ctx) = self.stream {
+            ctx.disconnected();
+        }
+    }
+
     fn prepare(&mut self, req: &mut HttpRequest, msg: HttpResponse)
     {
         trace!("Prepare message status={:?}", msg.status);
@@ -252,20 +291,26 @@ impl Task {
                 trace!("IO Frame: {:?}", frame);
                 match frame {
                     Frame::Message(response) => {
-                        self.prepare(req, response);
+                        if !self.disconnected {
+                            self.prepare(req, response);
+                        }
                     }
                     Frame::Payload(Some(chunk)) => {
-                        if self.prepared.is_some() {
-                            // TODO: add warning, write after EOF
-                            self.encoder.encode(&mut self.buffer, chunk.as_ref());
-                        } else {
-                            // might be response for EXCEPT
-                            self.buffer.extend_from_slice(chunk.as_ref())
+                        if !self.disconnected {
+                            if self.prepared.is_some() {
+                                // TODO: add warning, write after EOF
+                                self.encoder.encode(&mut self.buffer, chunk.as_ref());
+                            } else {
+                                // might be response for EXCEPT
+                                self.buffer.extend_from_slice(chunk.as_ref())
+                            }
                         }
                     },
                     Frame::Payload(None) => {
-                        // TODO: add error "not eof""
-                        if !self.encoder.encode(&mut self.buffer, [].as_ref()) {
+                        if !self.disconnected &&
+                            !self.encoder.encode(&mut self.buffer, [].as_ref())
+                        {
+                            // TODO: add error "not eof""
                             debug!("last payload item, but it is not EOF ");
                             return Err(())
                         }
@@ -276,15 +321,17 @@ impl Task {
         }
 
         // write bytes to TcpStream
-        while !self.buffer.is_empty() {
-            match io.write(self.buffer.as_ref()) {
-                Ok(n) => {
-                    self.buffer.split_to(n);
-                },
-                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                    break
+        if !self.disconnected {
+            while !self.buffer.is_empty() {
+                match io.write(self.buffer.as_ref()) {
+                    Ok(n) => {
+                        self.buffer.split_to(n);
+                    },
+                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        break
+                    }
+                    Err(_) => return Err(()),
                 }
-                Err(_) => return Err(()),
             }
         }
 
@@ -295,10 +342,13 @@ impl Task {
             } else if self.state == TaskRunningState::Paused {
                 self.state = TaskRunningState::Running;
             }
+        } else {
+            // at this point we wont get any more Frames
+            self.iostate = TaskIOState::Done;
         }
 
         // response is completed
-        if self.buffer.is_empty() && self.iostate.is_done() {
+        if (self.buffer.is_empty() || self.disconnected) && self.iostate.is_done() {
             // run middlewares
             if let Some(ref mut resp) = self.prepared {
                 if let Some(middlewares) = self.middlewares.take() {
@@ -313,6 +363,46 @@ impl Task {
             Ok(Async::NotReady)
         }
     }
+
+    fn poll_stream<S>(&mut self, stream: &mut S) -> Poll<(), ()>
+        where S: Stream<Item=Frame, Error=io::Error>
+    {
+        loop {
+            match stream.poll() {
+                Ok(Async::Ready(Some(frame))) => {
+                    match frame {
+                        Frame::Message(ref msg) => {
+                            if self.iostate != TaskIOState::ReadingMessage {
+                                error!("Non expected frame {:?}", frame);
+                                return Err(())
+                            }
+                            self.upgrade = msg.upgrade();
+                            if self.upgrade || msg.body().has_body() {
+                                self.iostate = TaskIOState::ReadingPayload;
+                            } else {
+                                self.iostate = TaskIOState::Done;
+                            }
+                        },
+                        Frame::Payload(ref chunk) => {
+                            if chunk.is_none() {
+                                self.iostate = TaskIOState::Done;
+                            } else if self.iostate != TaskIOState::ReadingPayload {
+                                error!("Non expected frame {:?}", self.iostate);
+                                return Err(())
+                            }
+                        },
+                    }
+                    self.frames.push_back(frame)
+                },
+                Ok(Async::Ready(None)) =>
+                    return Ok(Async::Ready(())),
+                Ok(Async::NotReady) =>
+                    return Ok(Async::NotReady),
+                Err(_) =>
+                    return Err(())
+            }
+        }
+    }
 }
 
 impl Future for Task {
@@ -320,45 +410,15 @@ impl Future for Task {
     type Error = ();
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        if let Some(ref mut stream) = self.stream {
-            loop {
-                match stream.poll() {
-                    Ok(Async::Ready(Some(frame))) => {
-                        match frame {
-                            Frame::Message(ref msg) => {
-                                if self.iostate != TaskIOState::ReadingMessage {
-                                    error!("Non expected frame {:?}", frame);
-                                    return Err(())
-                                }
-                                self.upgrade = msg.upgrade();
-                                if self.upgrade || msg.body().has_body() {
-                                    self.iostate = TaskIOState::ReadingPayload;
-                                } else {
-                                    self.iostate = TaskIOState::Done;
-                                }
-                            },
-                            Frame::Payload(ref chunk) => {
-                                if chunk.is_none() {
-                                    self.iostate = TaskIOState::Done;
-                                } else if self.iostate != TaskIOState::ReadingPayload {
-                                    error!("Non expected frame {:?}", self.iostate);
-                                    return Err(())
-                                }
-                            },
-                        }
-                        self.frames.push_back(frame)
-                    },
-                    Ok(Async::Ready(None)) =>
-                        return Ok(Async::Ready(())),
-                    Ok(Async::NotReady) =>
-                        return Ok(Async::NotReady),
-                    Err(_) =>
-                        return Err(())
-                }
-            }
-        } else {
-            Ok(Async::Ready(()))
-        }
+        let mut s = mem::replace(&mut self.stream, TaskStream::None);
+
+        let result = match s {
+            TaskStream::None => Ok(Async::Ready(())),
+            TaskStream::Stream(ref mut stream) => self.poll_stream(stream),
+            TaskStream::Context(ref mut context) => self.poll_stream(context),
+        };
+        self.stream = s;
+        result
     }
 }