use callback and shutdown queue

Author: Tomperez98

resonate/bridge.py

@@ -31,6 +31,7 @@
     ResolvePromiseRes,
     Resume,
     Retry,
+    Return,
 )
 from resonate.models.durable_promise import DurablePromise
 from resonate.models.result import Ko, Ok
@@ -84,7 +85,7 @@ def __init__(
             self._unicast,
             self._anycast,
         )
-        self._processor = Processor(self._cq)
+        self._processor = Processor()
 
         self._bridge_thread = threading.Thread(target=self._process_cq, name="bridge", daemon=True)
         self._message_source_thread = threading.Thread(target=self._process_msgs, name="message-source", daemon=True)
@@ -222,7 +223,7 @@ def _process_cq(self) -> None:
                                     return
 
                             case Function(id, cid, func):
-                                self._processor.enqueue(id, cid, func)
+                                self._processor.enqueue(func, lambda r, id=id, cid=cid: self._cq.put_nowait(Return(id, cid, r)))
                             case Delayed() as item:
                                 self._handle_delay(item)
 
@@ -345,7 +346,7 @@ def _process_delayed_events(self) -> None:
                 for item in events:
                     match item:
                         case Function(id, cid, func):
-                            self._processor.enqueue(id, cid, func)
+                            self._processor.enqueue(func, lambda r, id=id, cid=cid: self._cq.put_nowait(Return(id, cid, r)))
                         case retry:
                             self._cq.put_nowait(retry)
 

resonate/processor.py

@@ -4,47 +4,46 @@
 import queue
 from collections.abc import Callable
 from threading import Thread
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
-from resonate.models.commands import Command, Return
-from resonate.models.result import Ko, Ok
+from resonate.models.result import Ko, Ok, Result
 from resonate.utils import exit_on_exception
 
-if TYPE_CHECKING:
-    from concurrent.futures import Future
-
 
 class Processor:
-    def __init__(self, cq: queue.Queue[Command | tuple[Command, Future] | None]) -> None:
+    def __init__(self) -> None:
         self.threads = set[Thread]()
         for _ in range(min(32, (os.cpu_count() or 1))):
             self.threads.add(Thread(target=self._run, daemon=True))
 
-        self.sq = queue.Queue[tuple[str, str, Callable[[], Any]] | None]()
-        self.cq = cq
+        self.sq = queue.Queue[tuple[Callable[[], Any], Callable[[Result[Any]], None]]]()
 
     @exit_on_exception
     def _run(self) -> None:
-        while sqe := self.sq.get():
-            id, cid, func = sqe
+        while True:
+            try:
+                func, callback = self.sq.get()
+            except queue.ShutDown:
+                break
+
             try:
                 r = Ok(func())
             except Exception as e:
                 r = Ko(e)
-            self.cq.put_nowait(Return(id, cid, r))
-            self.sq.task_done()
 
-        self.sq.put(None)
+            callback(r)
+            self.sq.task_done()
 
-    def enqueue(self, id: str, cid: str, func: Callable[[], Any]) -> None:
-        self.sq.put((id, cid, func))
+    def enqueue(self, func: Callable[[], Any], callback: Callable[[Result[Any]], None]) -> None:
+        self.sq.put((func, callback))
 
     def start(self) -> None:
         for t in self.threads:
             if not t.is_alive():
                 t.start()
 
     def stop(self) -> None:
-        self.sq.put(None)
+        self.sq.shutdown()
         for t in self.threads:
             t.join()
+        self.sq.join()