Internal classes and RequestBatchManager Impelementation

services/sqs/pom.xml

@@ -91,5 +91,26 @@
             <version>${awsjavasdk.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>nl.jqno.equalsverifier</groupId>
+            <artifactId>equalsverifier</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-junit-jupiter</artifactId>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 </project>

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

@@ -0,0 +1,39 @@
+/*
+ * 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.services.sqs.internal.batchmanager;
+
+import java.util.concurrent.CompletableFuture;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+@SdkInternalApi
+public final class BatchingExecutionContext<RequestT, ResponseT> {
+
+    private final RequestT request;
+    private final CompletableFuture<ResponseT> response;
+
+    public BatchingExecutionContext(RequestT request, CompletableFuture<ResponseT> response) {
+        this.request = request;
+        this.response = response;
+    }
+
+    public RequestT request() {
+        return request;
+    }
+
+    public CompletableFuture<ResponseT> response() {
+        return response;
+    }
+}

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

@@ -0,0 +1,84 @@
+/*
+ * 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.services.sqs.internal.batchmanager;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledFuture;
+import java.util.function.BiConsumer;
+import java.util.function.Supplier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Outer map maps a batchKey (ex. queueUrl, overrideConfig etc.) to a {@link RequestBatchBuffer}
+ *
+ * @param <RequestT> the type of an outgoing response
+ */
+@SdkInternalApi
+public final class BatchingMap<RequestT, ResponseT> {
+
+    private final int maxBatchKeys;
+    private final int maxBufferSize;
+    private final Map<String, RequestBatchBuffer<RequestT, ResponseT>> batchContextMap;
+
+    public BatchingMap(int maxBatchKeys, int maxBufferSize) {
+        this.batchContextMap = new ConcurrentHashMap<>();
+        this.maxBatchKeys = maxBatchKeys;
+        this.maxBufferSize = maxBufferSize;
+    }
+
+    public void put(String batchKey, Supplier<ScheduledFuture<?>> scheduleFlush, RequestT request,
+                    CompletableFuture<ResponseT> response) throws IllegalStateException {
+        batchContextMap.computeIfAbsent(batchKey, k -> {
+            if (batchContextMap.size() == maxBatchKeys) {
+                throw new IllegalStateException("Reached MaxBatchKeys of: " + maxBatchKeys);
+            }
+            return new RequestBatchBuffer<>(maxBufferSize, scheduleFlush.get());
+        }).put(request, response);
+    }
+
+    public void putScheduledFlush(String batchKey, ScheduledFuture<?> scheduledFlush) {
+        batchContextMap.get(batchKey).putScheduledFlush(scheduledFlush);
+    }
+
+    public void forEach(BiConsumer<String, RequestBatchBuffer<RequestT, ResponseT>> action) {
+        batchContextMap.forEach(action);
+    }
+
+    public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableRequests(String batchKey,
+                                                                                        int maxBatchItems) {
+        return batchContextMap.get(batchKey).flushableRequests(maxBatchItems);
+    }
+
+    public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableScheduledRequests(String batchKey,
+                                                                                                 int maxBatchItems) {
+        return batchContextMap.get(batchKey).flushableScheduledRequests(maxBatchItems);
+    }
+
+    public void cancelScheduledFlush(String batchKey) {
+        batchContextMap.get(batchKey).cancelScheduledFlush();
+    }
+
+    public void clear() {
+        for (Map.Entry<String, RequestBatchBuffer<RequestT, ResponseT>> entry : batchContextMap.entrySet()) {
+            String key = entry.getKey();
+            entry.getValue().clear();
+            batchContextMap.remove(key);
+        }
+        batchContextMap.clear();
+    }
+}

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

@@ -0,0 +1,141 @@
+/*
+ * 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.services.sqs.internal.batchmanager;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ScheduledExecutorService;
+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;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResultEntry;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityRequest;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityResponse;
+import software.amazon.awssdk.services.sqs.model.SqsException;
+import software.amazon.awssdk.utils.Either;
+
+@SdkInternalApi
+public class ChangeMessageVisibilityBatchManager extends RequestBatchManager<ChangeMessageVisibilityRequest,
+    ChangeMessageVisibilityResponse,
+    ChangeMessageVisibilityBatchResponse> {
+
+    private final SqsAsyncClient sqsAsyncClient;
+
+    protected ChangeMessageVisibilityBatchManager(BatchOverrideConfiguration overrideConfiguration,
+                                                  ScheduledExecutorService scheduledExecutor,
+                                                  SqsAsyncClient sqsAsyncClient) {
+        super(overrideConfiguration, scheduledExecutor);
+        this.sqsAsyncClient = sqsAsyncClient;
+    }
+
+    private static ChangeMessageVisibilityBatchRequest createChangeMessageVisibilityBatchRequest(
+        List<IdentifiableMessage<ChangeMessageVisibilityRequest>> identifiedRequests, String batchKey) {
+        List<ChangeMessageVisibilityBatchRequestEntry> entries = identifiedRequests
+            .stream()
+            .map(identifiedRequest -> createChangeMessageVisibilityBatchRequestEntry(identifiedRequest.id(),
+                                                                                     identifiedRequest.message()))
+            .collect(Collectors.toList());
+        // Since requests are batched together according to a combination of their queueUrl and overrideConfiguration,
+        // all requests must have the same overrideConfiguration so it is sufficient to retrieve it from the first
+        // request.
+        Optional<AwsRequestOverrideConfiguration> overrideConfiguration = identifiedRequests.get(0).message()
+                                                                                            .overrideConfiguration();
+        return overrideConfiguration.map(
+                                        config -> ChangeMessageVisibilityBatchRequest.builder()
+                                                                                             .queueUrl(batchKey)
+                                                                                             .overrideConfiguration(config)
+                                                                                             .entries(entries)
+                                                                                             .build())
+                                    .orElseGet(() -> ChangeMessageVisibilityBatchRequest.builder()
+                                                                                        .queueUrl(batchKey)
+                                                                                        .entries(entries)
+                                                                                        .build());
+    }
+
+    private static ChangeMessageVisibilityBatchRequestEntry createChangeMessageVisibilityBatchRequestEntry(
+        String id,
+        ChangeMessageVisibilityRequest request) {
+        return ChangeMessageVisibilityBatchRequestEntry.builder().id(id).receiptHandle(request.receiptHandle())
+                                                       .visibilityTimeout(request.visibilityTimeout()).build();
+    }
+
+    private static IdentifiableMessage<ChangeMessageVisibilityResponse> createChangeMessageVisibilityResponse(
+        ChangeMessageVisibilityBatchResultEntry successfulEntry, ChangeMessageVisibilityBatchResponse batchResponse) {
+        String key = successfulEntry.id();
+        ChangeMessageVisibilityResponse.Builder builder = ChangeMessageVisibilityResponse.builder();
+        if (batchResponse.responseMetadata() != null) {
+            builder.responseMetadata(batchResponse.responseMetadata());
+        }
+        if (batchResponse.sdkHttpResponse() != null) {
+            builder.sdkHttpResponse(batchResponse.sdkHttpResponse());
+        }
+        ChangeMessageVisibilityResponse response = builder.build();
+        return new IdentifiableMessage<>(key, response);
+    }
+
+    private static IdentifiableMessage<Throwable> changeMessageVisibilityCreateThrowable(BatchResultErrorEntry failedEntry) {
+        String key = failedEntry.id();
+        AwsErrorDetails errorDetailsBuilder = AwsErrorDetails.builder().errorCode(failedEntry.code())
+                                                             .errorMessage(failedEntry.message()).build();
+        Throwable response = SqsException.builder().awsErrorDetails(errorDetailsBuilder).build();
+        return new IdentifiableMessage<>(key, response);
+    }
+
+
+
+    @Override
+    protected CompletableFuture<ChangeMessageVisibilityBatchResponse> batchAndSend(
+        List<IdentifiableMessage<ChangeMessageVisibilityRequest>> identifiedRequests, String batchKey) {
+        ChangeMessageVisibilityBatchRequest batchRequest = createChangeMessageVisibilityBatchRequest(identifiedRequests,
+                                                                                                     batchKey);
+        return sqsAsyncClient.changeMessageVisibilityBatch(batchRequest);
+    }
+
+    @Override
+    protected String getBatchKey(ChangeMessageVisibilityRequest request) {
+        return  request.overrideConfiguration().map(overrideConfig -> request.queueUrl() + overrideConfig.hashCode())
+                       .orElseGet(request::queueUrl);
+    }
+
+    @Override
+    protected List<Either<IdentifiableMessage<ChangeMessageVisibilityResponse>,
+        IdentifiableMessage<Throwable>>> mapBatchResponse(ChangeMessageVisibilityBatchResponse batchResponse) {
+
+        List<Either<IdentifiableMessage<ChangeMessageVisibilityResponse>, IdentifiableMessage<Throwable>>> mappedResponses =
+            new ArrayList<>();
+        batchResponse.successful().forEach(
+            batchResponseEntry -> {
+                IdentifiableMessage<ChangeMessageVisibilityResponse> response = createChangeMessageVisibilityResponse(
+                    batchResponseEntry, batchResponse);
+                mappedResponses.add(Either.left(response));
+            });
+        batchResponse.failed().forEach(batchResponseEntry -> {
+            IdentifiableMessage<Throwable> response = changeMessageVisibilityCreateThrowable(batchResponseEntry);
+            mappedResponses.add(Either.right(response));
+        });
+        return mappedResponses;
+
+    }
+}

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

@@ -15,25 +15,81 @@
 
 package software.amazon.awssdk.services.sqs.internal.batchmanager;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
 import software.amazon.awssdk.services.sqs.batchmanager.SqsAsyncBatchManager;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityRequest;
+import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityResponse;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageResponse;
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
+import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
+import software.amazon.awssdk.services.sqs.model.SendMessageResponse;
 import software.amazon.awssdk.utils.Validate;
 
 @SdkInternalApi
 public final class DefaultSqsAsyncBatchManager implements SqsAsyncBatchManager {
     // TODO : update the validation here while implementing this class in next PR
     private final SqsAsyncClient client;
-    private final  ScheduledExecutorService scheduledExecutor;
-    private final  BatchOverrideConfiguration overrideConfiguration;
+
+    private final SendMessageBatchManager sendMessageBatchManager;
+
+    private final DeleteMessageBatchManager deleteMessageBatchManager;
+
+    private final ChangeMessageVisibilityBatchManager changeMessageVisibilityBatchManager;
 
     private DefaultSqsAsyncBatchManager(DefaultBuilder builder) {
         this.client = Validate.notNull(builder.client, "client cannot be null");
-        this.scheduledExecutor = Validate.notNull(builder.scheduledExecutor, "scheduledExecutor cannot be null");
-        // TODO : create overrideConfiguration with Default values if null
-        this.overrideConfiguration = builder.overrideConfiguration;
+
+        ScheduledExecutorService scheduledExecutor = builder.scheduledExecutor;
+
+        this.sendMessageBatchManager = new SendMessageBatchManager(builder.overrideConfiguration,
+                                                                   scheduledExecutor,
+                                                                   client);
+        this.deleteMessageBatchManager = new DeleteMessageBatchManager(builder.overrideConfiguration,
+                                                                       scheduledExecutor,
+                                                                       client);
+        this.changeMessageVisibilityBatchManager = new ChangeMessageVisibilityBatchManager(builder.overrideConfiguration,
+                                                                                           scheduledExecutor,
+                                                                                           client);
+        //TODO : this will be updated while implementing the Receive Message Batch Manager
+    }
+
+    @SdkTestInternalApi
+    public DefaultSqsAsyncBatchManager(
+        SqsAsyncClient client,
+        SendMessageBatchManager sendMessageBatchManager,
+        DeleteMessageBatchManager deleteMessageBatchManager,
+        ChangeMessageVisibilityBatchManager changeMessageVisibilityBatchManager) {
+        this.sendMessageBatchManager = sendMessageBatchManager;
+        this.deleteMessageBatchManager = deleteMessageBatchManager;
+        this.changeMessageVisibilityBatchManager = changeMessageVisibilityBatchManager;
+        this.client = client;
+    }
+
+    @Override
+    public CompletableFuture<SendMessageResponse> sendMessage(SendMessageRequest request) {
+        return sendMessageBatchManager.batchRequest(request);
+    }
+
+    @Override
+    public CompletableFuture<DeleteMessageResponse> deleteMessage(DeleteMessageRequest request) {
+        return deleteMessageBatchManager.batchRequest(request);
+    }
+
+    @Override
+    public CompletableFuture<ChangeMessageVisibilityResponse> changeMessageVisibility(ChangeMessageVisibilityRequest request) {
+        return changeMessageVisibilityBatchManager.batchRequest(request);
+    }
+
+    @Override
+    public CompletableFuture<ReceiveMessageResponse> receiveMessage(ReceiveMessageRequest request) {
+        return null;
     }
 
     public static SqsAsyncBatchManager.Builder builder() {
@@ -42,6 +98,9 @@ public static SqsAsyncBatchManager.Builder builder() {
 
     @Override
     public void close() {
+        sendMessageBatchManager.close();
+        deleteMessageBatchManager.close();
+        changeMessageVisibilityBatchManager.close();
     }
 
     public static final class DefaultBuilder implements SqsAsyncBatchManager.Builder {