ConnectionPool.pop must ensure active connections are returned. (#178)

Co-authored-by: originmarkets-antoniodudarev <[email protected]>

Author: 2 people

channels_redis/core.py

@@ -69,6 +69,9 @@ async def pop(self, loop=None):
         if not conns:
             conns.append(await aioredis.create_redis(**self.host, loop=loop))
         conn = conns.pop()
+        if conn.closed:
+            conn = await self.pop(loop=loop)
+            return conn
         self.in_use[conn] = loop
         return conn
 
@@ -615,9 +618,11 @@ async def group_send(self, group, message):
                 x.decode("utf8") for x in await connection.zrange(key, 0, -1)
             ]
 
-        connection_to_channel_keys, channel_keys_to_message, channel_keys_to_capacity = self._map_channel_keys_to_connection(
-            channel_names, message
-        )
+        (
+            connection_to_channel_keys,
+            channel_keys_to_message,
+            channel_keys_to_capacity,
+        ) = self._map_channel_keys_to_connection(channel_names, message)
 
         for connection_index, channel_redis_keys in connection_to_channel_keys.items():
 

tests/test_core.py

@@ -5,7 +5,7 @@
 from async_generator import async_generator, yield_
 
 from asgiref.sync import async_to_sync
-from channels_redis.core import ChannelFull, RedisChannelLayer
+from channels_redis.core import ChannelFull, ConnectionPool, RedisChannelLayer
 
 TEST_HOSTS = [("localhost", 6379)]
 
@@ -323,6 +323,33 @@ async def test_group_send_capacity(channel_layer):
             await channel_layer.receive(channel)
 
 
+@pytest.mark.asyncio
+async def test_connection_pool_pop():
+    """
+    Makes sure that the connection pool does not return closed connections
+    """
+
+    # Setup scenario
+    connection_pool = ConnectionPool({"address": TEST_HOSTS[0]})
+    conn = await connection_pool.pop()
+
+    # Emualte a disconnect and return it to the pool
+    conn.close()
+    assert conn.closed == True
+    connection_pool.push(conn)
+
+    # Ensure the closed connection is inside the pool
+    conn_map_values = list(connection_pool.conn_map.values())
+    assert len(conn_map_values) == 1
+    conns = conn_map_values[0]
+    assert len(conns) == 1
+    assert conns[0].closed == True
+
+    # Retrieve new connection
+    conn = await connection_pool.pop()
+    assert conn.closed == False
+
+
 @pytest.mark.asyncio
 async def test_receive_cancel(channel_layer):
     """