refactor acceptor pipeline add client timeout

src/server/acceptor.rs

@@ -0,0 +1,315 @@
+use std::time::Duration;
+
+use actix_net::server::ServerMessage;
+use actix_net::service::{NewService, Service};
+use futures::future::{err, ok, Either, FutureResult};
+use futures::{Async, Future, Poll};
+use tokio_reactor::Handle;
+use tokio_tcp::TcpStream;
+use tokio_timer::{sleep, Delay};
+
+use super::handler::HttpHandler;
+use super::settings::WorkerSettings;
+use super::IoStream;
+
+/// This trait indicates types that can create acceptor service for http server.
+pub trait AcceptorServiceFactory: Send + Clone + 'static {
+    type Io: IoStream + Send;
+    type NewService: NewService<
+        Request = TcpStream,
+        Response = Self::Io,
+        Error = (),
+        InitError = (),
+    >;
+
+    fn create(&self) -> Self::NewService;
+}
+
+impl<F, T> AcceptorServiceFactory for F
+where
+    F: Fn() -> T + Send + Clone + 'static,
+    T::Response: IoStream + Send,
+    T: NewService<Request = TcpStream, Error = (), InitError = ()>,
+{
+    type Io = T::Response;
+    type NewService = T;
+
+    fn create(&self) -> T {
+        (self)()
+    }
+}
+
+#[derive(Clone)]
+/// Default acceptor service convert `TcpStream` to a `tokio_tcp::TcpStream`
+pub(crate) struct DefaultAcceptor;
+
+impl AcceptorServiceFactory for DefaultAcceptor {
+    type Io = TcpStream;
+    type NewService = DefaultAcceptor;
+
+    fn create(&self) -> Self::NewService {
+        DefaultAcceptor
+    }
+}
+
+impl NewService for DefaultAcceptor {
+    type Request = TcpStream;
+    type Response = TcpStream;
+    type Error = ();
+    type InitError = ();
+    type Service = DefaultAcceptor;
+    type Future = FutureResult<Self::Service, Self::InitError>;
+
+    fn new_service(&self) -> Self::Future {
+        ok(DefaultAcceptor)
+    }
+}
+
+impl Service for DefaultAcceptor {
+    type Request = TcpStream;
+    type Response = TcpStream;
+    type Error = ();
+    type Future = FutureResult<Self::Response, Self::Error>;
+
+    fn poll_ready(&mut self) -> Poll<(), Self::Error> {
+        Ok(Async::Ready(()))
+    }
+
+    fn call(&mut self, req: Self::Request) -> Self::Future {
+        ok(req)
+    }
+}
+
+pub(crate) struct TcpAcceptor<T, H: HttpHandler> {
+    inner: T,
+    settings: WorkerSettings<H>,
+}
+
+impl<T, H> TcpAcceptor<T, H>
+where
+    H: HttpHandler,
+    T: NewService<Request = TcpStream>,
+{
+    pub(crate) fn new(settings: WorkerSettings<H>, inner: T) -> Self {
+        TcpAcceptor { inner, settings }
+    }
+}
+
+impl<T, H> NewService for TcpAcceptor<T, H>
+where
+    H: HttpHandler,
+    T: NewService<Request = TcpStream>,
+{
+    type Request = ServerMessage;
+    type Response = ();
+    type Error = ();
+    type InitError = ();
+    type Service = TcpAcceptorService<T::Service, H>;
+    type Future = TcpAcceptorResponse<T, H>;
+
+    fn new_service(&self) -> Self::Future {
+        TcpAcceptorResponse {
+            fut: self.inner.new_service(),
+            settings: self.settings.clone(),
+        }
+    }
+}
+
+pub(crate) struct TcpAcceptorResponse<T, H>
+where
+    H: HttpHandler,
+    T: NewService<Request = TcpStream>,
+{
+    fut: T::Future,
+    settings: WorkerSettings<H>,
+}
+
+impl<T, H> Future for TcpAcceptorResponse<T, H>
+where
+    H: HttpHandler,
+    T: NewService<Request = TcpStream>,
+{
+    type Item = TcpAcceptorService<T::Service, H>;
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        match self.fut.poll() {
+            Err(_) => Err(()),
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            Ok(Async::Ready(service)) => Ok(Async::Ready(TcpAcceptorService {
+                inner: service,
+                settings: self.settings.clone(),
+            })),
+        }
+    }
+}
+
+pub(crate) struct TcpAcceptorService<T, H: HttpHandler> {
+    inner: T,
+    settings: WorkerSettings<H>,
+}
+
+impl<T, H> Service for TcpAcceptorService<T, H>
+where
+    H: HttpHandler,
+    T: Service<Request = TcpStream>,
+{
+    type Request = ServerMessage;
+    type Response = ();
+    type Error = ();
+    type Future = Either<TcpAcceptorServiceFut<T::Future>, FutureResult<(), ()>>;
+
+    fn poll_ready(&mut self) -> Poll<(), Self::Error> {
+        self.inner.poll_ready().map_err(|_| ())
+    }
+
+    fn call(&mut self, req: Self::Request) -> Self::Future {
+        match req {
+            ServerMessage::Connect(stream) => {
+                let stream =
+                    TcpStream::from_std(stream, &Handle::default()).map_err(|e| {
+                        error!("Can not convert to an async tcp stream: {}", e);
+                    });
+
+                if let Ok(stream) = stream {
+                    Either::A(TcpAcceptorServiceFut {
+                        fut: self.inner.call(stream),
+                    })
+                } else {
+                    Either::B(err(()))
+                }
+            }
+            ServerMessage::Shutdown(timeout) => Either::B(ok(())),
+            ServerMessage::ForceShutdown => {
+                // self.settings.head().traverse::<TcpStream, H>();
+                Either::B(ok(()))
+            }
+        }
+    }
+}
+
+pub(crate) struct TcpAcceptorServiceFut<T: Future> {
+    fut: T,
+}
+
+impl<T> Future for TcpAcceptorServiceFut<T>
+where
+    T: Future,
+{
+    type Item = ();
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        match self.fut.poll() {
+            Err(_) => Err(()),
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            Ok(Async::Ready(_)) => Ok(Async::Ready(())),
+        }
+    }
+}
+
+/// Errors produced by `AcceptorTimeout` service.
+#[derive(Debug)]
+pub enum TimeoutError<T> {
+    /// The inner service error
+    Service(T),
+
+    /// The request did not complete within the specified timeout.
+    Timeout,
+}
+
+/// Acceptor timeout middleware
+///
+/// Applies timeout to request prcoessing.
+pub(crate) struct AcceptorTimeout<T> {
+    inner: T,
+    timeout: usize,
+}
+
+impl<T: NewService> AcceptorTimeout<T> {
+    pub(crate) fn new(timeout: usize, inner: T) -> Self {
+        Self { inner, timeout }
+    }
+}
+
+impl<T: NewService> NewService for AcceptorTimeout<T> {
+    type Request = T::Request;
+    type Response = T::Response;
+    type Error = TimeoutError<T::Error>;
+    type InitError = T::InitError;
+    type Service = AcceptorTimeoutService<T::Service>;
+    type Future = AcceptorTimeoutFut<T>;
+
+    fn new_service(&self) -> Self::Future {
+        AcceptorTimeoutFut {
+            fut: self.inner.new_service(),
+            timeout: self.timeout,
+        }
+    }
+}
+
+#[doc(hidden)]
+pub(crate) struct AcceptorTimeoutFut<T: NewService> {
+    fut: T::Future,
+    timeout: usize,
+}
+
+impl<T: NewService> Future for AcceptorTimeoutFut<T> {
+    type Item = AcceptorTimeoutService<T::Service>;
+    type Error = T::InitError;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        let inner = try_ready!(self.fut.poll());
+        Ok(Async::Ready(AcceptorTimeoutService {
+            inner,
+            timeout: self.timeout as u64,
+        }))
+    }
+}
+
+/// Acceptor timeout service
+///
+/// Applies timeout to request prcoessing.
+pub(crate) struct AcceptorTimeoutService<T> {
+    inner: T,
+    timeout: u64,
+}
+
+impl<T: Service> Service for AcceptorTimeoutService<T> {
+    type Request = T::Request;
+    type Response = T::Response;
+    type Error = TimeoutError<T::Error>;
+    type Future = AcceptorTimeoutResponse<T>;
+
+    fn poll_ready(&mut self) -> Poll<(), Self::Error> {
+        self.inner.poll_ready().map_err(TimeoutError::Service)
+    }
+
+    fn call(&mut self, req: Self::Request) -> Self::Future {
+        AcceptorTimeoutResponse {
+            fut: self.inner.call(req),
+            sleep: sleep(Duration::from_millis(self.timeout)),
+        }
+    }
+}
+
+pub(crate) struct AcceptorTimeoutResponse<T: Service> {
+    fut: T::Future,
+    sleep: Delay,
+}
+impl<T: Service> Future for AcceptorTimeoutResponse<T> {
+    type Item = T::Response;
+    type Error = TimeoutError<T::Error>;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        match self.fut.poll() {
+            Ok(Async::NotReady) => match self.sleep.poll() {
+                Err(_) => Err(TimeoutError::Timeout),
+                Ok(Async::Ready(_)) => Err(TimeoutError::Timeout),
+                Ok(Async::NotReady) => Ok(Async::NotReady),
+            },
+            Ok(Async::Ready(resp)) => Ok(Async::Ready(resp)),
+            Err(err) => Err(TimeoutError::Service(err)),
+        }
+    }
+}