H3: Close body writer on any state

This commit is contained in:
stammw
2020-01-17 08:14:52 +01:00
committed by Dirkjan Ochtman
parent 384549faf2
commit ff4d736c06
+21 -18
View File
@@ -262,11 +262,9 @@ impl BodyWriter {
}
pub async fn close(mut self) -> Result<(), Error> {
let state = mem::replace(&mut self.state, BodyWriterState::Finished);
match state {
BodyWriterState::Idle(mut send) => send.finish().await.map_err(Into::into),
_ => panic!("cannot close while not in idle state"),
}
futures_util::future::poll_fn(|cx| self.poll_close_inner(cx))
.await
.map_err(Into::into)
}
pub fn cancel(mut self) {
@@ -281,6 +279,23 @@ impl BodyWriter {
_ => (),
}
}
fn poll_close_inner(&mut self, cx: &mut Context) -> Poll<Result<(), io::Error>> {
loop {
match self.state {
BodyWriterState::Finished => return Poll::Ready(Ok(())),
BodyWriterState::Idle(ref mut send) => {
ready!(Pin::new(send).poll_close(cx))?;
self.state = BodyWriterState::Finished;
return Poll::Ready(Ok(()));
}
BodyWriterState::Writing(ref mut write) => {
let send = ready!(Pin::new(write).poll(cx))?;
self.state = BodyWriterState::Idle(send);
}
}
}
}
}
enum BodyWriterState {
@@ -335,19 +350,7 @@ impl AsyncWrite for BodyWriter {
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), io::Error>> {
match self.state {
BodyWriterState::Finished => Poll::Ready(Ok(())),
BodyWriterState::Idle(ref mut send) => {
ready!(Pin::new(send).poll_close(cx))?;
self.state = BodyWriterState::Finished;
Poll::Ready(Ok(()))
}
BodyWriterState::Writing(ref mut write) => {
let send = ready!(Pin::new(write).poll(cx))?;
self.state = BodyWriterState::Idle(send);
Poll::Pending
}
}
self.poll_close_inner(cx)
}
}