//! Framed dispatcher service and related utilities #![allow(type_alias_bounds)] use std::collections::VecDeque; use std::pin::Pin; use std::task::{Context, Poll}; use std::{fmt, mem}; use actix_codec::{AsyncRead, AsyncWrite, Decoder, Encoder, Framed}; use actix_service::{IntoService, Service}; use futures::future::{ready, FutureExt}; use futures::{Future, Sink, Stream}; use log::debug; use crate::cell::Cell; use crate::mpsc; use crate::task::LocalWaker; type Request = ::Item; type Response = ::Item; /// Framed transport errors pub enum FramedTransportError { Service(E), Encoder(::Error), Decoder(::Error), } impl From for FramedTransportError { fn from(err: E) -> Self { FramedTransportError::Service(err) } } impl fmt::Debug for FramedTransportError where E: fmt::Debug, ::Error: fmt::Debug, ::Error: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { match *self { FramedTransportError::Service(ref e) => { write!(fmt, "FramedTransportError::Service({:?})", e) } FramedTransportError::Encoder(ref e) => { write!(fmt, "FramedTransportError::Encoder({:?})", e) } FramedTransportError::Decoder(ref e) => { write!(fmt, "FramedTransportError::Encoder({:?})", e) } } } } impl fmt::Display for FramedTransportError where E: fmt::Display, ::Error: fmt::Debug, ::Error: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { match *self { FramedTransportError::Service(ref e) => write!(fmt, "{}", e), FramedTransportError::Encoder(ref e) => write!(fmt, "{:?}", e), FramedTransportError::Decoder(ref e) => write!(fmt, "{:?}", e), } } } pub enum FramedMessage { Message(T), Close, } type Rx = Option::Item>>>; type Inner = Cell::Item, S::Error>>; /// FramedTransport - is a future that reads frames from Framed object /// and pass then to the service. #[pin_project::pin_project] pub struct FramedTransport where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Encoder + Decoder, ::Item: 'static, ::Error: std::fmt::Debug, { service: S, state: TransportState, framed: Framed, rx: Option::Item>>>, inner: Cell::Item, S::Error>>, } enum TransportState { Processing, Error(FramedTransportError), FramedError(FramedTransportError), FlushAndStop, Stopping, } struct FramedTransportInner { buf: VecDeque>, task: LocalWaker, } impl FramedTransport where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Decoder + Encoder, ::Item: 'static, ::Error: std::fmt::Debug, { pub fn new>(framed: Framed, service: F) -> Self { FramedTransport { framed, rx: None, service: service.into_service(), state: TransportState::Processing, inner: Cell::new(FramedTransportInner { buf: VecDeque::new(), task: LocalWaker::new(), }), } } /// Get Sender pub fn set_receiver( mut self, rx: mpsc::Receiver::Item>>, ) -> Self { self.rx = Some(rx); self } /// Get reference to a service wrapped by `FramedTransport` instance. pub fn get_ref(&self) -> &S { &self.service } /// Get mutable reference to a service wrapped by `FramedTransport` /// instance. pub fn get_mut(&mut self) -> &mut S { &mut self.service } /// Get reference to a framed instance wrapped by `FramedTransport` /// instance. pub fn get_framed(&self) -> &Framed { &self.framed } /// Get mutable reference to a framed instance wrapped by `FramedTransport` /// instance. pub fn get_framed_mut(&mut self) -> &mut Framed { &mut self.framed } } impl Future for FramedTransport where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Decoder + Encoder, ::Item: 'static, ::Error: std::fmt::Debug, ::Error: std::fmt::Debug, { type Output = Result<(), FramedTransportError>; fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { self.inner.get_ref().task.register(cx.waker()); let this = self.project(); poll( cx, this.service, this.state, this.framed, this.rx, this.inner, ) } } fn poll( cx: &mut Context, srv: &mut S, state: &mut TransportState, framed: &mut Framed, rx: &mut Rx, inner: &mut Inner, ) -> Poll>> where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Decoder + Encoder, ::Item: 'static, ::Error: std::fmt::Debug, { match mem::replace(state, TransportState::Processing) { TransportState::Processing => { if poll_read(cx, srv, state, framed, inner) || poll_write(cx, state, framed, rx, inner) { poll(cx, srv, state, framed, rx, inner) } else { Poll::Pending } } TransportState::Error(err) => { let is_empty = framed.is_write_buf_empty(); if is_empty || poll_write(cx, state, framed, rx, inner) { Poll::Ready(Err(err)) } else { *state = TransportState::Error(err); Poll::Pending } } TransportState::FlushAndStop => { if !framed.is_write_buf_empty() { match Pin::new(framed).poll_flush(cx) { Poll::Ready(Err(err)) => { debug!("Error sending data: {:?}", err); Poll::Ready(Ok(())) } Poll::Pending => Poll::Pending, Poll::Ready(Ok(_)) => Poll::Ready(Ok(())), } } else { Poll::Ready(Ok(())) } } TransportState::FramedError(err) => Poll::Ready(Err(err)), TransportState::Stopping => Poll::Ready(Ok(())), } } fn poll_read( cx: &mut Context, srv: &mut S, state: &mut TransportState, framed: &mut Framed, inner: &mut Inner, ) -> bool where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Decoder + Encoder, ::Item: 'static, ::Error: std::fmt::Debug, { loop { match srv.poll_ready(cx) { Poll::Ready(Ok(_)) => { let item = match framed.next_item(cx) { Poll::Ready(Some(Ok(el))) => el, Poll::Ready(Some(Err(err))) => { *state = TransportState::FramedError(FramedTransportError::Decoder(err)); return true; } Poll::Pending => return false, Poll::Ready(None) => { *state = TransportState::Stopping; return true; } }; let mut cell = inner.clone(); let fut = srv.call(item).then(move |item| { let inner = cell.get_mut(); inner.buf.push_back(item); inner.task.wake(); ready(()) }); tokio_executor::current_thread::spawn(fut); } Poll::Pending => return false, Poll::Ready(Err(err)) => { *state = TransportState::Error(FramedTransportError::Service(err)); return true; } } } } /// write to framed object fn poll_write( cx: &mut Context, state: &mut TransportState, framed: &mut Framed, rx: &mut Rx, inner: &mut Inner, ) -> bool where S: Service, Response = Response>, S::Error: 'static, S::Future: 'static, T: AsyncRead + AsyncWrite, U: Decoder + Encoder, ::Item: 'static, ::Error: std::fmt::Debug, { // let this = self.project(); let inner = inner.get_mut(); let mut rx_done = rx.is_none(); let mut buf_empty = inner.buf.is_empty(); loop { while !framed.is_write_buf_full() { if !buf_empty { match inner.buf.pop_front().unwrap() { Ok(msg) => { if let Err(err) = framed.write(msg) { *state = TransportState::FramedError(FramedTransportError::Encoder(err)); return true; } buf_empty = inner.buf.is_empty(); } Err(err) => { *state = TransportState::Error(FramedTransportError::Service(err)); return true; } } } if !rx_done && rx.is_some() { match Pin::new(rx.as_mut().unwrap()).poll_next(cx) { Poll::Ready(Some(FramedMessage::Message(msg))) => { if let Err(err) = framed.write(msg) { *state = TransportState::FramedError(FramedTransportError::Encoder(err)); return true; } } Poll::Ready(Some(FramedMessage::Close)) => { *state = TransportState::FlushAndStop; return true; } Poll::Ready(None) => { rx_done = true; let _ = rx.take(); } Poll::Pending => rx_done = true, } } if rx_done && buf_empty { break; } } if !framed.is_write_buf_empty() { match framed.flush(cx) { Poll::Pending => break, Poll::Ready(Err(err)) => { debug!("Error sending data: {:?}", err); *state = TransportState::FramedError(FramedTransportError::Encoder(err)); return true; } Poll::Ready(Ok(_)) => (), } } else { break; } } false }