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 std::marker::PhantomData;
use actix_net::codec::Framed; use actix_net::codec::Framed;
use actix_net::service::{NewService, Service}; use actix_net::service::{NewService, Service};
use futures::future::{ok, Either, FutureResult}; use futures::future::{ok, Either, FutureResult};
use futures::{Async, AsyncSink, Future, Poll, Sink}; 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 h1::{Codec, Message};
use response::Response; use response::Response;
@ -113,20 +113,43 @@ pub struct SendResponse<T>(PhantomData<(T,)>);
impl<T> Default for SendResponse<T> impl<T> Default for SendResponse<T>
where where
T: AsyncWrite, T: AsyncRead + AsyncWrite,
{ {
fn default() -> Self { fn default() -> Self {
SendResponse(PhantomData) 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> impl<T> NewService for SendResponse<T>
where where
T: AsyncWrite, T: AsyncRead + AsyncWrite,
{ {
type Request = (Response, Framed<T, Codec>); type Request = (Response, Framed<T, Codec>);
type Response = Framed<T, Codec>; type Response = Framed<T, Codec>;
type Error = io::Error; type Error = Error;
type InitError = (); type InitError = ();
type Service = SendResponse<T>; type Service = SendResponse<T>;
type Future = FutureResult<Self::Service, Self::InitError>; type Future = FutureResult<Self::Service, Self::InitError>;
@ -138,20 +161,23 @@ where
impl<T> Service for SendResponse<T> impl<T> Service for SendResponse<T>
where where
T: AsyncWrite, T: AsyncRead + AsyncWrite,
{ {
type Request = (Response, Framed<T, Codec>); type Request = (Response, Framed<T, Codec>);
type Response = Framed<T, Codec>; type Response = Framed<T, Codec>;
type Error = io::Error; type Error = Error;
type Future = SendResponseFut<T>; type Future = SendResponseFut<T>;
fn poll_ready(&mut self) -> Poll<(), Self::Error> { fn poll_ready(&mut self) -> Poll<(), Self::Error> {
Ok(Async::Ready(())) 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 { SendResponseFut {
res: Some(Message::Item(res)), res: Some(Message::Item(res)),
body: Some(body),
framed: Some(framed), framed: Some(framed),
} }
} }
@ -159,29 +185,72 @@ where
pub struct SendResponseFut<T> { pub struct SendResponseFut<T> {
res: Option<Message<Response>>, res: Option<Message<Response>>,
body: Option<Body>,
framed: Option<Framed<T, Codec>>, framed: Option<Framed<T, Codec>>,
} }
impl<T> Future for SendResponseFut<T> impl<T> Future for SendResponseFut<T>
where where
T: AsyncWrite, T: AsyncRead + AsyncWrite,
{ {
type Item = Framed<T, Codec>; type Item = Framed<T, Codec>;
type Error = io::Error; type Error = Error;
fn poll(&mut self) -> Poll<Self::Item, Self::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() { if let Some(res) = self.res.take() {
match self.framed.as_mut().unwrap().start_send(res)? { println!("SEND RESP: {:?}", res);
AsyncSink::Ready => (), framed.force_send(res)?;
AsyncSink::NotReady(res) => {
self.res = Some(res);
return Ok(Async::NotReady);
} }
} }
} }
// 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()? { match self.framed.as_mut().unwrap().poll_complete()? {
Async::Ready(_) => Ok(Async::Ready(self.framed.take().unwrap())), Async::Ready(_) => if self.res.is_some() || self.body.is_some() {
Async::NotReady => Ok(Async::NotReady), return self.poll();
},
Async::NotReady => return Ok(Async::NotReady),
} }
Ok(Async::Ready(self.framed.take().unwrap()))
} }
} }