Merge pull request #257 from zumalabs/pubsub-sentinel

Add sentinel support to pubsub layer

Author: acu192

channels_redis/pubsub.py

@@ -190,10 +190,12 @@ def on_close_noop(sender, exc=None):
 
 class RedisSingleShardConnection:
     def __init__(self, host, channel_layer):
-        self.host = host
+        self.host = host.copy() if type(host) is dict else {"address": host}
+        self.master_name = self.host.pop("master_name", None)
         self.channel_layer = channel_layer
         self._subscribed_to = set()
         self._lock = None
+        self._redis = None
         self._pub_conn = None
         self._sub_conn = None
         self._receiver = None
@@ -230,10 +232,12 @@ async def flush(self):
         if self._sub_conn is not None:
             self._sub_conn.close()
             await self._sub_conn.wait_closed()
+            self._put_redis_conn(self._sub_conn)
             self._sub_conn = None
         if self._pub_conn is not None:
             self._pub_conn.close()
             await self._pub_conn.wait_closed()
+            self._put_redis_conn(self._pub_conn)
             self._pub_conn = None
         self._subscribed_to = set()
 
@@ -247,11 +251,13 @@ async def _get_pub_conn(self):
             self._lock = asyncio.Lock()
         async with self._lock:
             if self._pub_conn is not None and self._pub_conn.closed:
+                self._put_redis_conn(self._pub_conn)
                 self._pub_conn = None
             while self._pub_conn is None:
                 try:
-                    self._pub_conn = await aioredis.create_redis(self.host)
+                    self._pub_conn = await self._get_redis_conn()
                 except BaseException:
+                    self._put_redis_conn(self._pub_conn)
                     logger.warning(
                         f"Failed to connect to Redis publish host: {self.host}; will try again in 1 second..."
                     )
@@ -270,6 +276,7 @@ async def _get_sub_conn(self):
             self._lock = asyncio.Lock()
         async with self._lock:
             if self._sub_conn is not None and self._sub_conn.closed:
+                self._put_redis_conn(self._sub_conn)
                 self._sub_conn = None
             if self._sub_conn is None:
                 if self._receive_task is not None:
@@ -287,8 +294,9 @@ async def _get_sub_conn(self):
                     self._receive_task = None
                 while self._sub_conn is None:
                     try:
-                        self._sub_conn = await aioredis.create_redis(self.host)
+                        self._sub_conn = await self._get_redis_conn()
                     except BaseException:
+                        self._put_redis_conn(self._sub_conn)
                         logger.warning(
                             f"Failed to connect to Redis subscribe host: {self.host}; will try again in 1 second..."
                         )
@@ -317,6 +325,31 @@ async def _do_receiving(self):
                     if channel_name in self.channel_layer.channels:
                         self.channel_layer.channels[channel_name].put_nowait(message)
 
+    async def _ensure_redis(self):
+        if self._redis is None:
+            if self.master_name is None:
+                self._redis = await aioredis.create_redis_pool(**self.host)
+            else:
+                # aioredis default timeout is way too low
+                self._redis = await aioredis.sentinel.create_sentinel(
+                    timeout=2, **self.host
+                )
+
+    def _get_aioredis_pool(self):
+        if self.master_name is None:
+            return self._redis._pool_or_conn
+        else:
+            return self._redis.master_for(self.master_name)._pool_or_conn
+
+    async def _get_redis_conn(self):
+        await self._ensure_redis()
+        conn = await self._get_aioredis_pool().acquire()
+        return aioredis.Redis(conn)
+
+    def _put_redis_conn(self, conn):
+        if conn:
+            self._get_aioredis_pool().release(conn._pool_or_conn)
+
     async def _do_keepalive(self):
         """
         This task's simple job is just to call `self._get_sub_conn()` periodically.

tests/test_pubsub_sentinel.py

@@ -0,0 +1,104 @@
+import asyncio
+import random
+
+import async_timeout
+import pytest
+from async_generator import async_generator, yield_
+
+from channels_redis.pubsub import RedisPubSubChannelLayer
+
+SENTINEL_MASTER = "sentinel"
+TEST_HOSTS = [{"sentinels": [("localhost", 26379)], "master_name": SENTINEL_MASTER}]
+
+
+@pytest.fixture()
+@async_generator
+async def channel_layer():
+    """
+    Channel layer fixture that flushes automatically.
+    """
+    channel_layer = RedisPubSubChannelLayer(hosts=TEST_HOSTS)
+    await yield_(channel_layer)
+    await channel_layer.flush()
+
+
+@pytest.mark.asyncio
+async def test_send_receive(channel_layer):
+    """
+    Makes sure we can send a message to a normal channel then receive it.
+    """
+    channel = await channel_layer.new_channel()
+    await channel_layer.send(channel, {"type": "test.message", "text": "Ahoy-hoy!"})
+    message = await channel_layer.receive(channel)
+    assert message["type"] == "test.message"
+    assert message["text"] == "Ahoy-hoy!"
+
+
+@pytest.mark.asyncio
+async def test_multi_send_receive(channel_layer):
+    """
+    Tests overlapping sends and receives, and ordering.
+    """
+    channel = await channel_layer.new_channel()
+    await channel_layer.send(channel, {"type": "message.1"})
+    await channel_layer.send(channel, {"type": "message.2"})
+    await channel_layer.send(channel, {"type": "message.3"})
+    assert (await channel_layer.receive(channel))["type"] == "message.1"
+    assert (await channel_layer.receive(channel))["type"] == "message.2"
+    assert (await channel_layer.receive(channel))["type"] == "message.3"
+
+
+@pytest.mark.asyncio
+async def test_groups_basic(channel_layer):
+    """
+    Tests basic group operation.
+    """
+    channel_name1 = await channel_layer.new_channel(prefix="test-gr-chan-1")
+    channel_name2 = await channel_layer.new_channel(prefix="test-gr-chan-2")
+    channel_name3 = await channel_layer.new_channel(prefix="test-gr-chan-3")
+    await channel_layer.group_add("test-group", channel_name1)
+    await channel_layer.group_add("test-group", channel_name2)
+    await channel_layer.group_add("test-group", channel_name3)
+    await channel_layer.group_discard("test-group", channel_name2)
+    await channel_layer.group_send("test-group", {"type": "message.1"})
+    # Make sure we get the message on the two channels that were in
+    async with async_timeout.timeout(1):
+        assert (await channel_layer.receive(channel_name1))["type"] == "message.1"
+        assert (await channel_layer.receive(channel_name3))["type"] == "message.1"
+    # Make sure the removed channel did not get the message
+    with pytest.raises(asyncio.TimeoutError):
+        async with async_timeout.timeout(1):
+            await channel_layer.receive(channel_name2)
+
+
+@pytest.mark.asyncio
+async def test_groups_same_prefix(channel_layer):
+    """
+    Tests group_send with multiple channels with same channel prefix
+    """
+    channel_name1 = await channel_layer.new_channel(prefix="test-gr-chan")
+    channel_name2 = await channel_layer.new_channel(prefix="test-gr-chan")
+    channel_name3 = await channel_layer.new_channel(prefix="test-gr-chan")
+    await channel_layer.group_add("test-group", channel_name1)
+    await channel_layer.group_add("test-group", channel_name2)
+    await channel_layer.group_add("test-group", channel_name3)
+    await channel_layer.group_send("test-group", {"type": "message.1"})
+
+    # Make sure we get the message on the channels that were in
+    async with async_timeout.timeout(1):
+        assert (await channel_layer.receive(channel_name1))["type"] == "message.1"
+        assert (await channel_layer.receive(channel_name2))["type"] == "message.1"
+        assert (await channel_layer.receive(channel_name3))["type"] == "message.1"
+
+
+@pytest.mark.asyncio
+async def test_random_reset__channel_name(channel_layer):
+    """
+    Makes sure resetting random seed does not make us reuse channel names.
+    """
+    random.seed(1)
+    channel_name_1 = await channel_layer.new_channel()
+    random.seed(1)
+    channel_name_2 = await channel_layer.new_channel()
+
+    assert channel_name_1 != channel_name_2