Add restate future combinators

Author: igalshilman

python/restate/__init__.py

@@ -20,6 +20,8 @@
 from .context import Context, ObjectContext, ObjectSharedContext
 from .context import WorkflowContext, WorkflowSharedContext
 from .context import DurablePromise, RestateDurableFuture, RestateDurableCallFuture, SendHandle
+from .combinators import wait, gather, as_completed, ALL_COMPLETED, FIRST_COMPLETED
+from .exceptions import TerminalError
 
 from .endpoint import app
 
@@ -46,6 +48,12 @@ def test_harness(app, follow_logs = False, restate_image = ""): # type: ignore
     "RestateDurableFuture",
     "RestateDurableCallFuture",
     "SendHandle",
+    "TerminalError",
     "app",
     "test_harness",
+    "wait",
+    "gather",
+    "as_completed",
+    "ALL_COMPLETED",
+    "FIRST_COMPLETED",
 ]

python/restate/combinators.py

@@ -0,0 +1,121 @@
+#
+#  Copyright (c) 2023-2024 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+# pylint: disable=R0913,C0301,R0917
+# pylint: disable=line-too-long
+"""combines multiple futures into a single future"""
+
+from typing import Any, AsyncGenerator, List, Tuple
+from restate.exceptions import TerminalError
+from restate.context import RestateDurableFuture
+from restate.server_context import ServerDurableFuture, ServerInvocationContext
+
+FIRST_COMPLETED = 1
+ALL_COMPLETED = 2
+
+async def wait(*futures: RestateDurableFuture[Any], mode: int = FIRST_COMPLETED) -> Tuple[List[RestateDurableFuture[Any]], List[RestateDurableFuture[Any]]]:
+    """
+    Blocks until at least one of the futures/all of the futures are completed. 
+
+    Returns a tuple of two lists: the first list contains the futures that are completed,
+    the second list contains the futures that are not completed.
+
+    The mode parameter can be either FIRST_COMPLETED or ALL_COMPLETED.
+    Using FIRST_COMPLETED will return as soon as one of the futures is completed.
+    Using ALL_COMPLETED will return only when all futures are completed.
+
+    examples:
+
+    completed, waiting = await wait(f1, f2, f3, mode=FIRST_COMPLETED)
+    for completed_future in completed:
+        # do something with the completed future
+        print(await completed_future)  # prints the result of the future
+
+    or 
+
+    completed, waiting = await wait(f1, f2, f3, mode=ALL_COMPLETED)
+    assert waiting == []
+
+
+    """
+    assert mode in (FIRST_COMPLETED, ALL_COMPLETED)
+
+    remaining = list(futures)
+    while remaining:
+        completed, waiting = await wait_completed(remaining)
+        if mode == FIRST_COMPLETED:
+            return completed, waiting
+        remaining = waiting
+
+    assert mode == ALL_COMPLETED
+    return list(futures), []
+
+async def gather(*futures: RestateDurableFuture[Any]) -> List[RestateDurableFuture[Any]]:
+    """
+    Blocks until all futures are completed.
+
+    Returns a list of all futures.
+    """
+    completed, _ = await wait(*futures, mode=ALL_COMPLETED)
+    return completed
+
+async def as_completed(*futures: RestateDurableFuture[Any]) -> AsyncGenerator[RestateDurableFuture[Any]]:
+    """
+    Returns an iterator that yields the futures as they are completed.
+    
+    example: 
+
+    async for future in as_completed(f1, f2, f3):
+        # do something with the completed future
+        print(await future)  # prints the result of the future
+
+    """
+    remaining = list(futures)
+    while remaining:
+        completed, waiting = await wait_completed(remaining)
+        for f in completed:
+            yield f
+        remaining = waiting
+
+async def wait_completed(futures: List[RestateDurableFuture[Any]]) -> Tuple[List[RestateDurableFuture[Any]], List[RestateDurableFuture[Any]]]:
+    """
+    Blocks until at least one of the futures is completed.
+
+    Returns a tuple of two lists: the first list contains the futures that are completed,
+    the second list contains the futures that are not completed.
+    """
+    if not futures:
+        return [], []
+    handles: List[int] = []
+    context: ServerInvocationContext | None = None
+    for f in futures:
+        if not isinstance(f, ServerDurableFuture):
+            raise TerminalError("All futures must SDK created futures.")
+        if context is None:
+            context = f.context
+        elif context is not f.context:
+            raise TerminalError("All futures must be created by the same SDK context.")
+        if f.is_completed():
+            return [f], []
+        handles.append(f.source_notification_handle)
+
+    assert context is not None
+    await context.create_poll_or_cancel_coroutine(handles)
+    completed = []
+    uncompleted = []
+    for index, handle in enumerate(handles):
+        future = futures[index]
+        if context.vm.is_completed(handle):
+            completed.append(future)
+        else:
+            uncompleted.append(future)
+    return completed, uncompleted
+
+

python/restate/context.py

@@ -33,6 +33,12 @@ class RestateDurableFuture(typing.Generic[T], Awaitable[T]):
     Represents a durable future.
     """
 
+    @abc.abstractmethod
+    def is_completed(self) -> bool:
+        """
+        Returns True if the future is completed, False otherwise.
+        """
+
     @abc.abstractmethod
     def __await__(self):
         pass

python/restate/server_context.py

@@ -36,27 +36,37 @@
 O = TypeVar('O')
 
 
-
 class ServerDurableFuture(RestateDurableFuture[T]):
     """This class implements a durable future API"""
     value: T | None = None
     error: TerminalError | None = None
     state: typing.Literal["pending", "fulfilled", "rejected"] = "pending"
 
-    def __init__(self, handle: int, factory, ctx: "ServerInvocationContext") -> None:
+    def __init__(self, context: "ServerInvocationContext", handle: int, awaitable_factory) -> None:
         super().__init__()
-        self.factory = factory
-        self.handle = handle
-        self.context = ctx
+        self.context = context
+        self.source_notification_handle = handle
+        self.awaitable_factory = awaitable_factory
         self.state = "pending"
 
+
+    def is_completed(self):
+        match self.state:
+            case "pending":
+                return self.context.vm.is_completed(self.source_notification_handle)
+            case "fulfilled":
+                return True
+            case "rejected":
+                return True
+
+
     def __await__(self):
 
         async def await_point():
             match self.state:
                 case "pending":
                     try:
-                        self.value = await self.factory()
+                        self.value = await self.awaitable_factory()
                         self.state = "fulfilled"
                         return self.value
                     except TerminalError as t:
@@ -71,31 +81,28 @@ async def await_point():
 
 
         return await_point().__await__()
-        #task = asyncio.create_task(self.factory())
-        #return task.__await__()
-
 
 class ServerCallDurableFuture(RestateDurableCallFuture[T], ServerDurableFuture[T]):
     """This class implements a durable future but for calls"""
     _invocation_id: typing.Optional[str] = None
 
     def __init__(self,
-                 ctx: "ServerInvocationContext",
+                 context: "ServerInvocationContext",
                  result_handle: int,
                  result_factory,
                  invocation_id_handle: int,
                  invocation_id_factory) -> None:
-        super().__init__(result_handle, result_factory, ctx)
+        super().__init__(context, result_handle, result_factory)
         self.invocation_id_handle = invocation_id_handle
         self.invocation_id_factory = invocation_id_factory
 
+
     async def invocation_id(self) -> str:
         """Get the invocation id."""
         if self._invocation_id is None:
             self._invocation_id  = await self.invocation_id_factory()
         return self._invocation_id
 
-
 class ServerSendHandle(SendHandle):
     """This class implements the send API"""
     _invocation_id: typing.Optional[str]
@@ -118,7 +125,6 @@ async def async_value(n: Callable[[], T]) -> T:
     """convert a simple value to a coroutine."""
     return n()
 
-
 class ServerDurablePromise(DurablePromise):
     """This class implements a durable promise API"""
 
@@ -314,7 +320,7 @@ async def transform():
                 return res
             return serde.deserialize(res)
 
-        return ServerDurableFuture(handle, transform, self)
+        return ServerDurableFuture(self, handle, transform)