Make sure to drain the request stream

src/endpoint/context.rs

@@ -752,7 +752,7 @@ impl ContextInternal {
         let _ = must_lock!(self.inner).vm.sys_end();
     }
 
-    pub(crate) fn consume_to_end(&self) {
+    pub(crate) fn consume_to_end(&mut self) {
         let mut inner_lock = must_lock!(self.inner);
 
         let out = inner_lock.vm.take_output();
@@ -763,6 +763,18 @@ impl ContextInternal {
         }
     }
 
+    pub(crate) fn drain_input(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        let mut inner_lock = must_lock!(self.inner);
+        loop {
+            let result = ready!(inner_lock.read.poll_recv(cx));
+            match result {
+                None => return Poll::Ready(Ok(())),
+                Some(Ok(_)) => continue,
+                Some(Err(err)) => return Poll::Ready(Err(ErrorInner::DrainError { err }.into())),
+            }
+        }
+    }
+
     pub(super) fn fail(&self, e: Error) {
         must_lock!(self.inner).fail(e)
     }

src/endpoint/futures/handler_state_aware.rs

@@ -1,66 +1,111 @@
-use crate::endpoint::{ContextInternal, Error};
+use crate::endpoint::{ContextInternal, Error, ErrorInner};
+use futures::{FutureExt, ready};
 use pin_project_lite::pin_project;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use std::time::Duration;
 use tokio::sync::oneshot;
 use tracing::warn;
 
 pin_project! {
     /// Future that will stop polling when handler is suspended/failed
-    pub struct HandlerStateAwareFuture<F> {
+    pub struct HandlerStateAwareFuture<F> where F: Future {
         #[pin]
-        fut: F,
+        state: HandlerStateAwareFutureState<F>,
         handler_state_rx: oneshot::Receiver<Error>,
         handler_context: ContextInternal,
     }
 }
 
-impl<F> HandlerStateAwareFuture<F> {
+impl<F> HandlerStateAwareFuture<F>
+where
+    F: Future,
+{
     pub fn new(
         handler_context: ContextInternal,
         handler_state_rx: oneshot::Receiver<Error>,
         fut: F,
     ) -> HandlerStateAwareFuture<F> {
         HandlerStateAwareFuture {
-            fut,
+            state: HandlerStateAwareFutureState::Running { fut },
             handler_state_rx,
             handler_context,
         }
     }
 }
 
+pin_project! {
+    #[project = HandlerStateAwareFutureStateProject]
+    enum HandlerStateAwareFutureState<F> where F: Future {
+        Running { #[pin] fut: F },
+        Draining {
+            output: Option<Result<F::Output, Error>>,
+            #[pin] sleep: tokio::time::Sleep
+        },
+    }
+}
+
 impl<F> Future for HandlerStateAwareFuture<F>
 where
     F: Future,
 {
     type Output = Result<F::Output, Error>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = self.project();
+        let mut this = self.project();
 
-        match this.handler_state_rx.try_recv() {
-            Ok(e) => {
-                warn!(
-                    rpc.system = "restate",
-                    rpc.service = %this.handler_context.service_name(),
-                    rpc.method = %this.handler_context.handler_name(),
-                    "Error while processing handler {e:#}"
-                );
-                this.handler_context.consume_to_end();
-                Poll::Ready(Err(e))
-            }
-            Err(oneshot::error::TryRecvError::Empty) => match this.fut.poll(cx) {
-                Poll::Ready(out) => {
-                    this.handler_context.consume_to_end();
-                    Poll::Ready(Ok(out))
+        loop {
+            match this.state.as_mut().project() {
+                HandlerStateAwareFutureStateProject::Running { fut } => {
+                    match this.handler_state_rx.try_recv() {
+                        Ok(e) => {
+                            warn!(
+                                rpc.system = "restate",
+                                rpc.service = %this.handler_context.service_name(),
+                                rpc.method = %this.handler_context.handler_name(),
+                                "Error while processing handler {e:#}"
+                            );
+                            this.handler_context.consume_to_end();
+                            this.state.set(HandlerStateAwareFutureState::Draining {
+                                output: Some(Err(e)),
+                                sleep: tokio::time::sleep(Duration::from_secs(60)),
+                            });
+                        }
+                        Err(oneshot::error::TryRecvError::Empty) => match fut.poll(cx) {
+                            Poll::Ready(output) => {
+                                this.handler_context.consume_to_end();
+                                this.state.set(HandlerStateAwareFutureState::Draining {
+                                    output: Some(Ok(output)),
+                                    sleep: tokio::time::sleep(Duration::from_secs(60)),
+                                });
+                                continue;
+                            }
+                            Poll::Pending => return Poll::Pending,
+                        },
+                        Err(oneshot::error::TryRecvError::Closed) => {
+                            panic!(
+                                "This is unexpected, this future is still being polled although the sender side was dropped. This should not be possible, because the sender is dropped when this future returns Poll:ready()."
+                            )
+                        }
+                    }
                 }
-                Poll::Pending => Poll::Pending,
-            },
-            Err(oneshot::error::TryRecvError::Closed) => {
-                panic!(
-                    "This is unexpected, this future is still being polled although the sender side was dropped. This should not be possible, because the sender is dropped when this future returns Poll:ready()."
-                )
+                HandlerStateAwareFutureStateProject::Draining {
+                    output,
+                    ref mut sleep,
+                } => match this.handler_context.drain_input(cx) {
+                    Poll::Ready(Ok(_)) => {
+                        return Poll::Ready(output.take().expect("Future polled after completion"));
+                    }
+                    Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+                    Poll::Pending => {
+                        ready!(sleep.poll_unpin(cx));
+                        return Poll::Ready(Err(ErrorInner::DrainError {
+                            err: "Timed out draining input stream after 60s".into(),
+                        }
+                        .into()));
+                    }
+                },
             }
         }
     }

src/endpoint/mod.rs

@@ -66,7 +66,8 @@ impl Error {
             | ErrorInner::UnexpectedValueVariantForSyscall { .. }
             | ErrorInner::Deserialization { .. }
             | ErrorInner::Serialization { .. }
-            | ErrorInner::HandlerResult { .. } => 500,
+            | ErrorInner::HandlerResult { .. }
+            | ErrorInner::DrainError { .. } => 500,
             ErrorInner::FieldRequiresMinimumVersion { .. } => 500,
             ErrorInner::BadDiscoveryVersion(_) => 415,
             ErrorInner::Header { .. } | ErrorInner::BadPath { .. } => 400,
@@ -123,6 +124,8 @@ pub(crate) enum ErrorInner {
         #[source]
         err: BoxError,
     },
+    #[error("Error while draining the input stream: {err}")]
+    DrainError { err: BoxError },
 }
 
 impl From<CoreError> for Error {
@@ -589,7 +592,7 @@ async fn handle_invocation(
         let user_code_fut = InterceptErrorFuture::new(ctx.clone(), svc.handle(ctx.clone()));
 
         // Wrap it in handler state aware future
-        HandlerStateAwareFuture::new(ctx.clone(), handler_state_rx, user_code_fut).await
+        HandlerStateAwareFuture::new(ctx, handler_state_rx, user_code_fut).await
     }
     .instrument(span)
     .await