Replace ConnectionPool with BlockingConnectionPool

Author: stinovlas

taskiq_redis/redis_broker.py

@@ -1,7 +1,7 @@
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -31,14 +31,16 @@ def __init__(
         :param result_backend: custom result backend.
         :param queue_name: name for a list in redis.
         :param max_connection_pool_size: maximum number of connections in pool.
-        :param connection_kwargs: additional arguments for aio-redis ConnectionPool.
+            Each worker opens its own connection. Therefore this value has to be
+            at least number of workers + 1.
+        :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
         """
         super().__init__(
             result_backend=result_backend,
             task_id_generator=task_id_generator,
         )
 
-        self.connection_pool: ConnectionPool = ConnectionPool.from_url(
+        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,

taskiq_redis/schedule_source.py

@@ -1,6 +1,6 @@
 from typing import Any, List, Optional
 
-from redis.asyncio import ConnectionPool, Redis, RedisCluster
+from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis, RedisCluster
 from taskiq import ScheduleSource
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
@@ -22,7 +22,7 @@ class RedisScheduleSource(ScheduleSource):
         This is how many keys will be fetched at once.
     :param max_connection_pool_size: maximum number of connections in pool.
     :param serializer: serializer for data.
-    :param connection_kwargs: additional arguments for aio-redis ConnectionPool.
+    :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
     """
 
     def __init__(
@@ -35,7 +35,7 @@ def __init__(
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.connection_pool: ConnectionPool = ConnectionPool.from_url(
+        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,

tests/test_broker.py

@@ -71,6 +71,29 @@ async def test_pub_sub_broker(
     await broker.shutdown()
 
 
+@pytest.mark.anyio
+async def test_pub_sub_broker_max_connections(
+    valid_broker_message: BrokerMessage,
+    redis_url: str,
+) -> None:
+    """Test PubSubBroker with connection limit set."""
+    broker = PubSubBroker(
+        url=redis_url,
+        queue_name=uuid.uuid4().hex,
+        max_connection_pool_size=4,
+        timeout=1,
+    )
+    worker_tasks = [asyncio.create_task(get_message(broker)) for _ in range(3)]
+    await asyncio.sleep(0.3)
+
+    await asyncio.gather(*[broker.kick(valid_broker_message) for _ in range(50)])
+    await asyncio.sleep(0.3)
+
+    for worker in worker_tasks:
+        worker.cancel()
+    await broker.shutdown()
+
+
 @pytest.mark.anyio
 async def test_list_queue_broker(
     valid_broker_message: BrokerMessage,
@@ -98,6 +121,29 @@ async def test_list_queue_broker(
     await broker.shutdown()
 
 
+@pytest.mark.anyio
+async def test_list_queue_broker_max_connections(
+    valid_broker_message: BrokerMessage,
+    redis_url: str,
+) -> None:
+    """Test ListQueueBroker with connection limit set."""
+    broker = ListQueueBroker(
+        url=redis_url,
+        queue_name=uuid.uuid4().hex,
+        max_connection_pool_size=4,
+        timeout=1,
+    )
+    worker_tasks = [asyncio.create_task(get_message(broker)) for _ in range(3)]
+    await asyncio.sleep(0.3)
+
+    await asyncio.gather(*[broker.kick(valid_broker_message) for _ in range(50)])
+    await asyncio.sleep(0.3)
+
+    for worker in worker_tasks:
+        worker.cancel()
+    await broker.shutdown()
+
+
 @pytest.mark.anyio
 async def test_list_queue_cluster_broker(
     valid_broker_message: BrokerMessage,

tests/test_schedule_source.py

@@ -1,3 +1,4 @@
+import asyncio
 import datetime as dt
 import uuid
 
@@ -108,6 +109,18 @@ async def test_buffer(redis_url: str) -> None:
     await source.shutdown()
 
 
+@pytest.mark.anyio
+async def test_max_connections(redis_url: str) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisScheduleSource(
+        redis_url,
+        prefix=prefix,
+        max_connection_pool_size=1,
+        timeout=3,
+    )
+    await asyncio.gather(*[source.get_schedules() for _ in range(10)])
+
+
 @pytest.mark.anyio
 async def test_cluster_set_schedule(redis_cluster_url: str) -> None:
     prefix = uuid.uuid4().hex