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/BatchAndSend.java

@@ -0,0 +1,32 @@
+/*
+ * 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.List;
+import java.util.concurrent.CompletableFuture;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Takes a list of identified requests in addition to a batchKey and batches the requests into a batch request.
+ * It then sends the batch request and returns a CompletableFuture of the response.
+ * @param <RequestT> the type of an outgoing request.
+ * @param <BatchResponseT> the type of an outgoing batch response.
+ */
+@FunctionalInterface
+@SdkInternalApi
+public interface BatchAndSend<RequestT, BatchResponseT> {
+    CompletableFuture<BatchResponseT> batchAndSend(List<IdentifiableMessage<RequestT>> identifiedRequests, String batchKey);
+}

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

@@ -0,0 +1,60 @@
+/*
+ * 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.time.Duration;
+import java.util.Optional;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
+
+@SdkInternalApi
+public final class BatchConfiguration {
+
+    public static final int DEFAULT_MAX_BATCH_ITEMS = 10;
+    public static final int DEFAULT_MAX_BATCH_KEYS = 100;
+    public static final int DEFAULT_MAX_BUFFER_SIZE = 500;
+    public static final Duration DEFAULT_MAX_BATCH_OPEN_IN_MS = Duration.ofMillis(200);
+
+    private final Integer maxBatchItems;
+    private final Integer maxBatchKeys;
+    private final Integer maxBufferSize;
+    private final Duration maxBatchOpenInMs;
+
+    public BatchConfiguration(BatchOverrideConfiguration overrideConfiguration) {
+        Optional<BatchOverrideConfiguration> configuration = Optional.ofNullable(overrideConfiguration);
+        this.maxBatchItems = configuration.flatMap(BatchOverrideConfiguration::maxBatchItems).orElse(DEFAULT_MAX_BATCH_ITEMS);
+        this.maxBatchKeys = configuration.flatMap(BatchOverrideConfiguration::maxBatchKeys).orElse(DEFAULT_MAX_BATCH_KEYS);
+        this.maxBufferSize = configuration.flatMap(BatchOverrideConfiguration::maxBufferSize).orElse(DEFAULT_MAX_BUFFER_SIZE);
+        this.maxBatchOpenInMs = configuration.flatMap(BatchOverrideConfiguration::maxBatchOpenInMs)
+                                             .orElse(DEFAULT_MAX_BATCH_OPEN_IN_MS);
+    }
+
+    public Duration maxBatchOpenInMs() {
+        return maxBatchOpenInMs;
+    }
+
+    public int maxBatchItems() {
+        return maxBatchItems;
+    }
+
+    public int maxBatchKeys() {
+        return maxBatchKeys;
+    }
+
+    public int maxBufferSize() {
+        return maxBufferSize;
+    }
+}

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

@@ -0,0 +1,28 @@
+/*
+ * 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 software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Takes a request and extracts a batchKey as determined by the caller.
+ * @param <RequestT> the request.
+ */
+@FunctionalInterface
+@SdkInternalApi
+public interface BatchKeyMapper<RequestT> {
+    String getBatchKey(RequestT request);
+}

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

@@ -0,0 +1,75 @@
+/*
+ * 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.ScheduledExecutorService;
+import java.util.function.Consumer;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
+
+@SdkInternalApi
+public interface BatchManagerBuilder<RequestT, ResponseT, BatchResponseT, B> {
+
+    /**
+     * Defines overrides to the default BatchManager configuration that should be used.
+     *
+     * @param overrideConfiguration the override configuration.
+     * @return a reference to this object so that method calls can be chained together.
+     */
+    B overrideConfiguration(BatchOverrideConfiguration overrideConfiguration);
+
+
+    B overrideConfiguration(Consumer<BatchOverrideConfiguration.Builder> overrideConfigurationConsumer);
+
+
+    /**
+     * Adds a {@link ScheduledExecutorService} to be used by the BatchManager to schedule periodic flushes of the underlying
+     * buffers.
+     * <p>
+     * Creating a BatchManager directly from a service client will use the service client's scheduled executor. If supplied by
+     * the user, this ScheduledExecutorService must be closed by the caller when it is ready to be shut down.
+     *
+     * @param scheduledExecutor the provided scheduled executor.
+     * @return a reference to this object so that method calls can be chained together.
+     */
+    B scheduledExecutor(ScheduledExecutorService scheduledExecutor);
+
+    /**
+     * Adds a function that defines how requests should be batched together into the appropriate batch response.
+     *
+     * @param batchFunction the provided function.
+     * @return a reference to this object so that method calls can be chained together.
+     */
+    B batchFunction(BatchAndSend<RequestT, BatchResponseT> batchFunction);
+
+    /**
+     * Adds a function that defines how a batch response should be extracted and transformed into its individual responses.
+     *
+     * @param responseMapper the provided function.
+     * @return a reference to this object so that method calls can be chained together.
+     */
+    B responseMapper(BatchResponseMapper<BatchResponseT, ResponseT> responseMapper);
+
+    /**
+     * Adds a function that calculates an appropriate batchKey from a given request.
+     *
+     * @param batchKeyMapper the provided function.
+     * @return a reference to this object so that method calls can be chained together.
+     */
+    B batchKeyMapper(BatchKeyMapper<RequestT> batchKeyMapper);
+
+
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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.List;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.utils.Either;
+
+/**
+ * Unpacks the batch response, then transforms individual entries to the appropriate response type. Each entry's batch ID
+ * is mapped to the individual response entry.
+ * @param <BatchResponseT> the type of an outgoing batch response.
+ * @param <ResponseT> the type of an outgoing response.
+ */
+@FunctionalInterface
+@SdkInternalApi
+public interface BatchResponseMapper<BatchResponseT, ResponseT> {
+    List<Either<IdentifiableMessage<ResponseT>, IdentifiableMessage<Throwable>>> mapBatchResponse(BatchResponseT batchResponse);
+}

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,89 @@
+/*
+ * 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.BiFunction;
+import java.util.function.Supplier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Outer map maps a batchKey (ex. queueUrl, overrideConfig etc.) to a nested BatchingGroupMap map.
+ *
+ * @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 BiFunction<Integer, ScheduledFuture<?>, RequestBatchBuffer<RequestT, ResponseT>> bufferSupplier;
+    private final Map<String, RequestBatchBuffer<RequestT, ResponseT>> batchContextMap;
+
+    public BatchingMap(int maxBatchKeys, int maxBufferSize,
+                       BiFunction<Integer, ScheduledFuture<?>, RequestBatchBuffer<RequestT, ResponseT>> bufferSupplier) {
+        this.batchContextMap = new ConcurrentHashMap<>();
+        this.maxBatchKeys = maxBatchKeys;
+        this.maxBufferSize = maxBufferSize;
+        this.bufferSupplier = bufferSupplier;
+    }
+
+    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 bufferSupplier.apply(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();
+    }
+}