dj-queue 0.1.0__py3-none-any.whl

dj_queue/routers.py

@@ -0,0 +1,26 @@
+from dj_queue.db import get_database_alias
+
+
+class DjQueueRouter:
+  app_label = "dj_queue"
+
+  def db_for_read(self, model, **hints):
+    if model._meta.app_label == self.app_label:
+      return get_database_alias()
+    return None
+
+  def db_for_write(self, model, **hints):
+    if model._meta.app_label == self.app_label:
+      return get_database_alias()
+    return None
+
+  def allow_relation(self, obj1, obj2, **hints):
+    labels = {obj1._meta.app_label, obj2._meta.app_label}
+    if self.app_label in labels:
+      return labels == {self.app_label}
+    return None
+
+  def allow_migrate(self, db, app_label, **hints):
+    if app_label == self.app_label:
+      return db == get_database_alias()
+    return None

dj_queue/runtime/__init__.py

@@ -0,0 +1 @@
+"""Runtime components for dj_queue."""

dj_queue/runtime/base.py

@@ -0,0 +1,198 @@
+from contextlib import contextmanager
+import threading
+
+from django.db import close_old_connections
+from django.utils import timezone
+
+from dj_queue.config import load_backend_config
+from dj_queue.db import get_database_alias
+from dj_queue.hooks import fire_hooks
+from dj_queue.models import Process
+from dj_queue.runtime.errors import handle_thread_error
+from dj_queue.runtime.interruptible import InterruptibleSleeper
+
+
+@contextmanager
+def app_executor():
+  close_old_connections()
+  try:
+    yield
+  finally:
+    close_old_connections()
+
+
+class BaseRunner:
+  process_kind = "Runner"
+  hook_prefix = "runner"
+
+  def __init__(
+    self,
+    config,
+    *,
+    backend_alias="default",
+    name,
+    pid,
+    hostname,
+    sleeper=None,
+    heartbeat_interval=None,
+    supervisor=None,
+  ):
+    self.config = config
+    self.backend_alias = backend_alias
+    self.name = name
+    self.pid = pid
+    self.hostname = hostname
+    self.sleeper = sleeper or InterruptibleSleeper()
+    self.supervisor = supervisor
+    self.process = None
+    self._stop_event = threading.Event()
+    self._heartbeat_thread = None
+    if heartbeat_interval is None:
+      heartbeat_interval = load_backend_config(backend_alias).process_heartbeat_interval
+    self._heartbeat_interval = heartbeat_interval
+    self._started = False
+    self._stopped = False
+
+  @property
+  def polling_interval(self):
+    return getattr(self.config, "polling_interval", 0)
+
+  def start(self):
+    if self.process is None:
+      self.process = self._register_process()
+      self._start_heartbeat_thread()
+      fire_hooks(f"{self.hook_prefix}.start", self.process, backend_alias=self.backend_alias)
+      self._started = True
+    return self.process
+
+  def run(self):
+    self.start()
+    try:
+      while self.should_continue():
+        try:
+          self.poll_once()
+        except Exception as error:
+          handle_thread_error(
+            error,
+            context=f"{self.hook_prefix}.run",
+            backend_alias=self.backend_alias,
+          )
+          break
+
+        if not self.should_continue():
+          break
+        self.sleeper.sleep(self.polling_interval)
+    finally:
+      self.stop()
+      close_old_connections()
+
+  def stop(self):
+    process = self._begin_stop()
+    if process is None:
+      return None
+    self._finish_stop(process)
+    return None
+
+  def request_stop(self):
+    self._stop_event.set()
+    wake_up = getattr(self.sleeper, "wake_up", None)
+    if callable(wake_up):
+      wake_up()
+
+  def should_continue(self):
+    if self._stop_event.is_set():
+      return False
+    if self.process is None:
+      return False
+
+    alias = get_database_alias(self.backend_alias)
+    if Process.objects.using(alias).filter(pk=self.process.pk).exists():
+      return True
+
+    self.request_stop()
+    return False
+
+  def poll_once(self):
+    raise NotImplementedError
+
+  def process_metadata(self):
+    return {}
+
+  def _begin_stop(self):
+    if self._stopped:
+      return None
+
+    self._stopped = True
+    self.request_stop()
+    process = self.process
+    if process is not None and self._started:
+      fire_hooks(f"{self.hook_prefix}.stop", process, backend_alias=self.backend_alias)
+    self._stop_heartbeat_thread()
+    return process
+
+  def _finish_stop(self, process):
+    self._deregister_process()
+    if process is not None and self._started:
+      fire_hooks(f"{self.hook_prefix}.exit", process, backend_alias=self.backend_alias)
+
+    close = getattr(self.sleeper, "close", None)
+    if callable(close):
+      close()
+
+  def _register_process(self):
+    alias = get_database_alias(self.backend_alias)
+    return Process.objects.using(alias).create(
+      kind=self.process_kind,
+      pid=self.pid,
+      hostname=self.hostname,
+      name=self.name,
+      metadata=self.process_metadata(),
+      supervisor=self.supervisor,
+      last_heartbeat_at=timezone.now(),
+    )
+
+  def _deregister_process(self):
+    if self.process is None:
+      return
+
+    alias = get_database_alias(self.backend_alias)
+    Process.objects.using(alias).filter(pk=self.process.pk).delete()
+    self.process = None
+
+  def _start_heartbeat_thread(self):
+    if self._heartbeat_interval <= 0:
+      return
+    self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
+    self._heartbeat_thread.start()
+
+  def _stop_heartbeat_thread(self):
+    thread = self._heartbeat_thread
+    if thread is None:
+      return
+    thread.join(timeout=max(self._heartbeat_interval, 0.1) + 0.1)
+    self._heartbeat_thread = None
+
+  def _heartbeat_loop(self):
+    alias = get_database_alias(self.backend_alias)
+    while not self._stop_event.wait(self._heartbeat_interval):
+      if self.process is None:
+        return
+      try:
+        with app_executor():
+          updated = (
+            Process.objects.using(alias)
+            .filter(pk=self.process.pk)
+            .update(last_heartbeat_at=timezone.now())
+          )
+      except Exception as error:
+        handle_thread_error(
+          error,
+          context=f"{self.hook_prefix}.heartbeat",
+          backend_alias=self.backend_alias,
+        )
+        self.request_stop()
+        return
+
+      if updated == 0:
+        self.request_stop()
+        return

dj_queue/runtime/dispatcher.py

@@ -0,0 +1,78 @@
+import os
+import socket
+
+from django.utils import timezone
+
+from dj_queue.operations.concurrency import (
+  cleanup_expired_semaphores,
+  promote_expired_blocked_jobs,
+)
+from dj_queue.operations.jobs import promote_scheduled_jobs
+from dj_queue.runtime.base import BaseRunner, app_executor
+
+
+class Dispatcher(BaseRunner):
+  process_kind = "Dispatcher"
+  hook_prefix = "dispatcher"
+
+  def __init__(
+    self,
+    config,
+    *,
+    backend_alias="default",
+    name=None,
+    pid=None,
+    hostname=None,
+    sleeper=None,
+    heartbeat_interval=None,
+    supervisor=None,
+  ):
+    super().__init__(
+      config,
+      backend_alias=backend_alias,
+      name=name or f"dispatcher-{os.getpid()}",
+      pid=pid or os.getpid(),
+      hostname=hostname or socket.gethostname(),
+      sleeper=sleeper,
+      heartbeat_interval=heartbeat_interval,
+      supervisor=supervisor,
+    )
+    self._last_maintenance_at = None
+
+  def poll_once(self):
+    if self.process is None:
+      self.start()
+
+    with app_executor():
+      promoted_jobs = promote_scheduled_jobs(
+        batch_size=self.config.batch_size,
+        backend_alias=self.backend_alias,
+      )
+      if self._maintenance_due():
+        cleanup_expired_semaphores(backend_alias=self.backend_alias)
+        promote_expired_blocked_jobs(
+          batch_size=self.config.batch_size,
+          backend_alias=self.backend_alias,
+        )
+        self._last_maintenance_at = timezone.now()
+    return promoted_jobs
+
+  def stop(self):
+    return super().stop()
+
+  def process_metadata(self):
+    return {
+      "batch_size": self.config.batch_size,
+      "polling_interval": self.config.polling_interval,
+      "concurrency_maintenance": self.config.concurrency_maintenance,
+      "concurrency_maintenance_interval": self.config.concurrency_maintenance_interval,
+    }
+
+  def _maintenance_due(self):
+    if not self.config.concurrency_maintenance:
+      return False
+    if self._last_maintenance_at is None:
+      return True
+    return (
+      timezone.now() - self._last_maintenance_at
+    ).total_seconds() >= self.config.concurrency_maintenance_interval

dj_queue/runtime/errors.py

@@ -0,0 +1,39 @@
+import logging
+
+from django.utils.module_loading import import_string
+
+from dj_queue.config import load_backend_config
+
+logger = logging.getLogger("dj_queue")
+
+
+def handle_thread_error(error, *, context="", backend_alias="default"):
+  callback_path = load_backend_config(backend_alias).on_thread_error
+  if callback_path:
+    try:
+      callback = import_string(callback_path)
+      callback(error)
+      return
+    except Exception:
+      logger.exception(
+        "on_thread_error callback raised",
+        extra={
+          "event": "dj_queue.thread_error_callback_failed",
+          "backend_alias": backend_alias,
+          "thread_error_context": context,
+          "on_thread_error": callback_path,
+          "thread_error_type": error.__class__.__name__,
+        },
+      )
+      return
+
+  logger.error(
+    "dj_queue infrastructure error",
+    exc_info=(error.__class__, error, error.__traceback__),
+    extra={
+      "event": "dj_queue.thread_error",
+      "backend_alias": backend_alias,
+      "thread_error_context": context,
+      "thread_error_type": error.__class__.__name__,
+    },
+  )

dj_queue/runtime/interruptible.py

@@ -0,0 +1,46 @@
+import os
+import select
+import threading
+
+
+class InterruptibleSleeper:
+  def __init__(self):
+    self._read_fd, self._write_fd = os.pipe()
+    os.set_blocking(self._read_fd, False)
+    os.set_blocking(self._write_fd, False)
+    self._closed = False
+    self._lock = threading.Lock()
+
+  def sleep(self, seconds):
+    if self._closed:
+      return
+
+    ready, _, _ = select.select([self._read_fd], [], [], max(seconds, 0))
+    if ready:
+      self._drain()
+
+  def wake_up(self):
+    with self._lock:
+      if self._closed:
+        return
+
+      try:
+        os.write(self._write_fd, b".")
+      except (BlockingIOError, OSError):
+        pass
+
+  def close(self):
+    with self._lock:
+      if self._closed:
+        return
+
+      self._closed = True
+      os.close(self._read_fd)
+      os.close(self._write_fd)
+
+  def _drain(self):
+    try:
+      while os.read(self._read_fd, 1024):
+        continue
+    except BlockingIOError:
+      return

dj_queue/runtime/notify.py

@@ -0,0 +1,119 @@
+import threading
+
+from django.db import connections
+
+from dj_queue.config import load_backend_config
+from dj_queue.db import get_database_alias, supports_listen_notify
+from dj_queue.runtime.errors import handle_thread_error
+
+READY_CHANNEL = "dj_queue_ready"
+
+
+class NoopWakeupBackend:
+  def start(self):
+    return None
+
+  def stop(self):
+    return None
+
+
+class NotifyWakeupBackend:
+  def __init__(self, *, backend_alias, wake_up):
+    self.backend_alias = backend_alias
+    self.wake_up = wake_up
+    self.failed = False
+    self._connection = None
+    self._watcher = None
+    self._stop_event = threading.Event()
+
+  def start(self):
+    if self._watcher is not None:
+      return None
+
+    try:
+      self._connection = self._open_connection()
+      self._start_watcher()
+    except Exception as error:
+      self.failed = True
+      self._close_connection()
+      handle_thread_error(error, context="worker.notify", backend_alias=self.backend_alias)
+    return None
+
+  def stop(self):
+    self._stop_event.set()
+    if self._watcher is not None:
+      self._watcher.join(timeout=1)
+      self._watcher = None
+    self._close_connection()
+    return None
+
+  def _start_watcher(self):
+    self._watcher = threading.Thread(target=self._watch, daemon=True, name="dj-queue-notify")
+    self._watcher.start()
+
+  def _watch(self):
+    connection = self._connection
+    if connection is None:
+      return
+
+    while not self._stop_event.is_set():
+      try:
+        notifications = connection.notifies(timeout=0.5, stop_after=1)
+      except Exception as error:
+        self.failed = True
+        handle_thread_error(error, context="worker.notify", backend_alias=self.backend_alias)
+        return
+
+      for _notify in notifications:
+        self.wake_up()
+
+  def _open_connection(self):
+    alias = get_database_alias(self.backend_alias)
+    wrapper = connections[alias]
+    wrapper.ensure_connection()
+    connection = wrapper.Database.connect(**wrapper.get_connection_params())
+    connection.autocommit = True
+    with connection.cursor() as cursor:
+      cursor.execute(f"LISTEN {READY_CHANNEL}")
+    return connection
+
+  def _close_connection(self):
+    connection = self._connection
+    if connection is None:
+      return
+    self._connection = None
+    try:
+      connection.close()
+    except Exception:
+      return
+
+
+def notify_ready_queues(queue_names, *, backend_alias="default"):
+  config = load_backend_config(backend_alias)
+  if not queue_names or not config.listen_notify:
+    return None
+
+  alias = get_database_alias(backend_alias)
+  if not supports_listen_notify(alias):
+    return None
+
+  for queue_name in queue_names:
+    _notify(READY_CHANNEL, queue_name, backend_alias=backend_alias)
+  return None
+
+
+def build_wakeup_backend(*, backend_alias="default", queues=(), wake_up=None):
+  config = load_backend_config(backend_alias)
+  alias = get_database_alias(backend_alias)
+  if wake_up is None or not config.listen_notify or not supports_listen_notify(alias):
+    return NoopWakeupBackend()
+  return NotifyWakeupBackend(backend_alias=backend_alias, wake_up=wake_up)
+
+
+def _notify(channel, payload, *, backend_alias):
+  try:
+    with connections[get_database_alias(backend_alias)].cursor() as cursor:
+      cursor.execute("SELECT pg_notify(%s, %s)", [channel, payload])
+  except Exception:
+    return None
+  return None

dj_queue/runtime/pidfile.py

@@ -0,0 +1,39 @@
+import os
+from pathlib import Path
+
+
+class PidFile:
+  def __init__(self, path, *, pid=None, probe=None):
+    self.path = Path(path)
+    self.pid = os.getpid() if pid is None else pid
+    self._probe = probe or _process_alive
+
+  def acquire(self):
+    if self.path.exists():
+      existing_pid = _read_pid(self.path)
+      if existing_pid is not None and self._probe(existing_pid):
+        raise RuntimeError(f"a dj_queue supervisor is already running (pid={existing_pid})")
+
+    self.path.parent.mkdir(parents=True, exist_ok=True)
+    self.path.write_text(str(self.pid))
+
+  def release(self):
+    try:
+      self.path.unlink(missing_ok=True)
+    except OSError:
+      pass
+
+
+def _read_pid(path):
+  try:
+    return int(path.read_text().strip())
+  except (OSError, TypeError, ValueError):
+    return None
+
+
+def _process_alive(pid):
+  try:
+    os.kill(pid, 0)
+  except (ProcessLookupError, PermissionError):
+    return False
+  return True

dj_queue/runtime/pool.py

@@ -0,0 +1,62 @@
+from concurrent.futures import ThreadPoolExecutor, wait
+import threading
+
+
+class WorkerPool:
+  def __init__(self, max_workers, *, wake_up):
+    self.max_workers = max_workers
+    self._wake_up = wake_up
+    self._executor = ThreadPoolExecutor(max_workers=max_workers)
+    self._lock = threading.Lock()
+    self._futures = set()
+    self._in_flight = 0
+    self._on_drained = None
+
+  @property
+  def idle_capacity(self):
+    with self._lock:
+      return max(0, self.max_workers - self._in_flight)
+
+  def submit(self, fn, *args, **kwargs):
+    with self._lock:
+      self._in_flight += 1
+
+    future = self._executor.submit(fn, *args, **kwargs)
+    with self._lock:
+      self._futures.add(future)
+    future.add_done_callback(self._complete)
+    return future
+
+  def shutdown(self, timeout, *, on_drained=None):
+    with self._lock:
+      self._on_drained = on_drained
+      futures = list(self._futures)
+    if not futures:
+      self._executor.shutdown(wait=False, cancel_futures=False)
+      self._notify_drained()
+      return True
+
+    _, not_done = wait(futures, timeout=timeout)
+    self._executor.shutdown(wait=False, cancel_futures=False)
+    if not not_done:
+      self._notify_drained()
+    return len(not_done) == 0
+
+  def _complete(self, future):
+    callback = None
+    with self._lock:
+      self._futures.discard(future)
+      self._in_flight -= 1
+      if self._in_flight == 0 and not self._futures:
+        callback = self._on_drained
+        self._on_drained = None
+    if callback is not None:
+      callback()
+    self._wake_up()
+
+  def _notify_drained(self):
+    with self._lock:
+      callback = self._on_drained
+      self._on_drained = None
+    if callback is not None:
+      callback()

dj_queue/runtime/procline.py

@@ -0,0 +1,11 @@
+import importlib
+
+
+def set_process_title(title):
+  try:
+    setproctitle = importlib.import_module("setproctitle")
+  except ModuleNotFoundError:
+    return False
+
+  setproctitle.setproctitle(title)
+  return True