Byte Based batching for SendMessage API

Author: joviegas

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/BatchOverrideConfiguration.java

@@ -51,7 +51,8 @@ private BatchOverrideConfiguration(Builder builder) {
         this.maxBatchItems = Validate.isPositiveOrNull(builder.maxBatchItems, "maxBatchItems");
         this.maxBatchKeys = Validate.isPositiveOrNull(builder.maxBatchKeys, "maxBatchKeys");
         this.maxBufferSize = Validate.isPositiveOrNull(builder.maxBufferSize, "maxBufferSize");
-        this.batchSendRequestFrequency = Validate.isPositiveOrNull(builder.batchSendRequestFrequency, "batchSendRequestFrequency");
+        this.batchSendRequestFrequency = Validate.isPositiveOrNull(builder.batchSendRequestFrequency,
+                                                                   "batchSendRequestFrequency");
         this.visibilityTimeout = Validate.isPositiveOrNull(builder.visibilityTimeout, "visibilityTimeout");
         this.longPollWaitTimeout = Validate.isPositiveOrNull(builder.longPollWaitTimeout, "longPollWaitTimeout");
         this.minReceiveWaitTime = Validate.isPositiveOrNull(builder.minReceiveWaitTime, "minReceiveWaitTime");

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/BatchingExecutionContext.java

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.services.sqs.internal.batchmanager;
 
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 
@@ -24,12 +25,12 @@ public final class BatchingExecutionContext<RequestT, ResponseT> {
     private final RequestT request;
     private final CompletableFuture<ResponseT> response;
 
-    private final long responsePayload;
+    private final Optional<Long> responsePayload;
 
     public BatchingExecutionContext(RequestT request, CompletableFuture<ResponseT> response) {
         this.request = request;
         this.response = response;
-        responsePayload = ResponsePayloadCalculator.calculateMessageSize(request);
+        responsePayload = RequestPayloadCalculator.calculateMessageSize(request);
     }
 
     public RequestT request() {
@@ -41,7 +42,7 @@ public CompletableFuture<ResponseT> response() {
     }
 
 
-    public long responsePayload() {
+    public Optional<Long> responsePayload() {
         return responsePayload;
     }
 }

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/BatchingMap.java

@@ -20,8 +20,6 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledFuture;
 import java.util.function.BiConsumer;
-import java.util.function.BiPredicate;
-import java.util.function.Predicate;
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 
@@ -34,15 +32,17 @@
 public final class BatchingMap<RequestT, ResponseT> {
 
     private final int maxBatchKeys;
+    private final int maxBatchBytesSize;
+    private final int maxBatchSize;
     private final Map<String, RequestBatchBuffer<RequestT, ResponseT>> batchContextMap;
-    private final BiPredicate<Map<String, BatchingExecutionContext<RequestT, ResponseT>>, RequestT> flushCondition ;
-
 
     public BatchingMap(int maxBatchKeys,
-                       BiPredicate<Map<String, BatchingExecutionContext<RequestT, ResponseT>>, RequestT> flushCondition) {
+                       int maxBatchBytesSize,
+                       int maxBatchSize) {
         this.batchContextMap = new ConcurrentHashMap<>();
         this.maxBatchKeys = maxBatchKeys;
-        this.flushCondition = flushCondition;
+        this.maxBatchBytesSize = maxBatchBytesSize;
+        this.maxBatchSize = maxBatchSize;
     }
 
     public void put(String batchKey, Supplier<ScheduledFuture<?>> scheduleFlush, RequestT request,
@@ -51,10 +51,14 @@ public void put(String batchKey, Supplier<ScheduledFuture<?>> scheduleFlush, Req
             if (batchContextMap.size() == maxBatchKeys) {
                 throw new IllegalStateException("Reached MaxBatchKeys of: " + maxBatchKeys);
             }
-            return new RequestBatchBuffer<>(scheduleFlush.get(), flushCondition);
+            return new RequestBatchBuffer<>(scheduleFlush.get(), maxBatchSize, maxBatchBytesSize);
         }).put(request, response);
     }
 
+    public boolean contains(String batchKey) {
+        return batchContextMap.containsKey(batchKey);
+    }
+
     public void putScheduledFlush(String batchKey, ScheduledFuture<?> scheduledFlush) {
         batchContextMap.get(batchKey).putScheduledFlush(scheduledFlush);
     }

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/ChangeMessageVisibilityBatchManager.java

@@ -17,17 +17,14 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
-import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
 import software.amazon.awssdk.services.sqs.model.BatchResultErrorEntry;
 import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest;
 import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry;
@@ -49,22 +46,10 @@ public class ChangeMessageVisibilityBatchManager extends RequestBatchManager<Cha
     protected ChangeMessageVisibilityBatchManager(RequestBatchConfiguration overrideConfiguration,
                                                   ScheduledExecutorService scheduledExecutor,
                                                   SqsAsyncClient sqsAsyncClient) {
-        super(overrideConfiguration,
-              scheduledExecutor, (stringBatchingExecutionContextMap, changeMessageVisibilityRequest)
-                  -> shouldFlush(stringBatchingExecutionContextMap, changeMessageVisibilityRequest, overrideConfiguration)
-        );
+        super(overrideConfiguration, scheduledExecutor);
         this.sqsAsyncClient = sqsAsyncClient;
     }
 
-    private static boolean shouldFlush(Map<String, BatchingExecutionContext<ChangeMessageVisibilityRequest,
-        ChangeMessageVisibilityResponse>> contextMap,
-                                       ChangeMessageVisibilityRequest request, RequestBatchConfiguration configuration) {
-        if (request != null) {
-            return false;
-        }
-        return contextMap.size() >= configuration.maxBatchItems();
-    }
-
     private static ChangeMessageVisibilityBatchRequest createChangeMessageVisibilityBatchRequest(
         List<IdentifiableMessage<ChangeMessageVisibilityRequest>> identifiedRequests, String batchKey) {
         List<ChangeMessageVisibilityBatchRequestEntry> entries = identifiedRequests
@@ -79,10 +64,10 @@ private static ChangeMessageVisibilityBatchRequest createChangeMessageVisibility
                                                                                             .overrideConfiguration();
         return overrideConfiguration.map(
                                         config -> ChangeMessageVisibilityBatchRequest.builder()
-                                                                                             .queueUrl(batchKey)
-                                                                                             .overrideConfiguration(config)
-                                                                                             .entries(entries)
-                                                                                             .build())
+                                                                                     .queueUrl(batchKey)
+                                                                                     .overrideConfiguration(config)
+                                                                                     .entries(entries)
+                                                                                     .build())
                                     .orElseGet(() -> ChangeMessageVisibilityBatchRequest.builder()
                                                                                         .queueUrl(batchKey)
                                                                                         .entries(entries)
@@ -119,7 +104,6 @@ private static IdentifiableMessage<Throwable> changeMessageVisibilityCreateThrow
     }
 
 
-
     @Override
     protected CompletableFuture<ChangeMessageVisibilityBatchResponse> batchAndSend(
         List<IdentifiableMessage<ChangeMessageVisibilityRequest>> identifiedRequests, String batchKey) {
@@ -130,8 +114,8 @@ protected CompletableFuture<ChangeMessageVisibilityBatchResponse> batchAndSend(
 
     @Override
     protected String getBatchKey(ChangeMessageVisibilityRequest request) {
-        return  request.overrideConfiguration().map(overrideConfig -> request.queueUrl() + overrideConfig.hashCode())
-                       .orElseGet(request::queueUrl);
+        return request.overrideConfiguration().map(overrideConfig -> request.queueUrl() + overrideConfig.hashCode())
+                      .orElseGet(request::queueUrl);
     }
 
     @Override

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.services.sqs.internal.batchmanager;
 
+import static software.amazon.awssdk.services.sqs.internal.batchmanager.SqsMessageDefault.MAX_PAYLOAD_SIZE_BYTES;
+
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -49,13 +51,24 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) {
 
         ScheduledExecutorService scheduledExecutor = builder.scheduledExecutor;
 
-        this.sendMessageBatchManager = new SendMessageBatchManager(new RequestBatchConfiguration(builder.overrideConfiguration),
+        RequestBatchConfiguration.Builder configBuilder =
+            builder.overrideConfiguration != null ?
+            RequestBatchConfiguration.builder()
+                                     .batchSendRequestFrequency(builder.overrideConfiguration.batchSendRequestFrequency())
+                                     .maxBatchItems(builder.overrideConfiguration.maxBatchItems())
+                                     .maxBufferSize(builder.overrideConfiguration.maxBufferSize())
+                                     .maxBatchKeys(builder.overrideConfiguration.maxBatchKeys())
+                                                  : RequestBatchConfiguration.builder();
+
+        this.sendMessageBatchManager = new SendMessageBatchManager(configBuilder
+                                                                       .maxBatchBytesSize(MAX_PAYLOAD_SIZE_BYTES)
+                                                                       .build(),
                                                                    scheduledExecutor,
                                                                    client);
-        this.deleteMessageBatchManager = new DeleteMessageBatchManager(new RequestBatchConfiguration(builder.overrideConfiguration),
+        this.deleteMessageBatchManager = new DeleteMessageBatchManager(configBuilder.build(),
                                                                        scheduledExecutor,
                                                                        client);
-        this.changeMessageVisibilityBatchManager = new ChangeMessageVisibilityBatchManager(new RequestBatchConfiguration(builder.overrideConfiguration),
+        this.changeMessageVisibilityBatchManager = new ChangeMessageVisibilityBatchManager(configBuilder.build(),
                                                                                            scheduledExecutor,
                                                                                            client);
 

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DeleteMessageBatchManager.java

@@ -27,8 +27,6 @@
 import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.model.BatchResultErrorEntry;
-import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityRequest;
-import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityResponse;
 import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
 import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
 import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
@@ -48,9 +46,7 @@ public class DeleteMessageBatchManager extends RequestBatchManager<DeleteMessage
     protected DeleteMessageBatchManager(RequestBatchConfiguration overrideConfiguration,
                                         ScheduledExecutorService scheduledExecutor,
                                         SqsAsyncClient sqsAsyncClient) {
-        super(overrideConfiguration, scheduledExecutor, (stringBatchingExecutionContextMap, changeMessageVisibilityRequest)
-            -> shouldFlush(stringBatchingExecutionContextMap, changeMessageVisibilityRequest, overrideConfiguration)
-        );
+        super(overrideConfiguration, scheduledExecutor);
         this.sqsAsyncClient = sqsAsyncClient;
     }
 
@@ -65,7 +61,6 @@ private static boolean shouldFlush(Map<String, BatchingExecutionContext<DeleteMe
     }
 
 
-
     private static DeleteMessageBatchRequest createDeleteMessageBatchRequest(
         List<IdentifiableMessage<DeleteMessageRequest>> identifiedRequests, String batchKey) {
         List<DeleteMessageBatchRequestEntry> entries = identifiedRequests

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchBuffer.java

@@ -23,32 +23,26 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledFuture;
-import java.util.function.BiPredicate;
-import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 
 @SdkInternalApi
-public final class RequestBatchBuffer<RequestT, ResponseT>  {
+public final class RequestBatchBuffer<RequestT, ResponseT> {
     private final Object flushLock = new Object();
 
     private final Map<String, BatchingExecutionContext<RequestT, ResponseT>> idToBatchContext;
-
-
-    private BiPredicate<Map<String, BatchingExecutionContext<RequestT, ResponseT>>, RequestT> flushCondition ;
-
+    private final int maxBatchItems;
+    private final int maxBatchSizeInBytes;
     /**
      * Batch entries in a batch request require a unique ID so nextId keeps track of the ID to assign to the next
      * BatchingExecutionContext. For simplicity, the ID is just an integer that is incremented everytime a new request and
      * response pair is received.
      */
     private int nextId;
-    private int maxBatchItems;
-
     /**
-     * Keeps track of the ID of the next entry to be added in a batch request. This ID does not necessarily correlate to a
-     * request that already exists in the idToBatchContext map since it refers to the next entry (ex. if the last entry added
-     * to idToBatchContext had an id of 22, nextBatchEntry will have a value of 23).
+     * Keeps track of the ID of the next entry to be added in a batch request. This ID does not necessarily correlate to a request
+     * that already exists in the idToBatchContext map since it refers to the next entry (ex. if the last entry added to
+     * idToBatchContext had an id of 22, nextBatchEntry will have a value of 23).
      */
     private int nextBatchEntry;
 
@@ -58,27 +52,50 @@ public final class RequestBatchBuffer<RequestT, ResponseT>  {
     private ScheduledFuture<?> scheduledFlush;
 
     public RequestBatchBuffer(ScheduledFuture<?> scheduledFlush,
-                              BiPredicate<Map<String, BatchingExecutionContext<RequestT, ResponseT>>, RequestT> flushCondition) {
+                              int maxBatchItems, int maxBatchSizeInBytes) {
         this.idToBatchContext = new ConcurrentHashMap<>();
         this.nextId = 0;
         this.nextBatchEntry = 0;
         this.scheduledFlush = scheduledFlush;
-        this.flushCondition = flushCondition;
+        this.maxBatchItems = maxBatchItems;
+        this.maxBatchSizeInBytes = maxBatchSizeInBytes;
     }
 
+    /**
+     * When request is null it checks if current contents in idToBatchContext are flushable. When request is passed it calculate
+     * if the new request can overflow of ByteSize if yes then it flushes the messages
+     */
     public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableRequests(RequestT request) {
         synchronized (flushLock) {
-            if (flushCondition.test(idToBatchContext, request)) {
+            if (isByteSizeThresholdCrossed(request) || isMaxBatchSizeLimitReached(request)) {
                 return extractFlushedEntries(maxBatchItems);
             }
+
             return new ConcurrentHashMap<>();
         }
     }
 
+    private boolean isMaxBatchSizeLimitReached(RequestT request) {
+        int batchSizeLimit = request != null ? this.maxBatchItems + 1 : this.maxBatchItems;
+        return idToBatchContext.size() >= batchSizeLimit;
+    }
+
+    private boolean isByteSizeThresholdCrossed(RequestT request) {
+        if (maxBatchSizeInBytes < 0) {
+            return false;
+        }
+        long incomingRequestBytes = RequestPayloadCalculator.calculateMessageSize(request).orElse(0L);
+        long totalPayloadSize = idToBatchContext.values().stream()
+                                                .map(BatchingExecutionContext::responsePayload)
+                                                .mapToLong(opt -> opt.orElse(0L))
+                                                .sum() + incomingRequestBytes;
+        return totalPayloadSize > maxBatchSizeInBytes;
+    }
+
 
     public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableScheduledRequests(int maxBatchItems) {
         synchronized (flushLock) {
-            if (idToBatchContext.size() > 0) {
+            if (!idToBatchContext.isEmpty()) {
                 return extractFlushedEntries(maxBatchItems);
             }
             return new ConcurrentHashMap<>();