Adding batchManager() method into code generated sync and async clients (#2666)

* Moving protected APIs out of internal package

* Refactoring test import

* Adding configurable option for users to specify batchKey and batchBuffer limit

* Modifying batchBuffer and batchingMap to handle limits

* Returning exception on batch entry failure

* Modifying test to expect exception on batch entry failure

* Refactoring core DefaultBatchManager

Allows for either an exception or response to be returned by the responseMapper. Also adding tests and option to specify batchKey and batchBuffer limits.

* Modifying SqsBatchManager and associated files to accept generated client's executor and ScheduledExecutor

* Adding batchManager configuration class

* Adding batchManager config into CustomizationConfig

* Adding batchManager() method

* Adding batchManager config into sqs customization.config

* Removed need for a hasExecutor and hasScheduledExecutor configuration

* Removed need to check and sepcify instance type

* Adding .batchManager() method to async client

* Removing create() method in AsyncBatchManager and adding async batch manager return type to customization config

* Adding expected generated test files

* Adding tests for codegen batchManager() method

* Modifying SqsBatchManager tests to use client's batchManager() method

* Modifying codegen to use a scheduledExecutor (sync clients don't support regular executor)

* Addressing PR comments

* Addressing more PR comments

* Only create default executor if builder.executor is null. Adding tests to check as well

* Fixing checkstyle issues

* Reverting to using Executor instead of ExecutorService

Co-authored-by: Michael Li <[email protected]>

Author: 16lim21

codegen/src/main/java/software/amazon/awssdk/codegen/model/config/customization/BatchManagerMethod.java

@@ -0,0 +1,47 @@
+/*
+ * 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.codegen.model.config.customization;
+
+/**
+ * Config required to generate a batchManager method that returns an instance of a BatchManager in addition to any required
+ * executors or scheduledExecutors.
+ */
+public class BatchManagerMethod {
+
+    public static final String METHOD_NAME = "batchManager";
+
+    /** Fqcn of the return type of the operation for the sync client */
+    private String returnType;
+
+    /** Fqcn of the return type of the operation for the async client */
+    private String asyncReturnType;
+
+    public String getReturnType() {
+        return returnType;
+    }
+
+    public void setReturnType(String returnType) {
+        this.returnType = returnType;
+    }
+
+    public String getAsyncReturnType() {
+        return asyncReturnType;
+    }
+
+    public void setAsyncReturnType(String asyncReturnType) {
+        this.asyncReturnType = asyncReturnType;
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/model/config/customization/CustomizationConfig.java

@@ -189,6 +189,8 @@ public class CustomizationConfig {
 
     private String userAgent;
 
+    private BatchManagerMethod batchManagerMethod;
+
     private CustomizationConfig() {
     }
 
@@ -485,4 +487,12 @@ public CustomizationConfig withUserAgent(String userAgent) {
         this.userAgent = userAgent;
         return this;
     }
+
+    public BatchManagerMethod getBatchManagerMethod() {
+        return batchManagerMethod;
+    }
+
+    public void setBatchManagerMethod(BatchManagerMethod batchManagerMethod) {
+        this.batchManagerMethod = batchManagerMethod;
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

@@ -23,6 +23,7 @@
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
+import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.batchMangerMethod;
 import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
 
 import com.squareup.javapoet.ClassName;
@@ -83,13 +84,15 @@ public final class AsyncClientClass extends AsyncClientInterface {
     private final PoetExtensions poetExtensions;
     private final ClassName className;
     private final ProtocolSpec protocolSpec;
+    private boolean hasScheduledExecutor;
 
     public AsyncClientClass(GeneratorTaskParams dependencies) {
         super(dependencies.getModel());
         this.model = dependencies.getModel();
         this.poetExtensions = dependencies.getPoetExtensions();
         this.className = poetExtensions.getClientClass(model.getMetadata().getAsyncClient());
         this.protocolSpec = getProtocolSpecs(poetExtensions, model);
+        this.hasScheduledExecutor = false;
     }
 
     @Override
@@ -141,12 +144,16 @@ public TypeSpec poetSpec() {
         protocolSpec.createErrorResponseHandler().ifPresent(classBuilder::addMethod);
 
         if (model.hasWaiters()) {
-            classBuilder.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
-                                                   .addModifiers(PRIVATE, FINAL)
-                                                   .build());
             classBuilder.addMethod(waiterImplMethod());
+            addScheduledExecutorIfNeeded(classBuilder);
+        }
+
+        if (model.getCustomizationConfig().getBatchManagerMethod() != null) {
+            classBuilder.addMethod(batchMangerMethod(model, false));
+            addScheduledExecutorIfNeeded(classBuilder);
         }
 
+
         return classBuilder.build();
     }
 
@@ -195,7 +202,7 @@ private MethodSpec constructor(Builder classBuilder) {
             builder.endControlFlow();
         }
 
-        if (model.hasWaiters()) {
+        if (model.hasWaiters() || model.getCustomizationConfig().getBatchManagerMethod() != null) {
             builder.addStatement("this.executorService = clientConfiguration.option($T.SCHEDULED_EXECUTOR_SERVICE)",
                                  SdkClientOption.class);
         }
@@ -485,4 +492,13 @@ private boolean hasStreamingV4AuthOperations() {
         return model.getOperations().values().stream()
                 .anyMatch(this::shouldUseAsyncWithBodySigner);
     }
+
+    private void addScheduledExecutorIfNeeded(Builder classBuilder) {
+        if (!hasScheduledExecutor) {
+            classBuilder.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
+                                           .addModifiers(PRIVATE, FINAL)
+                                           .build());
+            hasScheduledExecutor = true;
+        }
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientInterface.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.codegen.docs.DocConfiguration;
 import software.amazon.awssdk.codegen.docs.SimpleMethodOverload;
 import software.amazon.awssdk.codegen.docs.WaiterDocs;
+import software.amazon.awssdk.codegen.model.config.customization.BatchManagerMethod;
 import software.amazon.awssdk.codegen.model.config.customization.UtilitiesMethod;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -108,6 +109,10 @@ public TypeSpec poetSpec() {
             result.addMethod(waiterMethod());
         }
 
+        if (model.getCustomizationConfig().getBatchManagerMethod() != null) {
+            result.addMethod(batchManagerMethod());
+        }
+
         return result.build();
     }
 
@@ -467,4 +472,17 @@ private MethodSpec waiterMethod() {
                          .addJavadoc(WaiterDocs.waiterMethodInClient(poetExtensions.getAsyncWaiterInterface()))
                          .build();
     }
+
+    private MethodSpec batchManagerMethod() {
+        BatchManagerMethod config = model.getCustomizationConfig().getBatchManagerMethod();
+        ClassName returnType = PoetUtils.classNameFromFqcn(config.getAsyncReturnType());
+
+        return MethodSpec.methodBuilder("batchManager")
+                         .returns(returnType)
+                         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
+                         .addStatement("throw new $T()", UnsupportedOperationException.class)
+                         .addJavadoc("Creates an instance of {@link $T} object with the "
+                                     + "configuration set on this client.", returnType)
+                         .build();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java

@@ -32,6 +32,7 @@
 import software.amazon.awssdk.arns.Arn;
 import software.amazon.awssdk.auth.signer.EventStreamAws4Signer;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
+import software.amazon.awssdk.codegen.model.config.customization.BatchManagerMethod;
 import software.amazon.awssdk.codegen.model.config.customization.S3ArnableFieldConfig;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
@@ -268,6 +269,25 @@ static Optional<CodeBlock> addS3ArnableFieldCode(OperationModel opModel, Interme
         return Optional.empty();
     }
 
+    static MethodSpec batchMangerMethod(IntermediateModel model, boolean isSync) {
+        String scheduledExecutor = "executorService";
+        BatchManagerMethod config = model.getCustomizationConfig().getBatchManagerMethod();
+        ClassName returnType;
+        if (isSync) {
+            returnType = PoetUtils.classNameFromFqcn(config.getReturnType());
+        } else {
+            returnType = PoetUtils.classNameFromFqcn(config.getAsyncReturnType());
+        }
+
+        return MethodSpec.methodBuilder(BatchManagerMethod.METHOD_NAME)
+                         .addModifiers(Modifier.PUBLIC)
+                         .addAnnotation(Override.class)
+                         .returns(returnType)
+                         .addStatement("return $T.builder().client(this).scheduledExecutor($N).build()",
+                                       returnType, scheduledExecutor)
+                         .build();
+    }
+
     /**
      * Given operation and c2j name, returns the String that represents calling the
      * c2j member's getter method in the opmodel input shape.

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java

@@ -21,6 +21,7 @@
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.addS3ArnableFieldCode;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
+import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.batchMangerMethod;
 
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
@@ -32,6 +33,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -116,6 +118,13 @@ public TypeSpec poetSpec() {
             classBuilder.addMethod(utilitiesMethod());
         }
 
+        if (model.getCustomizationConfig().getBatchManagerMethod() != null) {
+            classBuilder.addMethod(batchMangerMethod(model, true));
+            classBuilder.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
+                                           .addModifiers(PRIVATE, FINAL)
+                                           .build());
+        }
+
         model.getEndpointOperation().ifPresent(
             o -> classBuilder.addField(EndpointDiscoveryRefreshCache.class, "endpointDiscoveryCache", PRIVATE));
 
@@ -181,6 +190,11 @@ private MethodSpec constructor() {
             builder.endControlFlow();
         }
 
+        if (model.getCustomizationConfig().getBatchManagerMethod() != null) {
+            builder.addStatement("this.executorService = clientConfiguration.option($T.SCHEDULED_EXECUTOR_SERVICE)",
+                                 SdkClientOption.class);
+        }
+
         return builder.build();
     }
 

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientInterface.java

@@ -42,6 +42,7 @@
 import software.amazon.awssdk.codegen.docs.DocConfiguration;
 import software.amazon.awssdk.codegen.docs.SimpleMethodOverload;
 import software.amazon.awssdk.codegen.docs.WaiterDocs;
+import software.amazon.awssdk.codegen.model.config.customization.BatchManagerMethod;
 import software.amazon.awssdk.codegen.model.config.customization.UtilitiesMethod;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -110,6 +111,10 @@ public TypeSpec poetSpec() {
             result.addMethod(waiterMethod());
         }
 
+        if (model.getCustomizationConfig().getBatchManagerMethod() != null) {
+            result.addMethod(batchManagerMethod());
+        }
+
         return result.build();
     }
 
@@ -498,4 +503,17 @@ private MethodSpec waiterMethod() {
                          .addJavadoc(WaiterDocs.waiterMethodInClient(poetExtensions.getSyncWaiterInterface()))
                          .build();
     }
+
+    private MethodSpec batchManagerMethod() {
+        BatchManagerMethod config = model.getCustomizationConfig().getBatchManagerMethod();
+        ClassName returnType = PoetUtils.classNameFromFqcn(config.getReturnType());
+
+        return MethodSpec.methodBuilder("batchManager")
+                         .returns(returnType)
+                         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
+                         .addStatement("throw new $T()", UnsupportedOperationException.class)
+                         .addJavadoc("Creates an instance of {@link $T} object with the "
+                                     + "configuration set on this client.", returnType)
+                         .build();
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java

@@ -109,6 +109,18 @@ public static IntermediateModel internalConfigModels() {
         return new IntermediateModelBuilder(models).build();
     }
 
+    public static IntermediateModel batchManagerModels() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/batchmanager/service-2.json").getFile());
+        File customizationModel = new File(ClientTestModels.class.getResource("client/c2j/batchmanager/customization.config").getFile());
+
+        C2jModels models = C2jModels.builder()
+                                    .serviceModel(getServiceModel(serviceModel))
+                                    .customizationConfig(getCustomizationConfig(customizationModel))
+                                    .build();
+
+        return new IntermediateModelBuilder(models).build();
+    }
+
     private static ServiceModel getServiceModel(File file) {
         return ModelLoaderUtils.loadModel(ServiceModel.class, file);
     }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/PoetClientFunctionalTests.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.ClientTestModels;
+import software.amazon.awssdk.codegen.poet.waiters.WaiterInterfaceSpec;
 
 public class PoetClientFunctionalTests {
 
@@ -117,4 +118,29 @@ public void syncClientCustomServiceMetaData() throws Exception {
         ClassSpec syncClientCustomServiceMetaData = createSyncClientClass(ClientTestModels.customContentTypeModels());
         assertThat(syncClientCustomServiceMetaData, generatesTo("test-customservicemetadata-sync.java"));
     }
+
+    @Test
+    public void syncClientBatchManager() throws Exception {
+        ClassSpec syncClientBatchManager = createSyncClientClass(ClientTestModels.batchManagerModels());
+        assertThat(syncClientBatchManager, generatesTo("test-batchmanager-sync-class.java"));
+    }
+
+    @Test
+    public void asyncClientBatchManager() throws Exception {
+        ClassSpec asyncClientBatchManager = new AsyncClientClass(
+            GeneratorTaskParams.create(ClientTestModels.batchManagerModels(), "sources/", "tests/"));
+        assertThat(asyncClientBatchManager, generatesTo("test-batchmanager-async-class.java"));
+    }
+
+    @Test
+    public void syncClientBatchManagerInterface() throws Exception {
+        ClassSpec syncClientBatchManagerInterface = new SyncClientInterface(ClientTestModels.batchManagerModels());
+        assertThat(syncClientBatchManagerInterface, generatesTo("test-batchmanager-sync-interface.java"));
+    }
+
+    @Test
+    public void asyncClientBatchManagerInterface() throws Exception {
+        ClassSpec asyncClientBatchManagerInterface = new AsyncClientInterface(ClientTestModels.batchManagerModels());
+        assertThat(asyncClientBatchManagerInterface, generatesTo("test-batchmanager-async-interface.java"));
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/batchmanager/customization.config

@@ -0,0 +1,6 @@
+{
+    "batchManagerMethod": {
+        "returnType": "software.amazon.awssdk.services.batchmanagertest.batchmanager.SyncBatchManagerTest",
+        "asyncReturnType": "software.amazon.awssdk.services.batchmanagertest.batchmanager.AsyncBatchManagerTest"
+    }
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/batchmanager/service-2.json

@@ -0,0 +1,20 @@
+{
+  "version":"2.0",
+  "metadata":{
+    "apiVersion":"2016-03-11",
+    "endpointPrefix":"batchmanager",
+    "jsonVersion":"1.1",
+    "protocol":"rest-json",
+    "serviceAbbreviation":"BatchManager",
+    "serviceFullName":"BatchManager",
+    "serviceId":"BatchManager",
+    "signatureVersion":"v4",
+    "uid":"batchmanager-2016-03-11"
+  },
+  "operations":{
+  },
+  "shapes": {
+    "String":{"type":"string"}
+  },
+  "documentation": "A service that implements the batchManager() method"
+}