More efficient completion callback serialization (#5)

As of now, running a group of workflows would lead to serializing the
`completion_callbacks` separately for each member of the group. With
large workflows, this could be quite slow. Additionally, there was a
long-standing TODO in the middleware to not unserialize the callbacks
since they just have to be serialized again immediately after.

This PR aims to solve both these issues. We never actually need an
unserialized version of the callbacks so we now only pass around the
serialized version. Only the workflow that needs to be run is now
unserialized by the middleware when needed.

Author: pencil

dramatiq_workflow/_base.py

@@ -8,8 +8,8 @@
 from ._constants import CALLBACK_BARRIER_TTL, OPTION_KEY_CALLBACKS
 from ._helpers import workflow_with_completion_callbacks
 from ._middleware import WorkflowMiddleware, workflow_noop
-from ._models import Barrier, Chain, CompletionCallbacks, Group, Message, WithDelay, WorkflowType
-from ._serialize import serialize_callbacks, serialize_workflow
+from ._models import Barrier, Chain, Group, Message, SerializedCompletionCallbacks, WithDelay, WorkflowType
+from ._serialize import serialize_workflow
 
 logger = logging.getLogger(__name__)
 
@@ -91,15 +91,15 @@ def __init__(
         self.broker = broker or dramatiq.get_broker()
 
         self._delay = None
-        self._completion_callbacks = []
+        self._completion_callbacks: SerializedCompletionCallbacks | None = None
 
         while isinstance(self.workflow, WithDelay):
             self._delay = (self._delay or 0) + self.workflow.delay
             self.workflow = self.workflow.task
 
     def run(self):
         current = self.workflow
-        completion_callbacks = self._completion_callbacks.copy()
+        completion_callbacks = self._completion_callbacks or []
 
         if isinstance(current, Message):
             current = self.__augment_message(current, completion_callbacks)
@@ -115,7 +115,10 @@ def run(self):
             task = tasks.pop(0)
             if tasks:
                 completion_id = self.__create_barrier(1)
-                completion_callbacks.append((completion_id, Chain(*tasks), False))
+                completion_callbacks = [
+                    *completion_callbacks,
+                    (completion_id, serialize_workflow(Chain(*tasks)), False),
+                ]
             self.__workflow_with_completion_callbacks(task, completion_callbacks).run()
             return
 
@@ -126,7 +129,7 @@ def run(self):
                 return
 
             completion_id = self.__create_barrier(len(tasks))
-            completion_callbacks.append((completion_id, None, True))
+            completion_callbacks = [*completion_callbacks, (completion_id, None, True)]
             for task in tasks:
                 self.__workflow_with_completion_callbacks(task, completion_callbacks).run()
             return
@@ -141,18 +144,22 @@ def __workflow_with_completion_callbacks(self, task, completion_callbacks) -> "W
             delay=self._delay,
         )
 
-    def __schedule_noop(self, completion_callbacks: CompletionCallbacks):
+    def __schedule_noop(self, completion_callbacks: SerializedCompletionCallbacks):
         noop_message = workflow_noop.message()
         noop_message = self.__augment_message(noop_message, completion_callbacks)
         self.broker.enqueue(noop_message, delay=self._delay)
 
-    def __augment_message(self, message: Message, completion_callbacks: CompletionCallbacks) -> Message:
+    def __augment_message(self, message: Message, completion_callbacks: SerializedCompletionCallbacks) -> Message:
+        options = {}
+        if completion_callbacks:
+            options = {OPTION_KEY_CALLBACKS: completion_callbacks}
+
         return message.copy(
             # We reset the message timestamp to better represent the time the
             # message was actually enqueued.  This is to avoid tripping the max_age
             # check in the broker.
             message_timestamp=time.time() * 1000,
-            options={OPTION_KEY_CALLBACKS: serialize_callbacks(completion_callbacks)},
+            options=options,
         )
 
     @property
@@ -170,7 +177,7 @@ def __rate_limiter_backend(self):
                 )
         return self.__cached_rate_limiter_backend
 
-    def __create_barrier(self, count: int):
+    def __create_barrier(self, count: int) -> str | None:
         if count == 1:
             # No need to create a distributed barrier if there is only one task
             return None

dramatiq_workflow/_helpers.py

@@ -1,12 +1,12 @@
 import dramatiq
 
-from ._models import CompletionCallbacks, WorkflowType
+from ._models import SerializedCompletionCallbacks, WorkflowType
 
 
 def workflow_with_completion_callbacks(
     workflow: WorkflowType,
     broker: dramatiq.Broker,
-    completion_callbacks: CompletionCallbacks,
+    completion_callbacks: SerializedCompletionCallbacks,
     delay: int | None = None,
 ):
     from ._base import Workflow

dramatiq_workflow/_middleware.py

@@ -5,8 +5,8 @@
 
 from ._constants import OPTION_KEY_CALLBACKS
 from ._helpers import workflow_with_completion_callbacks
-from ._models import Barrier
-from ._serialize import unserialize_callbacks, unserialize_workflow
+from ._models import Barrier, SerializedCompletionCallbacks
+from ._serialize import unserialize_workflow
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@ def after_process_message(
         if message.failed:
             return
 
-        completion_callbacks: list[dict] | None = message.options.get(OPTION_KEY_CALLBACKS)
+        completion_callbacks: SerializedCompletionCallbacks | None = message.options.get(OPTION_KEY_CALLBACKS)
         if completion_callbacks is None:
             return
 
@@ -48,9 +48,7 @@ def after_process_message(
             workflow_with_completion_callbacks(
                 unserialize_workflow(remaining_workflow),
                 broker,
-                # TODO: This is somewhat inefficient because we're unserializing all callbacks
-                # even though we are just going to serialize them again.
-                unserialize_callbacks(completion_callbacks),
+                completion_callbacks,
             ).run()
 
             if not propagate:

dramatiq_workflow/_models.py

@@ -42,6 +42,5 @@ def __eq__(self, other):
 Message = dramatiq.Message
 WorkflowType = Message | Chain | Group | WithDelay
 
-# CompletionCallback is a tuple containing: ID, WorkflowType, propagate
-CompletionCallback = tuple[str, WorkflowType | None, bool]
-CompletionCallbacks = list[CompletionCallback]
+SerializedCompletionCallback = tuple[str | None, dict | None, bool]
+SerializedCompletionCallbacks = list[SerializedCompletionCallback]

dramatiq_workflow/_serialize.py

@@ -1,9 +1,15 @@
 import typing
 
-from ._models import Chain, CompletionCallbacks, Group, Message, WithDelay, WorkflowType
+from ._models import (
+    Chain,
+    Group,
+    Message,
+    WithDelay,
+    WorkflowType,
+)
 
 
-def serialize_workflow(workflow: WorkflowType | None) -> typing.Any:
+def serialize_workflow(workflow: WorkflowType | None) -> dict | None:
     """
     Return a serialized version of the workflow that can be JSON-encoded.
     """
@@ -69,11 +75,3 @@ def unserialize_workflow_or_none(workflow: typing.Any) -> WorkflowType | None:
     if workflow is None:
         return None
     return unserialize_workflow(workflow)
-
-
-def serialize_callbacks(callbacks: CompletionCallbacks) -> list[tuple[str, typing.Any, bool]]:
-    return [(id, serialize_workflow(g), propagate) for id, g, propagate in callbacks]
-
-
-def unserialize_callbacks(callbacks: list[dict]) -> CompletionCallbacks:
-    return [(id, unserialize_workflow_or_none(g), propagate) for id, g, propagate in callbacks]

dramatiq_workflow/tests/test_middleware.py

@@ -0,0 +1,92 @@
+import unittest
+from unittest import mock
+
+import dramatiq
+from dramatiq.broker import Broker
+from dramatiq.rate_limits.backends import StubBackend
+
+from dramatiq_workflow import Chain, WorkflowMiddleware
+from dramatiq_workflow._barrier import AtMostOnceBarrier
+from dramatiq_workflow._constants import OPTION_KEY_CALLBACKS
+from dramatiq_workflow._serialize import serialize_workflow
+
+
+class WorkflowMiddlewareTests(unittest.TestCase):
+    def setUp(self):
+        # Initialize common mocks and the middleware instance for each test
+        self.rate_limiter_backend = StubBackend()
+        self.middleware = WorkflowMiddleware(self.rate_limiter_backend)
+
+        self.broker = mock.MagicMock(spec=Broker)
+
+    def _make_message(
+        self, message_options: dict | None = None, message_timestamp: int = 1717526084640
+    ) -> dramatiq.broker.MessageProxy:
+        """
+        Creates a dramatiq MessageProxy object with given options.
+        """
+        message_id = 1  # Simplistic message ID for testing
+        message = dramatiq.Message(
+            message_id=str(message_id),
+            message_timestamp=message_timestamp,
+            queue_name="default",
+            actor_name="test_task",
+            args=(),
+            kwargs={},
+            options=message_options or {},
+        )
+        return dramatiq.broker.MessageProxy(message)
+
+    def _create_serialized_workflow(self) -> dict | None:
+        """
+        Creates and serializes a simple workflow for testing.
+        """
+        # Define a simple workflow (Chain with a single task)
+        workflow = Chain(self._make_message()._message)
+        serialized = serialize_workflow(workflow)
+        return serialized
+
+    def test_after_process_message_without_callbacks(self):
+        message = self._make_message()
+
+        self.middleware.after_process_message(self.broker, message)
+
+        self.broker.enqueue.assert_not_called()
+
+    def test_after_process_message_with_exception(self):
+        message = self._make_message({OPTION_KEY_CALLBACKS: [(None, self._create_serialized_workflow(), True)]})
+
+        self.middleware.after_process_message(self.broker, message, exception=Exception("Test exception"))
+
+        self.broker.enqueue.assert_not_called()
+
+    def test_after_process_message_with_failed_message(self):
+        message = self._make_message({OPTION_KEY_CALLBACKS: [(None, self._create_serialized_workflow(), True)]})
+        message.failed = True
+
+        self.middleware.after_process_message(self.broker, message)
+
+        self.broker.enqueue.assert_not_called()
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_after_process_message_with_workflow(self, mock_time):
+        mock_time.return_value = 1337
+        message = self._make_message({OPTION_KEY_CALLBACKS: [(None, self._create_serialized_workflow(), True)]})
+
+        self.middleware.after_process_message(self.broker, message)
+
+        self.broker.enqueue.assert_called_once_with(self._make_message(message_timestamp=1337_000)._message, delay=None)
+
+    @mock.patch("dramatiq_workflow._base.time.time")
+    def test_after_process_message_with_barriered_workflow(self, mock_time):
+        mock_time.return_value = 1337
+        barrier = AtMostOnceBarrier(self.rate_limiter_backend, "barrier_1")
+        barrier.create(2)
+        message = self._make_message({OPTION_KEY_CALLBACKS: [(barrier.key, self._create_serialized_workflow(), True)]})
+
+        self.middleware.after_process_message(self.broker, message)
+        self.broker.enqueue.assert_not_called()
+
+        # Calling again, barrier should be completed now
+        self.middleware.after_process_message(self.broker, message)
+        self.broker.enqueue.assert_called_once_with(self._make_message(message_timestamp=1337_000)._message, delay=None)

dramatiq_workflow/tests/test_workflow.py

@@ -377,7 +377,6 @@ def test_nested_delays(self, time_mock):
             self.__make_message(
                 1,
                 message_timestamp=updated_timestamp,
-                message_options={"workflow_completion_callbacks": []},
             ),
             delay=20,
         )