SQS Automatic Request Batching (#5580)

* Codegenerate BatchManager API under AsyncClient and Initail Interfaces for BatchManager (#5321)

* Codegenerate BatchManager API under AsyncClient and Add empty initial Batchmanager interfaces and Implementations

* Addressed review comments

* Added Internal classes required for BatchManager Implementation

* Revert "Added Internal classes required for BatchManager Implementation"

This reverts commit 318969b.

* Internal classes and RequestBatchManager Impelementation (#5418)

* Added Internal classes required for BatchManager Implementation

* Added Batch Send Implementation

* Handled review comments

* Handled review comments

* Handled review comments

* Made RequestBatchManager class a Abstract class

* Checkstyle issues

* Removed unused methods

* New lines removed

* Made public static to private state for sqsBatch functions

* Constants added

* Sonar cloud issues fixed

* commit to check why test on codebuild

* Increased Timeouts for get

* Added abstract methods

* Handled comments to remove Builders

* Handled comments to take care when batchmanager closed while pending requests

* Handled comments

* Checkstyle issue

* Added Consumer builders args for existing APIs of BatchManager (#5514)

* Receive Batch Manager Implementation (#5488)

* Add Recieve Buffer Queue And its related configuration

* Update ReceiveBatch  manager

* Recieve Batch Manager Implementation

* Receive Batch Manager Implemetation

* Handled review comments

* Checkstyle failure

* Flsky test case fixed

* Flaky test case fixed

* Hamdled review comments

* Handled comments

* Removed ReceiveMessageCompletableFuture

* SdkClosable implemented

* Added ReceiveMessageBatchManager class for completeness

* Checkstyle issues

* Null checks

* Handled comments from Zoe

* Updated the defaults to 50ms same as V1 after surface area review

* Revert "Updated the defaults to 50ms same as V1 after surface area review"

This reverts commit e7d2295.

* Bytes Based batching for SendMessageRequest Batching (#5540)

* Initial changes

* Initial changes 2

* Byte Based batching for SendMessage API

* Byte Based batching for SendMessage API

* Handled comments

* Checkstyle issue

* Add User Agent for Sqs Calls made using Automatic Batching Manager  (#5546)

* Add User Agent for Sqs Calls made using Automatic Batching Manager as hll/abm

* Review comments

* Update comments from  PR 5488 (#5550)

* Update comments of PR 5488

* Update comments from  PR 5488

* Handled surface area review comments (#5563)

* Initial version

* Intermediate changes

* Update after internal poll

* ResponseCOnfiguration construction updated

* RequestOverride configuration check added to Bypass batch manager

* Handled review comments

* Removed TODO since validations are handled in BatchPverrideConfiguration

* Fix issue where the Scheduled Timeout was incorrectly completing the futures with empty messages (#5571)

* Fix issue where the Scheduled Timeout was incorrectly completing the futures with empty messages

* Handled review comments

* Integ test for Automatic Request Batching (#5576)

* feat(sqs): add BatchManager for client-side request batching to Amazon SQS

The new BatchManager allows for simple request batching using client-side buffering, improving cost efficiency and reducing the number of requests sent to Amazon SQS. The client-side buffering supports up to 10 requests per batch and is supported by the SqsAsyncClient. Batched requests, along with receive message polling, help to increase throughput.

* Add check for scheduledExecutor such that it not null when creating SqsAsyncBatchManager (#5582)

* Add check for scheduledExecutor such that it not null when creating SqsAsyncBatchManager

* Update services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManagerBuilderTest.java

Co-authored-by: David Ho <[email protected]>

* Update services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManagerBuilderTest.java

Co-authored-by: David Ho <[email protected]>

* Update services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManagerBuilderTest.java

Co-authored-by: David Ho <[email protected]>

---------

Co-authored-by: David Ho <[email protected]>

* Updating Timeouts in gets so that we dont wait infinitely

---------

Co-authored-by: David Ho <[email protected]>

Author: joviegas

.changes/next-release/feature-AmazonSimpleQueueService-97fc706.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "Amazon Simple Queue Service",
+    "contributor": "",
+    "description": "The AWS SDK for Java now supports a new `BatchManager` for Amazon Simple Queue Service (SQS), allowing for client-side request batching with `SqsAsyncClient`. This feature improves cost efficiency by buffering up to 10 requests before sending them as a batch to SQS. The implementation also supports receive message polling, which further enhances throughput by minimizing the number of individual requests sent. The batched requests help to optimize performance and reduce the costs associated with using Amazon SQS."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/AddMetadata.java

@@ -60,6 +60,7 @@ public static Metadata constructMetadata(ServiceModel serviceModel,
                 .withBaseBuilder(String.format(Constant.BASE_BUILDER_CLASS_NAME_PATTERN, serviceName))
                 .withDocumentation(serviceModel.getDocumentation())
                 .withServiceAbbreviation(serviceMetadata.getServiceAbbreviation())
+                .withBatchmanagerPackageName(namingStrategy.getBatchManagerPackageName(serviceName))
                 .withServiceFullName(serviceMetadata.getServiceFullName())
                 .withServiceName(serviceName)
                 .withSyncClient(String.format(Constant.SYNC_CLIENT_CLASS_NAME_PATTERN, serviceName))

codegen/src/main/java/software/amazon/awssdk/codegen/internal/Constant.java

@@ -76,6 +76,8 @@ public final class Constant {
 
     public static final String PACKAGE_NAME_SMOKE_TEST_PATTERN = "%s.smoketests";
 
+    public static final String PACKAGE_NAME_BATCHMANAGER_PATTERN = "%s.batchmanager";
+
     public static final String PACKAGE_NAME_CUSTOM_AUTH_PATTERN = "%s.auth";
 
     public static final String AUTH_POLICY_ENUM_CLASS_DIR = "software/amazon/awssdk/auth/policy/actions";

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

@@ -342,6 +342,11 @@ public class CustomizationConfig {
      */
     private Map<String, PreClientExecutionRequestCustomizer> preClientExecutionRequestCustomizer;
 
+    /**
+     * A boolean flag to indicate if Automatic Batch Request is supported.
+     */
+    private boolean batchManagerSupported;
+
     private CustomizationConfig() {
     }
 
@@ -901,4 +906,12 @@ public void setPreClientExecutionRequestCustomizer(Map<String, PreClientExecutio
         this.preClientExecutionRequestCustomizer = preClientExecutionRequestCustomizer;
     }
 
+    public boolean getBatchManagerSupported() {
+        return batchManagerSupported;
+    }
+
+    public void setBatchManagerSupported(boolean batchManagerSupported) {
+        this.batchManagerSupported = batchManagerSupported;
+    }
+
 }

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/Metadata.java

@@ -71,6 +71,8 @@ public class Metadata {
 
     private String waitersPackageName;
 
+    private String batchManagerPackageName;
+
     private String endpointRulesPackageName;
 
     private String authSchemePackageName;
@@ -789,4 +791,22 @@ public String getFullInternalJmesPathPackageName() {
         return joinPackageNames(getFullJmesPathPackageName(), "internal");
     }
 
+    public Metadata withBatchmanagerPackageName(String batchmanagerPackageName) {
+        setBatchManagerPackageName(batchmanagerPackageName);
+        return this;
+    }
+
+
+    public String getBatchManagerPackageName() {
+        return batchManagerPackageName;
+    }
+
+    public void setBatchManagerPackageName(String batchManagerPackageName) {
+        this.batchManagerPackageName = batchManagerPackageName;
+    }
+
+    public String getFullBatchManagerPackageName() {
+        return joinPackageNames(rootPackageName, getBatchManagerPackageName());
+    }
+
 }

codegen/src/main/java/software/amazon/awssdk/codegen/naming/DefaultNamingStrategy.java

@@ -189,6 +189,11 @@ public String getJmesPathPackageName(String serviceName) {
         return getCustomizedPackageName(concatServiceNameIfShareModel(serviceName), Constant.PACKAGE_NAME_JMESPATH_PATTERN);
     }
 
+    @Override
+    public String getBatchManagerPackageName(String serviceName) {
+        return getCustomizedPackageName(concatServiceNameIfShareModel(serviceName), Constant.PACKAGE_NAME_BATCHMANAGER_PATTERN);
+    }
+
     @Override
     public String getSmokeTestPackageName(String serviceName) {
 

codegen/src/main/java/software/amazon/awssdk/codegen/naming/NamingStrategy.java

@@ -74,6 +74,11 @@ public interface NamingStrategy {
      */
     String getJmesPathPackageName(String serviceName);
 
+    /**
+     * Retrieve the batchManager package name that should be used based on the service name.
+     */
+    String getBatchManagerPackageName(String serviceName);
+
     /**
      * Retrieve the smote test package name that should be used based on the service name.
      */

codegen/src/main/java/software/amazon/awssdk/codegen/poet/PoetExtension.java

@@ -196,4 +196,9 @@ public boolean isResponse(ShapeModel shapeModel) {
     public boolean isRequest(ShapeModel shapeModel) {
         return shapeModel.getShapeType() == ShapeType.Request;
     }
+
+    public ClassName getBatchManagerAsyncInterface() {
+        return ClassName.get(model.getMetadata().getFullBatchManagerPackageName(),
+                             model.getMetadata().getServiceName() + "AsyncBatchManager");
+    }
 }

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

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.codegen.poet.client;
 
+import static com.squareup.javapoet.TypeSpec.Builder;
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 import static javax.lang.model.element.Modifier.FINAL;
@@ -104,6 +105,7 @@ public final class AsyncClientClass extends AsyncClientInterface {
     private final ClassName serviceClientConfigurationClassName;
     private final ServiceClientConfigurationUtils configurationUtils;
     private final boolean useSraAuth;
+    private boolean hasScheduledExecutor;
 
     public AsyncClientClass(GeneratorTaskParams dependencies) {
         super(dependencies.getModel());
@@ -139,7 +141,7 @@ protected void addModifiers(TypeSpec.Builder type) {
     }
 
     @Override
-    protected void addFields(TypeSpec.Builder type) {
+    protected void addFields(Builder type) {
         type.addField(FieldSpec.builder(ClassName.get(Logger.class), "log")
                                .addModifiers(PRIVATE, STATIC, FINAL)
                                .initializer("$T.getLogger($T.class)", LoggerFactory.class,
@@ -155,6 +157,10 @@ protected void addFields(TypeSpec.Builder type) {
             type.addField(AwsJsonProtocolFactory.class, "jsonProtocolFactory", PRIVATE, FINAL);
         }
 
+        if (shouldAddScheduledExecutor()) {
+            addScheduledExecutorIfNeeded(type);
+        }
+
         model.getEndpointOperation().ifPresent(
             o -> type.addField(EndpointDiscoveryRefreshCache.class, "endpointDiscoveryCache", PRIVATE));
     }
@@ -180,9 +186,6 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
 
     @Override
     protected void addWaiterMethod(TypeSpec.Builder type) {
-        type.addField(FieldSpec.builder(ClassName.get(ScheduledExecutorService.class), "executorService")
-                                       .addModifiers(PRIVATE, FINAL)
-                                       .build());
 
         MethodSpec waiter = MethodSpec.methodBuilder("waiter")
                                       .addModifiers(PUBLIC)
@@ -263,14 +266,18 @@ private MethodSpec constructor(TypeSpec.Builder classBuilder) {
             builder.endControlFlow();
         }
 
-        if (model.hasWaiters()) {
+        if (shouldAddScheduledExecutor()) {
             builder.addStatement("this.executorService = clientConfiguration.option($T.SCHEDULED_EXECUTOR_SERVICE)",
                                  SdkClientOption.class);
         }
 
         return builder.build();
     }
 
+    private boolean shouldAddScheduledExecutor() {
+        return model.hasWaiters() || model.getCustomizationConfig().getBatchManagerSupported();
+    }
+
     private boolean hasOperationWithEventStreamOutput() {
         return model.getOperations().values().stream().anyMatch(OperationModel::hasEventStreamOutput);
     }
@@ -547,6 +554,26 @@ protected MethodSpec utilitiesMethod() {
                          .build();
     }
 
+    @Override
+    protected void addBatchManagerMethod(Builder type) {
+
+        String scheduledExecutor = "executorService";
+        ClassName returnType;
+
+        returnType = poetExtensions.getBatchManagerAsyncInterface();
+
+        MethodSpec batchManager = MethodSpec.methodBuilder("batchManager")
+                                            .addModifiers(PUBLIC)
+                                            .addAnnotation(Override.class)
+                                            .returns(returnType)
+                                            .addStatement("return $T.builder().client(this).scheduledExecutor($N).build()",
+                                                          returnType, scheduledExecutor)
+                                            .build();
+
+
+        type.addMethod(batchManager);
+    }
+
     private MethodSpec resolveMetricPublishersMethod() {
         String clientConfigName = "clientConfiguration";
         String requestOverrideConfigName = "requestOverrideConfiguration";
@@ -621,4 +648,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

@@ -96,6 +96,9 @@ public TypeSpec poetSpec() {
         if (model.hasWaiters()) {
             addWaiterMethod(result);
         }
+        if (model.getCustomizationConfig().getBatchManagerSupported()) {
+            addBatchManagerMethod(result);
+        }
         result.addMethod(serviceClientConfigMethod());
         addAdditionalMethods(result);
         addCloseMethod(result);
@@ -162,6 +165,16 @@ protected void addWaiterMethod(TypeSpec.Builder type) {
         type.addMethod(waiterOperationBody(builder).build());
     }
 
+    protected void addBatchManagerMethod(TypeSpec.Builder type) {
+        ClassName returnType = poetExtensions.getBatchManagerAsyncInterface();
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("batchManager")
+                                               .addModifiers(PUBLIC)
+                                               .returns(returnType)
+                                               .addJavadoc("Creates an instance of {@link $T} object with the "
+                                                           + "configuration set on this client.", returnType);
+        type.addMethod(batchManagerOperationBody(builder).build());
+    }
+
     @Override
     public ClassName className() {
         return className;
@@ -532,4 +545,10 @@ protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
         return builder.addModifiers(DEFAULT, PUBLIC)
                       .addStatement("throw new $T()", UnsupportedOperationException.class);
     }
+
+    protected MethodSpec.Builder batchManagerOperationBody(MethodSpec.Builder builder) {
+        return builder.addModifiers(DEFAULT, PUBLIC)
+                      .addStatement("throw new $T()", UnsupportedOperationException.class);
+    }
+
 }

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

@@ -217,4 +217,9 @@ protected MethodSpec.Builder utilitiesOperationBody(MethodSpec.Builder builder)
     protected MethodSpec.Builder waiterOperationBody(MethodSpec.Builder builder) {
         return builder.addAnnotation(Override.class).addStatement("return delegate.waiter()");
     }
+
+    @Override
+    protected MethodSpec.Builder batchManagerOperationBody(MethodSpec.Builder builder) {
+        return builder.addAnnotation(Override.class).addStatement("return delegate.batchManager()");
+    }
 }

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

@@ -398,6 +398,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/AsyncClientClassTest.java

@@ -18,6 +18,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.awsJsonServiceModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.awsQueryCompatibleJsonServiceModels;
+import static software.amazon.awssdk.codegen.poet.ClientTestModels.batchManagerModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.customContentTypeModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.customPackageModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.endpointDiscoveryModels;
@@ -92,6 +93,12 @@ public void asyncClientCustomPackageName() {
         assertThat(syncClientCustomServiceMetaData, generatesTo("test-custompackage-async.java"));
     }
 
+    @Test
+    public void asyncClientBatchManager() {
+        ClassSpec aSyncClientBatchManager = createAsyncClientClass(batchManagerModels());
+        assertThat(aSyncClientBatchManager, generatesTo("test-batchmanager-async.java"));
+    }
+
     private AsyncClientClass createAsyncClientClass(IntermediateModel model) {
         return new AsyncClientClass(GeneratorTaskParams.create(model, "sources/", "tests/", "resources/"));
     }

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

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.codegen.poet.client;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static software.amazon.awssdk.codegen.poet.ClientTestModels.batchManagerModels;
 import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels;
 import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
 
@@ -28,4 +29,10 @@ public void asyncClientInterface() {
         ClassSpec asyncClientInterface = new AsyncClientInterface(restJsonServiceModels());
         assertThat(asyncClientInterface, generatesTo("test-json-async-client-interface.java"));
     }
+
+    @Test
+    public void asyncClientInterfaceWithBatchManager() {
+        ClassSpec asyncClientInterface = new AsyncClientInterface(batchManagerModels());
+        assertThat(asyncClientInterface, generatesTo("test-json-async-client-interface-batchmanager.java"));
+    }
 }

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

@@ -0,0 +1,3 @@
+{
+    "batchManagerSupported": true
+}