Avoid per-job status requests on list views and limit concurrency (#9)

Co-authored-by: Slava Skvortsov <29122694+SlavaSkvortsov@users.noreply.github.com>

Author: fedtf

arq_admin/queue.py

@@ -1,17 +1,20 @@
 import asyncio
 from contextlib import suppress
 from dataclasses import dataclass, field
-from typing import Any, List, Optional, Set
+from typing import Any, Dict, List, Optional
 
 from arq import ArqRedis
 from arq.connections import RedisSettings, create_pool
-from arq.constants import result_key_prefix
 from arq.jobs import DeserializationError, Job as ArqJob, JobDef, JobStatus
+from arq.utils import timestamp_ms
 from django.utils import timezone
 
 from arq_admin import settings
 from arq_admin.job import JobInfo
 
+ARQ_PREFIX = 'arq:'
+PREFIX_PRIORITY = {prefix: i for i, prefix in enumerate(['job', 'in-progress', 'result'])}
+
 
 @dataclass
 class QueueStats:
@@ -31,6 +34,10 @@ class QueueStats:
 class Queue:
     redis_settings: RedisSettings
     name: str
+    concurrent_redis_access_sem: asyncio.Semaphore = field(
+        default_factory=lambda: asyncio.Semaphore(settings.ARQ_MAX_CONNECTIONS),
+    )
+    _cached_job_id_to_status_map: Optional[Dict[str, JobStatus]] = None
     _redis: ArqRedis = field(init=False, default=None)  # type: ignore
 
     async def __aenter__(self) -> 'Queue':
@@ -48,12 +55,12 @@ def from_name(cls, name: str) -> 'Queue':
         )
 
     async def get_jobs(self, status: Optional[JobStatus] = None) -> List[JobInfo]:
-        job_ids = await self._get_job_ids()
+        job_id_to_status_map = await self._get_job_id_to_status_map()
 
         if status:
-            job_ids_tuple = tuple(job_ids)
-            statuses = await asyncio.gather(*[self._get_job_status(job_id) for job_id in job_ids_tuple])
-            job_ids = {job_id for (job_id, job_status) in zip(job_ids_tuple, statuses) if job_status == status}
+            job_ids = {job_id for (job_id, job_status) in job_id_to_status_map.items() if job_status == status}
+        else:
+            job_ids = set(job_id_to_status_map.keys())
 
         jobs: List[JobInfo] = await asyncio.gather(*[self.get_job_by_id(job_id) for job_id in job_ids])
 
@@ -66,12 +73,14 @@ async def get_stats(self) -> QueueStats:
             port=self.redis_settings.port,
             database=self.redis_settings.database,
         )
+
         try:
-            job_ids = await self._get_job_ids()
-            statuses = await asyncio.gather(*[self._get_job_status(job_id) for job_id in job_ids])
+            job_id_to_status_map = await self._get_job_id_to_status_map()
         except Exception as ex:  # noqa: B902
             result.error = str(ex)
         else:
+            statuses = job_id_to_status_map.values()
+
             result.queued_jobs = len([status for status in statuses if status == JobStatus.queued])
             result.running_jobs = len([status for status in statuses if status == JobStatus.in_progress])
             result.deferred_jobs = len([status for status in statuses if status == JobStatus.deferred])
@@ -88,10 +97,11 @@ async def get_job_by_id(self, job_id: str) -> JobInfo:
 
         unknown_function_msg = "Can't find job"
         base_info = None
-        try:
-            base_info = await arq_job.info()
-        except DeserializationError:
-            unknown_function_msg = "Unknown, can't deserialize"
+        async with self.concurrent_redis_access_sem:
+            try:
+                base_info = await arq_job.info()
+            except DeserializationError:
+                unknown_function_msg = "Unknown, can't deserialize"
 
         if not base_info:
             base_info = JobDef(
@@ -104,7 +114,7 @@ async def get_job_by_id(self, job_id: str) -> JobInfo:
             )
 
         job_info = JobInfo.from_base(base_info, job_id)
-        job_info.status = await arq_job.status()
+        job_info.status = await self._get_job_status(job_id)
 
         return job_info
 
@@ -122,17 +132,53 @@ async def abort_job(self, job_id: str) -> Optional[bool]:
         return None
 
     async def _get_job_status(self, job_id: str) -> JobStatus:
+        if self._cached_job_id_to_status_map is not None:
+            return self._cached_job_id_to_status_map.get(job_id, JobStatus.not_found)
+
         arq_job = ArqJob(
             job_id=job_id,
             redis=self._redis,
             _queue_name=self.name,
             _deserializer=settings.ARQ_DESERIALIZER_BY_QUEUE.get(self.name),
         )
-        return await arq_job.status()
-
-    async def _get_job_ids(self) -> Set[str]:
-        raw_job_ids = set(await self._redis.zrangebyscore(self.name, '-inf', 'inf'))
-        result_keys = await self._redis.keys(f'{result_key_prefix}*')
-        raw_job_ids |= {key[len(result_key_prefix):] for key in result_keys}
+        async with self.concurrent_redis_access_sem:
+            return await arq_job.status()
+
+    async def _get_job_id_to_status_map(self) -> Dict[str, JobStatus]:
+        if self._cached_job_id_to_status_map is not None:
+            return self._cached_job_id_to_status_map
+
+        async with self._redis.pipeline(transaction=True) as pipe:
+            await pipe.keys(f'{ARQ_PREFIX}*:*')
+            await pipe.zrange(self.name, withscores=True, start=0, end=-1)
+            all_arq_keys, job_ids_with_scores = await pipe.execute()
+
+        # iter over lists of type [job_id, prefix];
+        # can't use dict here because we can have multiple keys for one job and need to use the more specific one
+        job_ids_with_prefixes = (
+            key.decode('utf-8')[len(ARQ_PREFIX):].split(':')[::-1] for key in all_arq_keys
+        )
 
-        return {job_id.decode('utf-8') if isinstance(job_id, bytes) else job_id for job_id in raw_job_ids}
+        job_ids_to_scores = {key[0].decode('utf-8'): key[1] for key in job_ids_with_scores}
+        job_ids_to_prefixes = dict(sorted(
+            # not only ensure that we don't get key error but also filter out stuff that's not a client job
+            ([job_id, prefix] for job_id, prefix in job_ids_with_prefixes if prefix in PREFIX_PRIORITY),
+            # make sure that more specific indices go after less specific ones
+            key=lambda job_id_with_prefix: PREFIX_PRIORITY[job_id_with_prefix[-1]],
+        ))
+
+        self._cached_job_id_to_status_map = {
+            job_id: self._get_job_status_from_raw_data(prefix, job_ids_to_scores.get(job_id))
+            for job_id, prefix in job_ids_to_prefixes.items()
+        }
+
+        return self._cached_job_id_to_status_map
+
+    def _get_job_status_from_raw_data(self, prefix: str, zscore: Optional[int]) -> JobStatus:  # noqa: CFQ004
+        if prefix == 'result':
+            return JobStatus.complete
+        if prefix == 'in-progress' and zscore:
+            return JobStatus.in_progress
+        if zscore:
+            return JobStatus.deferred if zscore > timestamp_ms() else JobStatus.queued
+        return JobStatus.not_found

arq_admin/settings.py

@@ -26,3 +26,5 @@
     ARQ_DESERIALIZER_BY_QUEUE = defaultdict(lambda: ARQ_DESERIALIZER)
 
 ARQ_JOB_ABORT_TIMEOUT = getattr(settings, 'ARQ_JOB_ABORT_TIMEOUT', 5)
+
+ARQ_MAX_CONNECTIONS = getattr(settings, 'ARQ_MAX_CONNECTIONS', 100)

requirements.txt

@@ -1,2 +1,2 @@
-arq==0.24.0
+arq==0.25.0
 Django==4.1.2

tests/test_queue.py

@@ -4,6 +4,7 @@
 
 import pytest
 import pytest_asyncio
+from arq import ArqRedis
 from arq.constants import default_queue_name
 from arq.jobs import DeserializationError, Job, JobStatus
 from django.conf import settings
@@ -49,7 +50,23 @@ async def test_stats(queue: Queue) -> None:
 
 
 @pytest.mark.asyncio()
-@patch.object(Queue, '_get_job_ids')
+@pytest.mark.usefixtures('all_jobs')
+async def test_stats_with_running_job_wo_zscore(redis: ArqRedis, queue: Queue) -> None:
+    await redis.zrem(queue.name, 'running_task')
+
+    assert await queue.get_stats() == QueueStats(
+        name=default_queue_name,
+        host=settings.REDIS_SETTINGS.host,
+        port=settings.REDIS_SETTINGS.port,
+        database=settings.REDIS_SETTINGS.database,
+        queued_jobs=1,
+        running_jobs=0,
+        deferred_jobs=1,
+    )
+
+
+@pytest.mark.asyncio()
+@patch.object(Queue, '_get_job_id_to_status_map')
 async def test_stats_with_error(mocked_get_job_ids: AsyncMock, queue: Queue) -> None:
     error_text = 'test error'
     mocked_get_job_ids.side_effect = Exception(error_text)