Merge branch 'bulk-enqueue'

PeterJCLaw · PeterJCLaw · commit 613533d8f924 · 2021-11-10T17:06:55.000Z
diff --git a/django_lightweight_queue/backends/base.py b/django_lightweight_queue/backends/base.py
@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import Tuple, TypeVar, Optional
+from typing import Tuple, TypeVar, Optional, Collection
 
 from ..job import Job
 from ..types import QueueName, WorkerNumber
@@ -18,6 +18,16 @@ def startup(self, queue: QueueName) -> None:
     def enqueue(self, job: Job, queue: QueueName) -> None:
         raise NotImplementedError()
 
+    def bulk_enqueue(self, jobs: Collection[Job], queue: QueueName) -> None:
+        """
+        Enqueue a number of tasks in one pass.
+
+        Backends are strongly encouraged to override this with a more efficient
+        implemenation if they can.
+        """
+        for job in jobs:
+            self.enqueue(job, queue)
+
     @abstractmethod
     def dequeue(self, queue: QueueName, worker_num: WorkerNumber, timeout: int) -> Optional[Job]:
         raise NotImplementedError()
diff --git a/django_lightweight_queue/backends/redis.py b/django_lightweight_queue/backends/redis.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Collection
 
 import redis
 
@@ -21,7 +21,13 @@ def __init__(self) -> None:
         )
 
     def enqueue(self, job: Job, queue: QueueName) -> None:
-        self.client.lpush(self._key(queue), job.to_json().encode('utf-8'))
+        return self.bulk_enqueue([job], queue)
+
+    def bulk_enqueue(self, jobs: Collection[Job], queue: QueueName) -> None:
+        self.client.lpush(
+            self._key(queue),
+            *(job.to_json().encode('utf-8') for job in jobs),
+        )
 
     def dequeue(self, queue: QueueName, worker_num: WorkerNumber, timeout: int) -> Optional[Job]:
         raw = self.client.brpop(self._key(queue), timeout)
diff --git a/django_lightweight_queue/backends/reliable_redis.py b/django_lightweight_queue/backends/reliable_redis.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple, TypeVar, Optional
+from typing import Dict, List, Tuple, TypeVar, Optional, Collection
 
 import redis
 
@@ -88,7 +88,13 @@ def move_processing_jobs_to_main(pipe: redis.client.Pipeline) -> None:
         )
 
     def enqueue(self, job: Job, queue: QueueName) -> None:
-        self.client.lpush(self._key(queue), job.to_json().encode('utf-8'))
+        return self.bulk_enqueue([job], queue)
+
+    def bulk_enqueue(self, jobs: Collection[Job], queue: QueueName) -> None:
+        self.client.lpush(
+            self._key(queue),
+            *(job.to_json().encode('utf-8') for job in jobs),
+        )
 
     def dequeue(self, queue: QueueName, worker_number: WorkerNumber, timeout: int) -> Optional[Job]:
         main_queue_key = self._key(queue)
diff --git a/django_lightweight_queue/task.py b/django_lightweight_queue/task.py
@@ -1,4 +1,16 @@
-from typing import Any, Generic, TypeVar, Callable, Optional
+from types import TracebackType
+from typing import (
+    Any,
+    cast,
+    Dict,
+    List,
+    Type,
+    Tuple,
+    Generic,
+    TypeVar,
+    Callable,
+    Optional,
+)
 
 from . import app_settings
 from .job import Job
@@ -83,6 +95,48 @@ def __call__(self, fn: TCallable) -> 'TaskWrapper[TCallable]':
         return TaskWrapper(fn, self.queue, self.timeout, self.sigkill_on_stop, self.atomic)
 
 
+class BulkEnqueueHelper(Generic[TCallable]):
+    def __init__(
+        self,
+        task_wrapper: 'TaskWrapper[TCallable]',
+        batch_size: int,
+        queue_override: Optional[QueueName],
+    ) -> None:
+        self._to_create: List[Job] = []
+        self._task_wrapper = task_wrapper
+        self.batch_size = batch_size
+        self.queue_override = queue_override
+
+    def __enter__(self) -> TCallable:
+        return cast(TCallable, self._create)
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self.flush()
+
+    def _create(self, *args: Any, **kwargs: Any) -> None:
+        self._to_create.append(
+            self._task_wrapper._build_job(args, kwargs),
+        )
+        if len(self._to_create) >= self.batch_size:
+            self.flush()
+
+    def flush(self) -> None:
+        if not self._to_create:
+            return
+
+        self._task_wrapper._enqueue_job_instances(
+            self._to_create,
+            queue_override=self.queue_override,
+        )
+
+        self._to_create = []
+
+
 class TaskWrapper(Generic[TCallable]):
     def __init__(
         self,
@@ -103,18 +157,55 @@ def __init__(
     def __repr__(self) -> str:
         return "<TaskWrapper: {}>".format(self.path)
 
-    def __call__(self, *args: Any, **kwargs: Any) -> None:
+    def _build_job(self, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Job:
         # Allow us to override the default values dynamically
         timeout = kwargs.pop('django_lightweight_queue_timeout', self.timeout)
         sigkill_on_stop = kwargs.pop(
             'django_lightweight_queue_sigkill_on_stop',
             self.sigkill_on_stop,
         )
 
-        # Allow queue overrides, but you must ensure that this queue will exist
-        queue = kwargs.pop('django_lightweight_queue_queue', self.queue)
-
         job = Job(self.path, args, kwargs, timeout, sigkill_on_stop)
         job.validate()
 
+        return job
+
+    def _enqueue_job_instances(
+        self,
+        new_jobs: List[Job],
+        queue_override: Optional[QueueName],
+    ) -> None:
+        queue = queue_override if queue_override is not None else self.queue
+        get_backend(queue).bulk_enqueue(new_jobs, queue)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> None:
+        job = self._build_job(args, kwargs)
+
+        # Allow queue overrides, but you must ensure that this queue will exist
+        queue = kwargs.pop('django_lightweight_queue_queue', self.queue)
+
         get_backend(queue).enqueue(job, queue)
+
+    def bulk_enqueue(
+        self,
+        batch_size: int = 1000,
+        queue_override: Optional[QueueName] = None,
+    ) -> BulkEnqueueHelper[TCallable]:
+        """
+        Enqueue jobs in bulk.
+
+        Use like:
+
+            with my_task.bulk_enqueue() as enqueue:
+                enqueue(the_ids=[42, 43])
+                enqueue(the_ids=[45, 46])
+
+        This is equivalent to:
+
+            my_task(the_ids=[42, 43])
+            my_task(the_ids=[45, 46])
+
+        The target queue for the whole batch may be overridden, however the
+        caller must ensure that the queue actually exists (i.e: has workers).
+        """
+        return BulkEnqueueHelper(self, batch_size, queue_override)
diff --git a/tests/test_task.py b/tests/test_task.py
@@ -0,0 +1,156 @@
+import unittest
+import contextlib
+from typing import Any, Mapping, Iterator
+from unittest import mock
+
+import fakeredis
+from django_lightweight_queue import task
+from django_lightweight_queue.types import QueueName, WorkerNumber
+from django_lightweight_queue.utils import get_path, get_backend
+from django_lightweight_queue.backends.redis import RedisBackend
+
+from . import settings
+
+QUEUE = QueueName('dummy-queue')
+
+
+@task(str(QUEUE))
+def dummy_task(num: int) -> None:
+    pass
+
+
+class TaskTests(unittest.TestCase):
+    longMessage = True
+    prefix = settings.LIGHTWEIGHT_QUEUE_REDIS_PREFIX
+
+    @contextlib.contextmanager
+    def mock_workers(self, workers: Mapping[str, int]) -> Iterator[None]:
+        with unittest.mock.patch(
+            'django_lightweight_queue.utils._accepting_implied_queues',
+            new=False,
+        ), unittest.mock.patch.dict(
+            'django_lightweight_queue.app_settings.WORKERS',
+            workers,
+        ):
+            yield
+
+    def setUp(self) -> None:
+        super().setUp()
+
+        get_backend.cache_clear()
+
+        with mock.patch('redis.StrictRedis', fakeredis.FakeStrictRedis):
+            self.backend = RedisBackend()
+
+        # Mock get_backend. Unfortunately due to the nameing of the 'task'
+        # decorator class being the same as its containing module and it being
+        # exposed as the symbol at django_lightweight_queue.task, we cannot mock
+        # this in the normal way. Instead we mock get_path (which get_backend
+        # calls) and intercept the our dummy value.
+        def mocked_get_path(path: str) -> Any:
+            if path == 'test-backend':
+                return lambda: self.backend
+            return get_path(path)
+
+        patch = mock.patch(
+            'django_lightweight_queue.app_settings.BACKEND',
+            new='test-backend',
+        )
+        patch.start()
+        self.addCleanup(patch.stop)
+        patch = mock.patch(
+            'django_lightweight_queue.utils.get_path',
+            side_effect=mocked_get_path,
+        )
+        patch.start()
+        self.addCleanup(patch.stop)
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        get_backend.cache_clear()
+
+    def test_enqueues_job(self) -> None:
+        self.assertEqual(0, self.backend.length(QUEUE))
+
+        dummy_task(42)
+
+        job = self.backend.dequeue(QUEUE, WorkerNumber(0), 5)
+        # Plain assert to placate mypy
+        assert job is not None, "Failed to get a job after enqueuing one"
+
+        self.assertEqual(
+            {
+                'path': 'tests.test_task.dummy_task',
+                'args': [42],
+                'kwargs': {},
+                'timeout': None,
+                'sigkill_on_stop': False,
+                'created_time': mock.ANY,
+            },
+            job.as_dict(),
+        )
+
+    def test_bulk_enqueues_jobs(self) -> None:
+        self.assertEqual(0, self.backend.length(QUEUE))
+
+        with dummy_task.bulk_enqueue() as enqueue:
+            enqueue(13)
+            enqueue(num=42)
+
+        job = self.backend.dequeue(QUEUE, WorkerNumber(0), 5)
+        # Plain assert to placate mypy
+        assert job is not None, "Failed to get a job after enqueuing one"
+
+        self.assertEqual(
+            {
+                'path': 'tests.test_task.dummy_task',
+                'args': [13],
+                'kwargs': {},
+                'timeout': None,
+                'sigkill_on_stop': False,
+                'created_time': mock.ANY,
+            },
+            job.as_dict(),
+            "First job",
+        )
+
+        job = self.backend.dequeue(QUEUE, WorkerNumber(0), 5)
+        # Plain assert to placate mypy
+        assert job is not None, "Failed to get a job after enqueuing one"
+
+        self.assertEqual(
+            {
+                'path': 'tests.test_task.dummy_task',
+                'args': [],
+                'kwargs': {'num': 42},
+                'timeout': None,
+                'sigkill_on_stop': False,
+                'created_time': mock.ANY,
+            },
+            job.as_dict(),
+            "Second job",
+        )
+
+    def test_bulk_enqueues_jobs_batch_size_boundary(self) -> None:
+        self.assertEqual(0, self.backend.length(QUEUE), "Should initially be empty")
+
+        with dummy_task.bulk_enqueue(batch_size=3) as enqueue:
+            enqueue(1)
+            enqueue(2)
+            enqueue(3)
+            enqueue(4)
+
+        jobs = [
+            self.backend.dequeue(QUEUE, WorkerNumber(0), 5)
+            for _ in range(4)
+        ]
+
+        self.assertEqual(0, self.backend.length(QUEUE), "Should be empty after dequeuing all jobs")
+
+        args = [x.args for x in jobs if x is not None]
+
+        self.assertEqual(
+            [[1], [2], [3], [4]],
+            args,
+            "Wrong jobs bulk enqueued",
+        )
