Fix HTTP/2 race condition waiting on stream events.

httpcore/_async/http2.py

@@ -231,19 +231,27 @@ async def _receive_stream_event(
         self, request: Request, stream_id: int
     ) -> h2.events.Event:
         while not self._events.get(stream_id):
-            await self._receive_events(request)
+            await self._receive_events(request, stream_id)
         return self._events[stream_id].pop(0)
 
-    async def _receive_events(self, request: Request) -> None:
-        events = await self._read_incoming_data(request)
-        for event in events:
-            event_stream_id = getattr(event, "stream_id", 0)
-
-            if hasattr(event, "error_code"):
-                raise RemoteProtocolError(event)
-
-            if event_stream_id in self._events:
-                self._events[event_stream_id].append(event)
+    async def _receive_events(self, request: Request, stream_id: int = None) -> None:
+        async with self._read_lock:
+            # This conditional is a bit icky. We don't want to block reading if we've
+            # actually got an event to return for a given stream. We need to do that
+            # check *within* the atomic read lock. Though it also need to be optional,
+            # because when we call it from `_wait_for_outgoing_flow` we *do* want to
+            # block until we've available flow control, event when we have events
+            # pending for the stream ID we're attempting to send on.
+            if stream_id is None or not self._events.get(stream_id):
+                events = await self._read_incoming_data(request)
+                for event in events:
+                    event_stream_id = getattr(event, "stream_id", 0)
+
+                    if hasattr(event, "error_code"):
+                        raise RemoteProtocolError(event)
+
+                    if event_stream_id in self._events:
+                        self._events[event_stream_id].append(event)
 
         await self._write_outgoing_data(request)
 
@@ -274,11 +282,12 @@ async def _read_incoming_data(
         timeouts = request.extensions.get("timeout", {})
         timeout = timeouts.get("read", None)
 
-        async with self._read_lock:
-            data = await self._network_stream.read(self.READ_NUM_BYTES, timeout)
-            if data == b"":
-                raise RemoteProtocolError("Server disconnected")
-            return self._h2_state.receive_data(data)
+        data = await self._network_stream.read(self.READ_NUM_BYTES, timeout)
+        if data == b"":
+            raise RemoteProtocolError("Server disconnected")
+        events = self._h2_state.receive_data(data)
+
+        return events
 
     async def _write_outgoing_data(self, request: Request) -> None:
         timeouts = request.extensions.get("timeout", {})

httpcore/_sync/http2.py

@@ -231,19 +231,27 @@ def _receive_stream_event(
         self, request: Request, stream_id: int
     ) -> h2.events.Event:
         while not self._events.get(stream_id):
-            self._receive_events(request)
+            self._receive_events(request, stream_id)
         return self._events[stream_id].pop(0)
 
-    def _receive_events(self, request: Request) -> None:
-        events = self._read_incoming_data(request)
-        for event in events:
-            event_stream_id = getattr(event, "stream_id", 0)
-
-            if hasattr(event, "error_code"):
-                raise RemoteProtocolError(event)
-
-            if event_stream_id in self._events:
-                self._events[event_stream_id].append(event)
+    def _receive_events(self, request: Request, stream_id: int = None) -> None:
+        with self._read_lock:
+            # This conditional is a bit icky. We don't want to block reading if we've
+            # actually got an event to return for a given stream. We need to do that
+            # check *within* the atomic read lock. Though it also need to be optional,
+            # because when we call it from `_wait_for_outgoing_flow` we *do* want to
+            # block until we've available flow control, event when we have events
+            # pending for the stream ID we're attempting to send on.
+            if stream_id is None or not self._events.get(stream_id):
+                events = self._read_incoming_data(request)
+                for event in events:
+                    event_stream_id = getattr(event, "stream_id", 0)
+
+                    if hasattr(event, "error_code"):
+                        raise RemoteProtocolError(event)
+
+                    if event_stream_id in self._events:
+                        self._events[event_stream_id].append(event)
 
         self._write_outgoing_data(request)
 
@@ -274,11 +282,12 @@ def _read_incoming_data(
         timeouts = request.extensions.get("timeout", {})
         timeout = timeouts.get("read", None)
 
-        with self._read_lock:
-            data = self._network_stream.read(self.READ_NUM_BYTES, timeout)
-            if data == b"":
-                raise RemoteProtocolError("Server disconnected")
-            return self._h2_state.receive_data(data)
+        data = self._network_stream.read(self.READ_NUM_BYTES, timeout)
+        if data == b"":
+            raise RemoteProtocolError("Server disconnected")
+        events = self._h2_state.receive_data(data)
+
+        return events
 
     def _write_outgoing_data(self, request: Request) -> None:
         timeouts = request.extensions.get("timeout", {})