Connection pool work

encode · Feb 7, 2024 · cda040d · cda040d
1 parent fcb0ca2
commit cda040d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 20 deletions.
diff --git a/httpcore/_async/connection_pool.py b/httpcore/_async/connection_pool.py
@@ -195,8 +195,7 @@ async def handle_async_request(self, request: Request) -> Response:
         try:
             while True:
                 with AsyncShieldCancellation():
-                    async with self._pool_lock:
-                        closing = self._assign_requests_to_connections()
+                    closing = self._assign_requests_to_connections()
                 await self._close_connections(closing)
                 connection = await pool_request.wait_for_connection(timeout=timeout)
 
@@ -211,9 +210,8 @@ async def handle_async_request(self, request: Request) -> Response:
 
         except BaseException as exc:
             with AsyncShieldCancellation():
-                async with self._pool_lock:
-                    self._requests.remove(pool_request)
-                    closing = self._assign_requests_to_connections()
+                self._requests.remove(pool_request)
+                closing = self._assign_requests_to_connections()
             await self._close_connections(closing)
             raise exc from None
 
@@ -351,10 +349,11 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
     async def aclose(self) -> None:
         if not self._closed:
             self._closed = True
+            if hasattr(self._stream, "aclose"):
+                await self._stream.aclose()
+
             with AsyncShieldCancellation():
-                if hasattr(self._stream, "aclose"):
-                    await self._stream.aclose()
-                async with self._pool._pool_lock:
-                    self._pool._requests.remove(self._pool_request)
-                    closing = self._pool._assign_requests_to_connections()
+                self._pool._requests.remove(self._pool_request)
+                closing = self._pool._assign_requests_to_connections()
+
             await self._pool._close_connections(closing)
diff --git a/httpcore/_sync/connection_pool.py b/httpcore/_sync/connection_pool.py
@@ -195,8 +195,7 @@ def handle_request(self, request: Request) -> Response:
         try:
             while True:
                 with ShieldCancellation():
-                    with self._pool_lock:
-                        closing = self._assign_requests_to_connections()
+                    closing = self._assign_requests_to_connections()
                 self._close_connections(closing)
                 connection = pool_request.wait_for_connection(timeout=timeout)
 
@@ -211,9 +210,8 @@ def handle_request(self, request: Request) -> Response:
 
         except BaseException as exc:
             with ShieldCancellation():
-                with self._pool_lock:
-                    self._requests.remove(pool_request)
-                    closing = self._assign_requests_to_connections()
+                self._requests.remove(pool_request)
+                closing = self._assign_requests_to_connections()
             self._close_connections(closing)
             raise exc from None
 
@@ -351,10 +349,11 @@ def __iter__(self) -> Iterator[bytes]:
     def close(self) -> None:
         if not self._closed:
             self._closed = True
+            if hasattr(self._stream, "close"):
+                self._stream.close()
+
             with ShieldCancellation():
-                if hasattr(self._stream, "close"):
-                    self._stream.close()
-                with self._pool._pool_lock:
-                    self._pool._requests.remove(self._pool_request)
-                    closing = self._pool._assign_requests_to_connections()
+                self._pool._requests.remove(self._pool_request)
+                closing = self._pool._assign_requests_to_connections()
+
             self._pool._close_connections(closing)