Generalize repo analysis (elastic#93899)

In elastic#93825 we introduced a CAS operation on snapshot repositories. In due
course we'll want to verify that this operation works as expected using
the repository analysis API. This commit reworks the implementation
slightly so as to simplify the change which will add this functionality.

Author: DaveCTurner

server/src/main/java/org/elasticsearch/common/util/concurrent/ThrottledIterator.java

@@ -0,0 +1,154 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util.concurrent;
+
+import org.elasticsearch.core.AbstractRefCounted;
+import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.core.Strings;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+
+import java.util.Iterator;
+import java.util.Objects;
+import java.util.concurrent.Semaphore;
+import java.util.function.BiConsumer;
+
+public class ThrottledIterator<T> implements Releasable {
+
+    private static final Logger logger = LogManager.getLogger(ThrottledIterator.class);
+
+    /**
+     * Iterate through the given collection, performing an operation on each item which may fork background tasks, but with a limit on the
+     * number of such background tasks running concurrently to avoid overwhelming the rest of the system (e.g. starving other work of access
+     * to an executor).
+     *
+     * @param iterator The items to iterate. May be accessed by multiple threads, but accesses are all protected by synchronizing on itself.
+     * @param itemConsumer The operation to perform on each item. Each operation receives a {@link RefCounted} which can be used to track
+     *                     the execution of any background tasks spawned for this item. This operation may run on the thread which
+     *                     originally called {@link #run}, if this method has not yet returned. Otherwise it will run on a thread on which a
+     *                     background task previously called {@link RefCounted#decRef()} on its ref count. This operation should not throw
+     *                     any exceptions.
+     * @param maxConcurrency The maximum number of ongoing operations at any time.
+     * @param onItemCompletion Executed when each item is completed, which can be used for instance to report on progress. Must not throw
+     *                         exceptions.
+     * @param onCompletion     Executed when all items are completed.
+     */
+    public static <T> void run(
+        Iterator<T> iterator,
+        BiConsumer<Releasable, T> itemConsumer,
+        int maxConcurrency,
+        Runnable onItemCompletion,
+        Runnable onCompletion
+    ) {
+        try (var throttledIterator = new ThrottledIterator<>(iterator, itemConsumer, maxConcurrency, onItemCompletion, onCompletion)) {
+            throttledIterator.run();
+        }
+    }
+
+    private final RefCounted refs; // one ref for each running item, plus one for the iterator if incomplete
+    private final Iterator<T> iterator;
+    private final BiConsumer<Releasable, T> itemConsumer;
+    private final Semaphore permits;
+    private final Runnable onItemCompletion;
+
+    private ThrottledIterator(
+        Iterator<T> iterator,
+        BiConsumer<Releasable, T> itemConsumer,
+        int maxConcurrency,
+        Runnable onItemCompletion,
+        Runnable onCompletion
+    ) {
+        this.iterator = Objects.requireNonNull(iterator);
+        this.itemConsumer = Objects.requireNonNull(itemConsumer);
+        if (maxConcurrency <= 0) {
+            throw new IllegalArgumentException("maxConcurrency must be positive");
+        }
+        this.permits = new Semaphore(maxConcurrency);
+        this.onItemCompletion = Objects.requireNonNull(onItemCompletion);
+        this.refs = AbstractRefCounted.of(onCompletion);
+    }
+
+    private void run() {
+        while (permits.tryAcquire()) {
+            final T item;
+            synchronized (iterator) {
+                if (iterator.hasNext()) {
+                    item = iterator.next();
+                } else {
+                    permits.release();
+                    return;
+                }
+            }
+            try (var itemRefs = new ItemRefCounted()) {
+                itemRefs.incRef();
+                itemConsumer.accept(Releasables.releaseOnce(itemRefs::decRef), item);
+            } catch (Exception e) {
+                logger.error(Strings.format("exception when processing [%s] with [%s]", item, itemConsumer), e);
+                assert false : e;
+            }
+        }
+    }
+
+    @Override
+    public void close() {
+        refs.decRef();
+    }
+
+    // A RefCounted for a single item, including protection against calling back into run() if it's created and closed within a single
+    // invocation of run().
+    private class ItemRefCounted extends AbstractRefCounted implements Releasable {
+        private boolean isRecursive = true;
+
+        ItemRefCounted() {
+            refs.incRef();
+        }
+
+        @Override
+        protected void closeInternal() {
+            try {
+                onItemCompletion.run();
+            } catch (Exception e) {
+                logger.error("exception in onItemCompletion", e);
+                assert false : e;
+            } finally {
+                permits.release();
+                try {
+                    // Someone must now pick up the next item. Here we might be called from the run() invocation which started processing
+                    // the just-completed item (via close() -> decRef()) if that item's processing didn't fork or all its forked tasks
+                    // finished first. If so, there's no need to call run() here, we can just return and the next iteration of the run()
+                    // loop will continue the processing; moreover calling run() in this situation could lead to a stack overflow. However
+                    // if we're not within that run() invocation then ...
+                    if (isRecursive() == false) {
+                        // ... we're not within any other run() invocation either, so it's safe (and necessary) to call run() here.
+                        run();
+                    }
+                } finally {
+                    refs.decRef();
+                }
+            }
+        }
+
+        // Note on blocking: we call both of these synchronized methods exactly once (and must enter close() before calling isRecursive()).
+        // If close() releases the last ref and calls closeInternal(), and hence isRecursive(), then there's no other threads involved and
+        // hence no blocking. In contrast if close() doesn't release the last ref then it exits immediately, so the call to isRecursive()
+        // will proceed without delay in this case too.
+
+        private synchronized boolean isRecursive() {
+            return isRecursive;
+        }
+
+        @Override
+        public synchronized void close() {
+            decRef();
+            isRecursive = false;
+        }
+    }
+}

server/src/test/java/org/elasticsearch/common/util/concurrent/ThrottledIteratorTests.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util.concurrent;
+
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.elasticsearch.action.support.RefCountingRunnable;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.FixedExecutorBuilder;
+import org.elasticsearch.threadpool.ScalingExecutorBuilder;
+import org.elasticsearch.threadpool.TestThreadPool;
+import org.elasticsearch.threadpool.ThreadPool;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BooleanSupplier;
+import java.util.stream.IntStream;
+
+public class ThrottledIteratorTests extends ESTestCase {
+    private static final String CONSTRAINED = "constrained";
+    private static final String RELAXED = "relaxed";
+
+    public void testConcurrency() throws InterruptedException {
+        final var maxConstrainedThreads = between(1, 3);
+        final var maxRelaxedThreads = between(1, 100);
+        final var constrainedQueue = between(3, 6);
+        final var threadPool = new TestThreadPool(
+            "test",
+            new FixedExecutorBuilder(Settings.EMPTY, CONSTRAINED, maxConstrainedThreads, constrainedQueue, CONSTRAINED, false),
+            new ScalingExecutorBuilder(RELAXED, 1, maxRelaxedThreads, TimeValue.timeValueSeconds(30), true)
+        );
+        try {
+            final var items = between(1, 10000); // large enough that inadvertent recursion will trigger a StackOverflowError
+            final var itemStartLatch = new CountDownLatch(items);
+            final var completedItems = new AtomicInteger();
+            final var maxConcurrency = between(1, (constrainedQueue + maxConstrainedThreads) * 2);
+            final var itemPermits = new Semaphore(maxConcurrency);
+            final var completionLatch = new CountDownLatch(1);
+            final BooleanSupplier forkSupplier = randomFrom(
+                () -> false,
+                ESTestCase::randomBoolean,
+                LuceneTestCase::rarely,
+                LuceneTestCase::usually,
+                () -> true
+            );
+            final var blockPermits = new Semaphore(between(0, Math.min(maxRelaxedThreads, maxConcurrency) - 1));
+
+            ThrottledIterator.run(IntStream.range(0, items).boxed().iterator(), (releasable, item) -> {
+                try (var refs = new RefCountingRunnable(releasable::close)) {
+                    assertTrue(itemPermits.tryAcquire());
+                    if (forkSupplier.getAsBoolean()) {
+                        var ref = refs.acquire();
+                        final var executor = randomFrom(CONSTRAINED, RELAXED);
+                        threadPool.executor(executor).execute(new AbstractRunnable() {
+
+                            @Override
+                            public void onRejection(Exception e) {
+                                assertEquals(CONSTRAINED, executor);
+                                itemStartLatch.countDown();
+                            }
+
+                            @Override
+                            protected void doRun() {
+                                itemStartLatch.countDown();
+                                if (RELAXED.equals(executor) && randomBoolean() && blockPermits.tryAcquire()) {
+                                    // simulate at most (maxConcurrency-1) long-running operations, to demonstrate that they don't
+                                    // hold up the processing of the other operations
+                                    try {
+                                        assertTrue(itemStartLatch.await(30, TimeUnit.SECONDS));
+                                    } catch (InterruptedException e) {
+                                        throw new AssertionError("unexpected", e);
+                                    } finally {
+                                        blockPermits.release();
+                                    }
+                                }
+                            }
+
+                            @Override
+                            public void onAfter() {
+                                itemPermits.release();
+                                ref.close();
+                            }
+
+                            @Override
+                            public void onFailure(Exception e) {
+                                throw new AssertionError("unexpected", e);
+                            }
+                        });
+                    } else {
+                        itemStartLatch.countDown();
+                        itemPermits.release();
+                    }
+                }
+            }, maxConcurrency, completedItems::incrementAndGet, completionLatch::countDown);
+
+            assertTrue(completionLatch.await(30, TimeUnit.SECONDS));
+            assertEquals(items, completedItems.get());
+            assertTrue(itemPermits.tryAcquire(maxConcurrency));
+            assertTrue(itemStartLatch.await(0, TimeUnit.SECONDS));
+        } finally {
+            ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
+        }
+    }
+}