mirror of
https://github.com/n0-computer/noq.git
synced 2026-09-19 01:36:08 +00:00
H3: Close body writer on any state
This commit is contained in:
+21
-18
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user