Detect correctly http.disconnect event

Author: igalshilman

python/restate/server.py

@@ -15,7 +15,7 @@
 import traceback
 from restate.discovery import compute_discovery_json
 from restate.endpoint import Endpoint
-from restate.server_context import ServerInvocationContext
+from restate.server_context import ServerInvocationContext, DisconnectedException
 from restate.server_types import Receive, Scope, Send, binary_to_header, header_to_binary
 from restate.vm import VMWrapper
 from restate._internal import PyIdentityVerifier, IdentityVerificationException # pylint: disable=import-error,no-name-in-module
@@ -81,6 +81,7 @@ async def send_health_check(send: Send):
         'more_body': False,
     })
 
+
 async def process_invocation_to_completion(vm: VMWrapper,
                                            handler,
                                            attempt_headers: Dict[str, str],
@@ -128,6 +129,10 @@ async def process_invocation_to_completion(vm: VMWrapper,
     except asyncio.exceptions.CancelledError:
         context.on_attempt_finished()
         raise
+    except DisconnectedException:
+        # The client disconnected before we could send the response
+        context.on_attempt_finished()
+        return
     # pylint: disable=W0718
     except Exception:
         traceback.print_exc()

python/restate/server_context.py

@@ -37,6 +37,15 @@
 I = TypeVar('I')
 O = TypeVar('O')
 
+class DisconnectedException(IOError):
+    """
+    This class implements the disconnected exception.
+    """
+
+    def __init__(self) -> None:
+        super().__init__("disconnected")
+
+
 
 class ServerTeardownEvent(AttemptFinishedEvent):
     """
@@ -202,7 +211,6 @@ def peek(self) -> Awaitable[Any | None]:
         handle = vm.sys_peek_promise(self.name)
         serde = self.serde
         assert serde is not None
-
         return self.server_context.create_future(handle, serde)
 
 
@@ -311,6 +319,17 @@ def on_attempt_finished(self):
         self.request_finished_event.set()
 
 
+    async def receive_and_notify_input(self):
+        """Receive input from the state machine."""
+        chunk = await self.receive()
+        if chunk.get('type') == 'http.request':
+            assert isinstance(chunk['body'], bytes)
+            self.vm.notify_input(chunk['body'])
+        if not chunk.get('more_body', False):
+            self.vm.notify_input_closed()
+        if chunk.get('type') == 'http.disconnect':
+            raise DisconnectedException()
+
     async def take_and_send_output(self):
         """Take output from state machine and send it"""
         output = self.vm.take_output()
@@ -344,12 +363,7 @@ async def create_poll_or_cancel_coroutine(self, handles: typing.List[int]) -> No
             if isinstance(do_progress_response, DoProgressCancelSignalReceived):
                 raise TerminalError("cancelled", 409)
             if isinstance(do_progress_response, DoProgressReadFromInput):
-                chunk = await self.receive()
-                if chunk.get('body', None) is not None:
-                    assert isinstance(chunk['body'], bytes)
-                    self.vm.notify_input(chunk['body'])
-                if not chunk.get('more_body', False):
-                    self.vm.notify_input_closed()
+                await self.receive_and_notify_input()
                 continue
             if isinstance(do_progress_response, DoProgressExecuteRun):
                 fn = self.run_coros_to_execute[do_progress_response.handle]
@@ -364,7 +378,14 @@ async def wrapper(f):
                 asyncio.create_task(wrapper(fn))
                 continue
             if isinstance(do_progress_response, DoWaitPendingRun):
-                await self.sync_point.wait()
+                sync_task = asyncio.create_task(self.sync_point.wait())
+                read_task = asyncio.create_task(self.receive_and_notify_input())
+                done, _ = await asyncio.wait([sync_task, read_task], return_when=asyncio.FIRST_COMPLETED)
+                if read_task in done:
+                    _ = read_task.result() # rethrow any exception
+                if sync_task in done:
+                    continue
+
 
     def _create_fetch_result_coroutine(self, handle: int, serde: Serde[T] | None = None):
         """Create a coroutine that fetches a result from a notification handle."""