Unset the keep-alive timer when upgrading to WebSocket (#3107)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

Author: kazuki-ma

tests/protocols/test_http.py

@@ -451,6 +451,21 @@ async def test_keepalive_timeout_with_pipelined_requests(http_protocol_cls: type
     assert protocol.timeout_keep_alive_task is not None
 
 
+async def test_keepalive_timeout_with_pipelined_websocket_upgrade(
+    http_protocol_cls: type[HTTPProtocol], ws_protocol_cls: type[WSProtocol]
+):
+    if http_protocol_cls.__name__ == "HttpToolsProtocol":
+        pytest.skip("httptools never upgrades with the keep-alive timer armed")
+
+    app = Response("Hello, world", media_type="text/plain")
+
+    protocol = get_connected_protocol(app, http_protocol_cls, ws=ws_protocol_cls)
+    protocol.data_received(SIMPLE_GET_REQUEST + UPGRADE_REQUEST)
+    await protocol.loop.run_one()
+
+    assert protocol.timeout_keep_alive_task is None
+
+
 async def test_close(http_protocol_cls: type[HTTPProtocol]):
     app = Response(b"", status_code=204, headers={"connection": "close"})
 

uvicorn/protocols/http/h11_impl.py

@@ -197,6 +197,9 @@ def handle_events(self) -> None:
                 break
 
             elif isinstance(event, h11.Request):
+                # Pipelined HTTP requests and WebSocket upgrades may be processed after the keep-alive timer is armed.
+                self._unset_keepalive_if_required()
+
                 self.headers = [(key.lower(), value) for key, value in event.headers]
                 raw_path, _, query_string = event.target.partition(b"?")
                 path = unquote(raw_path.decode("ascii"))
@@ -231,14 +234,6 @@ def handle_events(self) -> None:
                 else:
                     app = self.app
 
-                # When starting to process a request, disable the keep-alive
-                # timeout. Normally we disable this when receiving data from
-                # client and set back when finishing processing its request.
-                # However, for pipelined requests processing finishes after
-                # already receiving the next request and thus the timer may
-                # be set here, which we don't want.
-                self._unset_keepalive_if_required()
-
                 self.cycle = RequestResponseCycle(
                     scope=self.scope,
                     conn=self.conn,

uvicorn/protocols/http/zttp_impl.py

@@ -145,6 +145,9 @@ def data_received(self, data: bytes) -> None:
     def handle_events(self, event: zttp.Event) -> None:
         while event is not zttp.NEED_DATA:
             if isinstance(event, zttp.Request):
+                # Pipelined HTTP requests and WebSocket upgrades may be processed after the keep-alive timer is armed.
+                self._unset_keepalive_if_required()
+
                 assert isinstance(event.headers, zttp.HeaderBlock)
                 self.headers = event.headers.to_list(lowercase_names=True)
                 path = event.path.decode("ascii")
@@ -181,8 +184,6 @@ def handle_events(self, event: zttp.Event) -> None:
                 else:
                     app = self.app
 
-                self._unset_keepalive_if_required()
-
                 self.cycle = RequestResponseCycle(
                     scope=self.scope,
                     conn=self.conn,