Add ProgressCheck callbacks to end-to-end acknowledgements (#3565)

Add ProgressCheck callbacks to end-to-end acknowledgements

Signed-off-by: Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com>
Co-authored-by: Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com>

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

@@ -8,6 +8,9 @@
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.EventHandle;
 
+import java.time.Duration;
+import java.util.function.Consumer;
+
 /**
  * AcknowledgmentSet keeps track of set of events that
  * belong to the batch of events that a source creates.
@@ -58,4 +61,16 @@ public interface AcknowledgementSet {
       * initial events are going through the pipeline line.
      */
     public void complete();
+
+    /**
+     * adds progress check callback to the acknowledgement set. When added
+     * the callback is called every progressCheckInterval time with the 
+     * indication of current progress as a ratio of pending number of
+     * acknowledgements over total acknowledgements
+     *
+     * @param progressCheckCallback progress check callback to be called
+     * @param progressCheckInterval frequency of invocation of progress check callback
+     * @since 2.6
+     */
+    public void addProgressCheck(final Consumer<ProgressCheck> progressCheckCallback, final Duration progressCheckInterval);
 }

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

@@ -0,0 +1,17 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.acknowledgements;
+
+public interface ProgressCheck {
+    /**
+     * Returns the pending ratio
+     *
+     * @return returns the ratio of pending to the total acknowledgements
+     * @since 2.6
+     */
+    Double getRatio();
+}
+

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

@@ -8,7 +8,7 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 
@@ -23,7 +23,7 @@ CallbackTheadFactory callbackTheadFactory() {
     }
 
     @Bean(name = "acknowledgementCallbackExecutor")
-    ExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) {
-        return Executors.newFixedThreadPool(MAX_THREADS, callbackTheadFactory);
+    ScheduledExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) {
+        return Executors.newScheduledThreadPool(MAX_THREADS, callbackTheadFactory);
     }
 }

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

@@ -6,6 +6,7 @@
 package org.opensearch.dataprepper.acknowledgements;
 
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.EventHandle;
 import org.opensearch.dataprepper.model.event.InternalEventHandle;
@@ -18,37 +19,61 @@
 import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 public class DefaultAcknowledgementSet implements AcknowledgementSet {
     private static final Logger LOG = LoggerFactory.getLogger(DefaultAcknowledgementSet.class);
     private final Consumer<Boolean> callback;
+    private Consumer<ProgressCheck> progressCheckCallback;
     private final Instant expiryTime;
-    private final ExecutorService executor;
+    private final ScheduledExecutorService scheduledExecutor;
     // This lock protects all the non-final members
     private final ReentrantLock lock;
     private boolean result;
     private final Map<EventHandle, AtomicInteger> pendingAcknowledgments;
     private Future<?> callbackFuture;
     private final DefaultAcknowledgementSetMetrics metrics;
+    private ScheduledFuture<?> progressCheckFuture;
     private boolean completed;
+    private AtomicInteger totalEventsAdded;
 
-    public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer<Boolean> callback, final Duration expiryTime, final DefaultAcknowledgementSetMetrics metrics) {
+    public DefaultAcknowledgementSet(final ScheduledExecutorService scheduledExecutor,
+                                     final Consumer<Boolean> callback,
+                                     final Duration expiryTime,
+                                     final DefaultAcknowledgementSetMetrics metrics) {
         this.callback = callback;
         this.result = true;
-        this.executor = executor;
+        this.totalEventsAdded = new AtomicInteger(0);
+        this.scheduledExecutor = scheduledExecutor;
         this.expiryTime = Instant.now().plusMillis(expiryTime.toMillis());
         this.callbackFuture = null;
         this.metrics = metrics;
         this.completed = false;
+        this.progressCheckCallback = null;
         pendingAcknowledgments = new HashMap<>();
         lock = new ReentrantLock(true);
     }
 
+    public void addProgressCheck(final Consumer<ProgressCheck> progressCheckCallback, final Duration progressCheckInterval) {
+        this.progressCheckCallback = progressCheckCallback;
+        this.progressCheckFuture = scheduledExecutor.scheduleAtFixedRate(this::checkProgress, 0L, progressCheckInterval.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    public void checkProgress() {
+        lock.lock();
+        int numberOfEventsPending = pendingAcknowledgments.size();
+        lock.unlock();
+        if (progressCheckCallback != null) {
+            progressCheckCallback.accept(new DefaultProgressCheck((double)numberOfEventsPending/totalEventsAdded.get()));
+        }
+    }
+
     @Override
     public void add(Event event) {
         lock.lock();
@@ -59,6 +84,7 @@ public void add(Event event) {
                     InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
                     internalEventHandle.setAcknowledgementSet(this);
                     pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
+                    totalEventsAdded.incrementAndGet();
                 }
             }
         } finally {
@@ -88,6 +114,9 @@ public boolean isDone() {
                 return true;
             }
             if (Instant.now().isAfter(expiryTime)) {
+                if (progressCheckFuture != null) {
+                    progressCheckFuture.cancel(false);
+                }
                 if (callbackFuture != null) {
                     callbackFuture.cancel(true);
                     callbackFuture = null;
@@ -112,7 +141,10 @@ public void complete() {
         try {
             completed = true;
             if (pendingAcknowledgments.size() == 0) {
-                callbackFuture = executor.submit(() -> callback.accept(this.result));
+                if (progressCheckFuture != null) {
+                    progressCheckFuture.cancel(false);
+                }
+                callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result));
             }
         } finally {
             lock.unlock();
@@ -136,7 +168,10 @@ public boolean release(final EventHandle eventHandle, final boolean result) {
             if (pendingAcknowledgments.get(eventHandle).decrementAndGet() == 0) {
                 pendingAcknowledgments.remove(eventHandle);
                 if (completed && pendingAcknowledgments.size() == 0) {
-                    callbackFuture = executor.submit(() -> callback.accept(this.result));
+                    if (progressCheckFuture != null) {
+                        progressCheckFuture.cancel(false);
+                    }
+                    callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result));
                     return true;
                 } else if (pendingAcknowledgments.size() == 0) {
                     LOG.debug("Acknowledgement set is not completed. Delaying callback until it is completed");

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

@@ -15,35 +15,35 @@
 import javax.inject.Named;
 import java.time.Duration;
 import java.util.Objects;
-import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Consumer;
 
 @Named
 public class DefaultAcknowledgementSetManager implements AcknowledgementSetManager {
     private static final int DEFAULT_WAIT_TIME_MS = 15 * 1000;
     private final AcknowledgementSetMonitor acknowledgementSetMonitor;
-    private final ExecutorService executor;
+    private final ScheduledExecutorService scheduledExecutor;
     private final AcknowledgementSetMonitorThread acknowledgementSetMonitorThread;
     private PluginMetrics pluginMetrics;
     private DefaultAcknowledgementSetMetrics metrics;
 
     @Inject
     public DefaultAcknowledgementSetManager(
-            @Named("acknowledgementCallbackExecutor") final ExecutorService callbackExecutor) {
+            @Named("acknowledgementCallbackExecutor") final ScheduledExecutorService callbackExecutor) {
         this(callbackExecutor, Duration.ofMillis(DEFAULT_WAIT_TIME_MS));
     }
 
-    public DefaultAcknowledgementSetManager(final ExecutorService callbackExecutor, final Duration waitTime) {
+    public DefaultAcknowledgementSetManager(final ScheduledExecutorService callbackExecutor, final Duration waitTime) {
         this.acknowledgementSetMonitor = new AcknowledgementSetMonitor();
-        this.executor = Objects.requireNonNull(callbackExecutor);
+        this.scheduledExecutor = Objects.requireNonNull(callbackExecutor);
         acknowledgementSetMonitorThread = new AcknowledgementSetMonitorThread(acknowledgementSetMonitor, waitTime);
         acknowledgementSetMonitorThread.start();
         pluginMetrics = PluginMetrics.fromNames("acknowledgementSetManager", "acknowledgements");
         metrics = new DefaultAcknowledgementSetMetrics(pluginMetrics);
     }
 
     public AcknowledgementSet create(final Consumer<Boolean> callback, final Duration timeout) {
-        AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(executor, callback, timeout, metrics);
+        AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(scheduledExecutor, callback, timeout, metrics);
         acknowledgementSetMonitor.add(acknowledgementSet);
         metrics.increment(DefaultAcknowledgementSetMetrics.CREATED_METRIC_NAME);
         return acknowledgementSet;

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.acknowledgements;
+
+import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck;
+
+public class DefaultProgressCheck implements ProgressCheck {
+    double ratio;
+
+    public DefaultProgressCheck(double ratio) {
+        this.ratio = ratio;
+    }
+
+    @Override
+    public Double getRatio() {
+        return ratio;
+    }
+}