Release completed operations in the legacy graphql-ws handler (#4610)

Author: patrick91

RELEASE.md

@@ -0,0 +1,26 @@
+---
+release type: patch
+social_messages:
+  x: >-
+    {project_name} {version} is out! This release fixes the legacy graphql-ws
+    handler so completed subscriptions no longer count towards
+    max_subscriptions_per_connection. 🍓 https://strawberry.rocks/release/{version}
+  linkedin: >-
+    {project_name} {version} is out. This release fixes the legacy graphql-ws
+    handler so subscriptions that complete on their own release their slot and
+    no longer count towards max_subscriptions_per_connection.
+---
+
+This release fixes the legacy `graphql-ws` protocol handler so that
+subscriptions which complete on their own (or fail before execution) release
+their slot on the connection.
+
+Previously, completed operations were kept in the handler's bookkeeping until
+the client sent a `stop` message for them, reused their operation id, or
+disconnected. On connections with `max_subscriptions_per_connection`
+configured, a client using distinct operation ids could therefore hit
+`Subscription limit reached` even though none of its earlier subscriptions
+were still active. The `graphql-transport-ws` handler was not affected.
+
+Sending a `stop` message for an operation that has already completed is now a
+no-op instead of an error.

strawberry/subscriptions/protocols/graphql_ws/handlers.py

@@ -215,17 +215,28 @@ async def handle_async_results(
 
         except asyncio.CancelledError:
             await self.send_message(CompleteMessage(type="complete", id=operation_id))
+        finally:
+            # The operation is over, whether it completed on its own, failed
+            # before execution or was stopped. Release its bookkeeping so that
+            # it no longer counts towards ``max_subscriptions_per_connection``
+            # and its id can be reused right away.
+            self.subscriptions.pop(operation_id, None)
+            self.tasks.pop(operation_id, None)
 
     async def cleanup_operation(self, operation_id: str) -> None:
-        if operation_id in self.subscriptions:
+        result_source = self.subscriptions.pop(operation_id, None)
+        if result_source is not None:
             with suppress(RuntimeError):
-                await self.subscriptions[operation_id].aclose()
-            del self.subscriptions[operation_id]
+                await result_source.aclose()
+
+        task = self.tasks.pop(operation_id, None)
+        if task is None:
+            # Already finished (and released itself), or never existed
+            return
 
-        self.tasks[operation_id].cancel()
+        task.cancel()
         with suppress(BaseException):
-            await self.tasks[operation_id]
-        del self.tasks[operation_id]
+            await task
 
     async def cleanup(self) -> None:
         for operation_id in list(self.tasks.keys()):

tests/websockets/test_graphql_ws.py

@@ -1030,6 +1030,98 @@ async def test_reusing_operation_id_does_not_count_against_limit(
         await ws.close()
 
 
+async def test_completed_subscriptions_do_not_count_against_limit(
+    http_client_class: type[HttpClient],
+):
+    """Subscriptions that complete on their own must release their slot,
+    otherwise a client could hit the limit with operations that are no longer
+    active."""
+    test_client = http_client_class(schema, max_subscriptions_per_connection=2)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_legacy_message({"type": "connection_init"})
+        response: ConnectionAckMessage = await ws.receive_json()
+        assert response["type"] == "connection_ack"
+
+        # Each of these completes on its own after a single result, so the
+        # third one must not be rejected even though the limit is 2.
+        for operation_id in ("sub1", "sub2", "sub3"):
+            await ws.send_legacy_message(
+                {
+                    "type": "start",
+                    "id": operation_id,
+                    "payload": {
+                        "query": 'subscription { echo(message: "Hi") }',
+                    },
+                }
+            )
+
+            data_message: DataMessage = await ws.receive_json()
+            assert data_message["type"] == "data"
+            assert data_message["id"] == operation_id
+            assert data_message["payload"]["data"] == {"echo": "Hi"}
+
+            complete_message: CompleteMessage = await ws.receive_json()
+            assert complete_message["type"] == "complete"
+            assert complete_message["id"] == operation_id
+
+        # Stopping an operation that already completed is a no-op and must
+        # not affect the connection.
+        await ws.send_legacy_message({"type": "stop", "id": "sub1"})
+
+        await ws.send_legacy_message(
+            {
+                "type": "start",
+                "id": "sub4",
+                "payload": {
+                    "query": 'subscription { echo(message: "Hi") }',
+                },
+            }
+        )
+        data_message = await ws.receive_json()
+        assert data_message["type"] == "data"
+        assert data_message["id"] == "sub4"
+
+        complete_message = await ws.receive_json()
+        assert complete_message["type"] == "complete"
+        assert complete_message["id"] == "sub4"
+
+        await ws.close()
+
+
+async def test_failed_subscriptions_do_not_count_against_limit(
+    http_client_class: type[HttpClient],
+):
+    """Operations that fail before execution (e.g. validation errors) must
+    release their slot as well."""
+    test_client = http_client_class(schema, max_subscriptions_per_connection=2)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_legacy_message({"type": "connection_init"})
+        response: ConnectionAckMessage = await ws.receive_json()
+        assert response["type"] == "connection_ack"
+
+        for operation_id in ("sub1", "sub2", "sub3"):
+            await ws.send_legacy_message(
+                {
+                    "type": "start",
+                    "id": operation_id,
+                    "payload": {"query": "subscription { doesNotExist }"},
+                }
+            )
+
+            error_message: ErrorMessage = await ws.receive_json()
+            assert error_message["type"] == "error"
+            assert error_message["id"] == operation_id
+            assert error_message["payload"] != {"message": "Subscription limit reached"}
+
+        await ws.close()
+
+
 async def test_max_subscriptions_per_connection_disabled(
     http_client_class: type[HttpClient],
 ):