dj-queue 0.1.0__py3-none-any.whl

dj_queue/models/jobs.py

@@ -0,0 +1,328 @@
+import uuid
+
+from django.core.exceptions import ObjectDoesNotExist, ValidationError
+from django.db import models
+from django.db.models import Q
+from django.utils.module_loading import import_string
+
+from dj_queue.db import get_database_alias
+from dj_queue.exceptions import UndiscardableError
+
+JOB_STATUS_RELATIONS = (
+  ("ready", "ready_execution"),
+  ("scheduled", "scheduled_execution"),
+  ("claimed", "claimed_execution"),
+  ("blocked", "blocked_execution"),
+  ("failed", "failed_execution"),
+)
+
+
+class JobQuerySet(models.QuerySet):
+  def ready(self):
+    return self.filter(ready_execution__isnull=False)
+
+  def scheduled(self):
+    return self.filter(scheduled_execution__isnull=False)
+
+  def claimed(self):
+    return self.filter(claimed_execution__isnull=False)
+
+  def blocked(self):
+    return self.filter(blocked_execution__isnull=False)
+
+  def failed(self):
+    return self.filter(failed_execution__isnull=False)
+
+  def finished(self):
+    return self.filter(finished_at__isnull=False)
+
+
+class Job(models.Model):
+  objects = JobQuerySet.as_manager()
+
+  id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+  task_path = models.TextField()
+  queue_name = models.CharField(max_length=64, default="default")
+  priority = models.SmallIntegerField(default=0)
+  payload = models.JSONField(default=dict)
+  backend_name = models.CharField(max_length=64)
+  scheduled_at = models.DateTimeField(null=True, blank=True)
+  concurrency_key = models.CharField(max_length=255, null=True, blank=True)
+  finished_at = models.DateTimeField(null=True, blank=True)
+  return_value = models.JSONField(null=True, blank=True)
+  created_at = models.DateTimeField(auto_now_add=True)
+  updated_at = models.DateTimeField(auto_now=True)
+
+  class Meta:
+    db_table = "dj_queue_jobs"
+    constraints = [
+      models.CheckConstraint(
+        condition=Q(priority__gte=-100) & Q(priority__lte=100),
+        name="dj_queue_jobs_priority_range",
+      )
+    ]
+    indexes = [
+      models.Index(fields=["queue_name", "finished_at"]),
+      models.Index(fields=["scheduled_at", "finished_at"]),
+      models.Index(fields=["finished_at"]),
+    ]
+
+  @property
+  def status(self):
+    if self.finished_at is not None:
+      return "finished"
+
+    for status_name, relation_name in JOB_STATUS_RELATIONS:
+      if self._has_state_relation(relation_name):
+        return status_name
+    return None
+
+  @property
+  def ready(self):
+    return self.status == "ready"
+
+  @property
+  def scheduled(self):
+    return self.status == "scheduled"
+
+  @property
+  def claimed(self):
+    return self.status == "claimed"
+
+  @property
+  def blocked(self):
+    return self.status == "blocked"
+
+  @property
+  def failed(self):
+    return self.status == "failed"
+
+  @property
+  def finished(self):
+    return self.status == "finished"
+
+  def _has_state_relation(self, relation_name):
+    try:
+      getattr(self, relation_name)
+    except ObjectDoesNotExist:
+      return False
+    return True
+
+
+class ReadyExecution(models.Model):
+  job = models.OneToOneField(
+    Job,
+    on_delete=models.CASCADE,
+    related_name="ready_execution",
+  )
+  queue_name = models.CharField(max_length=64)
+  priority = models.SmallIntegerField()
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_ready_executions"
+    indexes = [
+      models.Index(fields=["priority", "id"]),
+      models.Index(fields=["queue_name", "priority", "id"]),
+    ]
+
+  def clean(self):
+    super().clean()
+    _validate_live_state(self)
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+  @classmethod
+  def discard_all_in_batches(cls, *, batch_size=500, backend_alias="default"):
+    operation = import_string("dj_queue.operations.jobs.discard_ready_jobs")
+    return _discard_jobs_for_state(
+      cls,
+      operation,
+      batch_size=batch_size,
+      backend_alias=backend_alias,
+    )
+
+
+class ScheduledExecution(models.Model):
+  job = models.OneToOneField(
+    Job,
+    on_delete=models.CASCADE,
+    related_name="scheduled_execution",
+  )
+  queue_name = models.CharField(max_length=64)
+  priority = models.SmallIntegerField()
+  scheduled_at = models.DateTimeField()
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_scheduled_executions"
+    indexes = [models.Index(fields=["scheduled_at", "priority", "id"])]
+
+  def clean(self):
+    super().clean()
+    _validate_live_state(self)
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+
+class ClaimedExecution(models.Model):
+  job = models.OneToOneField(
+    Job,
+    on_delete=models.CASCADE,
+    related_name="claimed_execution",
+  )
+  process = models.ForeignKey(
+    "dj_queue.Process",
+    null=True,
+    blank=True,
+    on_delete=models.SET_NULL,
+    related_name="claimed_executions",
+  )
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_claimed_executions"
+    indexes = [models.Index(fields=["process", "job"])]
+
+  def clean(self):
+    super().clean()
+    _validate_live_state(self)
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+  @classmethod
+  def discard_all_in_batches(cls, **_kwargs):
+    raise UndiscardableError("cannot discard in-progress jobs")
+
+
+class BlockedExecution(models.Model):
+  job = models.OneToOneField(
+    Job,
+    on_delete=models.CASCADE,
+    related_name="blocked_execution",
+  )
+  queue_name = models.CharField(max_length=64)
+  priority = models.SmallIntegerField()
+  concurrency_key = models.CharField(max_length=255)
+  expires_at = models.DateTimeField()
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_blocked_executions"
+    indexes = [
+      models.Index(fields=["concurrency_key", "priority", "id"]),
+      models.Index(fields=["expires_at", "concurrency_key"]),
+    ]
+
+  def clean(self):
+    super().clean()
+    _validate_live_state(self)
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+  @classmethod
+  def discard_all_in_batches(cls, *, batch_size=500, backend_alias="default"):
+    operation = import_string("dj_queue.operations.jobs.discard_blocked_jobs")
+    return _discard_jobs_for_state(
+      cls,
+      operation,
+      batch_size=batch_size,
+      backend_alias=backend_alias,
+    )
+
+
+class FailedExecution(models.Model):
+  job = models.OneToOneField(
+    Job,
+    on_delete=models.CASCADE,
+    related_name="failed_execution",
+  )
+  exception_class = models.CharField(max_length=255)
+  message = models.TextField(default="")
+  traceback = models.TextField(default="")
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_failed_executions"
+
+  def clean(self):
+    super().clean()
+    _validate_live_state(self)
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+  def retry(self):
+    return _retry_failed_job(self.job_id, backend_alias=self.job.backend_name)
+
+  def discard(self):
+    return _discard_failed_job(self.job_id, backend_alias=self.job.backend_name)
+
+  @classmethod
+  def retry_all(cls, queryset):
+    retried = 0
+    for execution in queryset.select_related("job"):
+      execution.retry()
+      retried += 1
+    return retried
+
+  @classmethod
+  def discard_all_in_batches(cls, *, batch_size=500, backend_alias="default"):
+    alias = get_database_alias(backend_alias)
+    deleted = 0
+    while True:
+      job_ids = list(cls.objects.using(alias).values_list("job_id", flat=True)[:batch_size])
+      if not job_ids:
+        return deleted
+      for job_id in job_ids:
+        deleted += _discard_failed_job(job_id, backend_alias=backend_alias)
+
+
+def _retry_failed_job(job_id, *, backend_alias):
+  operation = import_string("dj_queue.operations.jobs.retry_failed_job")
+  return operation(job_id, backend_alias=backend_alias)
+
+
+def _discard_failed_job(job_id, *, backend_alias):
+  operation = import_string("dj_queue.operations.jobs.discard_failed_job")
+  return operation(job_id, backend_alias=backend_alias)
+
+
+def _discard_jobs_for_state(model, operation, *, batch_size, backend_alias):
+  alias = get_database_alias(backend_alias)
+  deleted = 0
+  while True:
+    job_ids = list(model.objects.using(alias).values_list("job_id", flat=True)[:batch_size])
+    if not job_ids:
+      return deleted
+    deleted += operation(job_ids=job_ids, batch_size=batch_size, backend_alias=backend_alias)
+
+
+def _validate_live_state(instance):
+  if not instance.job_id:
+    return
+
+  for model in LIVE_STATE_MODELS:
+    queryset = model._default_manager.filter(job_id=instance.job_id)
+    if model is instance.__class__ and instance.pk is not None:
+      queryset = queryset.exclude(pk=instance.pk)
+    if queryset.exists():
+      raise ValidationError({"job": "job already has a live execution state"})
+
+
+LIVE_STATE_MODELS = (
+  ReadyExecution,
+  ScheduledExecution,
+  ClaimedExecution,
+  BlockedExecution,
+  FailedExecution,
+)

dj_queue/models/recurring.py

@@ -0,0 +1,51 @@
+from croniter import croniter
+from django.core.exceptions import ValidationError
+from django.db import models
+
+
+class RecurringTask(models.Model):
+  key = models.CharField(max_length=255, unique=True)
+  task_path = models.CharField(max_length=255)
+  payload = models.JSONField(null=True, blank=True)
+  schedule = models.CharField(max_length=255)
+  queue_name = models.CharField(max_length=64, default="default")
+  priority = models.SmallIntegerField(default=0)
+  description = models.TextField(default="", blank=True)
+  static = models.BooleanField(default=False)
+  created_at = models.DateTimeField(auto_now_add=True)
+  updated_at = models.DateTimeField(auto_now=True)
+
+  class Meta:
+    db_table = "dj_queue_recurring_tasks"
+
+  def clean(self):
+    super().clean()
+    if not croniter.is_valid(str(self.schedule)):
+      raise ValidationError({"schedule": "schedule must be a valid cron expression"})
+
+  def save(self, *args, **kwargs):
+    self.full_clean()
+    return super().save(*args, **kwargs)
+
+
+class RecurringExecution(models.Model):
+  job = models.OneToOneField(
+    "dj_queue.Job",
+    null=True,
+    blank=True,
+    on_delete=models.CASCADE,
+    related_name="recurring_execution",
+  )
+  task_key = models.CharField(max_length=255)
+  run_at = models.DateTimeField()
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_recurring_executions"
+    constraints = [
+      models.UniqueConstraint(
+        fields=["task_key", "run_at"],
+        name="dj_queue_recurring_executions_task_key_run_at_unique",
+      )
+    ]
+    indexes = [models.Index(fields=["task_key", "run_at"])]

dj_queue/models/runtime.py

@@ -0,0 +1,55 @@
+from django.db import models
+
+
+class Semaphore(models.Model):
+  key = models.CharField(max_length=255, unique=True)
+  value = models.IntegerField()
+  limit = models.IntegerField()
+  expires_at = models.DateTimeField()
+  created_at = models.DateTimeField(auto_now_add=True)
+  updated_at = models.DateTimeField(auto_now=True)
+
+  class Meta:
+    db_table = "dj_queue_semaphores"
+    indexes = [
+      models.Index(fields=["key", "value"]),
+      models.Index(fields=["expires_at"]),
+    ]
+
+
+class Process(models.Model):
+  kind = models.CharField(max_length=32)
+  pid = models.IntegerField()
+  hostname = models.CharField(max_length=255)
+  name = models.CharField(max_length=255)
+  metadata = models.JSONField(default=dict)
+  supervisor = models.ForeignKey(
+    "self",
+    null=True,
+    blank=True,
+    on_delete=models.SET_NULL,
+    related_name="children",
+  )
+  last_heartbeat_at = models.DateTimeField()
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_processes"
+    constraints = [
+      models.UniqueConstraint(
+        fields=["name", "supervisor"],
+        name="dj_queue_processes_name_supervisor_unique",
+      )
+    ]
+    indexes = [
+      models.Index(fields=["name", "supervisor"]),
+      models.Index(fields=["last_heartbeat_at"]),
+    ]
+
+
+class Pause(models.Model):
+  queue_name = models.CharField(max_length=64, unique=True)
+  created_at = models.DateTimeField(auto_now_add=True)
+
+  class Meta:
+    db_table = "dj_queue_pauses"

dj_queue/operations/__init__.py

@@ -0,0 +1 @@
+"""Transactional state transitions live here."""

dj_queue/operations/cleanup.py

@@ -0,0 +1,37 @@
+from datetime import timedelta
+
+from django.utils import timezone
+
+from dj_queue.config import load_backend_config
+from dj_queue.db import get_database_alias
+from dj_queue.models import Job
+
+
+def clear_finished_jobs(
+  *,
+  older_than=None,
+  task_path=None,
+  batch_size=500,
+  backend_alias="default",
+  now=None,
+):
+  config = load_backend_config(backend_alias)
+  if older_than is None:
+    older_than = config.clear_finished_jobs_after
+  if older_than is None:
+    return 0
+
+  alias = get_database_alias(backend_alias)
+  if now is None:
+    now = timezone.now()
+  cutoff = now - timedelta(seconds=older_than)
+  queryset = Job.objects.using(alias).filter(finished_at__lt=cutoff).order_by("finished_at", "id")
+  if task_path is not None:
+    queryset = queryset.filter(task_path=task_path)
+
+  job_ids = list(queryset.values_list("pk", flat=True)[:batch_size])
+  if not job_ids:
+    return 0
+
+  Job.objects.using(alias).filter(pk__in=job_ids).delete()
+  return len(job_ids)

dj_queue/operations/concurrency.py

@@ -0,0 +1,176 @@
+from datetime import timedelta
+
+from django.db import IntegrityError, transaction
+from django.utils import timezone
+from django.utils.module_loading import import_string
+
+from dj_queue.config import load_backend_config
+from dj_queue.db import get_database_alias, locked_queryset
+from dj_queue.log import log_event
+from dj_queue.models import BlockedExecution, ReadyExecution, Semaphore
+from dj_queue.runtime import notify as runtime_notify
+
+
+def semaphore_acquire(
+  key,
+  *,
+  limit,
+  duration_seconds,
+  backend_alias="default",
+):
+  alias = get_database_alias(backend_alias)
+  expires_at = timezone.now() + timedelta(seconds=duration_seconds)
+
+  for attempt in range(2):
+    try:
+      with transaction.atomic(using=alias):
+        semaphore = Semaphore.objects.using(alias).select_for_update().filter(key=key).first()
+        if semaphore is None:
+          Semaphore.objects.using(alias).create(
+            key=key,
+            value=limit - 1,
+            limit=limit,
+            expires_at=expires_at,
+          )
+          return True
+
+        if semaphore.value <= 0:
+          return False
+
+        semaphore.value -= 1
+        semaphore.expires_at = expires_at
+        semaphore.save(using=alias, update_fields=["value", "expires_at", "updated_at"])
+        return True
+    except IntegrityError:
+      # two workers can both miss the row, then race to create the unique key
+      # retry once so the loser can load the row created by the winner
+      if attempt == 0:
+        continue
+      continue
+
+  return False
+
+
+def semaphore_release(key, *, duration_seconds, backend_alias="default"):
+  alias = get_database_alias(backend_alias)
+  expires_at = timezone.now() + timedelta(seconds=duration_seconds)
+
+  with transaction.atomic(using=alias):
+    semaphore = Semaphore.objects.using(alias).select_for_update().filter(key=key).first()
+    if semaphore is None:
+      return False
+
+    semaphore.value = min(semaphore.limit, semaphore.value + 1)
+    semaphore.expires_at = expires_at
+    semaphore.save(using=alias, update_fields=["value", "expires_at", "updated_at"])
+    return True
+
+
+def unblock_next_blocked_job(
+  key,
+  *,
+  limit,
+  duration_seconds,
+  backend_alias="default",
+  use_skip_locked=True,
+):
+  alias = get_database_alias(backend_alias)
+
+  with transaction.atomic(using=alias):
+    queryset = (
+      BlockedExecution.objects.using(alias)
+      .select_related("job")
+      .filter(concurrency_key=key)
+      .order_by("-priority", "id")
+    )
+    blocked = locked_queryset(queryset, use_skip_locked=use_skip_locked).first()
+    if blocked is None:
+      return None
+
+    if not semaphore_acquire(
+      key,
+      limit=limit,
+      duration_seconds=duration_seconds,
+      backend_alias=backend_alias,
+    ):
+      return None
+
+    job = blocked.job
+    queue_name = blocked.queue_name
+    priority = blocked.priority
+    blocked.delete(using=alias)
+    ReadyExecution.objects.using(alias).create(
+      job=job,
+      queue_name=queue_name,
+      priority=priority,
+    )
+
+  log_event(
+    "job.unblocked",
+    job_id=str(job.id),
+    concurrency_key=key,
+  )
+  runtime_notify.notify_ready_queues((job.queue_name,), backend_alias=backend_alias)
+  return job
+
+
+def cleanup_expired_semaphores(*, backend_alias="default"):
+  alias = get_database_alias(backend_alias)
+  queryset = Semaphore.objects.using(alias).filter(expires_at__lte=timezone.now())
+  deleted = queryset.count()
+  if not deleted:
+    return 0
+
+  queryset.delete()
+  return deleted
+
+
+def promote_expired_blocked_jobs(*, batch_size=500, backend_alias="default", use_skip_locked=None):
+  alias = get_database_alias(backend_alias)
+  if use_skip_locked is None:
+    use_skip_locked = load_backend_config(backend_alias).use_skip_locked
+  now = timezone.now()
+  promoted_jobs = []
+  task_settings = {}
+
+  with transaction.atomic(using=alias):
+    queryset = (
+      BlockedExecution.objects.using(alias)
+      .select_related("job")
+      .filter(expires_at__lte=now)
+      .order_by("expires_at", "-priority", "id")
+    )
+    blocked_rows = list(locked_queryset(queryset, use_skip_locked=use_skip_locked)[:batch_size])
+
+    for blocked in blocked_rows:
+      limit, duration_seconds = task_settings.get(blocked.job.task_path, (None, None))
+      if limit is None:
+        task = import_string(blocked.job.task_path)
+        limit = int(getattr(task.func, "concurrency_limit"))
+        duration_seconds = int(getattr(task.func, "concurrency_duration", 60))
+        task_settings[blocked.job.task_path] = (limit, duration_seconds)
+
+      if semaphore_acquire(
+        blocked.concurrency_key,
+        limit=limit,
+        duration_seconds=duration_seconds,
+        backend_alias=backend_alias,
+      ):
+        job = blocked.job
+        queue_name = blocked.queue_name
+        priority = blocked.priority
+        blocked.delete(using=alias)
+        ReadyExecution.objects.using(alias).create(
+          job=job,
+          queue_name=queue_name,
+          priority=priority,
+        )
+        promoted_jobs.append(job)
+      else:
+        blocked.expires_at = now + timedelta(seconds=duration_seconds)
+        blocked.save(using=alias, update_fields=["expires_at"])
+
+  for job in promoted_jobs:
+    log_event("job.unblocked", job_id=str(job.id), concurrency_key=job.concurrency_key)
+    runtime_notify.notify_ready_queues((job.queue_name,), backend_alias=backend_alias)
+  return promoted_jobs