Adding BulkProcessor2.addWithBackpressure() (#94599)

masseyke · web-flow · commit b068ab2ee287 · 2023-03-21T13:44:20.000-05:00
This commit adds a method called addWithBackpressure() to BulkProcessor2. This method
is similar to the add() method that already exists, except that it blocks the calling thread until
there is room to add the request.
diff --git a/server/src/main/java/org/elasticsearch/action/bulk/BulkProcessor2.java b/server/src/main/java/org/elasticsearch/action/bulk/BulkProcessor2.java
@@ -26,7 +26,11 @@
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiConsumer;
+import java.util.function.Supplier;
 
 /**
  * A bulk processor is a thread safe bulk processing class, allowing to easily set when to "flush" a new bulk request
@@ -206,6 +210,13 @@ public static Builder builder(
      */
     private final Object mutex = new Object();
 
+    /*
+     * This Lock and Condition are used to throttle calls to bulkProcessor.addWithBackpressure when the bulkProcessor already has too many
+     * bytes in flight and cannot accept more data until already-in-flight requests complete.
+     */
+    private final Lock backpressureLock = new ReentrantLock();
+    private final Condition backpressureNotRequiredCondition = backpressureLock.newCondition();
+
     BulkProcessor2(
         BiConsumer<BulkRequest, ActionListener<BulkResponse>> consumer,
         int maxNumberOfRetries,
@@ -276,6 +287,21 @@ public BulkProcessor2 add(IndexRequest request) throws EsRejectedExecutionExcept
         return add((DocWriteRequest<?>) request);
     }
 
+    /**
+     * This method is similar to {@link #add(IndexRequest) add()}, except that if adding the approximate size in bytes of the request to
+     * totalBytesInFlight would exceed maxBytesInFlight then this method will block until the request can be added without exceeding
+     * maxBytesInFlight (or until shouldAbort returns false). This method should be used carefully, ideally from a single thread. This is
+     * because calling it with multiple threads would either hang up all the threads in the pool (in the case of a bounded thread pool) or
+     * effectively create an unbounded queue (in the case of an unbounded thread pool).
+     * @param request The request to add to a batch to be consumed
+     * @param shouldAbort If this returns true then this method bails out with an EsRejectedExecutionException
+     * @return this BulkProcessor2
+     * @throws EsRejectedExecutionException if shouldAbort returns true before the request has been added to a batch
+     */
+    public BulkProcessor2 addWithBackpressure(IndexRequest request, Supplier<Boolean> shouldAbort) throws EsRejectedExecutionException {
+        return addWithBackpressure((DocWriteRequest<?>) request, shouldAbort);
+    }
+
     /**
      * Adds an {@link DeleteRequest} to the list of actions to execute.
      * @throws EsRejectedExecutionException if adding the approximate size in bytes of the request to totalBytesInFlight would exceed
@@ -295,6 +321,44 @@ private BulkProcessor2 add(DocWriteRequest<?> request) throws EsRejectedExecutio
         return this;
     }
 
+    private BulkProcessor2 addWithBackpressure(DocWriteRequest<?> request, Supplier<Boolean> shouldAbort)
+        throws EsRejectedExecutionException {
+        /*
+         * We want this method to block until the bulkProcessor accepts the request. Otherwise, the subsequent calls to this method will
+         * likely continue rejecting. BulkProcessor2.add does not have the ability to exert backpressure.
+         * So we have to catch the EsRejectedExecutionException that is thrown when it already has too many bytes in flight. We then
+         * wait until some bulk has been completed, reducing the amount of data in flight and (probably) making room for this request.
+         */
+        boolean successfullyAdded = false;
+        while (successfullyAdded == false) {
+            if (shouldAbort.get()) {
+                throw new EsRejectedExecutionException("Rejecting request because bulk add has been cancelled by the caller");
+            }
+            try {
+                add(request);
+                successfullyAdded = true;
+            } catch (EsRejectedExecutionException e) {
+                logger.trace("Attempt to add request to batch rejected because too many bytes are in flight already. Will try again.");
+                /*
+                 * Note: It is possible that signalAll was called between the call to add above and acquiring this lock.
+                 * But in that case, either we wait 500ms, or another batch completes and another call to signalAll wakes us up.
+                 * Either way is preferable to requiring this lock for the whole try/catch block. That is why we ignore the
+                 * result of the call to await() -- either way we are going to try calling bulkProcessor.add() again.
+                 */
+                backpressureLock.lock();
+                try {
+                    backpressureNotRequiredCondition.await(500, TimeUnit.MILLISECONDS);
+                } catch (InterruptedException ex) {
+                    Thread.currentThread().interrupt();
+                    throw new RuntimeException(ex);
+                } finally {
+                    backpressureLock.unlock();
+                }
+            }
+        }
+        return this;
+    }
+
     /*
      * Exposed for unit testing
      */
@@ -393,22 +457,39 @@ private void execute(BulkRequest bulkRequest, long executionId) {
                 @Override
                 public void onResponse(BulkResponse response) {
                     totalBytesInFlight.addAndGet(-1 * bulkRequest.estimatedSizeInBytes());
+                    maybeNoLongerInExcessofMaxBytesInFlight();
                     listener.afterBulk(executionId, bulkRequest, response);
                 }
 
                 @Override
                 public void onFailure(Exception e) {
                     totalBytesInFlight.addAndGet(-1 * bulkRequest.estimatedSizeInBytes());
+                    maybeNoLongerInExcessofMaxBytesInFlight();
                     listener.afterBulk(executionId, bulkRequest, e);
                 }
             });
         } catch (Exception e) {
             logger.warn(() -> "Failed to execute bulk request " + executionId + ".", e);
             totalBytesInFlight.addAndGet(-1 * bulkRequest.estimatedSizeInBytes());
+            maybeNoLongerInExcessofMaxBytesInFlight();
             listener.afterBulk(executionId, bulkRequest, e);
         }
     }
 
+    /*
+     * This method is to be called whenever we deduct from totalBytesInFlight, signalling that it's worth retrying any requests blocked in
+     * addWithBackpressure.
+     */
+    private void maybeNoLongerInExcessofMaxBytesInFlight() {
+        // Signal in case any rejected docs were waiting for the space in bulkProcessor in addWithBackPressure
+        backpressureLock.lock();
+        try {
+            backpressureNotRequiredCondition.signalAll();
+        } finally {
+            backpressureLock.unlock();
+        }
+    }
+
     private void execute() {
         assert Thread.holdsLock(mutex);
         final BulkRequest bulkRequest = this.bulkRequestUnderConstruction;
diff --git a/server/src/test/java/org/elasticsearch/action/bulk/BulkProcessor2Tests.java b/server/src/test/java/org/elasticsearch/action/bulk/BulkProcessor2Tests.java
@@ -30,6 +30,7 @@
 import org.junit.Before;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.CountDownLatch;
@@ -38,6 +39,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
@@ -408,6 +410,128 @@ public void testRejections() throws Exception {
         consumerExecutor.shutdown();
     }
 
+    public void testAddWithNoBackpressureThrowsException() throws InterruptedException {
+        /*
+         * This tests that if we reduce the configured batch size and max bytes in flight then we can force add() to throw
+         * EsRejectedExecutionExceptions.
+         */
+        BulkResponse bulkResponse = new BulkResponse(
+            new BulkItemResponse[] { BulkItemResponse.success(0, randomFrom(DocWriteRequest.OpType.values()), mockResponse()) },
+            0
+        );
+        final BiConsumer<BulkRequest, ActionListener<BulkResponse>> consumer = (request, listener) -> {
+            threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
+                try {
+                    Thread.sleep(50);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                listener.onResponse(bulkResponse);
+            });
+        };
+        int numberOfRequests = 100;
+        final AtomicBoolean haveSeenRejections = new AtomicBoolean(false);
+        final BulkProcessor2.Listener listener = new BulkProcessor2.Listener() {
+
+            @Override
+            public void beforeBulk(long executionId, BulkRequest request) {}
+
+            @Override
+            public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {}
+
+            @Override
+            public void afterBulk(long executionId, BulkRequest request, Exception failure) {
+                fail("afterBulk should not return exception");
+            }
+        };
+        BulkProcessor2 bulkProcessor = BulkProcessor2.builder(consumer, listener, threadPool)
+            .setBulkSize(ByteSizeValue.ofBytes(512))
+            .setMaxBytesInFlight(ByteSizeValue.ofBytes(1024))
+            .setFlushInterval(TimeValue.timeValueMillis(1))
+            .build();
+        try {
+            for (int i = 0; i < numberOfRequests; i++) {
+                bulkProcessor.add(new IndexRequest().source(Collections.singletonMap("this_is_a_key" + i, "this_is_a_value" + i)));
+            }
+        } catch (EsRejectedExecutionException e) {
+            // We are throwing more data at the processor than the max bytes in flight setting can handle
+            haveSeenRejections.set(true);
+        } finally {
+            bulkProcessor.awaitClose(1, TimeUnit.SECONDS);
+        }
+        assertThat(haveSeenRejections.get(), equalTo(true));
+    }
+
+    public void testAddWithBackpressure() throws InterruptedException {
+        /*
+         * This test reduces the configured batch size and max bytes in flight so that using add() would throw
+         * EsRejectedExecutionExceptions (see testAddWithNoBackpressureThrowsException). But instead this test uses addWithBackpressure
+         * so that it blocks the calling thread until requests can be added.
+         */
+        BulkResponse bulkResponse = new BulkResponse(
+            new BulkItemResponse[] { BulkItemResponse.success(0, randomFrom(DocWriteRequest.OpType.values()), mockResponse()) },
+            0
+        );
+        final BiConsumer<BulkRequest, ActionListener<BulkResponse>> consumer = (request, listener) -> {
+            threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
+                try {
+                    Thread.sleep(50);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                listener.onResponse(bulkResponse);
+            });
+        };
+        int numberOfRequests = 100;
+        final CountDownLatch countDownLatch = new CountDownLatch(numberOfRequests);
+        final BulkProcessor2.Listener listener = new BulkProcessor2.Listener() {
+
+            @Override
+            public void beforeBulk(long executionId, BulkRequest request) {}
+
+            @Override
+            public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
+                for (int i = 0; i < request.requests.size(); i++) {
+                    countDownLatch.countDown();
+                }
+            }
+
+            @Override
+            public void afterBulk(long executionId, BulkRequest request, Exception failure) {
+                fail("afterBulk should not return exception");
+            }
+        };
+        BulkProcessor2 bulkProcessor = BulkProcessor2.builder(consumer, listener, threadPool)
+            .setBulkSize(ByteSizeValue.ofBytes(512))
+            .setMaxBytesInFlight(ByteSizeValue.ofBytes(1024))
+            .setFlushInterval(TimeValue.timeValueMillis(1))
+            .build();
+        AtomicBoolean abort = new AtomicBoolean(false);
+        final AtomicBoolean haveSeenRejections = new AtomicBoolean(false);
+        try {
+            for (int i = 0; i < numberOfRequests; i++) {
+                bulkProcessor.addWithBackpressure(
+                    new IndexRequest().source(Collections.singletonMap("this_is_a_key" + i, "this_is_a_value" + i)),
+                    abort::get
+                );
+            }
+            assertTrue(countDownLatch.await(5, TimeUnit.MINUTES));
+            assertThat(bulkProcessor.getTotalBytesInFlight(), equalTo(0L));
+            abort.set(true);
+            try {
+                bulkProcessor.addWithBackpressure(
+                    new IndexRequest().source(Collections.singletonMap("this_is_a_key", "this_is_a_value")),
+                    abort::get
+                );
+            } catch (EsRejectedExecutionException e) {
+                haveSeenRejections.set(true);
+            }
+        } finally {
+            bulkProcessor.awaitClose(1, TimeUnit.SECONDS);
+        }
+        assertThat(haveSeenRejections.get(), equalTo(true));
+    }
+
     private BulkProcessor2.Listener emptyListener() {
         return new BulkProcessor2.Listener() {
             @Override
