Merge pull request #3698 from aws/joviegas/stag_sqs_req_mem_leak

bugfix: Fixed memory leak issue in SqsBatchManager for RequestBatchManager sendRequest Calls in SqsBatchManager

Author: dagnir

.changes/2.30.31.json

@@ -55,6 +55,12 @@
             "category": "AWS SDK for Java v2",
             "contributor": "",
             "description": "Updated endpoint and partition metadata."
+        },
+        {
+            "type": "bugfix",
+            "category": "Amazon Simple Queue Service",
+            "contributor": "",
+            "description": "Fixed memory leak in SqsBatch Manager: Resolved an issue where pendingResponses and pendingBatchResponses collections in RequestBatchManager retained references to completed futures, causing memory accumulation over time."
         }
     ]
 }

CHANGELOG.md

@@ -36,6 +36,10 @@
   - ### Features
     - Added a mandatory parameter DataAutomationProfileArn to support for cross region inference for InvokeDataAutomationAsync API. Renamed DataAutomationArn to DataAutomationProjectArn. Added APIs to support tagging.
 
+## __Amazon Simple Queue Service__
+- ### Bugfixes
+    - Fixed memory leak in SqsBatchManager: Resolved an issue where pendingResponses and pendingBatchResponses collections in RequestBatchManager retained references to completed futures, causing memory accumulation over time.
+
 # __2.30.30__ __2025-02-27__
 ## __AWS SDK for Java v2__
   - ### Features

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

@@ -17,7 +17,6 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -58,15 +57,16 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration,
         this.maxBatchItems = batchConfiguration.maxBatchItems();
         this.sendRequestFrequency = batchConfiguration.sendRequestFrequency();
         this.scheduledExecutor = Validate.notNull(scheduledExecutor, "Null scheduledExecutor");
-        pendingBatchResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());
-        pendingResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());
+        pendingBatchResponses = ConcurrentHashMap.newKeySet();
+        pendingResponses = ConcurrentHashMap.newKeySet();
         this.requestsAndResponsesMaps = new BatchingMap<>(overrideConfiguration);
 
     }
 
     public CompletableFuture<ResponseT> batchRequest(RequestT request) {
         CompletableFuture<ResponseT> response = new CompletableFuture<>();
         pendingResponses.add(response);
+        response.whenComplete((r, t) -> pendingResponses.remove(response));
 
         try {
             String batchKey = getBatchKey(request);
@@ -120,16 +120,17 @@ private void flushBuffer(String batchKey, Map<String, BatchingExecutionContext<R
         flushableRequests.forEach((contextId, batchExecutionContext) ->
                                       requestEntries.add(new IdentifiableMessage<>(contextId, batchExecutionContext.request())));
         if (!requestEntries.isEmpty()) {
-            CompletableFuture<BatchResponseT> pendingBatchingRequest = batchAndSend(requestEntries, batchKey)
-                .whenComplete((result, ex) -> handleAndCompleteResponses(result, ex, flushableRequests));
-
+            CompletableFuture<BatchResponseT> pendingBatchingRequest = batchAndSend(requestEntries, batchKey);
             pendingBatchResponses.add(pendingBatchingRequest);
+            pendingBatchingRequest.whenComplete((result, ex) -> {
+                handleAndCompleteResponses(result, ex, flushableRequests);
+                pendingBatchResponses.remove(pendingBatchingRequest);
+            });
         }
     }
 
     private void handleAndCompleteResponses(BatchResponseT batchResult, Throwable exception,
                                             Map<String, BatchingExecutionContext<RequestT, ResponseT>> requests) {
-        requests.forEach((contextId, batchExecutionContext) ->  pendingResponses.add(batchExecutionContext.response()));
         if (exception != null) {
             requests.forEach((contextId, batchExecutionContext) -> batchExecutionContext.response()
                                                                                         .completeExceptionally(exception));