Repository: django/channels
Pull request #1976 (2023-02-09 22:00:25)

InMemoryChannelLayer improvements, test fixes
Improvements to InMemoryChannelLayer
- honor expiry / don't block for eternity
- honor per channel capacities
- threadsafer cleaning design with atomic operations
- parallel sending

Bugfixes tests
- channel_layer is sometimes overwritten and the tests become bogus

devkral: Note: I made an mistake with:
- honor expiry / don't block for eternity
This is the correct behavior for receive.
But maybe the receive method it should be put in a loop, in case the queue is expired
(see devkral/channels-multiprocessing for an example)

Patch:
diff --git a/channels/layers.py b/channels/layers.py
index 12bbd2b..48f7bac 100644
--- a/channels/layers.py
+++ b/channels/layers.py
@@ -198,13 +198,13 @@ class InMemoryChannelLayer(BaseChannelLayer):
         group_expiry=86400,
         capacity=100,
         channel_capacity=None,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(
             expiry=expiry,
             capacity=capacity,
             channel_capacity=channel_capacity,
-            **kwargs
+            **kwargs,
         )
         self.channels = {}
         self.groups = {}
@@ -225,13 +225,14 @@ class InMemoryChannelLayer(BaseChannelLayer):
         # name in message
         assert "__asgi_channel__" not in message
 
-        queue = self.channels.setdefault(channel, asyncio.Queue())
-        # Are we full
-        if queue.qsize() >= self.capacity:
-            raise ChannelFull(channel)
-
+        queue = self.channels.setdefault(
+            channel, asyncio.Queue(maxsize=self.get_capacity(channel))
+        )
         # Add message
-        await queue.put((time.time() + self.expiry, deepcopy(message)))
+        try:
+            queue.put_nowait((time.time() + self.expiry, deepcopy(message)))
+        except asyncio.queues.QueueFull:
+            raise ChannelFull(channel)
 
     async def receive(self, channel):
         """
@@ -242,14 +243,16 @@ class InMemoryChannelLayer(BaseChannelLayer):
         assert self.valid_channel_name(channel)
         self._clean_expired()
 
-        queue = self.channels.setdefault(channel, asyncio.Queue())
+        queue = self.channels.setdefault(
+            channel, asyncio.Queue(maxsize=self.get_capacity(channel))
+        )
 
         # Do a plain direct receive
         try:
             _, message = await queue.get()
         finally:
             if queue.empty():
-                del self.channels[channel]
+                self.channels.pop(channel, None)
 
         return message
 
@@ -279,19 +282,17 @@ class InMemoryChannelLayer(BaseChannelLayer):
                 self._remove_from_groups(channel)
                 # Is the channel now empty and needs deleting?
                 if queue.empty():
-                    del self.channels[channel]
+                    self.channels.pop(channel, None)
 
         # Group Expiration
         timeout = int(time.time()) - self.group_expiry
-        for group in self.groups:
-            for channel in list(self.groups.get(group, set())):
-                # If join time is older than group_expiry end the group membership
-                if (
-                    self.groups[group][channel]
-                    and int(self.groups[group][channel]) < timeout
-                ):
+        for channels in self.groups.values():
+            for name, timestamp in list(channels.items()):
+                # If join time is older than group_expiry
+                # end the group membership
+                if timestamp and timestamp < timeout:
                     # Delete from group
-                    del self.groups[group][channel]
+                    channels.pop(name, None)
 
     # Flush extension
 
@@ -308,8 +309,7 @@ class InMemoryChannelLayer(BaseChannelLayer):
         Removes a channel from all groups. Used when a message on it expires.
         """
         for channels in self.groups.values():
-            if channel in channels:
-                del channels[channel]
+            channels.pop(channel, None)
 
     # Groups extension
 
@@ -329,11 +329,13 @@ class InMemoryChannelLayer(BaseChannelLayer):
         assert self.valid_channel_name(channel), "Invalid channel name"
         assert self.valid_group_name(group), "Invalid group name"
         # Remove from group set
-        if group in self.groups:
-            if channel in self.groups[group]:
-                del self.groups[group][channel]
-            if not self.groups[group]:
-                del self.groups[group]
+        group_channels = self.groups.get(group, None)
+        if group_channels:
+            # remove channel if in group
+            group_channels.pop(channel, None)
+            # is group now empty? If yes remove it
+            if not group_channels:
+                self.groups.pop(group, None)
 
     async def group_send(self, group, message):
         # Check types
@@ -341,10 +343,15 @@ class InMemoryChannelLayer(BaseChannelLayer):
         assert self.valid_group_name(group), "Invalid group name"
         # Run clean
         self._clean_expired()
+
         # Send to each channel
-        for channel in self.groups.get(group, set()):
+        ops = []
+        if group in self.groups:
+            for channel in self.groups[group].keys():
+                ops.append(asyncio.create_task(self.send(channel, message)))
+        for send_result in asyncio.as_completed(ops):
             try:
-                await self.send(channel, message)
+                await send_result
             except ChannelFull:
                 pass
 
diff --git a/channels/sessions.py b/channels/sessions.py
index 0a4cadd..f1d51d4 100644
--- a/channels/sessions.py
+++ b/channels/sessions.py
@@ -2,6 +2,7 @@ import datetime
 import time
 from importlib import import_module
 
+import django
 from django.conf import settings
 from django.contrib.sessions.backends.base import UpdateError
 from django.core.exceptions import SuspiciousOperation
@@ -163,9 +164,7 @@ class InstanceSessionWrapper:
 
     async def resolve_session(self):
         session_key = self.scope["cookies"].get(self.cookie_name)
-        self.scope["session"]._wrapped = await database_sync_to_async(
-            self.session_store
-        )(session_key)
+        self.scope["session"]._wrapped = self.session_store(session_key)
 
     async def send(self, message):
         """
@@ -183,7 +182,7 @@ class InstanceSessionWrapper:
                 and message.get("status", 200) != 500
                 and (modified or settings.SESSION_SAVE_EVERY_REQUEST)
             ):
-                await database_sync_to_async(self.save_session)()
+                await self.save_session()
                 # If this is a message type that can transport cookies back to the
                 # client, then do so.
                 if message["type"] in self.cookie_response_message_types:
@@ -221,12 +220,15 @@ class InstanceSessionWrapper:
         # Pass up the send
         return await self.real_send(message)
 
-    def save_session(self):
+    async def save_session(self):
         """
         Saves the current session.
         """
         try:
-            self.scope["session"].save()
+            if django.VERSION >= (5, 1):
+                await self.scope["session"].asave()
+            else:
+                await database_sync_to_async(self.scope["session"].save)()
         except UpdateError:
             raise SuspiciousOperation(
                 "The request's session was deleted before the "
diff --git a/docs/topics/sessions.rst b/docs/topics/sessions.rst
index 29abb9c..871194c 100644
--- a/docs/topics/sessions.rst
+++ b/docs/topics/sessions.rst
@@ -73,7 +73,8 @@ whenever the session is modified.
 
 If you are in a WebSocket consumer, however, the session is populated
 **but will never be saved automatically** - you must call
-``scope["session"].save()`` yourself whenever you want to persist a session
+``scope["session"].save()`` (or the asynchronous version,
+``scope["session"].asave()``) yourself whenever you want to persist a session
 to your session store. If you don't save, the session will still work correctly
 inside the consumer (as it's stored as an instance variable), but other
 connections or HTTP views won't be able to see the changes.


Test patch:
diff --git a/tests/test_http.py b/tests/test_http.py
index aa1fbe5..bb55ba0 100644
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -1,6 +1,9 @@
 import re
+from importlib import import_module
 
+import django
 import pytest
+from django.conf import settings
 
 from channels.consumer import AsyncConsumer
 from channels.db import database_sync_to_async
@@ -93,15 +96,12 @@ async def test_session_samesite_invalid(samesite_invalid):
 
 @pytest.mark.django_db(transaction=True)
 @pytest.mark.asyncio
-async def test_muliple_sessions():
+async def test_multiple_sessions():
     """
     Create two application instances and test then out of order to verify that
     separate scopes are used.
     """
 
-    async def inner(scope, receive, send):
-        send(scope["path"])
-
     class SimpleHttpApp(AsyncConsumer):
         async def http_request(self, event):
             await database_sync_to_async(self.scope["session"].save)()
@@ -123,3 +123,84 @@ async def test_muliple_sessions():
 
     first_response = await first_communicator.get_response()
     assert first_response["body"] == b"/first/"
+
+
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.asyncio
+async def test_session_saves():
+    """
+    Saves information to a session and validates that it actually saves to the backend
+    """
+
+    class SimpleHttpApp(AsyncConsumer):
+        @database_sync_to_async
+        def set_fav_color(self):
+            self.scope["session"]["fav_color"] = "blue"
+
+        async def http_request(self, event):
+            if django.VERSION >= (5, 1):
+                await self.scope["session"].aset("fav_color", "blue")
+            else:
+                await self.set_fav_color()
+            await self.send(
+                {"type": "http.response.start", "status": 200, "headers": []}
+            )
+            await self.send(
+                {
+                    "type": "http.response.body",
+                    "body": self.scope["session"].session_key.encode(),
+                }
+            )
+
+    app = SessionMiddlewareStack(SimpleHttpApp.as_asgi())
+
+    communicator = HttpCommunicator(app, "GET", "/first/")
+
+    response = await communicator.get_response()
+    session_key = response["body"].decode()
+
+    SessionStore = import_module(settings.SESSION_ENGINE).SessionStore
+    session = SessionStore(session_key=session_key)
+    if django.VERSION >= (5, 1):
+        session_fav_color = await session.aget("fav_color")
+    else:
+        session_fav_color = await database_sync_to_async(session.get)("fav_color")
+
+    assert session_fav_color == "blue"
+
+
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.asyncio
+async def test_session_save_update_error():
+    """
+    Intentionally deletes the session to ensure that SuspiciousOperation is raised
+    """
+
+    async def inner(scope, receive, send):
+        send(scope["path"])
+
+    class SimpleHttpApp(AsyncConsumer):
+        @database_sync_to_async
+        def set_fav_color(self):
+            self.scope["session"]["fav_color"] = "blue"
+
+        async def http_request(self, event):
+            # Create a session as normal:
+            await database_sync_to_async(self.scope["session"].save)()
+
+            # Then simulate it's deletion from somewhere else:
+            # (e.g. logging out from another request)
+            SessionStore = import_module(settings.SESSION_ENGINE).SessionStore
+            session = SessionStore(session_key=self.scope["session"].session_key)
+            await database_sync_to_async(session.flush)()
+
+            await self.send(
+                {"type": "http.response.start", "status": 200, "headers": []}
+            )
+
+    app = SessionMiddlewareStack(SimpleHttpApp.as_asgi())
+
+    communicator = HttpCommunicator(app, "GET", "/first/")
+
+    with pytest.raises(django.core.exceptions.SuspiciousOperation):
+        await communicator.get_response()
diff --git a/tests/test_inmemorychannel.py b/tests/test_inmemorychannel.py
index 3f05ed7..4ba4bfa 100644
--- a/tests/test_inmemorychannel.py
+++ b/tests/test_inmemorychannel.py
@@ -26,9 +26,36 @@ async def test_send_receive(channel_layer):
     await channel_layer.send(
         "test-channel-1", {"type": "test.message", "text": "Ahoy-hoy!"}
     )
+    await channel_layer.send(
+        "test-channel-1", {"type": "test.message", "text": "Ahoy-hoy!"}
+    )
     message = await channel_layer.receive("test-channel-1")
     assert message["type"] == "test.message"
     assert message["text"] == "Ahoy-hoy!"
+    # not removed because not empty
+    assert "test-channel-1" in channel_layer.channels
+    message = await channel_layer.receive("test-channel-1")
+    assert message["type"] == "test.message"
+    assert message["text"] == "Ahoy-hoy!"
+    # removed because empty
+    assert "test-channel-1" not in channel_layer.channels
+
+
+@pytest.mark.asyncio
+async def test_race_empty(channel_layer):
+    """
+    Makes sure the race is handled gracefully.
+    """
+    receive_task = asyncio.create_task(channel_layer.receive("test-channel-1"))
+    await asyncio.sleep(0.1)
+    await channel_layer.send(
+        "test-channel-1", {"type": "test.message", "text": "Ahoy-hoy!"}
+    )
+    del channel_layer.channels["test-channel-1"]
+    await asyncio.sleep(0.1)
+    message = await receive_task
+    assert message["type"] == "test.message"
+    assert message["text"] == "Ahoy-hoy!"
 
 
 @pytest.mark.asyncio
@@ -62,7 +89,6 @@ async def test_multi_send_receive(channel_layer):
     """
     Tests overlapping sends and receives, and ordering.
     """
-    channel_layer = InMemoryChannelLayer()
     await channel_layer.send("test-channel-3", {"type": "message.1"})
     await channel_layer.send("test-channel-3", {"type": "message.2"})
     await channel_layer.send("test-channel-3", {"type": "message.3"})
@@ -76,7 +102,6 @@ async def test_groups_basic(channel_layer):
     """
     Tests basic group operation.
     """
-    channel_layer = InMemoryChannelLayer()
     await channel_layer.group_add("test-group", "test-gr-chan-1")
     await channel_layer.group_add("test-group", "test-gr-chan-2")
     await channel_layer.group_add("test-group", "test-gr-chan-3")
@@ -97,7 +122,6 @@ async def test_groups_channel_full(channel_layer):
     """
     Tests that group_send ignores ChannelFull
     """
-    channel_layer = InMemoryChannelLayer()
     await channel_layer.group_add("test-group", "test-gr-chan-1")
     await channel_layer.group_send("test-group", {"type": "message.1"})
     await channel_layer.group_send("test-group", {"type": "message.1"})