Add send_multiple method to core layer

Author: olzhasar

channels_redis/core.py

@@ -180,22 +180,39 @@ async def send(self, channel, message):
         """
         Send a message onto a (general or specific) channel.
         """
+        await self.send_multiple(channel, (message,))
+
+    async def send_multiple(self, channel, messages):
+        """
+        Send multiple messages at once onto a (general or specific) channel.
+        """
         # Typecheck
-        assert isinstance(message, dict), "message is not a dict"
         assert self.valid_channel_name(channel), "Channel name not valid"
-        # Make sure the message does not contain reserved keys
-        assert "__asgi_channel__" not in message
+        assert hasattr(messages, "__iter__"), "messages is not an iterable"
+
         # If it's a process-local channel, strip off local part and stick full name in message
         channel_non_local_name = channel
-        if "!" in channel:
-            message = dict(message.items())
-            message["__asgi_channel__"] = channel
+        process_local = "!" in channel
+        if process_local:
             channel_non_local_name = self.non_local_name(channel)
+
+        now = time.time()
+        mapping = {}
+        for message in messages:
+            assert isinstance(message, dict), "message is not a dict"
+            # Make sure the message does not contain reserved keys
+            assert "__asgi_channel__" not in message
+            if process_local:
+                message = dict(message.items())
+                message["__asgi_channel__"] = channel
+
+            mapping[self.serialize(message)] = now
+
         # Write out message into expiring key (avoids big items in list)
         channel_key = self.prefix + channel_non_local_name
         # Pick a connection to the right server - consistent for specific
         # channels, random for general channels
-        if "!" in channel:
+        if process_local:
             index = self.consistent_hash(channel)
         else:
             index = next(self._send_index_generator)
@@ -207,13 +224,13 @@ async def send(self, channel, message):
 
         # Check the length of the list before send
         # This can allow the list to leak slightly over capacity, but that's fine.
-        if await connection.zcount(channel_key, "-inf", "+inf") >= self.get_capacity(
-            channel
-        ):
+        current_length = await connection.zcount(channel_key, "-inf", "+inf")
+
+        if current_length + len(messages) > self.get_capacity(channel):
             raise ChannelFull()
 
         # Push onto the list then set it to expire in case it's not consumed
-        await connection.zadd(channel_key, {self.serialize(message): time.time()})
+        await connection.zadd(channel_key, mapping)
         await connection.expire(channel_key, int(self.expiry))
 
     def _backup_channel_name(self, channel):

tests/test_core.py

@@ -125,6 +125,25 @@ async def listen2():
     async_to_sync(channel_layer.flush)()
 
 
+@pytest.mark.asyncio
+async def test_send_multiple(channel_layer):
+    messsages = [
+        {"type": "test.message.1"},
+        {"type": "test.message.2"},
+        {"type": "test.message.3"},
+    ]
+
+    await channel_layer.send_multiple("test-channel-1", messsages)
+
+    expected = {"test.message.1", "test.message.2", "test.message.3"}
+    received = set()
+    for _ in range(3):
+        msg = await channel_layer.receive("test-channel-1")
+        received.add(msg["type"])
+
+    assert received == expected
+
+
 @pytest.mark.asyncio
 async def test_send_capacity(channel_layer):
     """