From d0af219b604ffdb90b18429833ce49a09458e4f9 Mon Sep 17 00:00:00 2001
From: Fabrizio Fortino <fabrizio.fortino@gmail.com>
Date: Tue, 20 Aug 2024 16:22:00 +0200
Subject: [PATCH] fix: bulk ingester might skip lister requests

---
 .../_helpers/bulk/BulkIngester.java           | 11 ++++----
 .../_helpers/bulk/BulkIngesterTest.java       | 25 +++++++++++++------
 2 files changed, 23 insertions(+), 13 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
index ad2920f4d..49b1b0e9a 100644
--- 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
@@ -311,12 +311,6 @@ 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) {
@@ -330,6 +324,11 @@ public void flush() {
                             exec.contexts, thr));
                     }
                 }
+
+                sendRequestCondition.signalIfReadyAfter(() -> {
+                    requestsInFlightCount--;
+                    closeCondition.signalAllIfReady();
+                });
                 return null;
             });
         }
diff --git a/java-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java b/java-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java
index a72473eda..a38569301 100644
--- a/java-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java
+++ b/java-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/bulk/BulkIngesterTest.java
@@ -89,26 +89,37 @@ 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 {
+    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)
@@ -138,7 +149,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);