diff --git a/Makefile b/Makefile index b6d493f..5da7162 100644 --- a/Makefile +++ b/Makefile @@ -14,4 +14,5 @@ lint: .PHONY: force-kill force-kill: - ps | grep steady_queue | cut -f 1 -d ' ' | xargs kill -9 && rm -f tmp/pids/steady_queue_supervisor.pid + ps | grep steady_queue | cut -f 1 -d ' ' | xargs kill -9 + rm -f tmp/pids/steady_queue_supervisor.pid diff --git a/README.md b/README.md index 2d4cc24..d9ae623 100644 --- a/README.md +++ b/README.md @@ -129,15 +129,15 @@ Steady Queue will try to find our configuration under the `STEADY_QUEUE` variabl from steady_queue.configuration import Configuration from datetime import timedelta -STEADY_QUEUE = Configuration.ConfigurationOptions( +STEADY_QUEUE = Configuration.Options( dispatchers=[ - Configuration.DispatcherConfiguration( + Configuration.Dispatcher( polling_interval=timedelta(seconds=1), batch_size=500 ) ], workers=[ - Configuration.WorkerConfiguration( + Configuration.Worker( queues=["*"], threads=3, polling_interval=timedelta(seconds=0.1) @@ -155,19 +155,19 @@ Here's an overview of the different options: - `batch_size`: the dispatcher will dispatch tasks in batches of this size. The default is 500. - `concurrency_maintenance_interval`: the time interval in seconds that the - dispatcher will wait before checking for blocked jobs that can be unblocked. + dispatcher will wait before checking for blocked tasks that can be unblocked. Read more about [concurrency controls](#concurrency-controls) to learn more about this setting. It defaults to `600` seconds. - `queues`: the list of queues that workers will pick tasks from. You can use `*` to indicate all queues (which is also the default and the behavior you'll get if you omit this). Tasks will be polled from those queues in order, so for - example, with `['real_time', 'background']`, no jobs will be taken from - `background` unless there aren't any more jobs waiting in `real_time`. + example, with `['real_time', 'background']`, no tasks will be taken from + `background` unless there aren't any more tasks waiting in `real_time`. You can also provide a prefix with a wildcard to match queues starting with a prefix. For example adding `staging*` to the queues list will create a worker - fetching jobs from all queues starting with `staging`. The wildcard `*` is + fetching tasks from all queues starting with `staging`. The wildcard `*` is only allowed on it's own or at the end of a queue name; you can't specify queue names such as `*_some_queue`. These will be ignored. @@ -388,9 +388,137 @@ TODO ## Concurrency controls -TODO +Steady Queue extends Django Tasks with concurrency controls, that allows you to limit how many tasks of a certain type or with certain arguments can run at the same time. When limited in this way, tasks will be blocked from running, and they'll stay blocked until another task finishes and unblocks them, or after the set expiry time (concurrency limit's _duration_) elapses. Tasks are never discarded or lost, just blocked. + + +```python +from django_tasks import task + +from steady_queue.concurrency import limits_concurrency + +@limits_concurrency( + key=lambda arg1, arg2, **kwargs: pass, + to=max_concurrent_executions, + duration=max_timedelta_to_guarantee_concurrency_limit, + group=concurrency_group +) +@task() +def my_task(arg1, arg2, **kwargs): + pass +``` + +- `key` is the only required parameter, and it can be a string or a callable + that receives the same arguments as the task and returns a string. It will be + used to identify the tasks that need to be limited together. +- `to` is `1` by default. +- `duration` is set to `steady_queue.default_concurrency_control_period` by + default, which itself defaults to `3 minutes`. +- `group` is used to control the concurrency of different tasks types together. + It defaults to the task's module path. + +When a task includes these controls, we'll ensure that, at most, the number of +tasks (indicated as `to`) that yield the same `key` will be performed +concurrently, and this guarantee will last for `duration` for each task +enqueued. Note that there is no guarantee about _the order of execution_, only +about tasks being performed at the same time (overlapping). + +The concurrency limits use the concept of semaphores when enqueueing, and work +as follows: when a task is enqueued, we check if it specifies concurrency +controls. If it does, we check the semaphore for the computed concurrency key. +If the semaphore is open, we claim it and we set the task as _ready_. Ready +means it can be picked up by workers for execution. When the task finishes +execution (be it successfully or unsuccessfully, resulting in a failed +execution), we signal the semaphore and try to unblock the next task with the +same key, if any. Unblocking the next task doesn't mean running that task right +away, but moving it from _blocked_ to _ready_. Since something can heppen that +prevents the first task from releasing the semaphore and unblocking the next +task (for example, someone pulling a plug in the machine where the worker is +running), we have the `duration` as a failsafe. Tasks that have been blocked for +more than duration are candidates to be released, but only as many of them as +the concurrency rules allow, as each one would need to go through the semaphore +dance check. This means that the `duration` is not really about the task that's +enqueued or being run, it's about the tasks that are blocked waiting. It's +important to note that after one or more candidate tasks are unblocked (either +because a task finishes or because `duration` expires and a semaphore is +released), the `duration` timer for the still blocked tasks is reset. This +happens indirectly via the expiration time of the semaphore, which is updated. + +For example + + +```python +@limits_concurrency( + to=2, + key=lambda contact: contact.account_id, + duration=timedelta(minutes=5) +) +@task() +def deliver_announcement(contact): + pass +``` + +In this case, we'll ensure that at most two tasks of the kind +`deliver_announcement` for the same account will run concurrently. If, for any +reason, one of those tasks takes longer than 5 minutes or doesn't release its +concurrency lock (signals the semaphore) within 5 minutes of acquiring it, a new +task with the same key might gain the lock. + +Let's see another example using `group`: + +```python +@limits_concurrency( + key=lambda contact: contact.pk, + duration=timedelta(minutes=15), + group='contact_tasks' +) +@task() +def contact_action(contact): + pass +``` + + +```python +@limits_concurrency( + key=lambda bundle: bundle.contact_id, + duration=timedelta(minutes=15), + group='contact_tasks' +) +@task() +def bundle_action(bundle): + pass +``` -## Failed jobs and retries +In this case, if we have a `contact_action` task enqueued for a contact record +with id `123` and another `bundle_action` task enqueued simultaneously for a +bundle record that references contact `123`, only one of them will be allowed to +proceed. The other one will stay blocked until the first one finishes (or 15 +minutes pass, whatever happens first). + +Note that the `duration` setting depends indirectly on the value for +`concurrency_maintenance_interval` that you set for your dispatcher(s), as +that'd be the frequency with which blocked tasks are checked and unblocked (at +which point, only one task per concurrency key, at most, is unblocked). In +general, you should set `duration` in a way that all your tasks would finish +well under that duration and think of the concurrency maintenance task as a +failsafe in case something goes wrong. + +Tasks are unblocked in order of priority but queue order is not taken into +account for unblocking tasks. That means that if you have a group of tasks that +share a concurrency group but are in different queues, or tasks of the same +class that you enqueue in different queues, the queue order you set for a worker +is not taken into account when unblocking blocked ones. The reason is that a +task that runs unblocks the next one, and the task itself doesn't know about a +particular worker's queue order (you could even have different workers with +different queue orders), it can only know about priority. Once blocked tasks are +unblocked and available for polling, they'll be picked up by a worker following +its queue order. + +Finally, failed tasks that are automatically or manually retried work in the +same way as new tasks that get enqueued: they get in the queue for getting an +open semaphore, and whenever they get it, they'll be run. It doesn't matter if +they had already gotten an open semaphore in the past. + +## Failed tasks and retries TODO @@ -492,7 +620,7 @@ It is possible to run multiple schedulers, for example, if you have multiple servers for redundancy and your run the `scheduler` in more than one of them. To avoid enqueueing duplicate tasks at the same time, an entry in the `steady_queue_recurringexecution` table is added in the same transaction as the -job is enqueued. This table has a unique index on `task_key` and `run_at`, +task is enqueued. This table has a unique index on `task_key` and `run_at`, ensuring only one entry per task per time will be created. This only works if you have `preserve_finished_tasks` set to `True` (the default), and the guarantee applies as long as you keep tasks around. diff --git a/TODOs.md b/TODOs.md index 310d208..6f80ae8 100644 --- a/TODOs.md +++ b/TODOs.md @@ -30,7 +30,10 @@ - [x] Remove demo app in favor of test dummy - [x] Configure package builds - [x] Review logging noisiness -- [ ] Concurrency controls +- [x] Cleanup configuration +- [x] Timer tasks that run immediately +- [x] Concurrency controls +- [ ] Support for multiple databases - [ ] Lifecycle hooks on processes - [ ] Tests - [ ] Class-based tasks @@ -40,5 +43,4 @@ - [ ] Signals on tasks: pre/post enqueue, pre/post perform - [ ] Contributing - [ ] Readme -- [ ] Support for multiple databases - [ ] Publish to PyPI diff --git a/steady_queue/__init__.py b/steady_queue/__init__.py index 7d64bd3..082c82e 100644 --- a/steady_queue/__init__.py +++ b/steady_queue/__init__.py @@ -1,7 +1,7 @@ from datetime import timedelta from typing import Optional -VERSION = (0, 1, "0b7") +VERSION = (0, 1, "0b8") __version__ = ".".join(map(str, VERSION)) diff --git a/steady_queue/backend.py b/steady_queue/backend.py index f259ba1..71a286e 100644 --- a/steady_queue/backend.py +++ b/steady_queue/backend.py @@ -2,7 +2,6 @@ from django_tasks.backends.base import BaseTaskBackend from django_tasks.task import P, T -from steady_queue.models import Job from steady_queue.task import SteadyQueueTask @@ -22,6 +21,8 @@ def validate_task(self, task: Task) -> None: def enqueue( self, task: Task[P, T], args: P.args, kwargs: P.kwargs ) -> TaskResult[T]: + from steady_queue.models import Job + if not isinstance(task, SteadyQueueTask): raise ValueError("Steady Queue only supports SteadyQueueTasks") @@ -35,7 +36,7 @@ def get_result(self, result_id: str) -> TaskResult: "This backend does not support retrieving or refreshing results." ) - def _to_task_result(self, task: SteadyQueueTask, job: Job) -> TaskResult: + def _to_task_result(self, task: SteadyQueueTask, job) -> TaskResult: return TaskResult( task=task, id=str(job.id), diff --git a/steady_queue/concurrency.py b/steady_queue/concurrency.py new file mode 100644 index 0000000..b3a09c5 --- /dev/null +++ b/steady_queue/concurrency.py @@ -0,0 +1,24 @@ +from datetime import timedelta +from typing import Optional + +import steady_queue +from steady_queue.task import SteadyQueueTask + + +def limits_concurrency( + key: str, + to: int = 1, + duration: Optional[timedelta] = None, + group: Optional[str] = None, +): + def wrapper(task: SteadyQueueTask): + task.concurrency_key = key + task.concurrency_limit = to + task.concurrency_duration = ( + duration or steady_queue.default_concurrency_control_period + ) + task.concurrency_group = group or task.module_path + + return task + + return wrapper diff --git a/steady_queue/configuration.py b/steady_queue/configuration.py index 7ccc3f9..d8ff311 100644 --- a/steady_queue/configuration.py +++ b/steady_queue/configuration.py @@ -7,21 +7,21 @@ class Configuration: @dataclass - class WorkerConfiguration: + class Worker: queues: list[str] = field(default_factory=lambda: ["*"]) threads: int = 3 processes: int = 1 polling_interval: timedelta = timedelta(seconds=1) @dataclass - class DispatcherConfiguration: + class Dispatcher: polling_interval: timedelta = timedelta(seconds=0.1) batch_size: int = 500 concurrency_maintenance: bool = True concurrency_maintenance_interval: timedelta = timedelta(minutes=5) @dataclass - class RecurringTaskConfiguration: + class RecurringTask: key: str class_name: Optional[str] = None command: Optional[str] = None @@ -32,36 +32,35 @@ class RecurringTaskConfiguration: description: Optional[str] = None @classmethod - def discover(cls) -> list["Configuration.RecurringTaskConfiguration"]: + def discover(cls) -> list["Configuration.RecurringTask"]: from steady_queue.recurring_task import configurations return configurations @dataclass - class ConfigurationOptions: - workers: list["Configuration.WorkerConfiguration"] - dispatchers: list["Configuration.DispatcherConfiguration"] - recurring_tasks: list["Configuration.RecurringTaskConfiguration"] + class Options: + workers: list["Configuration.Worker"] + dispatchers: list["Configuration.Dispatcher"] + recurring_tasks: list["Configuration.RecurringTask"] only_work: bool = False skip_recurring: bool = False def __init__( self, - workers: list["Configuration.WorkerConfiguration"] | None = None, - dispatchers: list["Configuration.DispatcherConfiguration"] | None = None, - recurring_tasks: list["Configuration.RecurringTaskConfiguration"] - | None = None, + workers: list["Configuration.Worker"] | None = None, + dispatchers: list["Configuration.Dispatcher"] | None = None, + recurring_tasks: list["Configuration.RecurringTask"] | None = None, only_work: bool = False, skip_recurring: bool = False, ): if workers is None: - workers = [Configuration.WorkerConfiguration()] + workers = [Configuration.Worker()] if dispatchers is None: - dispatchers = [Configuration.DispatcherConfiguration()] + dispatchers = [Configuration.Dispatcher()] if recurring_tasks is None: - recurring_tasks = Configuration.RecurringTaskConfiguration.discover() + recurring_tasks = Configuration.RecurringTask.discover() self.workers = workers self.dispatchers = dispatchers @@ -90,20 +89,20 @@ def instantiate(self) -> Base: raise ValueError(f"Invalid process kind: {self.kind}") - def __init__(self, options: Optional[ConfigurationOptions] = None): + def __init__(self, options: Optional[Options] = None): if options is None: - options = self.ConfigurationOptions() + options = self.Options() self.options = options @property - def configured_processes(self): + def configured_processes(self) -> list["Configuration.Process"]: if self.options.only_work: return self.workers return self.workers + self.dispatchers + self.schedulers @property - def workers(self): + def workers(self) -> list["Configuration.Process"]: workers = [] for worker_config in self.options.workers: workers += [ @@ -113,14 +112,14 @@ def workers(self): return workers @property - def dispatchers(self): + def dispatchers(self) -> list["Configuration.Process"]: return [ self.Process(kind="dispatcher", attributes=dispatcher_config) for dispatcher_config in self.options.dispatchers ] @property - def schedulers(self): + def schedulers(self) -> list["Configuration.Process"]: return [ self.Process( kind="scheduler", diff --git a/steady_queue/models/blocked_execution.py b/steady_queue/models/blocked_execution.py index f0c3e59..1796283 100644 --- a/steady_queue/models/blocked_execution.py +++ b/steady_queue/models/blocked_execution.py @@ -1,10 +1,49 @@ -from django.db import models +from django.db import models, transaction +from django.utils import timezone from steady_queue.models.execution import Execution, ExecutionQuerySet +from steady_queue.models.ready_execution import ReadyExecution +from steady_queue.models.semaphore import Semaphore class BlockedExecutionQuerySet(ExecutionQuerySet): - pass + def expired(self): + return self.filter(expires_at__lte=timezone.now()) + + def unblock(self, limit: int): + concurrency_keys = ( + self.expired() + .order_by("concurrency_key") + .distinct() + .values_list("concurrency_key", flat=True)[:limit] + ) + return self.release_many(concurrency_keys) + + def release_many(self, concurrency_keys: list[str]) -> int: + return sum(1 for key in concurrency_keys if self.release_one(key)) + + def release_one(self, concurrency_key: str): + with transaction.atomic(): + execution = ( + self.in_order() + .filter(concurrency_key=concurrency_key) + .select_for_update(skip_locked=True) + .first() + ) + if execution: + return execution.release() + + def releasable(self, concurrency_keys: list[str]) -> list[str]: + semaphores = dict( + Semaphore.objects.filter(key__in=concurrency_keys).values("key", "value") + ) + + # Concurrency keys without semaphore + concurrency keys with open semaphore + return [ + key + for key in concurrency_keys + if key not in semaphores or semaphores[key] > 0 + ] class BlockedExecution(Execution): @@ -35,6 +74,40 @@ class Meta: concurrency_key = models.CharField(max_length=255, verbose_name="concurrency key") expires_at = models.DateTimeField(verbose_name="expires at") + @property + def semaphore(self): + try: + return Semaphore.objects.get(key=self.concurrency_key) + except Semaphore.DoesNotExist: + return None + @property def type(self): return "blocked" + + def save(self, *args, **kwargs): + if self._state.adding: + self.concurrency_key = self.job.concurrency_key + self.set_expires_at() + + return super().save(*args, **kwargs) + + def set_expires_at(self): + self.expires_at = timezone.now() + self.job.concurrency_duration + + def release(self) -> bool: + with transaction.atomic(): + if self.acquire_concurrency_lock(): + self.promote_to_ready() + self.delete() + return True + + return False + + def acquire_concurrency_lock(self) -> bool: + return Semaphore.objects.wait(self.job) + + def promote_to_ready(self): + ReadyExecution.objects.create( + job=self.job, queue_name=self.queue_name, priority=self.priority + ) diff --git a/steady_queue/models/claimed_execution.py b/steady_queue/models/claimed_execution.py index a3b8b56..eba8681 100644 --- a/steady_queue/models/claimed_execution.py +++ b/steady_queue/models/claimed_execution.py @@ -76,6 +76,8 @@ def perform(self): except Exception as e: logger.exception("claimed execution failed", exc_info=e) self.failed_with(e) + finally: + self.unblock_next_job() def finished(self): logger.debug("claimed execution for job %s finished", self.job_id) diff --git a/steady_queue/models/concurrency_controls.py b/steady_queue/models/concurrency_controls.py index fe17bda..d336e5b 100644 --- a/steady_queue/models/concurrency_controls.py +++ b/steady_queue/models/concurrency_controls.py @@ -1,3 +1,65 @@ +from django.db import models +from django.utils.module_loading import import_string + +from steady_queue.models.blocked_execution import BlockedExecution + +from .semaphore import Semaphore + + +class ConcurrencyControlsQuerySet(models.QuerySet): + def release_all_concurrency_locks(self, jobs): + Semaphore.signal_all(filter(lambda job: job.is_concurrency_limited, jobs)) + + class ConcurrencyControls: + @property + def concurrency_limit(self): + return self.job_class.concurrency_limit + + @property + def concurrency_duration(self): + return self.job_class.concurrency_duration + def unblock_next_blocked_job(self): - pass + if self.release_concurrency_lock(): + self.release_next_blocked_job() + + @property + def is_concurrency_limited(self) -> bool: + return self.concurrency_key is not None + + @property + def is_blocked(self) -> bool: + return self.blocked_execution is not None + + def acquire_concurrency_lock(self) -> bool: + if not self.is_concurrency_limited: + return True + + return Semaphore.objects.wait(self) + + def release_concurrency_lock(self) -> bool: + if not self.is_concurrency_limited: + return False + + return Semaphore.objects.signal(self) + + def block(self): + BlockedExecution.objects.get_or_create(job=self) + + def release_next_blocked_job(self): + BlockedExecution.objects.release_one(self.concurrency_key) + + @property + def job_class(self): + return import_string(self.class_name) + + @property + def execution(self): + return super().execution or self.blocked_execution + + def delete(self, *args, **kwargs): + if self.is_concurrency_limited and self.is_ready: + self.unblock_next_blocked_job() + + return super().delete(*args, **kwargs) diff --git a/steady_queue/models/executable.py b/steady_queue/models/executable.py index 922eb63..8af5bbb 100644 --- a/steady_queue/models/executable.py +++ b/steady_queue/models/executable.py @@ -1,13 +1,19 @@ -import steady_queue from django.db import models from django.utils import timezone -from steady_queue.models.concurrency_controls import ConcurrencyControls + +import steady_queue +from steady_queue.models.concurrency_controls import ( + ConcurrencyControls, + ConcurrencyControlsQuerySet, +) from steady_queue.models.ready_execution import ReadyExecution from steady_queue.models.retryable import Retryable, RetryableQuerySet from steady_queue.models.schedulable import Schedulable, SchedulableQuerySet -class ExecutableQuerySet(RetryableQuerySet, SchedulableQuerySet, models.QuerySet): +class ExecutableQuerySet( + ConcurrencyControlsQuerySet, RetryableQuerySet, SchedulableQuerySet, models.QuerySet +): def ready(self): return self.filter(ready_execution__isnull=False) @@ -40,9 +46,8 @@ def prepare_all_for_execution(cls, jobs): @classmethod def dispatch_all(cls, jobs): - # TODO: concurrency limits - without_concurrency_limits = [j for j in jobs if not j.concurrency_key] - with_concurrency_limits = [j for j in jobs if j.concurrency_key] + without_concurrency_limits = [j for j in jobs if not j.is_concurrency_limited] + with_concurrency_limits = [j for j in jobs if j.is_concurrency_limited] cls.dispatch_all_at_once(without_concurrency_limits) cls.dispatch_all_one_by_one(with_concurrency_limits) @@ -77,8 +82,10 @@ def prepare_for_execution(self): return self.schedule() def dispatch(self): - # TODO: concurrency limits - return self.ready + if self.acquire_concurrency_lock(): + return self.ready + else: + return self.block() def dispatch_bypassing_concurrency_limits(self): return self.ready @@ -116,3 +123,10 @@ def execution(self): or getattr(self, "failed_execution", None) or getattr(self, "scheduled_execution", None) ) + + def save(self, *args, **kwargs): + creating = self._state.adding + super().save(*args, **kwargs) + + if creating: + self.prepare_for_execution() diff --git a/steady_queue/models/job.py b/steady_queue/models/job.py index 98f197e..bdc817f 100644 --- a/steady_queue/models/job.py +++ b/steady_queue/models/job.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Optional, Self +from typing import Optional from django.db import models from django.utils import timezone @@ -10,7 +10,14 @@ class JobQuerySet(ExecutableQuerySet, models.QuerySet): - pass + def enqueue(self, task: SteadyQueueTask, scheduled_at: Optional[datetime] = None): + try: + enqueued_job = self.create(**self.model.attributes_from_django_task(task)) + task.provider_task_id = enqueued_job.id + return enqueued_job + except Exception as e: + # TODO: enqueue error + raise e class Job(Executable, UpdatedAtMixin, BaseModel): @@ -53,24 +60,20 @@ class Meta: max_length=255, blank=True, null=True, verbose_name="concurrency key" ) - @classmethod - def enqueue( - cls, task: SteadyQueueTask, scheduled_at: Optional[datetime] = None - ) -> Self: - if scheduled_at is None: - scheduled_at = timezone.now() + DEFAULT_QUEUE_NAME = "default" + DEFAULT_PRIORITY = 0 - job = cls( - queue_name=task.queue_name, - class_name=task.module_path, - arguments=task.serialize(), - priority=task.priority, - scheduled_at=scheduled_at, - django_task_id=task.id, - ) - job.save() - job.prepare_for_execution() - return job + @classmethod + def attributes_from_django_task(cls, task: SteadyQueueTask): + return { + "queue_name": task.queue_name or cls.DEFAULT_QUEUE_NAME, + "django_task_id": task.id, + "priority": task.priority or cls.DEFAULT_PRIORITY, + "scheduled_at": task.run_after or timezone.now(), + "class_name": task.module_path, + "arguments": task.serialize(), + "concurrency_key": task.concurrency_key, + } def __str__(self): if isinstance(self.pk, int): diff --git a/steady_queue/models/recurring_task.py b/steady_queue/models/recurring_task.py index 53e704e..da6db5c 100644 --- a/steady_queue/models/recurring_task.py +++ b/steady_queue/models/recurring_task.py @@ -64,7 +64,7 @@ def wrap(cls, self_or_config): return cls.from_configuration(self_or_config) @classmethod - def from_configuration(cls, config: Configuration.RecurringTaskConfiguration): + def from_configuration(cls, config: Configuration.RecurringTask): return cls( key=config.key, static=True, diff --git a/steady_queue/models/semaphore.py b/steady_queue/models/semaphore.py index 2677499..4e49c1c 100644 --- a/steady_queue/models/semaphore.py +++ b/steady_queue/models/semaphore.py @@ -1,8 +1,26 @@ from django.db import models +from django.utils import timezone from steady_queue.models.base import BaseModel, UpdatedAtMixin +class SemaphoreQuerySet(models.QuerySet): + def wait(self, job) -> bool: + return Semaphore.Proxy(job).wait() + + def signal(self, job) -> bool: + return Semaphore.Proxy(job).signal() + + def signal_all(self, jobs) -> int: + return Semaphore.Proxy.signal_all(jobs) + + def available(self): + return self.filter(value__gt=0) + + def expired(self): + return self.filter(expires_at__lte=timezone.now()) + + class Semaphore(UpdatedAtMixin, BaseModel): class Meta: verbose_name = "semaphore" @@ -15,6 +33,69 @@ class Meta: models.Index(fields=("key", "value"), name="ix_sq_semaphores_key_value"), ) + objects = SemaphoreQuerySet.as_manager() + key = models.CharField(max_length=255, verbose_name="key") value = models.IntegerField(default=1, verbose_name="value") expires_at = models.DateTimeField(null=True, blank=True, verbose_name="expires at") + + class Proxy: + def __init__(self, job): + self.job = job + + def wait(self) -> bool: + try: + semaphore = Semaphore.objects.get(key=self.key) + return semaphore.value > 0 and self.attempt_decrement() + except Semaphore.DoesNotExist: + return self.attempt_creation() + + def signal(self) -> bool: + return self.attempt_increment() + + @classmethod + def signal_all(cls, jobs) -> int: + return Semaphore.objects.filter( + key=[job.concurrency_key for job in jobs] + ).update(value=models.F("value") + 1) + + def attempt_creation(self) -> bool: + semaphore, created = Semaphore.objects.get_or_create( + key=self.key, + defaults={"value": self.limit - 1, "expires_at": self.expires_at}, + ) + if created: + return True + + return self.check_limit_or_decrement() + + def check_limit_or_decrement(self) -> bool: + if self.limit == 1: + return False + return self.attempt_decrement() + + def attempt_decrement(self) -> bool: + updated = ( + Semaphore.objects.available() + .filter(key=self.key) + .update(value=models.F("value") - 1, expires_at=self.expires_at) + ) + return updated > 0 + + def attempt_increment(self) -> bool: + updated = Semaphore.objects.filter( + key=self.key, value__lte=self.limit + ).update(value=models.F("value") + 1, expires_at=self.expires_at) + return updated > 0 + + @property + def key(self) -> str: + return self.job.concurrency_key + + @property + def expires_at(self) -> timezone.datetime: + return timezone.now() + self.job.concurrency_duration + + @property + def limit(self) -> int: + return self.job.concurrency_limit or 1 diff --git a/steady_queue/processes/dispatcher.py b/steady_queue/processes/dispatcher.py index 9a3e5da..49aee24 100644 --- a/steady_queue/processes/dispatcher.py +++ b/steady_queue/processes/dispatcher.py @@ -1,19 +1,30 @@ import logging from datetime import timedelta -from typing import Any +from typing import Any, Optional +from steady_queue.app_executor import AppExecutor from steady_queue.configuration import Configuration +from steady_queue.models.blocked_execution import BlockedExecution from steady_queue.models.scheduled_execution import ScheduledExecution +from steady_queue.models.semaphore import Semaphore from steady_queue.processes.poller import Poller +from steady_queue.processes.timer import TimerTask logger = logging.getLogger("steady_queue") class Dispatcher(Poller): batch_size: int + concurrency_maintenance: Optional["ConcurrencyMaintenance"] = None - def __init__(self, options: Configuration.DispatcherConfiguration): + def __init__(self, options: Configuration.Dispatcher): self.batch_size = options.batch_size + if options.concurrency_maintenance: + self.concurrency_maintenance = self.ConcurrencyMaintenance( + interval=options.concurrency_maintenance_interval, + batch_size=options.batch_size, + ) + super().__init__(polling_interval=options.polling_interval) @property @@ -21,8 +32,19 @@ def metadata(self) -> dict[str, Any]: return { **super().metadata, "batch_size": self.batch_size, + "concurrency_maintenance_interval": self.concurrency_maintenance.interval + if self.concurrency_maintenance + else None, } + def boot(self): + super().boot() + self.start_concurrency_maintenance() + + def shutdown(self): + self.stop_concurrency_maintenance() + super().shutdown() + def poll(self) -> timedelta: batch = self.dispatch_next_batch() if batch > 0: @@ -32,6 +54,41 @@ def poll(self) -> timedelta: def dispatch_next_batch(self) -> int: return ScheduledExecution.dispatch_next_batch(self.batch_size) + def start_concurrency_maintenance(self): + if self.concurrency_maintenance: + self.concurrency_maintenance.start() + + def stop_concurrency_maintenance(self): + if self.concurrency_maintenance: + self.concurrency_maintenance.stop() + @property def is_all_work_completed(self) -> bool: return ScheduledExecution.objects.count() == 0 + + class ConcurrencyMaintenance: + def __init__(self, interval: timedelta, batch_size: int): + self.interval = interval + self.batch_size = batch_size + + def start(self): + self.concurrency_maintenance_task = TimerTask( + interval=self.interval, callable=self.run, run_now=True + ) + self.concurrency_maintenance_task.start() + + def stop(self): + self.concurrency_maintenance_task.stop() + + def run(self): + self.expire_semaphores() + self.unblock_blocked_executions() + + def expire_semaphores(self): + with AppExecutor.wrap_in_app_executor(): + # TODO batch deletes + Semaphore.objects.expired().delete() + + def unblock_blocked_executions(self): + with AppExecutor.wrap_in_app_executor(): + BlockedExecution.objects.unblock(self.batch_size) diff --git a/steady_queue/processes/scheduler.py b/steady_queue/processes/scheduler.py index 596080e..2e69fef 100644 --- a/steady_queue/processes/scheduler.py +++ b/steady_queue/processes/scheduler.py @@ -19,7 +19,7 @@ class Scheduler(Runnable, Interruptible, Registrable, Base): def __init__( self, recurring_tasks: Optional[ - list[RecurringTask | Configuration.RecurringTaskConfiguration] + list[RecurringTask | Configuration.RecurringTask] ] = None, **kwargs, ): diff --git a/steady_queue/processes/supervisor.py b/steady_queue/processes/supervisor.py index dc4a8e2..d7e6968 100644 --- a/steady_queue/processes/supervisor.py +++ b/steady_queue/processes/supervisor.py @@ -21,9 +21,7 @@ class Supervisor(Maintenance, Signals, Pidfiled, Registrable, Interruptible, Base): @classmethod - def launch( - cls, options: Optional[Configuration.ConfigurationOptions] = None - ) -> None: + def launch(cls, options: Optional[Configuration.Options] = None) -> None: configuration = Configuration(options) if not configuration.is_valid: raise ValueError("Invalid configuration") diff --git a/steady_queue/processes/timer.py b/steady_queue/processes/timer.py index 3728ebe..860b407 100644 --- a/steady_queue/processes/timer.py +++ b/steady_queue/processes/timer.py @@ -27,9 +27,15 @@ class TimerTask: isolation as possible, inspired by Ruby's Concurrent::TimerTask. """ - def __init__(self, interval: timedelta, callable: Callable): + def __init__( + self, + interval: timedelta, + callable: Callable, + run_now: bool = False, + ): self.interval = interval self.callable = callable + self.run_now = run_now self._stop_event = threading.Event() def start(self): @@ -49,6 +55,9 @@ def stop(self): logger.debug("timer task stopped") def run(self): + if self.run_now: + self.perform_task() + while not self._stop_event.is_set(): # Use Event.wait() for efficient interruptible sleep logger.debug("timer task waiting for %s", self.interval) @@ -56,14 +65,16 @@ def run(self): # Event was set (stop was called), break out of loop break - # Run the callable in a separate thread to isolate crashes - work_thread = threading.Thread(target=self.wrapped_callable) - work_thread.start() - work_thread.join() + self.perform_task() + + def perform_task(self): + # Run the callable in a separate thread to isolate crashes + work_thread = threading.Thread(target=self.wrapped_callable) + work_thread.start() + work_thread.join() def wrapped_callable(self): try: - time.sleep(10) self.callable() except Exception as e: logger.exception( diff --git a/steady_queue/processes/worker.py b/steady_queue/processes/worker.py index c53b48c..9b57dbc 100644 --- a/steady_queue/processes/worker.py +++ b/steady_queue/processes/worker.py @@ -14,7 +14,7 @@ class Worker(Poller): pool: Pool - def __init__(self, options: Configuration.WorkerConfiguration): + def __init__(self, options: Configuration.Worker): self.queues = options.queues self.pool = Pool(options.threads, on_idle=lambda: self.wake_up()) diff --git a/steady_queue/recurring_schedule.py b/steady_queue/recurring_schedule.py index 564a42f..901f53a 100644 --- a/steady_queue/recurring_schedule.py +++ b/steady_queue/recurring_schedule.py @@ -11,9 +11,7 @@ class RecurringSchedule: - def __init__( - self, tasks: list[RecurringTask | Configuration.RecurringTaskConfiguration] - ): + def __init__(self, tasks: list[RecurringTask | Configuration.RecurringTask]): self.configured_tasks: list[RecurringTask] = [ RecurringTask.wrap(t) for t in tasks ] diff --git a/steady_queue/recurring_task.py b/steady_queue/recurring_task.py index 47b289f..0ba3852 100644 --- a/steady_queue/recurring_task.py +++ b/steady_queue/recurring_task.py @@ -38,7 +38,7 @@ def wrapper(task: SteadyQueueTask): task.args = args task.kwargs = kwargs - configuration = Configuration.RecurringTaskConfiguration( + configuration = Configuration.RecurringTask( key=key, class_name=class_name, schedule=schedule, diff --git a/steady_queue/task.py b/steady_queue/task.py index 267b94d..e703085 100644 --- a/steady_queue/task.py +++ b/steady_queue/task.py @@ -18,6 +18,11 @@ class SteadyQueueTask(Task[P, T]): args: P.args kwargs: P.kwargs + concurrency_key: Optional[str] = None + concurrency_limit: Optional[int] = None + concurrency_duration: Optional[timezone.timedelta] = None + concurrency_group: Optional[str] = None + def __post_init__(self): super().__post_init__() self.arguments = {} diff --git a/tests/dummy/tasks.py b/tests/dummy/tasks.py index d049cfb..2fe48d2 100644 --- a/tests/dummy/tasks.py +++ b/tests/dummy/tasks.py @@ -2,6 +2,7 @@ from django_tasks import task +from steady_queue.concurrency import limits_concurrency from steady_queue.recurring_task import recurring @@ -17,6 +18,14 @@ def long_running_task(): print("long running task finished") +@limits_concurrency(key="limited_task") +@task() +def limited_task(duration: int = 10): + print(f"limited task for {duration} seconds") + time.sleep(duration) + print("limited task finished") + + @recurring(schedule="*/1 * * * *", key="dummy_recurring_task") @task() def dummy_recurring_task(): diff --git a/tests/settings.py b/tests/settings.py index e2b2c40..a816a58 100644 --- a/tests/settings.py +++ b/tests/settings.py @@ -91,14 +91,12 @@ ] -STEADY_QUEUE = Configuration.ConfigurationOptions( +STEADY_QUEUE = Configuration.Options( dispatchers=[ - Configuration.DispatcherConfiguration( - polling_interval=timedelta(seconds=1), batch_size=500 - ) + Configuration.Dispatcher(polling_interval=timedelta(seconds=1), batch_size=500) ], workers=[ - Configuration.WorkerConfiguration( + Configuration.Worker( queues=["*"], threads=2, polling_interval=timedelta(seconds=0.1), diff --git a/tests/test_queue_selector.py b/tests/test_queue_selector.py index 9f7c00b..eafc983 100644 --- a/tests/test_queue_selector.py +++ b/tests/test_queue_selector.py @@ -74,7 +74,4 @@ def create_dummy_job_in_queue(queue_name: str) -> Job: job = Job.objects.create( queue_name=queue_name, class_name="test.dummy", arguments={} ) - return ReadyExecution.objects.create( - job=job, - **ReadyExecution.attributes_from_job(job), - ) + return job.ready_execution