From c061e4ffcac7f2926ead8b708eb3520a980bf2dc Mon Sep 17 00:00:00 2001 From: MarkLux Date: Mon, 19 Aug 2024 21:17:47 +0800 Subject: [PATCH] fix(h2_connection_reset): add stream reset when client cancel the request --- httpcore/_async/http2.py | 9 +++++++++ httpcore/_sync/http2.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/httpcore/_async/http2.py b/httpcore/_async/http2.py index c201ee4cb..da0eacc41 100644 --- a/httpcore/_async/http2.py +++ b/httpcore/_async/http2.py @@ -401,6 +401,9 @@ async def _receive_remote_settings_change(self, event: h2.events.Event) -> None: await self._max_streams_semaphore.acquire() self._max_streams -= 1 + async def _reset_steam(self, stream_id: int, error_code: int) -> None: + self._h2_state.reset_stream(stream_id=stream_id, error_code=error_code) + async def _response_closed(self, stream_id: int) -> None: await self._max_streams_semaphore.release() del self._events[stream_id] @@ -578,6 +581,12 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]: # we want to close the response (and possibly the connection) # before raising that exception. with AsyncShieldCancellation(): + # need send cancel frame when the exception is from remote peer. + if not isinstance(exc, RemoteProtocolError): + await self._connection._reset_steam( + stream_id=self._stream_id, + error_code=h2.settings.ErrorCodes.CANCEL, + ) await self.aclose() raise exc diff --git a/httpcore/_sync/http2.py b/httpcore/_sync/http2.py index 1ee4bbb34..00117e579 100644 --- a/httpcore/_sync/http2.py +++ b/httpcore/_sync/http2.py @@ -401,6 +401,9 @@ def _receive_remote_settings_change(self, event: h2.events.Event) -> None: self._max_streams_semaphore.acquire() self._max_streams -= 1 + def _reset_steam(self, stream_id: int, error_code: int) -> None: + self._h2_state.reset_stream(stream_id=stream_id, error_code=error_code) + def _response_closed(self, stream_id: int) -> None: self._max_streams_semaphore.release() del self._events[stream_id] @@ -578,6 +581,12 @@ def __iter__(self) -> typing.Iterator[bytes]: # we want to close the response (and possibly the connection) # before raising that exception. with ShieldCancellation(): + # need send cancel frame when the exception is from remote peer. + if not isinstance(exc, RemoteProtocolError): + self._connection._reset_steam( + stream_id=self._stream_id, + error_code=h2.settings.ErrorCodes.CANCEL, + ) self.close() raise exc