(feat:prime/crac) Invoke the selected warm-up operation in the generated SdkWarmUpProvider

codegen/src/main/java/software/amazon/awssdk/codegen/poet/crac/WarmUpProviderSpec.java

@@ -22,9 +22,11 @@
 import com.squareup.javapoet.TypeSpec;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.util.Optional;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.model.intermediate.Protocol;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
@@ -38,10 +40,10 @@
  * <p>The {@code warmUpClient(ClientType)} body instantiates and closes a synthetic sync client (when the service
  * generates one) and a synthetic async client, each guarded by the requested {@link ClientType} and wired to an
  * in-memory canned HTTP client, dummy credentials, a fixed region and a local {@code endpointOverride}. Building the
- * clients JIT-compiles the client construction and configuration-resolution path before a CRaC checkpoint. The
- * {@code syncClientClassName()}/{@code asyncClientClassName()} strings let {@code SdkWarmUp.prime(Class...)} match a
- * requested client class to this provider without loading excluded client interfaces. The operation call that
- * exercises the marshal/unmarshal pipeline is added in a later stage.
+ * clients JIT-compiles the client construction and configuration-resolution path before a CRaC checkpoint. Each client
+ * also invokes the operation chosen by {@link WarmUpOperationSelector}, if any, to prime the marshalling, signing and
+ * unmarshalling paths. The {@code syncClientClassName()}/{@code asyncClientClassName()} strings let
+ * {@code SdkWarmUp.prime(Class...)} match a requested client class to this provider.
  */
 public class WarmUpProviderSpec implements ClassSpec {
 
@@ -71,10 +73,12 @@ public class WarmUpProviderSpec implements ClassSpec {
 
     private final IntermediateModel model;
     private final PoetExtension poetExtensions;
+    private final Optional<OperationModel> warmUpOperation;
 
     public WarmUpProviderSpec(IntermediateModel model) {
         this.model = model;
         this.poetExtensions = new PoetExtension(model);
+        this.warmUpOperation = WarmUpOperationSelector.selectWarmUpOperation(model);
     }
 
     @Override
@@ -123,12 +127,12 @@ private MethodSpec warmUpClientMethod() {
         if (!model.getCustomizationConfig().isSkipSyncClientGeneration()) {
             body.beginControlFlow("if (clientType == $T.SYNC)", ClientType.class)
                 .add(clientBlock(syncClientClass(),
-                                 CANNED_RESPONSE_HTTP_CLIENT, SDK_HTTP_CLIENT, "httpClient", "client"))
+                                 CANNED_RESPONSE_HTTP_CLIENT, SDK_HTTP_CLIENT, "httpClient", "client", false))
                 .endControlFlow();
         }
         body.beginControlFlow("if (clientType == $T.ASYNC)", ClientType.class)
             .add(clientBlock(asyncClientClass(),
-                             CANNED_RESPONSE_ASYNC_HTTP_CLIENT, SDK_ASYNC_HTTP_CLIENT, "asyncHttpClient", "asyncClient"))
+                             CANNED_RESPONSE_ASYNC_HTTP_CLIENT, SDK_ASYNC_HTTP_CLIENT, "asyncHttpClient", "asyncClient", true))
             .endControlFlow();
 
         return MethodSpec.methodBuilder("warmUpClient")
@@ -148,31 +152,52 @@ private ClassName asyncClientClass() {
     }
 
     /**
-     * Emits a canned HTTP client plus a try-with-resources that builds and closes {@code clientType}. The sync and async
-     * paths differ only in these types and variable names, so they share this emitter.
+     * Emits a canned HTTP client plus a try-with-resources that builds {@code clientType}, calls the selected
+     * warm-up operation (if any) and closes the client. The sync and async paths differ only in these types, variable
+     * names and joining the async call, so they share this emitter.
      */
     private CodeBlock clientBlock(ClassName clientType, ClassName cannedHttpClientType, ClassName httpClientType,
-                                  String httpClientVar, String clientVar) {
+                                  String httpClientVar, String clientVar, boolean async) {
+        CodeBlock.Builder block = CodeBlock.builder()
+            .addStatement("$T $N = $T.builder().responseBody($L).statusCode($L).build()",
+                          httpClientType, httpClientVar, cannedHttpClientType, CANNED_RESPONSE_FIELD,
+                          SUCCESS_STATUS_CODE)
+            .beginControlFlow("try ($1T $2N = $1T.builder()\n"
+                              + ".httpClient($3N)\n"
+                              + ".credentialsProvider($4T.create($5T.create($6S, $7S)))\n"
+                              + ".region($8T.US_EAST_1)\n"
+                              + ".endpointOverride($9T.create($10S))\n"
+                              + ".build())",
+                              clientType,
+                              clientVar,
+                              httpClientVar,
+                              STATIC_CREDENTIALS_PROVIDER,
+                              AWS_BASIC_CREDENTIALS,
+                              DUMMY_ACCESS_KEY_ID, DUMMY_SECRET_ACCESS_KEY,
+                              REGION,
+                              URI.class,
+                              LOCAL_ENDPOINT);
+
+        warmUpOperation.ifPresent(op -> block.add(warmUpOperationCall(op, clientVar, async)));
+
+        return block.endControlFlow()
+                    .build();
+    }
+
+    /**
+     * Verified simple methods generate a no-arg overload on both clients, so the call uses it. Other operations pass
+     * an empty request.
+     */
+    private CodeBlock warmUpOperationCall(OperationModel operation, String clientVar, boolean async) {
+        String join = async ? ".join()" : "";
+        if (operation.getInputShape().isSimpleMethod()) {
+            return CodeBlock.builder()
+                            .addStatement("$N.$N()" + join, clientVar, operation.getMethodName())
+                            .build();
+        }
+        ClassName requestType = poetExtensions.getModelClass(operation.getInputShape().getShapeName());
         return CodeBlock.builder()
-                        .addStatement("$T $N = $T.builder().responseBody($L).statusCode($L).build()",
-                                      httpClientType, httpClientVar, cannedHttpClientType, CANNED_RESPONSE_FIELD,
-                                      SUCCESS_STATUS_CODE)
-                        .beginControlFlow("try ($1T $2N = $1T.builder()\n"
-                                          + ".httpClient($3N)\n"
-                                          + ".credentialsProvider($4T.create($5T.create($6S, $7S)))\n"
-                                          + ".region($8T.US_EAST_1)\n"
-                                          + ".endpointOverride($9T.create($10S))\n"
-                                          + ".build())",
-                                          clientType,
-                                          clientVar,
-                                          httpClientVar,
-                                          STATIC_CREDENTIALS_PROVIDER,
-                                          AWS_BASIC_CREDENTIALS,
-                                          DUMMY_ACCESS_KEY_ID, DUMMY_SECRET_ACCESS_KEY,
-                                          REGION,
-                                          URI.class,
-                                          LOCAL_ENDPOINT)
-                        .endControlFlow()
+                        .addStatement("$N.$N($T.builder().build())" + join, clientVar, operation.getMethodName(), requestType)
                         .build();
     }
 

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-async-only.java

@@ -12,6 +12,7 @@
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.query.QueryAsyncClient;
+import software.amazon.awssdk.services.query.model.GetOperationWithChecksumRequest;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
@@ -36,6 +37,7 @@ public void warmUpClient(ClientType clientType) {
             try (QueryAsyncClient asyncClient = QueryAsyncClient.builder().httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.getOperationWithChecksum(GetOperationWithChecksumRequest.builder().build()).join();
             }
         }
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-cbor.java

@@ -37,6 +37,7 @@ public void warmUpClient(ClientType clientType) {
             try (JsonClient client = JsonClient.builder().httpClient(httpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                client.paginatedOperationWithResultKey();
             }
         }
         if (clientType == ClientType.ASYNC) {
@@ -45,6 +46,7 @@ public void warmUpClient(ClientType clientType) {
             try (JsonAsyncClient asyncClient = JsonAsyncClient.builder().httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.paginatedOperationWithResultKey().join();
             }
         }
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-query.java

@@ -15,6 +15,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.query.QueryAsyncClient;
 import software.amazon.awssdk.services.query.QueryClient;
+import software.amazon.awssdk.services.query.model.GetOperationWithChecksumRequest;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
@@ -38,6 +39,7 @@ public void warmUpClient(ClientType clientType) {
             try (QueryClient client = QueryClient.builder().httpClient(httpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                client.getOperationWithChecksum(GetOperationWithChecksumRequest.builder().build());
             }
         }
         if (clientType == ClientType.ASYNC) {
@@ -46,6 +48,7 @@ public void warmUpClient(ClientType clientType) {
             try (QueryAsyncClient asyncClient = QueryAsyncClient.builder().httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.getOperationWithChecksum(GetOperationWithChecksumRequest.builder().build()).join();
             }
         }
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-rest-json.java

@@ -38,6 +38,7 @@ public void warmUpClient(ClientType clientType) {
             try (JsonClient client = JsonClient.builder().httpClient(httpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                client.paginatedOperationWithResultKey();
             }
         }
         if (clientType == ClientType.ASYNC) {
@@ -46,6 +47,7 @@ public void warmUpClient(ClientType clientType) {
             try (JsonAsyncClient asyncClient = JsonAsyncClient.builder().httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.paginatedOperationWithResultKey().join();
             }
         }
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-rpcv2.java

@@ -14,6 +14,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.smithyrpcv2protocol.SmithyRpcV2ProtocolAsyncClient;
 import software.amazon.awssdk.services.smithyrpcv2protocol.SmithyRpcV2ProtocolClient;
+import software.amazon.awssdk.services.smithyrpcv2protocol.model.EmptyInputOutputRequest;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
@@ -37,6 +38,7 @@ public void warmUpClient(ClientType clientType) {
             try (SmithyRpcV2ProtocolClient client = SmithyRpcV2ProtocolClient.builder().httpClient(httpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                client.emptyInputOutput(EmptyInputOutputRequest.builder().build());
             }
         }
         if (clientType == ClientType.ASYNC) {
@@ -46,6 +48,7 @@ public void warmUpClient(ClientType clientType) {
                     .httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.emptyInputOutput(EmptyInputOutputRequest.builder().build()).join();
             }
         }
     }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/crac/warmup-provider-xml.java

@@ -15,6 +15,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.xml.XmlAsyncClient;
 import software.amazon.awssdk.services.xml.XmlClient;
+import software.amazon.awssdk.services.xml.model.GetOperationWithChecksumRequest;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
@@ -38,6 +39,7 @@ public void warmUpClient(ClientType clientType) {
             try (XmlClient client = XmlClient.builder().httpClient(httpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                client.getOperationWithChecksum(GetOperationWithChecksumRequest.builder().build());
             }
         }
         if (clientType == ClientType.ASYNC) {
@@ -46,6 +48,7 @@ public void warmUpClient(ClientType clientType) {
             try (XmlAsyncClient asyncClient = XmlAsyncClient.builder().httpClient(asyncHttpClient)
                     .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")))
                     .region(Region.US_EAST_1).endpointOverride(URI.create("http://localhost")).build()) {
+                asyncClient.getOperationWithChecksum(GetOperationWithChecksumRequest.builder().build()).join();
             }
         }
     }

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/crac/OperationRecordingInterceptor.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.crac;
+
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
+
+/**
+ * Records invoked operation names. Registered as a global interceptor via
+ * {@code software/amazon/awssdk/global/handlers/execution.interceptors} so tests can observe calls made by clients
+ * they cannot configure, such as the client a generated {@code SdkWarmUpProvider} builds internally.
+ */
+public class OperationRecordingInterceptor implements ExecutionInterceptor {
+
+    private static final List<String> OPERATION_NAMES = new CopyOnWriteArrayList<>();
+
+    public static List<String> operationNames() {
+        return OPERATION_NAMES;
+    }
+
+    public static void reset() {
+        OPERATION_NAMES.clear();
+    }
+
+    @Override
+    public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
+        OPERATION_NAMES.add(executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME));
+    }
+}

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/crac/WarmUpProviderBindingTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.crac;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.ClientType;
+import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient;
+import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient;
+import software.amazon.awssdk.services.protocolrestjson.internal.crac.ProtocolRestJsonWarmUpProvider;
+
+/**
+ * Verifies the generated {@code SdkWarmUpProvider} invokes the operation selected by
+ * {@code WarmUpOperationSelector}. {@link OperationRecordingInterceptor} is registered as a global interceptor to
+ * observe calls made by the clients the provider builds internally.
+ */
+class WarmUpProviderBindingTest {
+
+    /**
+     * The operation WarmUpOperationSelector picks for the ProtocolRestJson test model.
+     */
+    private static final String EXPECTED_OPERATION = "AllTypes";
+
+    private final ProtocolRestJsonWarmUpProvider provider = new ProtocolRestJsonWarmUpProvider();
+
+    @BeforeEach
+    void resetRecorder() {
+        OperationRecordingInterceptor.reset();
+    }
+
+    @Test
+    void warmUpClient_sync_invokesSelectedOperation() {
+        assertThatCode(() -> provider.warmUpClient(ClientType.SYNC)).doesNotThrowAnyException();
+
+        assertThat(OperationRecordingInterceptor.operationNames()).containsExactly(EXPECTED_OPERATION);
+    }
+
+    @Test
+    void warmUpClient_async_invokesSelectedOperation() {
+        assertThatCode(() -> provider.warmUpClient(ClientType.ASYNC)).doesNotThrowAnyException();
+
+        assertThat(OperationRecordingInterceptor.operationNames()).containsExactly(EXPECTED_OPERATION);
+    }
+
+    @Test
+    void warmUp_invokesSelectedOperationOnBothClients() {
+        assertThatCode(provider::warmUp).doesNotThrowAnyException();
+
+        assertThat(OperationRecordingInterceptor.operationNames())
+            .containsExactly(EXPECTED_OPERATION, EXPECTED_OPERATION);
+    }
+
+    @Test
+    void clientClassNames_matchGeneratedClients() {
+        assertThat(provider.syncClientClassName()).isEqualTo(ProtocolRestJsonClient.class.getName());
+        assertThat(provider.asyncClientClassName()).isEqualTo(ProtocolRestJsonAsyncClient.class.getName());
+    }
+}

test/codegen-generated-classes-test/src/test/resources/software/amazon/awssdk/global/handlers/execution.interceptors

@@ -0,0 +1 @@
+software.amazon.awssdk.crac.OperationRecordingInterceptor