Fix visibility timeout errors (#4812) (#4831)

Fix visibility timeout errors (#4812) Signed-off-by: Daniel Li <lhouqua@amazon.com> (cherry picked from commit 910533a)
opensearch-project · Aug 23, 2024 · a966bac · a966bac
1 parent 476b9a8
commit a966bac
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/...ugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java b/...ugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java
@@ -192,6 +192,8 @@ private ReceiveMessageRequest createReceiveMessageRequest() {
     private List<DeleteMessageBatchRequestEntry> processS3EventNotificationRecords(final Collection<ParsedMessage> s3EventNotificationRecords) {
         final List<DeleteMessageBatchRequestEntry> deleteMessageBatchRequestEntryCollection = new ArrayList<>();
         final List<ParsedMessage> parsedMessagesToRead = new ArrayList<>();
+        final Map<ParsedMessage, AcknowledgementSet> messageAcknowledgementSetMap = new HashMap<>();
+        final Map<ParsedMessage, List<DeleteMessageBatchRequestEntry>> messageWaitingForAcknowledgementsMap = new HashMap<>();
 
         for (ParsedMessage parsedMessage : s3EventNotificationRecords) {
             if (parsedMessage.isFailedParsing()) {
@@ -224,7 +226,7 @@ && isEventBridgeEventTypeCreated(parsedMessage)) {
         }
 
         LOG.info("Received {} messages from SQS. Processing {} messages.", s3EventNotificationRecords.size(), parsedMessagesToRead.size());
-
+        
         for (ParsedMessage parsedMessage : parsedMessagesToRead) {
             List<DeleteMessageBatchRequestEntry> waitingForAcknowledgements = new ArrayList<>();
             AcknowledgementSet acknowledgementSet = null;
@@ -262,7 +264,19 @@ && isEventBridgeEventTypeCreated(parsedMessage)) {
                         },
                         Duration.ofSeconds(progressCheckInterval));
                 }
+                messageAcknowledgementSetMap.put(parsedMessage, acknowledgementSet);
+                messageWaitingForAcknowledgementsMap.put(parsedMessage, waitingForAcknowledgements);
             }
+        }
+
+        if (endToEndAcknowledgementsEnabled) {
+            LOG.debug("Created acknowledgement sets for {} messages.", parsedMessagesToRead.size());
+        }
+
+        // Use a separate loop for processing the S3 objects
+        for (ParsedMessage parsedMessage : parsedMessagesToRead) {
+            final AcknowledgementSet acknowledgementSet = messageAcknowledgementSetMap.get(parsedMessage);
+            final List<DeleteMessageBatchRequestEntry> waitingForAcknowledgements = messageWaitingForAcknowledgementsMap.get(parsedMessage);
             final S3ObjectReference s3ObjectReference = populateS3Reference(parsedMessage.getBucketName(), parsedMessage.getObjectKey());
             final Optional<DeleteMessageBatchRequestEntry> deleteMessageBatchRequestEntry = processS3Object(parsedMessage, s3ObjectReference, acknowledgementSet);
             if (endToEndAcknowledgementsEnabled) {
@@ -271,7 +285,7 @@ && isEventBridgeEventTypeCreated(parsedMessage)) {
             } else {
                 deleteMessageBatchRequestEntry.ifPresent(deleteMessageBatchRequestEntryCollection::add);
             }
-        }
+        }         
 
         return deleteMessageBatchRequestEntryCollection;
     }