From ad218679d9e0ce2862d5dec71c39839f23426c0f Mon Sep 17 00:00:00 2001 From: jonathan343 Date: Sat, 8 Aug 2026 23:05:50 -0400 Subject: [PATCH] Add async client transport cleanup Generate close and async context manager methods for clients, and add matching cleanup support to aiohttp and CRT transports. Preserve shared transports when copying operation configs to avoid duplicating sessions and connection pools. --- README.md | 6 ++--- .../codegen/test/PythonCodegenTest.java | 10 ++++++- .../python/codegen/ClientGenerator.java | 26 ++++++++++++++++++- .../python/codegen/PythonSymbolProvider.java | 4 +++ .../smithy/python/codegen/RuntimeTypes.java | 1 + .../codegen/PythonSymbolProviderTest.java | 24 +++++++++++++++++ ...gfix-9b22c3201ef34610abb155ba32d0b097.json | 4 +++ ...ture-4fe36219987843b19f163600ae3a25a2.json | 4 +++ .../src/smithy_http/aio/aiohttp.py | 18 ++++++++----- .../smithy-http/src/smithy_http/aio/crt.py | 21 ++++++++++----- .../tests/unit/aio/test_aiohttp.py | 26 +++++++++++++++++++ .../smithy-http/tests/unit/aio/test_crt.py | 20 ++++++++++++-- 12 files changed, 145 insertions(+), 19 deletions(-) create mode 100644 packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json create mode 100644 packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json diff --git a/README.md b/README.md index a8422154d..cbda5db4c 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,9 @@ from echo.models import EchoMessageInput async def main() -> None: - client = EchoService(Config(endpoint_uri="https://example.com/")) - response = await client.echo_message(EchoMessageInput(message="spam")) - print(response.message) + async with EchoService(Config(endpoint_uri="https://example.com/")) as client: + response = await client.echo_message(EchoMessageInput(message="spam")) + print(response.message) if __name__ == "__main__": diff --git a/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java b/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java index 929aeb602..68f5a1d91 100644 --- a/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java +++ b/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java @@ -4,6 +4,10 @@ */ package software.amazon.smithy.python.codegen.test; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.file.Files; import java.nio.file.Path; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -20,7 +24,7 @@ public class PythonCodegenTest { @Test - public void testCodegen(@TempDir Path tempDir) { + public void testCodegen(@TempDir Path tempDir) throws IOException { // TODO: Move this to its own package once client codegen is in its own package PythonClientCodegenPlugin plugin = new PythonClientCodegenPlugin(); Model model = Model.assembler(PythonCodegenTest.class.getClassLoader()) @@ -38,5 +42,9 @@ public void testCodegen(@TempDir Path tempDir) { .model(model) .build(); plugin.execute(context); + + String client = Files.readString(tempDir.resolve("src/weather/client.py")); + assertTrue(client.contains("async def close(self) -> None:")); + assertTrue(client.contains("{id(self._config.transport): self._config.transport}")); } } diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index e9f5d7a35..993aeacb8 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -94,6 +94,27 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)), RuntimeTypes.RETRY_STRATEGY_RESOLVER); + writer.addStdlibImport("typing", "Any"); + writer.write(""" + + async def close(self) -> None: + \"\"\"Close any resources held by this client's transport.\"\"\" + await $1T(self._config.transport) + + async def __aenter__(self) -> "$2L": + return self + + async def __aexit__( + self, + exc_type: Any, + exc_value: Any, + traceback: Any, + ) -> None: + await self.close() + """, + RuntimeTypes.ASYNC_CLOSE, + serviceSymbol.getName()); + var topDownIndex = TopDownIndex.of(model); var eventStreamIndex = EventStreamIndex.of(model); for (OperationShape operation : topDownIndex.getContainedOperations(service)) { @@ -242,7 +263,10 @@ private void writeSharedOperationInit( ] if plugins: operation_plugins.extend(plugins) - config = deepcopy(self._config) + config = deepcopy( + self._config, + {id(self._config.transport): self._config.transport}, + ) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java index b738779c5..c18cd679c 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java @@ -66,6 +66,7 @@ public final class PythonSymbolProvider implements SymbolProvider, ShapeVisitor< private static final Logger LOGGER = Logger.getLogger(PythonSymbolProvider.class.getName()); private static final String SHAPES_FILE = "models"; private static final String SCHEMAS_FILE = "_private/schemas"; + private static final Set CLIENT_RESERVED_METHOD_NAMES = Set.of("close"); private final Model model; private final ReservedWordSymbolProvider.Escaper escaper; @@ -297,6 +298,9 @@ public Symbol operationShape(OperationShape shape) { // Operation names are escaped like members because ultimately they're // properties on an object too. var methodName = escaper.escapeMemberName(CaseUtils.toSnakeCase(shape.getId().getName(service))); + if (CLIENT_RESERVED_METHOD_NAMES.contains(methodName)) { + methodName = escapeWord(methodName); + } var methodSymbol = createGeneratedSymbolBuilder(shape, methodName, "client", false) .putProperty(SymbolProperties.IMPORTABLE, false) .build(); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java index 4a959e48f..7c5aed8e5 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java @@ -117,6 +117,7 @@ public final class RuntimeTypes { // smithy_core.aio.utils public static final Symbol ASYNC_LIST = createSymbol("aio.utils", "async_list", SmithyPythonDependency.SMITHY_CORE); + public static final Symbol ASYNC_CLOSE = createSymbol("aio.utils", "close", SmithyPythonDependency.SMITHY_CORE); // smithy_http public static final Symbol TUPLES_TO_FIELDS = diff --git a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java index 3b8c184f2..dcf1f46e2 100644 --- a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java +++ b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java @@ -9,6 +9,7 @@ import org.junit.jupiter.api.Test; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.shapes.MemberShape; +import software.amazon.smithy.model.shapes.OperationShape; import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.shapes.UnionShape; @@ -84,6 +85,29 @@ public void testUnionUnknownVariantNameCollidingWithShapeUsesUnderscoreSeparator provider.toSymbol(union).expectProperty(SymbolProperties.UNION_UNKNOWN).getName()); } + @Test + public void testOperationNameCollidingWithClientMethodIsEscaped() { + Model model = loadModel(""" + $version: "2" + namespace smithy.example + + service TestService { + version: "2024-01-01" + operations: [Close] + } + + operation Close {} + """); + PythonSymbolProvider provider = createProvider(model); + var operation = model.expectShape(ShapeId.from(NS + "#Close"), OperationShape.class); + + assertEquals( + "close_", + provider.toSymbol(operation) + .expectProperty(SymbolProperties.OPERATION_METHOD) + .getName()); + } + private static Model loadModel(String smithyIdl) { return Model.assembler().addUnparsedModel("test.smithy", smithyIdl).assemble().unwrap(); } diff --git a/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json new file mode 100644 index 000000000..4e478b341 --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "Preserved HTTP clients across operation configuration copies to avoid duplicating sessions and discarding connection pools." +} \ No newline at end of file diff --git a/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json b/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json new file mode 100644 index 000000000..2e1c8d504 --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json @@ -0,0 +1,4 @@ +{ + "type": "feature", + "description": "Added deterministic connection-pool cleanup and async context manager support to HTTP clients." +} diff --git a/packages/smithy-http/src/smithy_http/aio/aiohttp.py b/packages/smithy-http/src/smithy_http/aio/aiohttp.py index c5f08a1b4..37e5b13b0 100644 --- a/packages/smithy-http/src/smithy_http/aio/aiohttp.py +++ b/packages/smithy-http/src/smithy_http/aio/aiohttp.py @@ -1,8 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from copy import copy, deepcopy from itertools import chain -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from urllib.parse import parse_qs import yarl @@ -111,6 +110,16 @@ async def send( ) as resp: return await self._marshal_response(resp) + async def close(self) -> None: + """Close the underlying aiohttp session and its connection pool.""" + await self._session.close() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + async def _prepare_body(self, body: StreamingBlob) -> AsyncBytesReader | None: """Convert a body for aiohttp, omitting seekable bodies with no data.""" if not isinstance(body, AsyncBytesReader): @@ -159,7 +168,4 @@ async def _marshal_response( ) def __deepcopy__(self, memo: Any) -> "AIOHTTPClient": - return AIOHTTPClient( - client_config=deepcopy(self._config), - _session=copy(self._session), - ) + return self diff --git a/packages/smithy-http/src/smithy_http/aio/crt.py b/packages/smithy-http/src/smithy_http/aio/crt.py index 6e8a3f5b8..02182d40c 100644 --- a/packages/smithy-http/src/smithy_http/aio/crt.py +++ b/packages/smithy-http/src/smithy_http/aio/crt.py @@ -1,12 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 # pyright: reportMissingTypeStubs=false,reportUnknownMemberType=false +from asyncio import gather from collections.abc import AsyncGenerator, AsyncIterable -from copy import deepcopy from dataclasses import dataclass from inspect import iscoroutinefunction from io import BytesIO -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from awscrt.exceptions import AwsCrtError @@ -195,6 +195,18 @@ async def send( raise _CRTTimeoutError(f"CRT {e.name}: {e.message}") from e raise + async def close(self) -> None: + """Close all pooled HTTP connections.""" + connections = tuple(self._connections.values()) + self._connections.clear() + await gather(*(connection.close() for connection in connections)) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + async def _await_response( self, stream: "AIOHttpClientStreamUnified" ) -> AWSCRTHTTPResponse: @@ -364,7 +376,4 @@ async def _create_body_generator( yield chunk def __deepcopy__(self, memo: Any) -> "AWSCRTHTTPClient": - return AWSCRTHTTPClient( - eventloop=self._eventloop, - client_config=deepcopy(self._config), - ) + return self diff --git a/packages/smithy-http/tests/unit/aio/test_aiohttp.py b/packages/smithy-http/tests/unit/aio/test_aiohttp.py index 4f4548faf..d55846c7f 100644 --- a/packages/smithy-http/tests/unit/aio/test_aiohttp.py +++ b/packages/smithy-http/tests/unit/aio/test_aiohttp.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # pyright: reportPrivateUsage=false from collections.abc import AsyncIterator +from copy import deepcopy from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -17,11 +18,36 @@ def _create_client() -> tuple[AIOHTTPClient, MagicMock]: response.read = AsyncMock(return_value=b"") session = MagicMock() + session.close = AsyncMock() session.request.return_value.__aenter__ = AsyncMock(return_value=response) session.request.return_value.__aexit__ = AsyncMock(return_value=None) return AIOHTTPClient(_session=cast(Any, session)), session +def test_deepcopy_returns_same_client() -> None: + client, _ = _create_client() + + assert deepcopy(client) is client + + +async def test_close_closes_session() -> None: + client, session = _create_client() + + await client.close() + await client.close() + + assert session.close.await_count == 2 + + +async def test_context_manager_closes_session() -> None: + client, session = _create_client() + + async with client as entered: + assert entered is client + + session.close.assert_awaited_once() + + async def test_send_omits_empty_async_reader_body() -> None: client, session = _create_client() request = HTTPRequest( diff --git a/packages/smithy-http/tests/unit/aio/test_crt.py b/packages/smithy-http/tests/unit/aio/test_crt.py index 9718b3a2a..89f6845a0 100644 --- a/packages/smithy-http/tests/unit/aio/test_crt.py +++ b/packages/smithy-http/tests/unit/aio/test_crt.py @@ -22,9 +22,25 @@ def test_deepcopy_client() -> None: - """Test that AWSCRTHTTPClient can be deep copied.""" + """Test that config copies share the stateful HTTP client.""" client = AWSCRTHTTPClient() - deepcopy(client) + assert deepcopy(client) is client + + +async def test_close_closes_and_clears_pooled_connections() -> None: + client = AWSCRTHTTPClient() + connections = [AsyncMock(), AsyncMock()] + client._connections = { + ("https", "one.example.com", None): connections[0], + ("https", "two.example.com", None): connections[1], + } + + await client.close() + await client.close() + + assert client._connections == {} + for connection in connections: + connection.close.assert_awaited_once() def test_client_marshal_request() -> None: