lx/boitalettres
1
0
Fork 0
forked from KokaKiwi/boitalettres
Code Issues Pull requests Projects Releases Wiki Activity

Support closing connection after a Response

Browse source 
Closes #11
This commit is contained in:
Jill 2022-07-25 13:12:39 +02:00
parent c8e0aefc42
commit ce339cf89b
No known key found for this signature in database
GPG key ID: 09A5A2688F13FAC1
4 changed files with 26 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/proto/res/body.rs
View file

@ -70,6 +70,7 @@ impl fmt::Debug for Body {
pub enum Data { pub enum Data {
Data(ImapData), Data(ImapData),
Status(ImapStatus), Status(ImapStatus),
Close,
} }
impl Encode for Data { impl Encode for Data {
@ -77,6 +78,7 @@ impl Encode for Data {
match self { match self {
Data::Data(ref data) => data.encode(writer), Data::Data(ref data) => data.encode(writer),
Data::Status(ref status) => status.encode(writer), Data::Status(ref status) => status.encode(writer),
Data::Close => Ok(()),
} }
} }
} }

7
src/proto/res/mod.rs
View file

@ -13,6 +13,7 @@ pub(crate) mod stream;
pub struct Response { pub struct Response {
pub(crate) status: Status, pub(crate) status: Status,
pub(crate) body: Option<Body>, pub(crate) body: Option<Body>,
pub(crate) close: bool,
} }
impl Response { impl Response {
@ -20,6 +21,7 @@ impl Response {
Ok(Response { Ok(Response {
status: Status::new(code, msg)?, status: Status::new(code, msg)?,
body: None, body: None,
close: false,
}) })
} }
@ -50,6 +52,11 @@ impl Response {
self.body = Some(body.into()); self.body = Some(body.into());
self self
} }
pub fn close(mut self, close: bool) -> Self {
self.close = close;
self
}
} }
impl Response { impl Response {

5
src/proto/res/stream.rs
View file

@ -5,6 +5,7 @@ use super::body::Data;
use super::Response; use super::Response;
pub fn response_stream(res: Response, tag: Option<Tag>) -> impl Stream<Item = Data> { pub fn response_stream(res: Response, tag: Option<Tag>) -> impl Stream<Item = Data> {
let close = res.close;
let (body, status) = res.split(); let (body, status) = res.split();
let body = body.map(|body| body.into_stream()); let body = body.map(|body| body.into_stream());
@ -19,5 +20,9 @@ pub fn response_stream(res: Response, tag: Option<Tag>) -> impl Stream<Item = Da
let item = status.into_imap(tag); let item = status.into_imap(tag);
tracing::trace!(?item, "response_stream.yield"); tracing::trace!(?item, "response_stream.yield");
yield item.into(); yield item.into();
if close {
yield Data::Close;
}
} }
} }

12
src/server/pipeline.rs
View file

@ -23,6 +23,8 @@ pub struct Connection<C> {
#[pin] #[pin]
outbox: ConcatAll<BoxStream<'static, Data>>, outbox: ConcatAll<BoxStream<'static, Data>>,
write_buf: BytesMut, write_buf: BytesMut,
close: bool,
} }
impl<C> Connection<C> { impl<C> Connection<C> {
@ -34,6 +36,8 @@ impl<C> Connection<C> {
outbox: ConcatAll::new(), outbox: ConcatAll::new(),
write_buf: BytesMut::new(), write_buf: BytesMut::new(),
close: false,
} }
} }
} }
@ -105,6 +109,10 @@ where
while let Poll::Ready(Some(data)) = this.outbox.as_mut().poll_next(cx) { while let Poll::Ready(Some(data)) = this.outbox.as_mut().poll_next(cx) {
tracing::trace!(?data, "transport.write_buf"); tracing::trace!(?data, "transport.write_buf");
if let Data::Close = data {
*this.close = true;
}
if let Err(err) = data.encode(&mut writer) { if let Err(err) = data.encode(&mut writer) {
tracing::error!(?err, "transport.encode_error"); tracing::error!(?err, "transport.encode_error");
return Poll::Ready(Err(Box::new(err))); return Poll::Ready(Err(Box::new(err)));
@ -118,6 +126,10 @@ where
} }
this.write_buf.clear(); this.write_buf.clear();
if *this.close {
futures::ready!(this.conn.as_mut().poll_close(cx))?;
}
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
} }
} }