From a3dc9d8cf1f1dd4d17c43ae4efcf614be93d0e58 Mon Sep 17 00:00:00 2001
From: adnanhd <adnanharundogan@gmail.com>
Date: Wed, 24 Jun 2026 09:44:20 +0300
Subject: [PATCH] feat(dispatcher): add LeastLoadedDispatcher.route() returning
 the result

Factors the atomic select-and-reserve into _select(), shared by the existing
fire-and-forget dispatch() and a new route() that invokes the chosen
subscriber and returns its result (propagating errors). Lets a scheduler
reuse the load-aware routing instead of reimplementing it.
---
 eventforge/observers.py | 57 +++++++++++++++++++++++++++++------------
 tests/test_observers.py | 18 +++++++++++++
 2 files changed, 58 insertions(+), 17 deletions(-)

diff --git a/eventforge/observers.py b/eventforge/observers.py
index 0f779b1..9746079 100644
--- a/eventforge/observers.py
+++ b/eventforge/observers.py
@@ -228,6 +228,45 @@ class LeastLoadedDispatcher(Dispatcher):
     def __init__(self) -> None:
         self._lock = threading.Lock()
 
+    def _select(self, subscribers: List[Callable[..., Any]]) -> "Tuple[Any, bool]":
+        """Atomically pick and reserve the least-loaded subscriber.
+
+        Returns ``(subscriber, reserved)`` -- ``reserved`` is True when it's a
+        :class:`Node` whose slot was taken (the caller must ``release()``).
+        Selection runs under one lock so the pick can't go stale between
+        reading ``load()`` and committing. Raises if every subscriber is full.
+        """
+        with self._lock:
+            for sub in sorted(subscribers, key=lambda s: cast("LoadAware", s).load()):
+                if isinstance(sub, Node):
+                    if sub.try_acquire():
+                        return sub, True
+                elif cast("LoadAware", sub).load() < 1.0:
+                    return sub, False
+        raise RuntimeError("all subscribers saturated")
+
+    def route(
+        self,
+        subscribers: List[Callable[..., Any]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Select the least-loaded subscriber, invoke it, and return its result.
+
+        The result-returning counterpart of :meth:`dispatch` (which is
+        fire-and-forget): use this when a caller needs the value back -- e.g. a
+        scheduler running a job to completion. Reservation is held for the call
+        and released after; the subscriber's own exceptions propagate.
+        """
+        chosen, reserved = self._select(subscribers)
+        try:
+            # A reserved Node already counts the slot, so invoke its handler
+            # directly (its __call__ would double-count); others via __call__.
+            return (chosen.handler if reserved else chosen)(*args, **kwargs)
+        finally:
+            if reserved:
+                chosen.release()
+
     def dispatch(
         self,
         subscribers: List[Callable[..., Any]],
@@ -236,24 +275,8 @@ def dispatch(
     ) -> None:
         if not subscribers:
             return
-        chosen: Any = None
-        reserved = False
-        # Select-and-reserve under one lock so the pick can't go stale between
-        # reading load() and committing to a node.
-        with self._lock:
-            for sub in sorted(subscribers, key=lambda s: cast("LoadAware", s).load()):
-                if isinstance(sub, Node):
-                    if sub.try_acquire():
-                        chosen, reserved = sub, True
-                        break
-                elif cast("LoadAware", sub).load() < 1.0:
-                    chosen = sub
-                    break
-            if chosen is None:
-                raise RuntimeError("all subscribers saturated")
+        chosen, reserved = self._select(subscribers)  # saturated propagates
         try:
-            # A reserved Node already counts the slot, so invoke its handler
-            # directly (its __call__ would double-count); others via __call__.
             (chosen.handler if reserved else chosen)(*args, **kwargs)
         except Exception:
             logger.exception("least-loaded subscriber failed: %r", chosen)
diff --git a/tests/test_observers.py b/tests/test_observers.py
index 1da6141..10d6098 100644
--- a/tests/test_observers.py
+++ b/tests/test_observers.py
@@ -172,6 +172,24 @@ def __call__(self, *a, **kw):
         with pytest.raises(RuntimeError, match="saturated"):
             LeastLoadedDispatcher().dispatch([Sat(), Sat()], (), {})
 
+    def test_least_loaded_route_returns_result(self):
+        node = Node(
+            "w", cpus=4, memory_gb=8, gpus=[0], handler=lambda x: ("ran", x * 2)
+        )
+        disp = LeastLoadedDispatcher()
+        assert disp.route([node], 21) == ("ran", 42)
+        assert node.load() == 0.0  # released after the call
+
+    def test_least_loaded_route_propagates_and_releases_on_error(self):
+        def boom(_):
+            raise ValueError("nope")
+
+        node = Node("w", cpus=4, memory_gb=8, gpus=[0], handler=boom)
+        disp = LeastLoadedDispatcher()
+        with pytest.raises(ValueError, match="nope"):
+            disp.route([node], 1)
+        assert node.load() == 0.0  # released even though the call raised
+
     def test_least_loaded_reserves_node_atomically(self):
         # A capacity-1 Node under concurrent dispatch is reserved for the call's
         # duration: a second dispatch sees it saturated, and it never runs the