Fixed issue relating to cancelling scheduled flush and removing requests.

core/sdk-core/src/main/java/software/amazon/awssdk/core/BatchOverrideConfiguration.java

@@ -18,7 +18,7 @@
 import java.time.Duration;
 import java.util.concurrent.ScheduledExecutorService;
 import software.amazon.awssdk.annotations.SdkPublicApi;
-import software.amazon.awssdk.core.internal.batchutilities.BatchManager;
+import software.amazon.awssdk.core.internal.batchmanager.BatchManager;
 import software.amazon.awssdk.utils.ToString;
 import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
@@ -78,10 +78,38 @@ public Builder toBuilder() {
     public String toString() {
         return ToString.builder("BatchOverrideConfiguration")
                        .add("maxBatchItems", maxBatchItems)
-                       .add("maxBatchOpenInMs", maxBatchOpenInMs)
+                       .add("maxBatchOpenInMs", maxBatchOpenInMs.toMillis())
                        .build();
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        BatchOverrideConfiguration that = (BatchOverrideConfiguration) o;
+
+        if (maxBatchItems != null ? !maxBatchItems.equals(that.maxBatchItems) : that.maxBatchItems != null) {
+            return false;
+        }
+        if (maxBatchOpenInMs != null ? !maxBatchOpenInMs.equals(that.maxBatchOpenInMs) : that.maxBatchOpenInMs != null) {
+            return false;
+        }
+        return scheduledExecutor.equals(that.scheduledExecutor);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = maxBatchItems != null ? maxBatchItems.hashCode() : 0;
+        result = 31 * result + (maxBatchOpenInMs != null ? maxBatchOpenInMs.hashCode() : 0);
+        result = 31 * result + scheduledExecutor.hashCode();
+        return result;
+    }
+
     public static final class Builder implements CopyableBuilder<Builder, BatchOverrideConfiguration> {
 
         private Integer maxBatchItems;

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/batchutilities/BatchAndSend.java → core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/batchmanager/BatchAndSend.java

@@ -13,7 +13,7 @@
  * permissions and limitations under the License.
  */
 
-package software.amazon.awssdk.core.internal.batchutilities;
+package software.amazon.awssdk.core.internal.batchmanager;
 
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -28,5 +28,5 @@
 @FunctionalInterface
 @SdkProtectedApi
 public interface BatchAndSend<RequestT, BatchResponseT> {
-    CompletableFuture<BatchResponseT> batchAndSend(List<IdentifiableRequest<RequestT>> identifiedRequests, String batchGroupId);
+    CompletableFuture<BatchResponseT> batchAndSend(List<IdentifiableMessage<RequestT>> identifiedRequests, String batchGroupId);
 }

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/batchmanager/BatchBuffer.java

@@ -0,0 +1,136 @@
+/*
+ * 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.core.internal.batchmanager;
+
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledFuture;
+import java.util.stream.Collectors;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+@SdkInternalApi
+public final class BatchBuffer<RequestT, ResponseT> {
+    private final Object flushLock = new Object();
+
+    private final Map<String, BatchingExecutionContext<RequestT, ResponseT>> idToBatchContext;
+
+    // TODO: Figure out better name for nextId and nextBatchEntry.
+    /**
+     * 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;
+
+    /**
+     * 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;
+
+    /**
+     * The scheduled flush tasks associated with this batchBuffer.
+     */
+    private ScheduledFuture<?> scheduledFlush;
+
+    public BatchBuffer(ScheduledFuture<?> scheduledFlush) {
+        this.idToBatchContext = new ConcurrentHashMap<>();
+        this.nextId = 0;
+        this.nextBatchEntry = 0;
+        this.scheduledFlush = scheduledFlush;
+    }
+
+    public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableRequests(int maxBatchItems) {
+        synchronized (flushLock) {
+            if (idToBatchContext.size() >= maxBatchItems) {
+                return extractFlushedEntries(maxBatchItems);
+            }
+            return new ConcurrentHashMap<>();
+        }
+    }
+
+    public Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableScheduledRequests(int maxBatchItems) {
+        synchronized (flushLock) {
+            if (idToBatchContext.size() > 0) {
+                return extractFlushedEntries(maxBatchItems);
+            }
+            return new ConcurrentHashMap<>();
+        }
+    }
+
+    private Map<String, BatchingExecutionContext<RequestT, ResponseT>> extractFlushedEntries(int maxBatchItems) {
+        LinkedHashMap<String, BatchingExecutionContext<RequestT, ResponseT>> requestEntries = new LinkedHashMap<>();
+        String nextEntry;
+        while (requestEntries.size() < maxBatchItems && hasNextBatchEntry()) {
+            nextEntry = nextBatchEntry();
+            requestEntries.put(nextEntry, idToBatchContext.get(nextEntry));
+            idToBatchContext.remove(nextEntry);
+        }
+        return requestEntries;
+    }
+
+    public RequestT getRequest(String key) {
+        return idToBatchContext.get(key).request();
+    }
+
+    public CompletableFuture<ResponseT> getResponse(String key) {
+        return idToBatchContext.get(key).response();
+    }
+
+    public BatchingExecutionContext<RequestT, ResponseT> put(RequestT request, CompletableFuture<ResponseT> response) {
+        synchronized (this) {
+            if (nextId == Integer.MAX_VALUE) {
+                nextId = 0;
+            }
+            String id = Integer.toString(nextId++);
+            return idToBatchContext.put(id, new BatchingExecutionContext<>(request, response));
+        }
+    }
+
+    private boolean hasNextBatchEntry() {
+        return idToBatchContext.containsKey(Integer.toString(nextBatchEntry));
+    }
+
+    private String nextBatchEntry() {
+        if (nextBatchEntry == Integer.MAX_VALUE) {
+            nextBatchEntry = 0;
+        }
+        return Integer.toString(nextBatchEntry++);
+    }
+
+    public void putScheduledFlush(ScheduledFuture<?> scheduledFlush) {
+        this.scheduledFlush = scheduledFlush;
+    }
+
+    public void cancelScheduledFlush() {
+        scheduledFlush.cancel(false);
+    }
+
+    public Collection<CompletableFuture<ResponseT>> responses() {
+        return idToBatchContext.values()
+                               .stream()
+                               .map(BatchingExecutionContext::response)
+                               .collect(Collectors.toList());
+    }
+
+    public void clear() {
+        idToBatchContext.clear();
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/batchutilities/BatchKeyMapper.java → core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/batchmanager/BatchKeyMapper.java

@@ -13,7 +13,7 @@
  * permissions and limitations under the License.
  */
 
-package software.amazon.awssdk.core.internal.batchutilities;
+package software.amazon.awssdk.core.internal.batchmanager;
 
 import software.amazon.awssdk.annotations.SdkProtectedApi;