Modified to do the synchronization in the acknowledgement set framework

Signed-off-by: Krishna Kondaka <[email protected]>

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java

@@ -50,4 +50,12 @@ public interface AcknowledgementSet {
      * @since 2.2
      */
     public boolean release(final EventHandle eventHandle, final boolean result);
+
+    /**
+      * Indicates that the addition of initial set of events to
+      * the acknowledgement set is completed.
+      * It is possible that more events are added to the set as the 
+      * initial events are going through the pipeline line.
+     */
+    public void complete();
 }

data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java

@@ -34,6 +34,7 @@ public class DefaultAcknowledgementSet implements AcknowledgementSet {
     private final Map<EventHandle, AtomicInteger> pendingAcknowledgments;
     private Future<?> callbackFuture;
     private final DefaultAcknowledgementSetMetrics metrics;
+    private boolean closed;
 
     public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer<Boolean> callback, final Duration expiryTime, final DefaultAcknowledgementSetMetrics metrics) {
         this.callback = callback;
@@ -42,6 +43,7 @@ public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer<
         this.expiryTime = Instant.now().plusMillis(expiryTime.toMillis());
         this.callbackFuture = null;
         this.metrics = metrics;
+        this.closed = false;
         pendingAcknowledgments = new HashMap<>();
         lock = new ReentrantLock(true);
     }
@@ -84,6 +86,8 @@ public boolean isDone() {
             if (Instant.now().isAfter(expiryTime)) {
                 if (callbackFuture != null) {
                     callbackFuture.cancel(true);
+                    callbackFuture = null;
+                    LOG.warn("AcknowledgementSet expired");
                 }
                 metrics.increment(DefaultAcknowledgementSetMetrics.EXPIRED_METRIC_NAME);
                 return true;
@@ -98,6 +102,19 @@ public Instant getExpiryTime() {
         return expiryTime;
     }
 
+    @Override
+    public void complete() {
+        lock.lock();
+        try {
+            closed = true;
+            if (pendingAcknowledgments.size() == 0) {
+                callbackFuture = executor.submit(() -> callback.accept(this.result));
+            }
+        } finally {
+            lock.unlock();
+        }
+    }
+
     @Override
     public boolean release(final EventHandle eventHandle, final boolean result) {
         lock.lock();
@@ -114,9 +131,11 @@ public boolean release(final EventHandle eventHandle, final boolean result) {
             }
             if (pendingAcknowledgments.get(eventHandle).decrementAndGet() == 0) {
                 pendingAcknowledgments.remove(eventHandle);
-                if (pendingAcknowledgments.size() == 0) {
+                if (closed && pendingAcknowledgments.size() == 0) {
                     callbackFuture = executor.submit(() -> callback.accept(this.result));
                     return true;
+                } else if (pendingAcknowledgments.size() == 0) {
+                    LOG.warn("Acknowledgement set is not closed. Delaying callback until it is closed");
                 }
             }
         } finally {

data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/SqsWorker.java

@@ -232,28 +232,18 @@ && isEventBridgeEventTypeCreated(parsedMessage)) {
                 // Acknowledgement Set timeout is slightly smaller than the visibility timeout;
                 int timeout = (int) sqsOptions.getVisibilityTimeout().getSeconds() - 2;
                 acknowledgementSet = acknowledgementSetManager.create((result) -> {
-                   synchronized (waitingForAcknowledgements) {
-                        while (!acknowledgementSetReady.get()) {
-                            try {
-                                waitingForAcknowledgements.wait();
-                            } catch (InterruptedException e){}
-                        }
-                        acknowledgementSetCallbackCounter.increment();
-                        // Delete only if this is positive acknowledgement
-                        if (result == true) {
-                            deleteSqsMessages(waitingForAcknowledgements);
-                        }
+                    acknowledgementSetCallbackCounter.increment();
+                    // Delete only if this is positive acknowledgement
+                    if (result == true) {
+                        deleteSqsMessages(waitingForAcknowledgements);
                     }
                 }, Duration.ofSeconds(timeout));
             }
             final S3ObjectReference s3ObjectReference = populateS3Reference(parsedMessage.getBucketName(), parsedMessage.getObjectKey());
             final Optional<DeleteMessageBatchRequestEntry> deleteMessageBatchRequestEntry = processS3Object(parsedMessage, s3ObjectReference, acknowledgementSet);
             if (endToEndAcknowledgementsEnabled) {
                 deleteMessageBatchRequestEntry.ifPresent(waitingForAcknowledgements::add);
-                synchronized (waitingForAcknowledgements) {
-                    acknowledgementSetReady.set(true);
-                    waitingForAcknowledgements.notify();
-                }
+                acknowledgementSet.complete();
             } else {
                 deleteMessageBatchRequestEntry.ifPresent(deleteMessageBatchRequestEntryCollection::add);
             }