Avoid stack of StepListeners in SharedBlobCacheService (#94393)

DaveCTurner · web-flow · commit d875dc29f33e · 2023-03-08T07:38:43.000-05:00
Today `CacheService.CacheFile#populateAndRead` constructs a stack of
`StepListener` objects, one per region, which chain their results
together via `Math::addExact`. There's no need to do this recursively,
we can just accumulate the results in an `AtomicInteger` and return its
value when all the reads are complete.
diff --git a/server/src/main/java/org/elasticsearch/action/support/RefCountingListener.java b/server/src/main/java/org/elasticsearch/action/support/RefCountingListener.java
@@ -167,7 +167,8 @@ public String toString() {
 
     /**
      * Acquire a reference to this object and return a listener which consumes a response and releases the reference. The delegate {@link
-     * ActionListener} is called when all its references have been released. The consumer must not throw any exception.
+     * ActionListener} is called when all its references have been released. If the consumer throws an exception, the exception is passed
+     * to the final listener as if the returned listener was completed exceptionally.
      *
      * It is invalid to call this method once all references are released. Doing so will trip an assertion if assertions are enabled, and
      * will throw an {@link IllegalStateException} otherwise.
@@ -189,8 +190,7 @@ public void onResponse(Response response) {
                         acquiredConsumer.accept(response);
                     }
                 } catch (Exception e) {
-                    assert false : e;
-                    throw e;
+                    addException(e);
                 }
             }
 
diff --git a/server/src/test/java/org/elasticsearch/action/support/RefCountingListenerTests.java b/server/src/test/java/org/elasticsearch/action/support/RefCountingListenerTests.java
@@ -182,6 +182,25 @@ public void testValidation() {
         }
     }
 
+    public void testConsumerFailure() {
+        final var executed = new AtomicBoolean();
+        try (var refs = new RefCountingListener(new ActionListener<Void>() {
+            @Override
+            public void onResponse(Void unused) {
+                fail("unexpected success");
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                assertThat(e.getMessage(), equalTo("simulated"));
+                executed.set(true);
+            }
+        })) {
+            refs.acquire(ignored -> { throw new ElasticsearchException("simulated"); }).onResponse(null);
+        }
+        assertTrue(executed.get());
+    }
+
     public void testJavaDocExample() {
         final var flag = new AtomicBoolean();
         runExample(ActionListener.running(() -> assertTrue(flag.compareAndSet(false, true))));
diff --git a/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java b/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java
@@ -12,7 +12,8 @@
 import org.apache.lucene.store.AlreadyClosedException;
 import org.elasticsearch.Assertions;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.StepListener;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.RefCountingListener;
 import org.elasticsearch.blobcache.BlobCacheUtils;
 import org.elasticsearch.blobcache.common.ByteRange;
 import org.elasticsearch.blobcache.common.SparseFileTracker;
@@ -50,6 +51,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.LongConsumer;
@@ -703,26 +705,31 @@ private static void throwAlreadyEvicted() {
             throw new AlreadyClosedException("File chunk is evicted");
         }
 
-        StepListener<Integer> populateAndRead(
+        void populateAndRead(
             final ByteRange rangeToWrite,
             final ByteRange rangeToRead,
             final RangeAvailableHandler reader,
             final RangeMissingHandler writer,
-            final Executor executor
+            final Executor executor,
+            final ActionListener<Integer> listener
         ) {
             assert rangeToRead.length() > 0;
-            final StepListener<Integer> listener = new StepListener<>();
-            Releasable decrementRef = null;
+            final Releasable[] resources = new Releasable[2];
             try {
                 ensureOpen();
                 incRef();
-                decrementRef = Releasables.releaseOnce(this::decRef);
+                resources[1] = Releasables.releaseOnce(this::decRef);
+
                 ensureOpen();
-                Releasable finalDecrementRef = decrementRef;
-                listener.whenComplete(integer -> finalDecrementRef.close(), throwable -> finalDecrementRef.close());
                 final SharedBytes.IO fileChannel = sharedBytes.getFileChannel(sharedBytesPos);
-                listener.whenComplete(integer -> fileChannel.decRef(), e -> fileChannel.decRef());
-                final ActionListener<Void> rangeListener = rangeListener(rangeToRead, reader, listener, fileChannel);
+                resources[0] = Releasables.releaseOnce(fileChannel::decRef);
+
+                final ActionListener<Void> rangeListener = rangeListener(
+                    rangeToRead,
+                    reader,
+                    ActionListener.runBefore(listener, () -> Releasables.close(resources)),
+                    fileChannel
+                );
                 final List<SparseFileTracker.Gap> gaps = tracker.waitForRange(rangeToWrite, rangeToRead, rangeListener);
 
                 for (SparseFileTracker.Gap gap : gaps) {
@@ -758,9 +765,8 @@ public void onFailure(Exception e) {
                     });
                 }
             } catch (Exception e) {
-                releaseAndFail(listener, decrementRef, e);
+                releaseAndFail(listener, Releasables.wrap(resources), e);
             }
-            return listener;
         }
 
         private ActionListener<Void> rangeListener(
@@ -831,48 +837,35 @@ public int populateAndRead(
             final RangeMissingHandler writer,
             final String executor
         ) throws Exception {
-            StepListener<Integer> stepListener = null;
-            final long writeStart = rangeToWrite.start();
-            final long readStart = rangeToRead.start();
-            for (int region = getRegion(rangeToWrite.start()); region <= getEndingRegion(rangeToWrite.end()); region++) {
-                final ByteRange subRangeToWrite = mapSubRangeToRegion(rangeToWrite, region);
-                final ByteRange subRangeToRead = mapSubRangeToRegion(rangeToRead, region);
-                if (subRangeToRead.length() == 0L) {
-                    // nothing to read, skip
-                    if (stepListener == null) {
-                        stepListener = new StepListener<>();
-                        stepListener.onResponse(0);
+            final PlainActionFuture<Void> readsComplete = new PlainActionFuture<>();
+            final AtomicInteger bytesRead = new AtomicInteger();
+            try (var listeners = new RefCountingListener(1, readsComplete)) {
+                final long writeStart = rangeToWrite.start();
+                final long readStart = rangeToRead.start();
+                for (int region = getRegion(rangeToWrite.start()); region <= getEndingRegion(rangeToWrite.end()); region++) {
+                    final ByteRange subRangeToWrite = mapSubRangeToRegion(rangeToWrite, region);
+                    final ByteRange subRangeToRead = mapSubRangeToRegion(rangeToRead, region);
+                    if (subRangeToRead.length() == 0L) {
+                        // nothing to read, skip
+                        continue;
                     }
-                    continue;
-                }
-                final CacheFileRegion fileRegion = get(cacheKey, length, region);
-                final long regionStart = getRegionStart(region);
-                final long writeOffset = writeStart - regionStart;
-                final long readOffset = readStart - regionStart;
-                final StepListener<Integer> lis = fileRegion.populateAndRead(
-                    subRangeToWrite,
-                    subRangeToRead,
-                    (channel, channelPos, relativePos, len) -> {
+                    final CacheFileRegion fileRegion = get(cacheKey, length, region);
+                    final long regionStart = getRegionStart(region);
+                    final long writeOffset = writeStart - regionStart;
+                    final long readOffset = readStart - regionStart;
+                    fileRegion.populateAndRead(subRangeToWrite, subRangeToRead, (channel, channelPos, relativePos, len) -> {
                         assert regionOwners[fileRegion.sharedBytesPos].get() == fileRegion;
                         assert channelPos >= fileRegion.physicalStartOffset() && channelPos + len <= fileRegion.physicalEndOffset();
                         return reader.onRangeAvailable(channel, channelPos, relativePos - readOffset, len);
-                    },
-                    (channel, channelPos, relativePos, len, progressUpdater) -> {
+                    }, (channel, channelPos, relativePos, len, progressUpdater) -> {
                         assert regionOwners[fileRegion.sharedBytesPos].get() == fileRegion;
                         assert channelPos >= fileRegion.physicalStartOffset() && channelPos + len <= fileRegion.physicalEndOffset();
                         writer.fillCacheRange(channel, channelPos, relativePos - writeOffset, len, progressUpdater);
-                    },
-                    threadPool.executor(executor)
-                );
-                assert lis != null;
-                if (stepListener == null) {
-                    stepListener = lis;
-                } else {
-                    stepListener = stepListener.thenCombine(lis, Math::addExact);
+                    }, threadPool.executor(executor), listeners.acquire(i -> bytesRead.updateAndGet(j -> Math.addExact(i, j))));
                 }
-
             }
-            return stepListener.asFuture().get();
+            readsComplete.get();
+            return bytesRead.get();
         }
 
         @Override
diff --git a/x-pack/plugin/blob-cache/src/test/java/org/elasticsearch/blobcache/shared/SharedBlobCacheServiceTests.java b/x-pack/plugin/blob-cache/src/test/java/org/elasticsearch/blobcache/shared/SharedBlobCacheServiceTests.java
@@ -7,6 +7,7 @@
 
 package org.elasticsearch.blobcache.shared;
 
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.blobcache.common.ByteRange;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.common.settings.Setting;
@@ -71,20 +72,25 @@ public void testBasicEviction() throws IOException {
             assertEquals(3, cacheService.freeRegionCount());
             assertFalse(region1.tryEvict());
             assertEquals(3, cacheService.freeRegionCount());
+            final var bytesReadFuture = new PlainActionFuture<Integer>();
             region0.populateAndRead(
                 ByteRange.of(0L, 1L),
                 ByteRange.of(0L, 1L),
                 (channel, channelPos, relativePos, length) -> 1,
                 (channel, channelPos, relativePos, length, progressUpdater) -> progressUpdater.accept(length),
-                taskQueue.getThreadPool().executor(ThreadPool.Names.GENERIC)
+                taskQueue.getThreadPool().executor(ThreadPool.Names.GENERIC),
+                bytesReadFuture
             );
             assertFalse(region0.tryEvict());
             assertEquals(3, cacheService.freeRegionCount());
+            assertFalse(bytesReadFuture.isDone());
             taskQueue.runAllRunnableTasks();
             assertTrue(region0.tryEvict());
             assertEquals(4, cacheService.freeRegionCount());
             assertTrue(region2.tryEvict());
             assertEquals(5, cacheService.freeRegionCount());
+            assertTrue(bytesReadFuture.isDone());
+            assertEquals(Integer.valueOf(1), bytesReadFuture.actionGet());
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,8 @@ public String toString() {`
`167`	`167`
`168`	`168`	`/**`
`169`	`169`	`* Acquire a reference to this object and return a listener which consumes a response and releases the reference. The delegate {@link`
`170`		`- * ActionListener} is called when all its references have been released. The consumer must not throw any exception.`
	`170`	`+ * ActionListener} is called when all its references have been released. If the consumer throws an exception, the exception is passed`
	`171`	`+ * to the final listener as if the returned listener was completed exceptionally.`
`171`	`172`	`*`
`172`	`173`	`* It is invalid to call this method once all references are released. Doing so will trip an assertion if assertions are enabled, and`
`173`	`174`	`* will throw an {@link IllegalStateException} otherwise.`
`@@ -189,8 +190,7 @@ public void onResponse(Response response) {`
`189`	`190`	`acquiredConsumer.accept(response);`
`190`	`191`	`}`
`191`	`192`	`} catch (Exception e) {`
`192`		`- assert false : e;`
`193`		`- throw e;`
	`193`	`+ addException(e);`
`194`	`194`	`}`
`195`	`195`	`}`
`196`	`196`