Merge branch 'pause-queue'

Author: PeterJCLaw

django_lightweight_queue/backends/base.py

@@ -1,3 +1,4 @@
+import datetime
 from abc import ABCMeta, abstractmethod
 from typing import Tuple, TypeVar, Optional, Collection
 
@@ -51,3 +52,19 @@ def deduplicate(
         progress_logger: ProgressLogger = NULL_PROGRESS_LOGGER
     ) -> Tuple[int, int]:
         raise NotImplementedError()
+
+
+class BackendWithPause(BaseBackend, metaclass=ABCMeta):
+    @abstractmethod
+    def pause(self, queue: QueueName, until: datetime.datetime) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_paused(self, queue: QueueName) -> bool:
+        raise NotImplementedError()
+
+
+class BackendWithPauseResume(BackendWithPause, metaclass=ABCMeta):
+    @abstractmethod
+    def resume(self, queue: QueueName) -> None:
+        raise NotImplementedError()

django_lightweight_queue/backends/redis.py

@@ -1,14 +1,16 @@
+import datetime
 from typing import Optional, Collection
 
 import redis
 
 from .. import app_settings
 from ..job import Job
-from .base import BaseBackend
+from .base import BackendWithPauseResume
 from ..types import QueueName, WorkerNumber
+from ..utils import block_for_time
 
 
-class RedisBackend(BaseBackend):
+class RedisBackend(BackendWithPauseResume):
     """
     This backend has at-most-once semantics.
     """
@@ -30,6 +32,15 @@ def bulk_enqueue(self, jobs: Collection[Job], queue: QueueName) -> None:
         )
 
     def dequeue(self, queue: QueueName, worker_num: WorkerNumber, timeout: int) -> Optional[Job]:
+        if self.is_paused(queue):
+            # Block for a while to avoid constant polling ...
+            block_for_time(
+                lambda: self.is_paused(queue),
+                timeout=datetime.timedelta(seconds=timeout),
+            )
+            # ... but always indicate that we did no work
+            return None
+
         raw = self.client.brpop(self._key(queue), timeout)
         if raw is None:
             return None
@@ -40,6 +51,33 @@ def dequeue(self, queue: QueueName, worker_num: WorkerNumber, timeout: int) -> O
     def length(self, queue: QueueName) -> int:
         return self.client.llen(self._key(queue))
 
+    def pause(self, queue: QueueName, until: datetime.datetime) -> None:
+        """
+        Pause the given queue by setting a pause marker.
+        """
+
+        pause_key = self._pause_key(queue)
+
+        now = datetime.datetime.now(datetime.timezone.utc)
+        delta = until - now
+
+        self.client.setex(
+            pause_key,
+            time=int(delta.total_seconds()),
+            # Store the value for debugging, we rely on setex behaviour for
+            # implementation.
+            value=until.isoformat(' '),
+        )
+
+    def resume(self, queue: QueueName) -> None:
+        """
+        Resume the given queue by deleting the pause marker (if present).
+        """
+        self.client.delete(self._pause_key(queue))
+
+    def is_paused(self, queue: QueueName) -> bool:
+        return bool(self.client.exists(self._pause_key(queue)))
+
     def _key(self, queue: QueueName) -> str:
         if app_settings.REDIS_PREFIX:
             return '{}:django_lightweight_queue:{}'.format(
@@ -48,3 +86,6 @@ def _key(self, queue: QueueName) -> str:
             )
 
         return 'django_lightweight_queue:{}'.format(queue)
+
+    def _pause_key(self, queue: QueueName) -> str:
+        return self._key(queue) + ':pause'

django_lightweight_queue/backends/reliable_redis.py

@@ -1,20 +1,21 @@
+import datetime
 from typing import Dict, List, Tuple, TypeVar, Optional, Collection
 
 import redis
 
 from .. import app_settings
 from ..job import Job
-from .base import BackendWithDeduplicate
+from .base import BackendWithDeduplicate, BackendWithPauseResume
 from ..types import QueueName, WorkerNumber
-from ..utils import get_worker_numbers
+from ..utils import block_for_time, get_worker_numbers
 from ..progress_logger import ProgressLogger, NULL_PROGRESS_LOGGER
 
 # Work around https://github.com/python/mypy/issues/9914. Name needs to match
 # that in progress_logger.py.
 T = TypeVar('T')
 
 
-class ReliableRedisBackend(BackendWithDeduplicate):
+class ReliableRedisBackend(BackendWithDeduplicate, BackendWithPauseResume):
     """
     This backend manages a per-queue-per-worker 'processing' queue. E.g. if we
     had a queue called 'django_lightweight_queue:things', and two workers, we
@@ -107,6 +108,15 @@ def dequeue(self, queue: QueueName, worker_number: WorkerNumber, timeout: int) -
         main_queue_key = self._key(queue)
         processing_queue_key = self._processing_key(queue, worker_number)
 
+        if self.is_paused(queue):
+            # Block for a while to avoid constant polling ...
+            block_for_time(
+                lambda: self.is_paused(queue),
+                timeout=datetime.timedelta(seconds=timeout),
+            )
+            # ... but always indicate that we did no work
+            return None
+
         # Get any jobs off our 'processing' queue - but do not block doing so -
         # this is to catch the fact there may be a job already in our
         # processing queue if this worker crashed and has just been restarted.
@@ -191,11 +201,41 @@ def deduplicate(
 
         return original_size, self.client.llen(main_queue_key)
 
+    def pause(self, queue: QueueName, until: datetime.datetime) -> None:
+        """
+        Pause the given queue by setting a pause marker.
+        """
+
+        pause_key = self._pause_key(queue)
+
+        now = datetime.datetime.now(datetime.timezone.utc)
+        delta = until - now
+
+        self.client.setex(
+            pause_key,
+            time=int(delta.total_seconds()),
+            # Store the value for debugging, we rely on setex behaviour for
+            # implementation.
+            value=until.isoformat(' '),
+        )
+
+    def resume(self, queue: QueueName) -> None:
+        """
+        Resume the given queue by deleting the pause marker (if present).
+        """
+        self.client.delete(self._pause_key(queue))
+
+    def is_paused(self, queue: QueueName) -> bool:
+        return bool(self.client.exists(self._pause_key(queue)))
+
     def _key(self, queue: QueueName) -> str:
         key = 'django_lightweight_queue:{}'.format(queue)
 
         return self._prefix_key(key)
 
+    def _pause_key(self, queue: QueueName) -> str:
+        return self._key(queue) + ':pause'
+
     def _processing_key(self, queue: QueueName, worker_number: WorkerNumber) -> str:
         key = 'django_lightweight_queue:{}:processing:{}'.format(
             queue,

django_lightweight_queue/management/commands/queue_pause.py

@@ -0,0 +1,92 @@
+import re
+import argparse
+import datetime
+
+from django.core.management.base import BaseCommand, CommandError
+
+from ...types import QueueName
+from ...utils import get_backend
+from ...backends.base import BackendWithPause
+
+DURATION_PATTERN = r'^((?P<hours>\d+)h)?((?P<minutes>\d+)m)?((?P<seconds>\d+)s)?$'
+TIME_FORMAT = r'%Y-%m-%dT%H:%M:%S%z'
+
+
+def utcnow() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
+def parse_duration_to_time(duration: str) -> datetime.datetime:
+    match = re.match(DURATION_PATTERN, duration)
+    if match is None:
+        raise ValueError(
+            f"Unknown duration format {duration!r}. Try something like '1h2m3s'.",
+        )
+
+    delta = datetime.timedelta(
+        hours=int(match['hours'] or 0),
+        minutes=int(match['minutes'] or 0),
+        seconds=int(match['seconds'] or 0),
+    )
+
+    return utcnow() + delta
+
+
+def parse_time(date_string: str) -> datetime.datetime:
+    return datetime.datetime.strptime(date_string, TIME_FORMAT)
+
+
+class Command(BaseCommand):
+    help = """
+    Command to pause work on a redis-backed queue.
+
+    New jobs can still be added to the queue, however no jobs will be pulled off
+    the queue for processing.
+    """  # noqa:A003 # inherited name
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.add_argument(
+            'queue',
+            action='store',
+            help="The queue to pause.",
+        )
+
+        group = parser.add_mutually_exclusive_group(required=True)
+        group.add_argument(
+            '--for',
+            dest='until',
+            action='store',
+            type=parse_duration_to_time,
+            help=(
+                "The duration to pause the queue for. Specify a value like 1h2m3s, "
+                "all levels of precision are optional, so 5m and 1h are both valid."
+            ),
+        )
+        group.add_argument(
+            '--until',
+            action='store',
+            type=parse_time,
+            help=(
+                "The time at which the queue should reactivate. Specify as an "
+                "ISO 8601 value, specifically one parsable via datetime.strptime "
+                f"using {TIME_FORMAT.replace('%', r'%%')!r}, such as {utcnow():{TIME_FORMAT}}."
+            ),
+        )
+
+    def handle(self, queue: QueueName, until: datetime.datetime, **options: object) -> None:
+        if until < utcnow():
+            raise CommandError("Refusing to pause until a time in the past.")
+
+        backend = get_backend(queue)
+
+        if not isinstance(backend, BackendWithPause):
+            raise CommandError(
+                "Configured backend '{}.{}' doesn't support pausing".format(
+                    type(backend).__module__,
+                    type(backend).__name__,
+                ),
+            )
+
+        backend.pause(queue, until)
+
+        self.stdout.write(f"Paused queue {queue} until {until.isoformat(' ')}.")

django_lightweight_queue/management/commands/queue_resume.py

@@ -0,0 +1,36 @@
+import argparse
+
+from django.core.management.base import BaseCommand, CommandError
+
+from ...types import QueueName
+from ...utils import get_backend
+from ...backends.base import BackendWithPauseResume
+
+
+class Command(BaseCommand):
+    help = """
+    Command to resume work immediately on a redis-backed queue.
+
+    This removes a pause which may be in place for the given queue, though it
+    may not cause workers to resume work immediately.
+    """  # noqa:A003 # inherited name
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.add_argument(
+            'queue',
+            action='store',
+            help="The queue to resume.",
+        )
+
+    def handle(self, queue: QueueName, **options: object) -> None:
+        backend = get_backend(queue)
+
+        if not isinstance(backend, BackendWithPauseResume):
+            raise CommandError(
+                "Configured backend '{}.{}' doesn't support resuming from paused".format(
+                    type(backend).__module__,
+                    type(backend).__name__,
+                ),
+            )
+
+        backend.resume(queue)

django_lightweight_queue/utils.py

@@ -1,4 +1,6 @@
 import imp
+import time
+import datetime
 import warnings
 import importlib
 from typing import (
@@ -27,6 +29,8 @@
 
 _accepting_implied_queues = True
 
+FIVE_SECONDS = datetime.timedelta(seconds=5)
+
 
 def load_extra_config(file_path: str) -> None:
     extra_settings = imp.load_source('extra_settings', file_path)
@@ -139,6 +143,35 @@ def load_all_tasks() -> None:
     import_all_submodules('tasks', app_settings.IGNORE_APPS)
 
 
+def block_for_time(
+    should_continue_blocking: Callable[[], bool],
+    timeout: datetime.timedelta,
+    check_frequency: datetime.timedelta = FIVE_SECONDS,
+) -> bool:
+    """
+    Block until a cancellation function or timeout indicates otherwise.
+
+    Returns whether or not the timeout was encountered.
+    """
+    if not should_continue_blocking():
+        return False
+
+    end = time.time() + timeout.total_seconds()
+
+    while should_continue_blocking():
+        now = time.time()
+        if now > end:
+            # timed out
+            return True
+
+        time.sleep(min(
+            check_frequency.total_seconds(),
+            end - now,
+        ))
+
+    return False
+
+
 try:
     import setproctitle