actix/actix-extras
1
0
Fork 0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-28 09:42:40 +01:00
Code Issues Releases Wiki Activity

SendResponse service sends body as well

Browse Source
This commit is contained in:
Nikolay Kim 2018-11-14 18:57:58 -08:00
parent 03ad9a3105
commit 6e7560e287
1 changed files with 89 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

105
src/service.rs
View File

@ -1,13 +1,13 @@
use std::io;
use std::marker::PhantomData;
use actix_net::codec::Framed;
use actix_net::service::{NewService, Service};
use futures::future::{ok, Either, FutureResult};
use futures::{Async, AsyncSink, Future, Poll, Sink};
use tokio_io::AsyncWrite;
use tokio_io::{AsyncRead, AsyncWrite};
use error::ResponseError;
use body::Body;
use error::{Error, ResponseError};
use h1::{Codec, Message};
use response::Response;
@ -113,20 +113,43 @@ pub struct SendResponse<T>(PhantomData<(T,)>);
impl<T> Default for SendResponse<T>
where
T: AsyncWrite,
T: AsyncRead + AsyncWrite,
{
fn default() -> Self {
SendResponse(PhantomData)
}
}
impl<T> SendResponse<T>
where
T: AsyncRead + AsyncWrite,
{
pub fn send(
mut framed: Framed<T, Codec>,
mut res: Response,
) -> impl Future<Item = Framed<T, Codec>, Error = Error> {
// init codec
framed.get_codec_mut().prepare_te(&mut res);
// extract body from response
let body = res.replace_body(Body::Empty);
// write response
SendResponseFut {
res: Some(Message::Item(res)),
body: Some(body),
framed: Some(framed),
}
}
}
impl<T> NewService for SendResponse<T>
where
T: AsyncWrite,
T: AsyncRead + AsyncWrite,
{
type Request = (Response, Framed<T, Codec>);
type Response = Framed<T, Codec>;
type Error = io::Error;
type Error = Error;
type InitError = ();
type Service = SendResponse<T>;
type Future = FutureResult<Self::Service, Self::InitError>;
@ -138,20 +161,23 @@ where
impl<T> Service for SendResponse<T>
where
T: AsyncWrite,
T: AsyncRead + AsyncWrite,
{
type Request = (Response, Framed<T, Codec>);
type Response = Framed<T, Codec>;
type Error = io::Error;
type Error = Error;
type Future = SendResponseFut<T>;
fn poll_ready(&mut self) -> Poll<(), Self::Error> {
Ok(Async::Ready(()))
}
fn call(&mut self, (res, framed): Self::Request) -> Self::Future {
fn call(&mut self, (mut res, mut framed): Self::Request) -> Self::Future {
framed.get_codec_mut().prepare_te(&mut res);
let body = res.replace_body(Body::Empty);
SendResponseFut {
res: Some(Message::Item(res)),
body: Some(body),
framed: Some(framed),
}
}
@ -159,29 +185,72 @@ where
pub struct SendResponseFut<T> {
res: Option<Message<Response>>,
body: Option<Body>,
framed: Option<Framed<T, Codec>>,
}
impl<T> Future for SendResponseFut<T>
where
T: AsyncWrite,
T: AsyncRead + AsyncWrite,
{
type Item = Framed<T, Codec>;
type Error = io::Error;
type Error = Error;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
// send response
if self.res.is_some() {
let framed = self.framed.as_mut().unwrap();
if !framed.is_write_buf_full() {
if let Some(res) = self.res.take() {
match self.framed.as_mut().unwrap().start_send(res)? {
AsyncSink::Ready => (),
AsyncSink::NotReady(res) => {
self.res = Some(res);
return Ok(Async::NotReady);
println!("SEND RESP: {:?}", res);
framed.force_send(res)?;
}
}
}
// send body
if self.res.is_none() && self.body.is_some() {
let framed = self.framed.as_mut().unwrap();
if !framed.is_write_buf_full() {
let body = self.body.take().unwrap();
match body {
Body::Empty => (),
Body::Streaming(mut stream) => loop {
match stream.poll()? {
Async::Ready(item) => {
let done = item.is_none();
framed.force_send(Message::Chunk(item.into()))?;
if !done {
if !framed.is_write_buf_full() {
continue;
} else {
self.body = Some(Body::Streaming(stream));
break;
}
}
}
Async::NotReady => {
self.body = Some(Body::Streaming(stream));
break;
}
}
},
Body::Binary(mut bin) => {
framed.force_send(Message::Chunk(Some(bin.take())))?;
framed.force_send(Message::Chunk(None))?;
}
}
}
}
// flush
match self.framed.as_mut().unwrap().poll_complete()? {
Async::Ready(_) => Ok(Async::Ready(self.framed.take().unwrap())),
Async::NotReady => Ok(Async::NotReady),
Async::Ready(_) => if self.res.is_some() || self.body.is_some() {
return self.poll();
},
Async::NotReady => return Ok(Async::NotReady),
}
Ok(Async::Ready(self.framed.take().unwrap()))
}
}