Unify cases

Author: igalshilman

python/restate/server_context.py

@@ -225,15 +225,34 @@ class SyncPoint:
 
     def __init__(self):
         self._cond = asyncio.Condition()
+        self._running = set[asyncio.Task]()
 
-    async def wait(self):
+    async def wait(self, snapshot: set[asyncio.Task]):
         """Wait for the sync point."""
+
+        def wait_fn():
+            """
+            Check if any task had finished.
+            """
+            assert self._cond.locked()
+            for task in snapshot:
+                if task not in self._running:
+                    return True
+            return False
+
         async with self._cond:
-            await self._cond.wait()
+            await self._cond.wait_for(wait_fn)
 
-    async def arrive(self):
+    async def enter(self):
         """Arrive at the sync point."""
         async with self._cond:
+            self._running.add(asyncio.current_task())
+            await self._cond.notify_all()
+
+    async def leave(self):
+        """leave the sync point."""
+        async with self._cond:
+            self._running.remove(asyncio.current_task())
             self._cond.notify_all()
 
 class Tasks:
@@ -398,33 +417,31 @@ async def create_poll_or_cancel_coroutine(self, handles: typing.List[int]) -> No
                 return
             if isinstance(do_progress_response, DoProgressCancelSignalReceived):
                 raise TerminalError("cancelled", 409)
-            if isinstance(do_progress_response, DoProgressReadFromInput):
-                await self.receive_and_notify_input()
-                continue
             if isinstance(do_progress_response, DoProgressExecuteRun):
                 fn = self.run_coros_to_execute[do_progress_response.handle]
                 del self.run_coros_to_execute[do_progress_response.handle]
                 assert fn is not None
 
                 async def wrapper(f):
+                    await self.sync_point.enter()
                     await f()
                     await self.take_and_send_output()
-                    await self.sync_point.arrive()
+                    await self.sync_point.leave()
 
                 task = asyncio.create_task(wrapper(fn))
                 self.tasks.add(task)
+
                 continue
-            if isinstance(do_progress_response, DoWaitPendingRun):
-                sync_task = asyncio.create_task(self.sync_point.wait())
+            if isinstance(do_progress_response, (DoWaitPendingRun, DoProgressReadFromInput)):
+                snapshot = self.sync_point._running.copy()
+                sync_task = asyncio.create_task(self.sync_point.wait(snapshot))
                 read_task = asyncio.create_task(self.receive_and_notify_input())
                 self.tasks.add(sync_task)
                 self.tasks.add(read_task)
                 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."""
@@ -520,6 +537,8 @@ async def create_run_coroutine(self,
         except TerminalError as t:
             failure = Failure(code=t.status_code, message=t.message)
             self.vm.propose_run_completion_failure(handle, failure)
+        except asyncio.CancelledError as e:
+            raise e from None
         # pylint: disable=W0718
         except Exception as e:
             if max_attempts is None and max_retry_duration is None: