Use asyncio for synchronization

MarkusSintonen · MarkusSintonen · commit 393035a0445f · 2024-06-10T11:58:56.000+03:00
diff --git a/httpcore/__init__.py b/httpcore/__init__.py
@@ -130,7 +130,7 @@ def __init__(self, *args, **kwargs):  # type: ignore
     "WriteError",
 ]
 
-__version__ = "1.0.5"
+__version__ = "1.0.6"
 
 
 __locals = locals()
diff --git a/httpcore/_async/connection_pool.py b/httpcore/_async/connection_pool.py
@@ -299,11 +299,16 @@ def _assign_requests_to_connections(self) -> List[AsyncConnectionInterface]:
         return closing_connections
 
     async def _close_connections(self, closing: List[AsyncConnectionInterface]) -> None:
+        if not closing:
+            return
+
         # Close connections which have been removed from the pool.
-        with AsyncShieldCancellation():
+        async def close() -> None:
             for connection in closing:
                 await connection.aclose()
 
+        await AsyncShieldCancellation.shield(close)
+
     async def aclose(self) -> None:
         # Explicitly close the connection pool.
         # Clears all existing requests and connections.
@@ -369,9 +374,9 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
     async def aclose(self) -> None:
         if not self._closed:
             self._closed = True
-            with AsyncShieldCancellation():
-                if hasattr(self._stream, "aclose"):
-                    await self._stream.aclose()
+
+            if hasattr(self._stream, "aclose"):
+                await AsyncShieldCancellation.shield(self._stream.aclose)
 
             with self._pool._optional_thread_lock:
                 self._pool._requests.remove(self._pool_request)
diff --git a/httpcore/_async/http11.py b/httpcore/_async/http11.py
@@ -137,9 +137,8 @@ async def handle_async_request(self, request: Request) -> Response:
                 },
             )
         except BaseException as exc:
-            with AsyncShieldCancellation():
-                async with Trace("response_closed", logger, request) as trace:
-                    await self._response_closed()
+            async with Trace("response_closed", logger, request) as trace:
+                await AsyncShieldCancellation.shield(self._response_closed)
             raise exc
 
     # Sending the request...
@@ -344,8 +343,7 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
             # If we get an exception while streaming the response,
             # we want to close the response (and possibly the connection)
             # before raising that exception.
-            with AsyncShieldCancellation():
-                await self.aclose()
+            await AsyncShieldCancellation.shield(self.aclose)
             raise exc
 
     async def aclose(self) -> None:
diff --git a/httpcore/_async/http2.py b/httpcore/_async/http2.py
@@ -108,8 +108,7 @@ async def handle_async_request(self, request: Request) -> Response:
                     async with Trace("send_connection_init", logger, request, kwargs):
                         await self._send_connection_init(**kwargs)
                 except BaseException as exc:
-                    with AsyncShieldCancellation():
-                        await self.aclose()
+                    await AsyncShieldCancellation.shield(self.aclose)
                     raise exc
 
                 self._sent_connection_init = True
@@ -160,12 +159,15 @@ async def handle_async_request(self, request: Request) -> Response:
                     "stream_id": stream_id,
                 },
             )
-        except BaseException as exc:  # noqa: PIE786
-            with AsyncShieldCancellation():
+        except BaseException as exc:
+
+            async def close() -> None:
                 kwargs = {"stream_id": stream_id}
                 async with Trace("response_closed", logger, request, kwargs):
                     await self._response_closed(stream_id=stream_id)
 
+            await AsyncShieldCancellation.shield(close)
+
             if isinstance(exc, h2.exceptions.ProtocolError):
                 # One case where h2 can raise a protocol error is when a
                 # closed frame has been seen by the state machine.
@@ -577,8 +579,7 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
             # If we get an exception while streaming the response,
             # we want to close the response (and possibly the connection)
             # before raising that exception.
-            with AsyncShieldCancellation():
-                await self.aclose()
+            await AsyncShieldCancellation.shield(self.aclose)
             raise exc
 
     async def aclose(self) -> None:
diff --git a/httpcore/_sync/connection_pool.py b/httpcore/_sync/connection_pool.py
@@ -299,11 +299,16 @@ def _assign_requests_to_connections(self) -> List[ConnectionInterface]:
         return closing_connections
 
     def _close_connections(self, closing: List[ConnectionInterface]) -> None:
+        if not closing:
+            return
+
         # Close connections which have been removed from the pool.
-        with ShieldCancellation():
+        def close() -> None:
             for connection in closing:
                 connection.close()
 
+        ShieldCancellation.shield(close)
+
     def close(self) -> None:
         # Explicitly close the connection pool.
         # Clears all existing requests and connections.
@@ -369,9 +374,9 @@ def __iter__(self) -> Iterator[bytes]:
     def close(self) -> None:
         if not self._closed:
             self._closed = True
-            with ShieldCancellation():
-                if hasattr(self._stream, "close"):
-                    self._stream.close()
+
+            if hasattr(self._stream, "close"):
+                ShieldCancellation.shield(self._stream.close)
 
             with self._pool._optional_thread_lock:
                 self._pool._requests.remove(self._pool_request)
diff --git a/httpcore/_sync/http11.py b/httpcore/_sync/http11.py
@@ -137,9 +137,8 @@ def handle_request(self, request: Request) -> Response:
                 },
             )
         except BaseException as exc:
-            with ShieldCancellation():
-                with Trace("response_closed", logger, request) as trace:
-                    self._response_closed()
+            with Trace("response_closed", logger, request) as trace:
+                ShieldCancellation.shield(self._response_closed)
             raise exc
 
     # Sending the request...
@@ -344,8 +343,7 @@ def __iter__(self) -> Iterator[bytes]:
             # If we get an exception while streaming the response,
             # we want to close the response (and possibly the connection)
             # before raising that exception.
-            with ShieldCancellation():
-                self.close()
+            ShieldCancellation.shield(self.close)
             raise exc
 
     def close(self) -> None:
diff --git a/httpcore/_sync/http2.py b/httpcore/_sync/http2.py
@@ -108,8 +108,7 @@ def handle_request(self, request: Request) -> Response:
                     with Trace("send_connection_init", logger, request, kwargs):
                         self._send_connection_init(**kwargs)
                 except BaseException as exc:
-                    with ShieldCancellation():
-                        self.close()
+                    ShieldCancellation.shield(self.close)
                     raise exc
 
                 self._sent_connection_init = True
@@ -160,12 +159,15 @@ def handle_request(self, request: Request) -> Response:
                     "stream_id": stream_id,
                 },
             )
-        except BaseException as exc:  # noqa: PIE786
-            with ShieldCancellation():
+        except BaseException as exc:
+
+            def close() -> None:
                 kwargs = {"stream_id": stream_id}
                 with Trace("response_closed", logger, request, kwargs):
                     self._response_closed(stream_id=stream_id)
 
+            ShieldCancellation.shield(close)
+
             if isinstance(exc, h2.exceptions.ProtocolError):
                 # One case where h2 can raise a protocol error is when a
                 # closed frame has been seen by the state machine.
@@ -577,8 +579,7 @@ def __iter__(self) -> typing.Iterator[bytes]:
             # If we get an exception while streaming the response,
             # we want to close the response (and possibly the connection)
             # before raising that exception.
-            with ShieldCancellation():
-                self.close()
+            ShieldCancellation.shield(self.close)
             raise exc
 
     def close(self) -> None:
diff --git a/httpcore/_synchronization.py b/httpcore/_synchronization.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_synchronization.py b/tests/test_synchronization.py

Original file line number	Diff line number	Diff line change
`@@ -130,7 +130,7 @@ def __init__(self, args, *kwargs): # type: ignore`
`130`	`130`	`"WriteError",`
`131`	`131`	`]`
`132`	`132`
`133`		`-__version__ = "1.0.5"`
	`133`	`+__version__ = "1.0.6"`
`134`	`134`
`135`	`135`
`136`	`136`	`__locals = locals()`