support aioredis v2 (#259)

* support aioredis 2

* use psetex for float ttl

* health_check_interval + 1 sec

* fix health_check_interval psetex

* encoder context to encode/decode on per-request basis

* fix hiredis parser

* remove excessive comments

* update docstrings, pipe commands in run_job, gather in enqueue_job, expire in ms

* fix mypy, tweak tests

* using latest aioredis

* switch how latest aioredis is installed

* fix install

* fixing worker tests removing encoder_context usage

* removing usage of encoder_context

* removing parser and encoder_context

* tweak test_abort_job

* tests at last passing!

* switch to redis from aioredis

* cleanup

Co-authored-by: Samuel Colvin <[email protected]>

Author: Yolley

.gitignore

@@ -17,3 +17,4 @@ __pycache__/
 .vscode/
 .venv/
 /.auto-format
+/scratch/

arq/connections.py

@@ -9,9 +9,10 @@
 from urllib.parse import urlparse
 from uuid import uuid4
 
-import aioredis
-from aioredis import MultiExecError, Redis
 from pydantic.validators import make_arbitrary_type_validator
+from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio.sentinel import Sentinel
+from redis.exceptions import RedisError, WatchError
 
 from .constants import default_queue_name, job_key_prefix, result_key_prefix
 from .jobs import Deserializer, Job, JobDef, JobResult, Serializer, deserialize_job, serialize_job
@@ -70,20 +71,20 @@ def __repr__(self) -> str:
 expires_extra_ms = 86_400_000
 
 
-class ArqRedis(Redis):  # type: ignore
+class ArqRedis(Redis):  # type: ignore[misc]
     """
-    Thin subclass of ``aioredis.Redis`` which adds :func:`arq.connections.enqueue_job`.
+    Thin subclass of ``redis.asyncio.Redis`` which adds :func:`arq.connections.enqueue_job`.
 
     :param redis_settings: an instance of ``arq.connections.RedisSettings``.
     :param job_serializer: a function that serializes Python objects to bytes, defaults to pickle.dumps
     :param job_deserializer: a function that deserializes bytes into Python objects, defaults to pickle.loads
     :param default_queue_name: the default queue name to use, defaults to ``arq.queue``.
-    :param kwargs: keyword arguments directly passed to ``aioredis.Redis``.
+    :param kwargs: keyword arguments directly passed to ``redis.asyncio.Redis``.
     """
 
     def __init__(
         self,
-        pool_or_conn: Any,
+        pool_or_conn: Optional[ConnectionPool] = None,
         job_serializer: Optional[Serializer] = None,
         job_deserializer: Optional[Deserializer] = None,
         default_queue_name: str = default_queue_name,
@@ -92,7 +93,9 @@ def __init__(
         self.job_serializer = job_serializer
         self.job_deserializer = job_deserializer
         self.default_queue_name = default_queue_name
-        super().__init__(pool_or_conn, **kwargs)
+        if pool_or_conn:
+            kwargs['connection_pool'] = pool_or_conn
+        super().__init__(**kwargs)
 
     async def enqueue_job(
         self,
@@ -129,14 +132,10 @@ async def enqueue_job(
         defer_by_ms = to_ms(_defer_by)
         expires_ms = to_ms(_expires)
 
-        with await self as conn:
-            pipe = conn.pipeline()
-            pipe.unwatch()
-            pipe.watch(job_key)
-            job_exists = pipe.exists(job_key)
-            job_result_exists = pipe.exists(result_key_prefix + job_id)
-            await pipe.execute()
-            if await job_exists or await job_result_exists:
+        async with self.pipeline(transaction=True) as pipe:
+            await pipe.watch(job_key)
+            if any(await asyncio.gather(pipe.exists(job_key), pipe.exists(result_key_prefix + job_id))):
+                await pipe.reset()
                 return None
 
             enqueue_time_ms = timestamp_ms()
@@ -150,24 +149,22 @@ async def enqueue_job(
             expires_ms = expires_ms or score - enqueue_time_ms + expires_extra_ms
 
             job = serialize_job(function, args, kwargs, _job_try, enqueue_time_ms, serializer=self.job_serializer)
-            tr = conn.multi_exec()
-            tr.psetex(job_key, expires_ms, job)
-            tr.zadd(_queue_name, score, job_id)
+            pipe.multi()
+            pipe.psetex(job_key, expires_ms, job)
+            pipe.zadd(_queue_name, {job_id: score})
             try:
-                await tr.execute()
-            except MultiExecError:
+                await pipe.execute()
+            except WatchError:
                 # job got enqueued since we checked 'job_exists'
-                # https://github.com/samuelcolvin/arq/issues/131, avoid warnings in log
-                await asyncio.gather(*tr._results, return_exceptions=True)
                 return None
         return Job(job_id, redis=self, _queue_name=_queue_name, _deserializer=self.job_deserializer)
 
-    async def _get_job_result(self, key: str) -> JobResult:
-        job_id = key[len(result_key_prefix) :]
+    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)
         r = await job.result_info()
         if r is None:
-            raise KeyError(f'job "{key}" not found')
+            raise KeyError(f'job "{key.decode()}" not found')
         r.job_id = job_id
         return r
 
@@ -179,8 +176,8 @@ async def all_job_results(self) -> List[JobResult]:
         results = await asyncio.gather(*[self._get_job_result(k) for k in keys])
         return sorted(results, key=attrgetter('enqueue_time'))
 
-    async def _get_job_def(self, job_id: str, score: int) -> JobDef:
-        v = await self.get(job_key_prefix + job_id, encoding=None)
+    async def _get_job_def(self, job_id: bytes, score: int) -> JobDef:
+        v = await self.get(job_key_prefix + job_id.decode())
         jd = deserialize_job(v, deserializer=self.job_deserializer)
         jd.score = score
         return jd
@@ -189,8 +186,8 @@ async def queued_jobs(self, *, queue_name: str = default_queue_name) -> List[Job
         """
         Get information about queued, mostly useful when testing.
         """
-        jobs = await self.zrange(queue_name, withscores=True)
-        return await asyncio.gather(*[self._get_job_def(job_id, score) for job_id, score in jobs])
+        jobs = await self.zrange(queue_name, withscores=True, start=0, end=-1)
+        return await asyncio.gather(*[self._get_job_def(job_id, int(score)) for job_id, score in jobs])
 
 
 async def create_pool(
@@ -204,8 +201,7 @@ async def create_pool(
     """
     Create a new redis pool, retrying up to ``conn_retries`` times if the connection fails.
 
-    Similar to ``aioredis.create_redis_pool`` except it returns a :class:`arq.connections.ArqRedis` instance,
-    thus allowing job enqueuing.
+    Returns a :class:`arq.connections.ArqRedis` instance, thus allowing job enqueuing.
     """
     settings: RedisSettings = RedisSettings() if settings_ is None else settings_
 
@@ -214,32 +210,33 @@ async def create_pool(
     ), "str provided for 'host' but 'sentinel' is true; list of sentinels expected"
 
     if settings.sentinel:
-        addr: Any = settings.host
 
-        async def pool_factory(*args: Any, **kwargs: Any) -> Redis:
-            client = await aioredis.sentinel.create_sentinel_pool(*args, ssl=settings.ssl, **kwargs)
-            return client.master_for(settings.sentinel_master)
+        def pool_factory(*args: Any, **kwargs: Any) -> ArqRedis:
+            client = Sentinel(*args, sentinels=settings.host, ssl=settings.ssl, **kwargs)
+            return client.master_for(settings.sentinel_master, redis_class=ArqRedis)
 
     else:
         pool_factory = functools.partial(
-            aioredis.create_pool, create_connection_timeout=settings.conn_timeout, ssl=settings.ssl
+            ArqRedis,
+            host=settings.host,
+            port=settings.port,
+            socket_connect_timeout=settings.conn_timeout,
+            ssl=settings.ssl,
         )
-        addr = settings.host, settings.port
 
     try:
-        pool = await pool_factory(addr, db=settings.database, password=settings.password, encoding='utf8')
-        pool = ArqRedis(
-            pool,
-            job_serializer=job_serializer,
-            job_deserializer=job_deserializer,
-            default_queue_name=default_queue_name,
-        )
+        pool = pool_factory(db=settings.database, password=settings.password, encoding='utf8')
+        pool.job_serializer = job_serializer
+        pool.job_deserializer = job_deserializer
+        pool.default_queue_name = default_queue_name
+        await pool.ping()
 
-    except (ConnectionError, OSError, aioredis.RedisError, asyncio.TimeoutError) as e:
+    except (ConnectionError, OSError, RedisError, asyncio.TimeoutError) as e:
         if retry < settings.conn_retries:
             logger.warning(
-                'redis connection error %s %s %s, %d retries remaining...',
-                addr,
+                'redis connection error %s:%s %s %s, %d retries remaining...',
+                settings.host,
+                settings.port,
                 e.__class__.__name__,
                 e,
                 settings.conn_retries - retry,
@@ -264,17 +261,16 @@ async def pool_factory(*args: Any, **kwargs: Any) -> Redis:
 
 
 async def log_redis_info(redis: Redis, log_func: Callable[[str], Any]) -> None:
-    with await redis as r:
-        info_server, info_memory, info_clients, key_count = await asyncio.gather(
-            r.info(section='Server'),
-            r.info(section='Memory'),
-            r.info(section='Clients'),
-            r.dbsize(),
-        )
-
-    redis_version = info_server.get('server', {}).get('redis_version', '?')
-    mem_usage = info_memory.get('memory', {}).get('used_memory_human', '?')
-    clients_connected = info_clients.get('clients', {}).get('connected_clients', '?')
+    async with redis.pipeline(transaction=True) as pipe:
+        pipe.info(section='Server')
+        pipe.info(section='Memory')
+        pipe.info(section='Clients')
+        pipe.dbsize()
+        info_server, info_memory, info_clients, key_count = await pipe.execute()
+
+    redis_version = info_server.get('redis_version', '?')
+    mem_usage = info_memory.get('used_memory_human', '?')
+    clients_connected = info_clients.get('connected_clients', '?')
 
     log_func(
         f'redis_version={redis_version} '

arq/jobs.py

@@ -7,7 +7,7 @@
 from enum import Enum
 from typing import Any, Callable, Dict, Optional, Tuple
 
-from aioredis import Redis
+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 .utils import ms_to_datetime, poll, timestamp_ms
@@ -44,6 +44,10 @@ class JobDef:
     enqueue_time: datetime
     score: Optional[int]
 
+    def __post_init__(self) -> None:
+        if isinstance(self.score, float):
+            self.score = int(self.score)
+
 
 @dataclass
 class JobResult(JobDef):
@@ -110,7 +114,7 @@ async def info(self) -> Optional[JobDef]:
         """
         info: Optional[JobDef] = await self.result_info()
         if not info:
-            v = await self._redis.get(job_key_prefix + self.job_id, encoding=None)
+            v = await self._redis.get(job_key_prefix + self.job_id)
             if v:
                 info = deserialize_job(v, deserializer=self._deserializer)
         if info:
@@ -122,7 +126,7 @@ 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
         even if the job raised one.
         """
-        v = await self._redis.get(result_key_prefix + self.job_id, encoding=None)
+        v = await self._redis.get(result_key_prefix + self.job_id)
         if v:
             return deserialize_result(v, deserializer=self._deserializer)
         else:
@@ -151,7 +155,7 @@ 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
         """
-        await self._redis.zadd(abort_jobs_ss, timestamp_ms(), self.job_id)
+        await self._redis.zadd(abort_jobs_ss, {self.job_id: timestamp_ms()})
         try:
             await self.result(timeout=timeout, poll_delay=poll_delay)
         except asyncio.CancelledError:
@@ -179,7 +183,7 @@ def serialize_job(
     enqueue_time_ms: int,
     *,
     serializer: Optional[Serializer] = None,
-) -> Optional[bytes]:
+) -> bytes:
     data = {'t': job_try, 'f': function_name, 'a': args, 'k': kwargs, 'et': enqueue_time_ms}
     if serializer is None:
         serializer = pickle.dumps