fix: bulk ingester might skip listener requests (#867)

* fix: bulk ingester might skip lister requests * minor: fix style * always waiting for listener to be done before closing --------- Co-authored-by: Laura Trotta <[email protected]> Co-authored-by: Laura Trotta <[email protected]>
elastic · Sep 12, 2024 · 27f87e6 · 27f87e6
1 parent 67b90b7
commit 27f87e6
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 18 deletions.
diff --git a/java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngester.java b/java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngester.java
@@ -76,6 +76,7 @@ public class BulkIngester<Context> implements AutoCloseable {
     private final FnCondition addCondition = new FnCondition(lock, this::canAddOperation);
     private final FnCondition sendRequestCondition = new FnCondition(lock, this::canSendRequest);
     private final FnCondition closeCondition = new FnCondition(lock, this::closedAndFlushed);
+    private AtomicInteger listenerInProgressCount = new AtomicInteger();
 
     private static class RequestExecution<Context> {
         public final long id;
@@ -235,7 +236,7 @@ private boolean canAddOperation() {
     }
 
     private boolean closedAndFlushed() {
-        return isClosed && operations.isEmpty() && requestsInFlightCount == 0;
+        return isClosed && operations.isEmpty() && requestsInFlightCount == 0 && listenerInProgressCount.get() == 0;
     }
 
     //----- Ingester logic
@@ -311,25 +312,42 @@ public void flush() {
         if (exec != null) {
             // A request was actually sent
             exec.futureResponse.handle((resp, thr) -> {
-
-                sendRequestCondition.signalIfReadyAfter(() -> {
-                    requestsInFlightCount--;
-                    closeCondition.signalAllIfReady();
-                });
-
                 if (resp != null) {
                     // Success
                     if (listener != null) {
-                        scheduler.submit(() -> listener.afterBulk(exec.id, exec.request,
-                            exec.contexts, resp));
+                        listenerInProgressCount.incrementAndGet();
+                        scheduler.submit(() -> {
+                            try {
+                                listener.afterBulk(exec.id, exec.request, exec.contexts, resp);
+                            }
+                            finally {
+                                if(listenerInProgressCount.decrementAndGet() == 0){
+                                    closeCondition.signalIfReady();
+                                }
+                            }
+                        });
                     }
                 } else {
                     // Failure
                     if (listener != null) {
-                        scheduler.submit(() -> listener.afterBulk(exec.id, exec.request,
-                            exec.contexts, thr));
+                        listenerInProgressCount.incrementAndGet();
+                        scheduler.submit(() -> {
+                            try {
+                                listener.afterBulk(exec.id, exec.request, exec.contexts, thr);
+                            }
+                            finally {
+                                if(listenerInProgressCount.decrementAndGet() == 0){
+                                    closeCondition.signalIfReady();
+                                }
+                            }
+                        });
                     }
                 }
+
+                sendRequestCondition.signalIfReadyAfter(() -> {
+                    requestsInFlightCount--;
+                    closeCondition.signalAllIfReady();
+                });
                 return null;
             });
         }

diff --git a/...client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java b/...client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java
@@ -90,26 +90,44 @@ private void printStats(TestTransport transport) {
     @Test
     public void basicTestFlush() throws Exception {
         // Prime numbers, so that we have leftovers to flush before shutting down
-        multiThreadTest(7, 3, 5, 101);
+        multiThreadTest(7, 3, 5, 101, true);
+    }
+
+    @Test
+    public void basicTestFlushWithInternalScheduler() throws Exception {
+        // Prime numbers, so that we have leftovers to flush before shutting down
+        multiThreadTest(7, 3, 5, 101, false);
     }
 
     @Test
     public void basicTestNoFlush() throws Exception {
         // Will have nothing to flush on close.
-        multiThreadTest(10, 3, 5, 100);
+        multiThreadTest(10, 3, 5, 100, true);
     }
 
-    private void multiThreadTest(int maxOperations, int maxRequests, int numThreads, int numOperations) throws Exception {
+    @Test
+    public void basicTestNoFlushWithInternalScheduler() throws Exception {
+        // Will have nothing to flush on close.
+        multiThreadTest(10, 3, 5, 100, false);
+    }
+
+    private void multiThreadTest(int maxOperations, int maxRequests, int numThreads, int numOperations,
+                                 boolean externalScheduler) throws Exception {
 
         CountingListener listener = new CountingListener();
         TestTransport transport = new TestTransport();
         ElasticsearchAsyncClient client = new ElasticsearchAsyncClient(transport);
-        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
+        ScheduledExecutorService scheduler;
+        if (externalScheduler) {
+            scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
                 Thread t = Executors.defaultThreadFactory().newThread(r);
-                t.setName("my-bulk-ingester-executor#" );
+                t.setName("my-bulk-ingester-executor#");
                 t.setDaemon(true);
                 return t;
-        });
+            });
+        } else {
+            scheduler = null;
+        }
 
         BulkIngester<Void> ingester = BulkIngester.of(b -> b
             .client(client)
@@ -139,7 +157,7 @@ private void multiThreadTest(int maxOperations, int maxRequests, int numThreads,
 
         ingester.close();
         transport.close();
-        scheduler.shutdownNow();
+        if (scheduler != null) scheduler.shutdownNow();
 
         printStats(ingester);
         printStats(listener);