43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
import asyncio
|
|
from typing import Any, Dict, List
|
|
|
|
import aioredis
|
|
from channels_redis.core import ConnectionPool as ChannelRedisConnectionPool
|
|
from django.conf import settings
|
|
|
|
from . import logging
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
connection_pool_limit = getattr(settings, "CONNECTION_POOL_LIMIT", 100)
|
|
logger.info(f"CONNECTION_POOL_LIMIT={connection_pool_limit}")
|
|
|
|
|
|
class ConnectionPool(ChannelRedisConnectionPool):
|
|
""" Adds a trivial, soft limit for the pool """
|
|
|
|
def __init__(self, host: Any) -> None:
|
|
self.counter = 0
|
|
super().__init__(host)
|
|
|
|
async def pop(
|
|
self, *args: List[Any], **kwargs: Dict[str, Any]
|
|
) -> aioredis.commands.Redis:
|
|
while self.counter > connection_pool_limit:
|
|
await asyncio.sleep(0.1)
|
|
self.counter += 1
|
|
|
|
return await super().pop(*args, **kwargs)
|
|
|
|
def push(self, conn: aioredis.commands.Redis) -> None:
|
|
super().push(conn)
|
|
self.counter -= 1
|
|
|
|
def conn_error(self, conn: aioredis.commands.Redis) -> None:
|
|
super().conn_error(conn)
|
|
self.counter -= 1
|
|
|
|
def reset(self) -> None:
|
|
super().reset()
|
|
self.counter = 0
|