Merge branch 'main' into always-use-barrier

Author: pencil

README.md

@@ -153,6 +153,50 @@ workflow = Workflow(
 In this example, Task 2 will run roughly 1 second after Task 1 finishes, and
 Task 3 and will run 2 seconds after Task 2 finishes.
 
+### Large Workflows
+
+Because of how `dramatiq-workflow` is implemented, each task in a workflow has
+to know about the remaining tasks in the workflow that could potentially run
+after it. When a workflow has a large number of tasks, this can lead to an
+increase of memory usage in the broker and increased network traffic between
+the broker and the workers, especially when using `Group` tasks: Each task in a
+`Group` can potentially be the last one to finish, so each task has to retain a
+copy of the remaining tasks that run after the `Group`.
+
+There are a few things you can do to alleviate this issue:
+
+- Minimize the usage of parameters in the `message` method. Instead, consider
+  using a database to store data that is required by your tasks.
+- Limit the size of groups to a reasonable number of tasks. Instead of
+  scheduling one task with 1000 tasks in a group, consider scheduling 10 groups
+  with 100 tasks each and chaining them together.
+- Consider breaking down large workflows into smaller partial workflows that
+  then schedule a subsequent workflow at the very end of the outermost `Chain`.
+
+Lastly, you can use compression to reduce the size of the messages in your
+queue. While dramatiq does not provide a compression implementation by default,
+one can be added with just a few lines of code. For example:
+
+```python
+import dramatiq
+from dramatiq.encoder import JSONEncoder, MessageData
+import lz4.frame
+
+class DramatiqLz4JSONEncoder(JSONEncoder):
+    def encode(self, data: MessageData) -> bytes:
+        return lz4.frame.compress(super().encode(data))
+
+    def decode(self, data: bytes) -> MessageData:
+        try:
+            decompressed = lz4.frame.decompress(data)
+        except RuntimeError:
+            # Uncompressed data from before the switch to lz4
+            decompressed = data
+        return super().decode(decompressed)
+
+dramatiq.set_encoder(DramatiqLz4JSONEncoder())
+```
+
 ### Barrier
 
 `dramatiq-workflow` uses a barrier mechanism to keep track of the current state

dramatiq_workflow/_barrier.py

@@ -20,7 +20,7 @@ def __init__(self, backend, key, *args, ttl=900000):
         self.ran_key = f"{key}_ran"
 
     def create(self, parties):
-        self.backend.add(self.ran_key, 0, self.ttl)
+        self.backend.add(self.ran_key, -1, self.ttl)
         return super().create(parties)
 
     def wait(self, *args, block=True, timeout=None):
@@ -32,7 +32,7 @@ def wait(self, *args, block=True, timeout=None):
 
         released = super().wait(*args, block=False)
         if released:
-            never_released = self.backend.incr(self.ran_key, 1, 1, self.ttl)
+            never_released = self.backend.incr(self.ran_key, 1, 0, self.ttl)
             return never_released
 
         return False

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 Chain, CompletionCallbacks, Group, Message, WithDelay, WorkflowType
-from ._serialize import serialize_callbacks, serialize_workflow
+from ._models import 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
@@ -178,7 +185,7 @@ def __rate_limiter_backend(self) -> dramatiq.rate_limits.RateLimiterBackend:
     def __barrier(self) -> type[dramatiq.rate_limits.Barrier]:
         return self.__middleware.barrier
 
-    def __create_barrier(self, count: int):
+    def __create_barrier(self, count: int) -> str:
         completion_uuid = str(uuid4())
         completion_barrier = self.__barrier(self.__rate_limiter_backend, completion_uuid, ttl=CALLBACK_BARRIER_TTL)
         completion_barrier.create(count)

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

@@ -6,7 +6,8 @@
 from ._barrier import AtMostOnceBarrier
 from ._constants import OPTION_KEY_CALLBACKS
 from ._helpers import workflow_with_completion_callbacks
-from ._serialize import unserialize_callbacks, unserialize_workflow
+from ._models import SerializedCompletionCallbacks
+from ._serialize import unserialize_workflow
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +34,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
 
@@ -53,9 +54,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

@@ -39,6 +39,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_barrier.py

@@ -0,0 +1,39 @@
+import unittest
+
+from dramatiq.rate_limits.backends import StubBackend
+
+from .._barrier import AtMostOnceBarrier
+
+
+class AtMostOnceBarrierTests(unittest.TestCase):
+    def setUp(self):
+        self.backend = StubBackend()
+        self.key = "test_barrier"
+        self.parties = 3
+        self.ttl = 900000
+        self.barrier = AtMostOnceBarrier(self.backend, self.key, ttl=self.ttl)
+
+    def test_wait_block_true_raises(self):
+        with self.assertRaises(ValueError) as context:
+            self.barrier.wait(block=True)
+        self.assertEqual(str(context.exception), "Blocking is not supported by AtMostOnceBarrier")
+
+    def test_wait_releases_once(self):
+        self.barrier.create(self.parties)
+        for _ in range(self.parties - 1):
+            result = self.barrier.wait(block=False)
+            self.assertFalse(result)
+        result = self.barrier.wait(block=False)
+        self.assertTrue(result)
+        result = self.barrier.wait(block=False)
+        self.assertFalse(result)
+
+    def test_wait_does_not_release_when_db_emptied(self):
+        """
+        If the store is emptied, the barrier should not be released.
+        """
+        self.barrier.create(self.parties)
+        self.backend.db = {}
+        for _ in range(self.parties):
+            result = self.barrier.wait(block=False)
+            self.assertFalse(result)

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

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