Add async client transport cleanup

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__":

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}"));
     }
 }

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:

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<String> 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();

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 =

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();
     }

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."
+}

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."
+}

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

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

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(

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: