[feat] priority-based scheduling

Author: sanchit-arora-glean

arq/connections.py

@@ -13,7 +13,18 @@
 from redis.asyncio.sentinel import Sentinel
 from redis.exceptions import RedisError, WatchError
 
-from .constants import default_queue_name, expires_extra_ms, job_key_prefix, result_key_prefix
+from .constants import (
+    DEFAULT_PRIORITY,
+    MAX_PRIORITY,
+    MIN_PRIORITY,
+    compute_ready_score,
+    default_queue_name,
+    deferred_queue_key,
+    expires_extra_ms,
+    job_key_prefix,
+    ready_queue_key,
+    result_key_prefix,
+)
 from .jobs import Deserializer, Job, JobDef, JobResult, Serializer, deserialize_job, serialize_job
 from .utils import timestamp_ms, to_ms, to_unix_ms
 
@@ -126,6 +137,7 @@ async def enqueue_job(
         _defer_by: Union[None, int, float, timedelta] = None,
         _expires: Union[None, int, float, timedelta] = None,
         _job_try: Optional[int] = None,
+        _priority: int = DEFAULT_PRIORITY,
         **kwargs: Any,
     ) -> Optional[Job]:
         """
@@ -140,9 +152,13 @@ async def enqueue_job(
         :param _expires: do not start or retry a job after this duration;
             defaults to 24 hours plus deferring time, if any
         :param _job_try: useful when re-enqueueing jobs within a job
+        :param _priority: integer priority in the range 1..10 (1 = highest, 10 = lowest);
+            defaults to 5. Within the same priority, jobs run FIFO by enqueue time.
         :param kwargs: any keyword arguments to pass to the function
         :return: :class:`arq.jobs.Job` instance or ``None`` if a job with this ID already exists
         """
+        if not (MIN_PRIORITY <= _priority <= MAX_PRIORITY):
+            raise ValueError(f'_priority must be between {MIN_PRIORITY} and {MAX_PRIORITY}, got {_priority}')
         if _queue_name is None:
             _queue_name = self.default_queue_name
         job_id = _job_id or uuid4().hex
@@ -161,25 +177,80 @@ async def enqueue_job(
 
             enqueue_time_ms = timestamp_ms()
             if _defer_until is not None:
-                score = to_unix_ms(_defer_until)
+                run_at_ms = to_unix_ms(_defer_until)
+                deferred = True
             elif defer_by_ms:
-                score = enqueue_time_ms + defer_by_ms
+                run_at_ms = enqueue_time_ms + defer_by_ms
+                deferred = True
             else:
-                score = enqueue_time_ms
-
-            expires_ms = expires_ms or score - enqueue_time_ms + self.expires_extra_ms
-
-            job = serialize_job(function, args, kwargs, _job_try, enqueue_time_ms, serializer=self.job_serializer)
+                run_at_ms = enqueue_time_ms
+                deferred = False
+
+            expires_ms = expires_ms or run_at_ms - enqueue_time_ms + self.expires_extra_ms
+
+            job = serialize_job(
+                function,
+                args,
+                kwargs,
+                _job_try,
+                enqueue_time_ms,
+                priority=_priority,
+                serializer=self.job_serializer,
+            )
             pipe.multi()
             pipe.psetex(job_key, expires_ms, job)
-            pipe.zadd(_queue_name, {job_id: score})
+            if deferred:
+                pipe.zadd(deferred_queue_key(_queue_name), {job_id: run_at_ms})
+            else:
+                pipe.zadd(
+                    ready_queue_key(_queue_name),
+                    {job_id: compute_ready_score(_priority, enqueue_time_ms)},
+                )
             try:
                 await pipe.execute()
             except WatchError:
                 # job got enqueued since we checked 'job_exists'
                 return None
         return Job(job_id, redis=self, _queue_name=_queue_name, _deserializer=self.job_deserializer)
 
+    async def queue_size(self, queue_name: Optional[str] = None) -> int:
+        """
+        Total number of jobs in this queue across both the ready and deferred sub-zsets.
+        """
+        if queue_name is None:
+            queue_name = self.default_queue_name
+        async with self.pipeline(transaction=False) as pipe:
+            pipe.zcard(ready_queue_key(queue_name))
+            pipe.zcard(deferred_queue_key(queue_name))
+            ready, deferred = await pipe.execute()
+        return int(ready) + int(deferred)
+
+    async def job_queue_score(
+        self, queue_name: str, job_id: str
+    ) -> tuple[Optional[int], bool]:
+        """
+        Look up a job's score in the queue. Returns ``(score, is_deferred)``;
+        ``score`` is ``None`` when the job is in neither sub-zset.
+        """
+        async with self.pipeline(transaction=False) as pipe:
+            pipe.zscore(ready_queue_key(queue_name), job_id)
+            pipe.zscore(deferred_queue_key(queue_name), job_id)
+            ready_score, deferred_score = await pipe.execute()
+        if ready_score is not None:
+            return int(ready_score), False
+        if deferred_score is not None:
+            return int(deferred_score), True
+        return None, False
+
+    @staticmethod
+    def zrem_from_queue(pipe: Any, queue_name: str, job_id: str) -> None:
+        """
+        Stage ``ZREM`` against both sub-zsets on an existing pipeline.
+        Idempotent — safe even if the job only lives in one of them.
+        """
+        pipe.zrem(ready_queue_key(queue_name), job_id)
+        pipe.zrem(deferred_queue_key(queue_name), job_id)
+
     async def _get_job_result(self, key: bytes) -> JobResult:
         job_id = key[len(result_key_prefix) :].decode()
         job = Job(job_id, self, _deserializer=self.job_deserializer)
@@ -209,11 +280,16 @@ async def _get_job_def(self, job_id: bytes, score: int) -> JobDef:
 
     async def queued_jobs(self, *, queue_name: Optional[str] = None) -> list[JobDef]:
         """
-        Get information about queued, mostly useful when testing.
+        Get information about queued jobs across both the ready and deferred sub-zsets.
+        Mostly useful when testing.
         """
         if queue_name is None:
             queue_name = self.default_queue_name
-        jobs = await self.zrange(queue_name, withscores=True, start=0, end=-1)
+        async with self.pipeline(transaction=False) as pipe:
+            pipe.zrange(ready_queue_key(queue_name), withscores=True, start=0, end=-1)
+            pipe.zrange(deferred_queue_key(queue_name), withscores=True, start=0, end=-1)
+            ready, deferred = await pipe.execute()
+        jobs = list(ready) + list(deferred)
         return await asyncio.gather(*[self._get_job_def(job_id, int(score)) for job_id, score in jobs])
 
 

arq/constants.py

@@ -16,3 +16,24 @@
 
 # extra time after the job is expected to start when the job key should expire, 1 day in ms
 expires_extra_ms = 86_400_000
+
+# priority constants. ready-zset score = PRIORITY_FACTOR * priority + timestamp_ms,
+# so lower priority value runs first; FIFO is preserved within a tie via the timestamp.
+# 1e15 keeps the priority bucket strictly above any realistic timestamp_ms; doubles still
+# resolve ms-level differences at this magnitude (gap ~2.0 at ~1e16).
+PRIORITY_FACTOR = 10**15
+DEFAULT_PRIORITY = 5
+MIN_PRIORITY = 1
+MAX_PRIORITY = 10
+
+
+def ready_queue_key(queue_name: str) -> str:
+    return f'{queue_name}:ready'
+
+
+def deferred_queue_key(queue_name: str) -> str:
+    return f'{queue_name}:deferred'
+
+
+def compute_ready_score(priority: int, timestamp_ms: int) -> int:
+    return PRIORITY_FACTOR * priority + timestamp_ms

arq/jobs.py

@@ -2,14 +2,25 @@
 import logging
 import pickle
 import warnings
-from dataclasses import dataclass
-from datetime import datetime
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Callable, Optional
 
 from redis.asyncio import Redis
 
-from .constants import abort_jobs_ss, default_queue_name, in_progress_key_prefix, job_key_prefix, result_key_prefix
+from .constants import (
+    DEFAULT_PRIORITY,
+    MIN_PRIORITY,
+    abort_jobs_ss,
+    compute_ready_score,
+    default_queue_name,
+    deferred_queue_key,
+    in_progress_key_prefix,
+    job_key_prefix,
+    ready_queue_key,
+    result_key_prefix,
+)
 from .utils import ms_to_datetime, poll, timestamp_ms
 
 logger = logging.getLogger('arq.jobs')
@@ -39,6 +50,9 @@ class JobStatus(str, Enum):
     not_found = 'not_found'
 
 
+_EPOCH = datetime.fromtimestamp(0, tz=timezone.utc)
+
+
 @dataclass
 class JobDef:
     function: str
@@ -48,6 +62,9 @@ class JobDef:
     enqueue_time: datetime
     score: Optional[int]
     job_id: Optional[str]
+    # priority defaults so existing constructor sites that predate the priority feature keep working;
+    # dataclass inheritance forces every JobResult-only field below to also carry a default.
+    priority: int = DEFAULT_PRIORITY
 
     def __post_init__(self) -> None:
         if isinstance(self.score, float):
@@ -56,11 +73,11 @@ def __post_init__(self) -> None:
 
 @dataclass
 class JobResult(JobDef):
-    success: bool
-    result: Any
-    start_time: datetime
-    finish_time: datetime
-    queue_name: str
+    success: bool = False
+    result: Any = None
+    start_time: datetime = field(default_factory=lambda: _EPOCH)
+    finish_time: datetime = field(default_factory=lambda: _EPOCH)
+    queue_name: str = ''
 
 
 class Job:
@@ -104,8 +121,9 @@ async def result(
         async for delay in poll(poll_delay):
             async with self._redis.pipeline(transaction=True) as tr:
                 tr.get(result_key_prefix + self.job_id)
-                tr.zscore(self._queue_name, self.job_id)
-                v, s = await tr.execute()
+                tr.zscore(ready_queue_key(self._queue_name), self.job_id)
+                tr.zscore(deferred_queue_key(self._queue_name), self.job_id)
+                v, ready_score, deferred_score = await tr.execute()
 
             if v:
                 info = deserialize_result(v, deserializer=self._deserializer)
@@ -115,7 +133,7 @@ async def result(
                     raise info.result
                 else:
                     raise SerializationError(info.result)
-            elif s is None:
+            elif ready_score is None and deferred_score is None:
                 raise ResultNotFound(
                     'Not waiting for job result because the job is not in queue. '
                     'Is the worker function configured to keep result?'
@@ -134,10 +152,25 @@ async def info(self) -> Optional[JobDef]:
             if v:
                 info = deserialize_job(v, deserializer=self._deserializer)
         if info:
-            s = await self._redis.zscore(self._queue_name, self.job_id)
-            info.score = None if s is None else int(s)
+            score, _ = await self._lookup_queue_score()
+            info.score = score
         return info
 
+    async def _lookup_queue_score(self) -> tuple[Optional[int], bool]:
+        """
+        Return ``(score, is_deferred)`` for this job. ``score`` is ``None`` when the job
+        isn't in either sub-zset.
+        """
+        async with self._redis.pipeline(transaction=False) as pipe:
+            pipe.zscore(ready_queue_key(self._queue_name), self.job_id)
+            pipe.zscore(deferred_queue_key(self._queue_name), self.job_id)
+            ready_score, deferred_score = await pipe.execute()
+        if ready_score is not None:
+            return int(ready_score), False
+        if deferred_score is not None:
+            return int(deferred_score), True
+        return None, False
+
     async def result_info(self) -> Optional[JobResult]:
         """
         Information about the job result if available, does not wait for the result. Does not raise an exception
@@ -156,15 +189,18 @@ async def status(self) -> JobStatus:
         async with self._redis.pipeline(transaction=True) as tr:
             tr.exists(result_key_prefix + self.job_id)
             tr.exists(in_progress_key_prefix + self.job_id)
-            tr.zscore(self._queue_name, self.job_id)
-            is_complete, is_in_progress, score = await tr.execute()
+            tr.zscore(ready_queue_key(self._queue_name), self.job_id)
+            tr.zscore(deferred_queue_key(self._queue_name), self.job_id)
+            is_complete, is_in_progress, ready_score, deferred_score = await tr.execute()
 
         if is_complete:
             return JobStatus.complete
         elif is_in_progress:
             return JobStatus.in_progress
-        elif score:
-            return JobStatus.deferred if score > timestamp_ms() else JobStatus.queued
+        elif ready_score is not None:
+            return JobStatus.queued
+        elif deferred_score is not None:
+            return JobStatus.deferred
         else:
             return JobStatus.not_found
 
@@ -177,11 +213,16 @@ async def abort(self, *, timeout: Optional[float] = None, poll_delay: float = 0.
         :param poll_delay: how often to poll redis for the job result
         :return: True if the job aborted properly, False otherwise
         """
-        job_info = await self.info()
-        if job_info and job_info.score and job_info.score > timestamp_ms():
+        # if the job is currently sitting in the deferred sub-zset, hoist it into ready at
+        # top priority so the worker picks it up on the next poll and observes the abort flag.
+        _, is_deferred = await self._lookup_queue_score()
+        if is_deferred:
             async with self._redis.pipeline(transaction=True) as tr:
-                tr.zrem(self._queue_name, self.job_id)
-                tr.zadd(self._queue_name, {self.job_id: 1})
+                tr.zrem(deferred_queue_key(self._queue_name), self.job_id)
+                tr.zadd(
+                    ready_queue_key(self._queue_name),
+                    {self.job_id: compute_ready_score(MIN_PRIORITY, timestamp_ms())},
+                )
                 await tr.execute()
 
         await self._redis.zadd(abort_jobs_ss, {self.job_id: timestamp_ms()})
@@ -215,9 +256,17 @@ def serialize_job(
     job_try: Optional[int],
     enqueue_time_ms: int,
     *,
+    priority: int = DEFAULT_PRIORITY,
     serializer: Optional[Serializer] = None,
 ) -> bytes:
-    data = {'t': job_try, 'f': function_name, 'a': args, 'k': kwargs, 'et': enqueue_time_ms}
+    data = {
+        't': job_try,
+        'f': function_name,
+        'a': args,
+        'k': kwargs,
+        'et': enqueue_time_ms,
+        'priority': priority,
+    }
     if serializer is None:
         serializer = pickle.dumps
     try:
@@ -240,6 +289,7 @@ def serialize_result(
     queue_name: str,
     job_id: str,
     *,
+    priority: int = DEFAULT_PRIORITY,
     serializer: Optional[Serializer] = None,
 ) -> Optional[bytes]:
     data = {
@@ -254,6 +304,7 @@ def serialize_result(
         'ft': finished_ms,
         'q': queue_name,
         'id': job_id,
+        'priority': priority,
     }
     if serializer is None:
         serializer = pickle.dumps
@@ -284,19 +335,20 @@ def deserialize_job(r: bytes, *, deserializer: Optional[Deserializer] = None) ->
             enqueue_time=ms_to_datetime(d['et']),
             score=None,
             job_id=None,
+            priority=d.get('priority', DEFAULT_PRIORITY),
         )
     except Exception as e:
         raise DeserializationError('unable to deserialize job') from e
 
 
 def deserialize_job_raw(
     r: bytes, *, deserializer: Optional[Deserializer] = None
-) -> tuple[str, tuple[Any, ...], dict[str, Any], int, int]:
+) -> tuple[str, tuple[Any, ...], dict[str, Any], int, int, int]:
     if deserializer is None:
         deserializer = pickle.loads
     try:
         d = deserializer(r)
-        return d['f'], d['a'], d['k'], d['t'], d['et']
+        return d['f'], d['a'], d['k'], d['t'], d['et'], d.get('priority', DEFAULT_PRIORITY)
     except Exception as e:
         raise DeserializationError('unable to deserialize job') from e
 
@@ -319,6 +371,7 @@ def deserialize_result(r: bytes, *, deserializer: Optional[Deserializer] = None)
             finish_time=ms_to_datetime(d['ft']),
             queue_name=d.get('q', '<unknown>'),
             job_id=d.get('id', '<unknown>'),
+            priority=d.get('priority', DEFAULT_PRIORITY),
         )
     except Exception as e:
         raise DeserializationError('unable to deserialize job result') from e