Add race()

Author: chriso

src/dispatch/__init__.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import dispatch.integrations
-from dispatch.coroutine import all, call, gather
+from dispatch.coroutine import all, any, call, gather, race
 from dispatch.function import DEFAULT_API_URL, Client
 from dispatch.id import DispatchID
 from dispatch.proto import Call, Error, Input, Output
@@ -21,4 +21,6 @@
     "call",
     "gather",
     "all",
+    "any",
+    "race",
 ]

src/dispatch/coroutine.py

@@ -25,7 +25,7 @@ def gather(*awaitables: Awaitable[Any]) -> list[Any]:  # type: ignore[misc]
 @durable
 def all(*awaitables: Awaitable[Any]) -> list[Any]:  # type: ignore[misc]
     """Concurrently run a set of coroutines, blocking until all coroutines
-    return or any coroutine raises an error. If a coroutine fails with an
+    return or any coroutine raises an error. If any coroutine fails with an
     uncaught exception, the exception will be re-raised here."""
     return (yield AllDirective(awaitables))
 
@@ -35,10 +35,19 @@ def all(*awaitables: Awaitable[Any]) -> list[Any]:  # type: ignore[misc]
 def any(*awaitables: Awaitable[Any]) -> list[Any]:  # type: ignore[misc]
     """Concurrently run a set of coroutines, blocking until any coroutine
     returns or all coroutines raises an error. If all coroutines fail with
-    uncaught exceptions, an AnyException will be re-raised here."""
+    uncaught exceptions, the exception(s) will be re-raised here."""
     return (yield AnyDirective(awaitables))
 
 
+@coroutine
+@durable
+def race(*awaitables: Awaitable[Any]) -> list[Any]:  # type: ignore[misc]
+    """Concurrently run a set of coroutines, blocking until any coroutine
+    returns or raises an error. If any coroutine fails with an uncaught
+    exception, the exception will be re-raised here."""
+    return (yield RaceDirective(awaitables))
+
+
 @dataclass(slots=True)
 class AllDirective:
     awaitables: tuple[Awaitable[Any], ...]
@@ -49,6 +58,11 @@ class AnyDirective:
     awaitables: tuple[Awaitable[Any], ...]
 
 
+@dataclass(slots=True)
+class RaceDirective:
+    awaitables: tuple[Awaitable[Any], ...]
+
+
 class AnyException(RuntimeError):
     """Error indicating that all coroutines passed to any() failed
     with an exception."""

src/dispatch/scheduler.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Awaitable, Callable, Protocol, TypeAlias
 
-from dispatch.coroutine import AllDirective, AnyDirective, AnyException
+from dispatch.coroutine import AllDirective, AnyDirective, AnyException, RaceDirective
 from dispatch.error import IncompatibleStateError
 from dispatch.experimental.durable.function import DurableCoroutine, DurableGenerator
 from dispatch.proto import Call, Error, Input, Output
@@ -178,6 +178,46 @@ def value(self) -> Any:
         return self.first_result.value
 
 
+@dataclass(slots=True)
+class RaceFuture:
+    """A future result of a dispatch.coroutine.race() operation."""
+
+    waiting: set[CoroutineID] = field(default_factory=set)
+    first_result: CoroutineResult | None = None
+    first_error: Exception | None = None
+
+    def add_result(self, result: CallResult | CoroutineResult):
+        assert isinstance(result, CoroutineResult)
+
+        if result.error is not None:
+            if self.first_error is None:
+                self.first_error = result.error
+        else:
+            if self.first_result is None:
+                self.first_result = result
+
+        self.waiting.remove(result.coroutine_id)
+
+    def add_error(self, error: Exception):
+        if self.first_error is None:
+            self.first_error = error
+
+    def ready(self) -> bool:
+        return (
+            self.first_error is not None
+            or self.first_result is not None
+            or len(self.waiting) == 0
+        )
+
+    def error(self) -> Exception | None:
+        assert self.ready()
+        return self.first_error
+
+    def value(self) -> Any:
+        assert self.first_error is None
+        return self.first_result.value if self.first_result else None
+
+
 @dataclass(slots=True)
 class Coroutine:
     """An in-flight coroutine."""
@@ -470,6 +510,16 @@ def _run(self, input: Input) -> Output:
                     )
                     state.suspended[coroutine.id] = coroutine
 
+                case RaceDirective():
+                    children = spawn_children(
+                        state, coroutine, coroutine_yield.awaitables
+                    )
+
+                    coroutine.result = RaceFuture(
+                        waiting={child.id for child in children}
+                    )
+                    state.suspended[coroutine.id] = coroutine
+
                 case _:
                     raise RuntimeError(
                         f"coroutine unexpectedly yielded '{coroutine_yield}'"

tests/dispatch/test_scheduler.py

@@ -1,11 +1,17 @@
 import unittest
 from typing import Any, Callable
 
-from dispatch.coroutine import AnyException, any, call, gather
+from dispatch.coroutine import AnyException, any, call, gather, race
 from dispatch.experimental.durable import durable
 from dispatch.proto import Call, CallResult, Error, Input, Output
 from dispatch.proto import _any_unpickle as any_unpickle
-from dispatch.scheduler import AllFuture, AnyFuture, CoroutineResult, OneShotScheduler
+from dispatch.scheduler import (
+    AllFuture,
+    AnyFuture,
+    CoroutineResult,
+    OneShotScheduler,
+    RaceFuture,
+)
 from dispatch.sdk.v1 import call_pb2 as call_pb
 from dispatch.sdk.v1 import exit_pb2 as exit_pb
 from dispatch.sdk.v1 import poll_pb2 as poll_pb
@@ -21,6 +27,11 @@ async def call_any(*functions):
     return await any(*[call_one(function) for function in functions])
 
 
+@durable
+async def call_race(*functions):
+    return await race(*[call_one(function) for function in functions])
+
+
 @durable
 async def call_concurrently(*functions):
     return await gather(*[call_one(function) for function in functions])
@@ -201,6 +212,37 @@ async def main():
             output, AnyException, "4 coroutine(s) failed with an exception"
         )
 
+    def test_resume_after_race_result(self):
+        @durable
+        async def main():
+            return await call_race("a", "b", "c", "d")
+
+        output = self.start(main)
+        calls = self.assert_poll_call_functions(output, ["a", "b", "c", "d"])
+
+        output = self.resume(
+            main,
+            output,
+            [CallResult.from_value(23, correlation_id=calls[1].correlation_id)],
+        )
+        self.assert_exit_result_value(output, 23)
+
+    def test_resume_after_race_error(self):
+        @durable
+        async def main():
+            return await call_race("a", "b", "c", "d")
+
+        output = self.start(main)
+        calls = self.assert_poll_call_functions(output, ["a", "b", "c", "d"])
+
+        error = Error.from_exception(RuntimeError("oops"))
+        output = self.resume(
+            main,
+            output,
+            [CallResult.from_error(error, correlation_id=calls[2].correlation_id)],
+        )
+        self.assert_exit_result_error(output, RuntimeError, "oops")
+
     def test_dag(self):
         @durable
         async def main():
@@ -600,3 +642,83 @@ def test_two_result_errors(self):
 
         with self.assertRaises(AssertionError):
             future.value()
+
+
+class TestRaceFuture(unittest.TestCase):
+    def test_empty(self):
+        future = RaceFuture()
+
+        self.assertTrue(future.ready())
+        self.assertIsNone(future.value())
+        self.assertIsNone(future.error())
+
+    def test_one_result_value(self):
+        future = RaceFuture(waiting={10})
+
+        self.assertFalse(future.ready())
+        future.add_result(CoroutineResult(coroutine_id=10, value="foobar"))
+
+        self.assertTrue(future.ready())
+        self.assertIsNone(future.error())
+        self.assertEqual(future.value(), "foobar")
+
+    def test_one_result_error(self):
+        future = RaceFuture(waiting={10})
+
+        self.assertFalse(future.ready())
+        error = RuntimeError("oops")
+        future.add_result(CoroutineResult(coroutine_id=10, error=error))
+
+        self.assertTrue(future.ready())
+        self.assertIs(future.error(), error)
+
+        with self.assertRaises(AssertionError):
+            future.value()
+
+    def test_one_generic_error(self):
+        future = RaceFuture(waiting={10})
+
+        self.assertFalse(future.ready())
+        error = RuntimeError("oops")
+        future.add_error(error)
+
+        self.assertTrue(future.ready())
+        self.assertIs(future.error(), error)
+
+        with self.assertRaises(AssertionError):
+            future.value()
+
+    def test_two_result_values(self):
+        future = RaceFuture(waiting={10, 20})
+
+        self.assertFalse(future.ready())
+
+        future.add_result(CoroutineResult(coroutine_id=20, value="bar"))
+        self.assertTrue(future.ready())
+        self.assertIsNone(future.error())
+        self.assertEqual(future.value(), "bar")
+
+        future.add_result(CoroutineResult(coroutine_id=10, value="foo"))
+        self.assertTrue(future.ready())
+        self.assertIsNone(future.error())
+        self.assertEqual(future.value(), "bar")
+
+    def test_two_result_errors(self):
+        future = RaceFuture(waiting={10, 20})
+
+        self.assertFalse(future.ready())
+        error1 = RuntimeError("oops")
+        future.add_result(CoroutineResult(coroutine_id=10, error=error1))
+
+        self.assertTrue(future.ready())
+        self.assertIs(future.error(), error1)
+
+        error2 = RuntimeError("oops2")
+        future.add_result(CoroutineResult(coroutine_id=20, error=error2))
+        self.assertIs(future.error(), error1)
+
+        future.add_error(error2)
+        self.assertIs(future.error(), error1)
+
+        with self.assertRaises(AssertionError):
+            future.value()