Java class to run CompletableFuture<T> instances in batches

Question

This is a utility class for running CompletableFuture<T>s in batches (e.g. to avoid using too much bandwidth making HTTP requests). It seems to work, but I'm a little iffy on the use of my CompletableFutureReferenceWrapper class— I need it so that wrappedFuture can remove itself from running after it completes. Is there a more idiomatic way to do this? Other code suggestions are also appreciated.

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.function.Supplier;

/**
 * Run {@link CompletableFuture}s in batches.
 *
 * @param <T> The type of the {@link CompletableFuture}.
 */
public final class JobsQueue<T> {

    private final Queue<Supplier<CompletableFutureReferenceWrapper>> waiting = new ConcurrentLinkedQueue<>();
    private final Set<CompletableFutureReferenceWrapper> running = new HashSet<>();
    private final Queue<T> output = new ArrayDeque<>();
    private final int maxJobsRunningAtOnce;

    /**
     * Initialize a new {@link JobsQueue} instance.
     *
     * @param maxJobsRunningAtOnce The maxmimum number of {@link CompletableFuture}s that will run simultaneously.
     */
    public JobsQueue(int maxJobsRunningAtOnce) {
        this.maxJobsRunningAtOnce = maxJobsRunningAtOnce;
    }

    /**
     * Enqueue a future.
     *
     * @param completableFutureSupplier Supplier for the future to be run.
     */
    public void add(Supplier<CompletableFuture<T>> completableFutureSupplier) {
        Supplier<CompletableFutureReferenceWrapper> wrappedSupplier = () -> {
            CompletableFutureReferenceWrapper wrapper = new CompletableFutureReferenceWrapper(null);
            CompletableFuture<T> originalFuture = completableFutureSupplier.get();
            CompletableFuture<Void> wrappedFuture = originalFuture.thenAccept(result -> {
                synchronized (this) {
                    if (result != null) {
                        this.output.add(result);
                    }

                    this.running.remove(wrapper);
                    if (!this.waiting.isEmpty()) {
                        this.running.add(this.waiting.remove().get());
                    }
                }
            });
            wrapper.future = wrappedFuture;
            return wrapper;
        };
        synchronized (this) {
            if (this.running.size() < this.maxJobsRunningAtOnce) {
                this.running.add(wrappedSupplier.get());
            } else {
                this.waiting.add(wrappedSupplier);
            }
        }
    }

    /**
     * Wait for all running and waiting futures to complete before returning.
     */
    public void joinAll() {
        // not only must we join all futures in running, we must wait for those in waiting to trickle out as well
        while (true) {
            final boolean breakFlag;
            synchronized (this.waiting) {
                breakFlag = this.waiting.isEmpty();
            }

            final List<CompletableFutureReferenceWrapper> wrappers;
            synchronized (this.running) {
                wrappers = new ArrayList<>(this.running);
            }

            for (CompletableFutureReferenceWrapper wrapper : wrappers) {
                assert wrapper.future != null;
                wrapper.future.join();
            }

            if (breakFlag) {
                break;
            }
        }
    }

    /**
     * Get the results of all the completed futures so far.
     * Empties the results queue.
     *
     * @return a list of all the results of the completed futures.
     * If <code>T</code> is {@link Void}, returns an empty list.
     */
    public @NotNull List<T> getAll() {
        synchronized (this.output) {
            List<T> out = new ArrayList<>(this.output);
            this.output.clear();
            return out;
        }
    }

    private static class CompletableFutureReferenceWrapper {
        @Nullable CompletableFuture<Void> future;

        CompletableFutureReferenceWrapper(@Nullable CompletableFuture<Void> future) {
            this.future = future;
        }
    }
}

Answer 1

If I understood it correctly, the main idea is to control the number of parallel executions using maxJobsRunningAtOnce?

This is a violation of single responsibility principle as JobsQueue acts like both queue and executor.

If to separate the queue from the executor, you could simply use something like running.forEach(CompletableFuture::join) and control the parallel thread by an external executor.

Ok, let's forget about the initial design.

You do use an explicit synchronization on both waiting and running collections. The first one is already ConcurrentLinkedQueue and it's safe to check its size. Additionally, you could avoid the second synchronized block by making running as thread-safe also, for example as ConcurrentHashMap.newKeySet()

By the way, output is declared as ArrayDeque but in the end, is used just to copy the values to the plain array list. Plain ArrayList would do the same trick.