add per request waiterOverrideConfiguration overloaded waiter operation

codegen/src/main/java/software/amazon/awssdk/codegen/docs/WaiterDocs.java

@@ -66,7 +66,9 @@ public static CodeBlock waiterOperationConsumerBuilderJavadoc(ClassName clientCl
                                                                  + "<p>This is a convenience method to create an instance of "
                                                                  + "the request builder without the need "
                                                                  + "to create one manually using {@link $T.builder()} ")
-                                                    .param(operationModel.getInput().getVariableName(), "the request to be used"
+                                                    .param(operationModel.getInput().getVariableName(), "The consumer that will"
+                                                                                                        + " configure the "
+                                                                                                        + "request to be used"
                                                                                                         + " for polling")
                                                     .returns(clientClassName.simpleName().contains("Async") ?
                                                              "CompletableFuture of the WaiterResponse containing either a "
@@ -79,7 +81,6 @@ public static CodeBlock waiterOperationConsumerBuilderJavadoc(ClassName clientCl
                         .add(javadocs, clientClassName, operationModel.getMethodName(), waiterDefinition.getKey(),
                              requestClassName)
                         .build();
-
     }
 
     public static CodeBlock waiterBuilderMethodJavadoc(ClassName className) {
@@ -170,4 +171,49 @@ public static CodeBlock waiterBuilderBuildJavadoc(ClassName className) {
                         .add(javadocs, className, className)
                         .build();
     }
+
+    public static CodeBlock waiterOperationWithOverrideConfigConsumerBuilder(ClassName clientClassName,
+                                                                             ClassName requestClassName,
+                                                                             Map.Entry<String, WaiterDefinition> waiterDefinition,
+                                                                             OperationModel opModel) {
+        String javadocs = new DocumentationBuilder().description("Polls {@link $T#$N} API until the desired condition "
+                                                                 + "{@code $N} is met, "
+                                                                 + "or until it is determined that the resource will never "
+                                                                 + "enter into the desired state. \n "
+                                                                 + "<p>This is a convenience method to create an instance of "
+                                                                 + "the request builder and instance of the override config "
+                                                                 + "builder")
+                                                    .param(opModel.getInput().getVariableName(),
+                                                           "The consumer that will configure the request to be used for polling")
+                                                    .param("overrideConfig",
+                                                           "The consumer that will configure the per request override "
+                                                           + "configuration for waiters")
+                                                    .returns("WaiterResponse containing either a response or an exception that "
+                                                             + "has matched with the waiter success condition")
+                                                    .build();
+        return CodeBlock.builder()
+                        .add(javadocs, clientClassName, opModel.getMethodName(), waiterDefinition.getKey())
+                        .build();
+
+    }
+
+    public static CodeBlock waiterOperationWithOverrideConfig(ClassName clientClassName,
+                                                              Map.Entry<String, WaiterDefinition> waiterDefinition,
+                                                              OperationModel opModel) {
+        String javadocs = new DocumentationBuilder().description("Polls {@link $T#$N} API until the desired condition "
+                                                                 + "{@code $N} is met, "
+                                                                 + "or until it is determined that the resource will never "
+                                                                 + "enter into the desired state")
+                                                    .param(opModel.getInput().getVariableName(), "The request to be"
+                                                                                                 + " used"
+                                                                                                 + " for polling")
+                                                    .param("overrideConfig", "Per request "
+                                                                             + "override configuration for waiters")
+                                                    .returns("WaiterResponse containing either a response or an exception that "
+                                                             + "has matched with the waiter success condition")
+                                                    .build();
+        return CodeBlock.builder()
+                        .add(javadocs, clientClassName, opModel.getMethodName(), waiterDefinition.getKey())
+                        .build();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/waiters/BaseWaiterClassSpec.java

@@ -56,6 +56,7 @@
 import software.amazon.awssdk.codegen.poet.PoetUtils;
 import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
+import software.amazon.awssdk.core.retry.backoff.BackoffStrategy;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.core.waiters.WaiterAcceptor;
 import software.amazon.awssdk.core.waiters.WaiterOverrideConfiguration;
@@ -102,6 +103,7 @@ public TypeSpec poetSpec() {
         typeSpecBuilder.addMethod(staticErrorCodeMethod());
         typeSpecBuilder.addMethods(waiterOperations());
         typeSpecBuilder.addMethods(waiterAcceptorInitializers());
+        typeSpecBuilder.addMethods(waiterConfigInitializers());
         typeSpecBuilder.addFields(waitersFields());
         additionalTypeSpecModification(typeSpecBuilder);
 
@@ -173,31 +175,55 @@ private MethodSpec constructor() {
         return ctor.build();
     }
 
+    private List<MethodSpec> waiterConfigInitializers() {
+        List<MethodSpec> initializers = new ArrayList<>();
+        waiters.forEach((k, v) -> initializers.add(waiterConfigInitializer(k, v)));
+        return initializers;
+    }
+
+    private MethodSpec waiterConfigInitializer(String waiterKey, WaiterDefinition waiterDefinition) {
+        ClassName overrideConfig = ClassName.get(WaiterOverrideConfiguration.class);
+        MethodSpec.Builder configMethod =
+            MethodSpec.methodBuilder(waiterFieldName(waiterKey) + "Config")
+                      .addModifiers(PRIVATE, STATIC)
+                      .addParameter(overrideConfig, "overrideConfig")
+                      .returns(overrideConfig);
+
+        configMethod.addStatement("$T<$T> optionalOverrideConfig = Optional.ofNullable(overrideConfig)",
+                                  Optional.class,
+                                  WaiterOverrideConfiguration.class);
+        configMethod.addStatement("int maxAttempts = optionalOverrideConfig.flatMap(WaiterOverrideConfiguration::maxAttempts)"
+                                  + ".orElse($L)",
+                                  waiterDefinition.getMaxAttempts());
+        configMethod.addStatement("$T backoffStrategy = optionalOverrideConfig."
+                                  + "flatMap(WaiterOverrideConfiguration::backoffStrategy).orElse($T.create($T.ofSeconds($L)))",
+                                  BackoffStrategy.class,
+                                  FixedDelayBackoffStrategy.class,
+                                  Duration.class,
+                                  waiterDefinition.getDelay());
+        configMethod.addStatement("$T waitTimeout = optionalOverrideConfig.flatMap(WaiterOverrideConfiguration::waitTimeout)"
+                                  + ".orElse(null)",
+                                  Duration.class);
+
+        configMethod.addStatement("return WaiterOverrideConfiguration.builder().maxAttempts(maxAttempts).backoffStrategy"
+                                  + "(backoffStrategy).waitTimeout(waitTimeout).build()");
+        return configMethod.build();
+    }
+
     private CodeBlock waiterFieldInitialization(Map.Entry<String, WaiterDefinition> waiterDefinition) {
         String waiterKey = waiterDefinition.getKey();
-        String waiterName = lowercaseFirstChar(waiterKey);
         WaiterDefinition waiter = waiterDefinition.getValue();
-        String overrideConfigurationVarName = waiterName + "Strategy";
         OperationModel opModel = operationModel(waiter);
         CodeBlock.Builder codeBlockBuilder = CodeBlock
-            .builder()
-            .addStatement("$T $N = builder.overrideConfiguration == null ? $T.builder().maxAttempts($L)"
-                          + ".backoffStrategy($T.create($T.ofSeconds($L))).build() : builder.overrideConfiguration",
-                          WaiterOverrideConfiguration.class,
-                          overrideConfigurationVarName,
-                          WaiterOverrideConfiguration.class,
-                          waiter.getMaxAttempts(),
-                          FixedDelayBackoffStrategy.class,
-                          Duration.class,
-                          waiter.getDelay());
-
+            .builder();
 
         String waiterFieldName = waiterFieldName(waiterKey);
-        codeBlockBuilder.add("this.$L = $T.builder($T.class).overrideConfiguration($L).acceptors($LAcceptors())",
+        codeBlockBuilder.add("this.$L = $T.builder($T.class)"
+                             + ".acceptors($LAcceptors()).overrideConfiguration($LConfig(builder.overrideConfiguration))",
                              waiterFieldName,
                              waiterClassName,
                              ClassName.get(modelPackage, opModel.getReturnType().getReturnType()),
-                             overrideConfigurationVarName,
+                             waiterFieldName,
                              waiterFieldName);
 
         additionalWaiterConfig().ifPresent(codeBlockBuilder::add);
@@ -275,9 +301,33 @@ private List<MethodSpec> waiterOperations() {
     private Stream<MethodSpec> waiterOperations(Map.Entry<String, WaiterDefinition> waiterDefinition) {
         List<MethodSpec> methods = new ArrayList<>();
         methods.add(waiterOperation(waiterDefinition));
+        methods.add(waiterOperationWithOverrideConfig(waiterDefinition));
         return methods.stream();
     }
 
+    private MethodSpec waiterOperationWithOverrideConfig(Map.Entry<String, WaiterDefinition> waiterDefinition) {
+        String waiterMethodName = waiterDefinition.getKey();
+        OperationModel opModel = operationModel(waiterDefinition.getValue());
+
+        ClassName overrideConfig = ClassName.get(WaiterOverrideConfiguration.class);
+        ClassName requestType = ClassName.get(modelPackage, opModel.getInput().getVariableType());
+
+        String waiterFieldName = waiterFieldName(waiterDefinition.getKey());
+        MethodSpec.Builder builder = methodSignatureWithReturnType(waiterMethodName, opModel)
+            .addParameter(requestType, opModel.getInput().getVariableName())
+            .addParameter(overrideConfig, "overrideConfig")
+            .addModifiers(PUBLIC)
+            .addAnnotation(Override.class)
+            .addStatement("return $L.$L(() -> client.$N(applyWaitersUserAgent($N)), $LConfig(overrideConfig))",
+                          waiterFieldName,
+                          waiterClassName.simpleName().equals("Waiter") ? "run" : "runAsync",
+                          lowercaseFirstChar(waiterDefinition.getValue().getOperation()),
+                          opModel.getInput().getVariableName(),
+                          waiterFieldName);
+
+        return builder.build();
+    }
+
     private MethodSpec waiterOperation(Map.Entry<String, WaiterDefinition> waiterDefinition) {
         String waiterMethodName = waiterDefinition.getKey();
         OperationModel opModel = operationModel(waiterDefinition.getValue());

codegen/src/main/java/software/amazon/awssdk/codegen/poet/waiters/BaseWaiterInterfaceSpec.java

@@ -102,6 +102,8 @@ private Stream<MethodSpec> waiterOperations(Map.Entry<String, WaiterDefinition>
         List<MethodSpec> methods = new ArrayList<>();
         methods.add(waiterOperation(waiterDefinition));
         methods.add(waiterConsumerBuilderOperation(waiterDefinition));
+        methods.add(waiterOperationWithOverrideConfig(waiterDefinition));
+        methods.add(waiterConsumerBuilderOperationWithOverrideConfig(waiterDefinition));
         return methods.stream();
     }
 
@@ -118,6 +120,54 @@ private MethodSpec waiterOperation(Map.Entry<String, WaiterDefinition> waiterDef
         return unsupportedOperation(builder).build();
     }
 
+    private MethodSpec waiterOperationWithOverrideConfig(Map.Entry<String, WaiterDefinition> waiterDefinition) {
+        String waiterMethodName = waiterDefinition.getKey();
+        OperationModel opModel = model.getOperation(waiterDefinition.getValue().getOperation());
+        ClassName requestClass = ClassName.get(modelPackage,
+                                               opModel.getInput().getVariableType());
+        CodeBlock javadoc = WaiterDocs.waiterOperationWithOverrideConfig(
+            clientClassName(), waiterDefinition, opModel);
+
+        MethodSpec.Builder builder = methodSignatureWithReturnType(waiterMethodName, opModel)
+            .addParameter(requestClass, opModel.getInput().getVariableName())
+            .addParameter(ClassName.get(WaiterOverrideConfiguration.class), "overrideConfig")
+            .addJavadoc(javadoc);
+        return unsupportedOperation(builder).build();
+    }
+
+    private MethodSpec waiterConsumerBuilderOperationWithOverrideConfig(Map.Entry<String, WaiterDefinition> waiterDefinition) {
+        String waiterMethodName = waiterDefinition.getKey();
+        OperationModel opModel = model.getOperation(waiterDefinition.getValue().getOperation());
+        ClassName requestClass = ClassName.get(modelPackage,
+                                               opModel.getInput().getVariableType());
+        ParameterizedTypeName requestType = ParameterizedTypeName.get(ClassName.get(Consumer.class),
+                                                                      requestClass.nestedClass("Builder"));
+        ParameterizedTypeName overrideConfigType =
+            ParameterizedTypeName.get(ClassName.get(Consumer.class),
+                                      ClassName.get(WaiterOverrideConfiguration.class).nestedClass("Builder"));
+
+        CodeBlock javadoc = WaiterDocs.waiterOperationWithOverrideConfigConsumerBuilder(
+            clientClassName(), requestClass, waiterDefinition, opModel);
+
+        String inputVariable = opModel.getInput().getVariableName();
+
+        MethodSpec.Builder builder = methodSignatureWithReturnType(waiterMethodName, opModel)
+            .addParameter(requestType, inputVariable)
+            .addParameter(overrideConfigType, "overrideConfig")
+            .addJavadoc(javadoc);
+
+        builder.addModifiers(Modifier.DEFAULT, Modifier.PUBLIC)
+               .addStatement("return $L($T.builder().applyMutation($L).build(),"
+                             + "$T.builder().applyMutation($L).build())",
+                             getWaiterMethodName(waiterMethodName),
+                             requestClass,
+                             inputVariable,
+                             ClassName.get(WaiterOverrideConfiguration.class),
+                             "overrideConfig");
+
+        return builder.build();
+    }
+
     private MethodSpec waiterConsumerBuilderOperation(Map.Entry<String, WaiterDefinition> waiterDefinition) {
         String waiterMethodName = waiterDefinition.getKey();
         OperationModel opModel = model.getOperation(waiterDefinition.getValue().getOperation());

codegen/src/main/resources/software/amazon/awssdk/codegen/waiters/WaitersRuntime.java

@@ -33,17 +33,7 @@ private WaitersRuntime() {
     }
 
     private static List<WaiterAcceptor<Object>> defaultAcceptors() {
-        return Arrays.asList(throwOnUnmatchedExceptionWaiter(), retryOnUnmatchedResponseWaiter());
-    }
-
-    private static WaiterAcceptor<Object> throwOnUnmatchedExceptionWaiter() {
-        return WaiterAcceptor.errorOnExceptionAcceptor(t -> {
-            if (t instanceof RuntimeException) {
-                throw (RuntimeException) t;
-            }
-
-            throw SdkClientException.create("Encountered unexpected exception.", t);
-        });
+        return Collections.singletonList(retryOnUnmatchedResponseWaiter());
     }
 
     private static WaiterAcceptor<Object> retryOnUnmatchedResponseWaiter() {
@@ -502,7 +492,8 @@ public WaiterState waiterState() {
 
         @Override
         public boolean matches(SdkResponse response) {
-            return response.sdkHttpResponse().statusCode() == statusCode;
+            return response.sdkHttpResponse() != null &&
+                   response.sdkHttpResponse().statusCode() == statusCode;
         }
 
         @Override