com.pivovarit.collectors.Dispatcher Maven / Gradle / Ivy

Go to download
package com.pivovarit.collectors;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.FutureTask;
import java.util.concurrent.Semaphore;
import java.util.function.BiConsumer;
import java.util.function.Supplier;

/**
 * @author Grzegorz Piwowarek
 */
final class Dispatcher {

    private final CompletableFuture completionSignaller = new CompletableFuture<>();
    private final Executor executor;
    private final Semaphore limiter;

    private Dispatcher() {
        this.executor = Executors.newVirtualThreadPerTaskExecutor();
        this.limiter = null;
    }

    private Dispatcher(Executor executor, int permits) {
        this.executor = executor;
        this.limiter = new Semaphore(permits);
    }

    static  Dispatcher from(Executor executor, int permits) {
        return new Dispatcher<>(executor, permits);
    }

    static  Dispatcher virtual() {
        return new Dispatcher<>();
    }

    CompletableFuture enqueue(Supplier supplier) {
        InterruptibleCompletableFuture future = new InterruptibleCompletableFuture<>();
        completionSignaller.whenComplete(shortcircuit(future));
        try {
            executor.execute(completionTask(supplier, future));
        } catch (Throwable e) {
            completionSignaller.completeExceptionally(e);
            return CompletableFuture.failedFuture(e);
        }
        return future;
    }

    private FutureTask completionTask(Supplier supplier, InterruptibleCompletableFuture future) {
        FutureTask task = new FutureTask<>(() -> {
            if (!completionSignaller.isCompletedExceptionally()) {
                try {
                    if (limiter == null) {
                        future.complete(supplier.get());
                    } else {
                        try {
                            limiter.acquire();
                            future.complete(supplier.get());
                        } finally {
                            limiter.release();
                        }
                    }
                } catch (Throwable e) {
                    completionSignaller.completeExceptionally(e);
                }
            }
        }, null);
        future.completedBy(task);
        return task;
    }

    private static  BiConsumer shortcircuit(InterruptibleCompletableFuture future) {
        return (__, throwable) -> {
            if (throwable != null) {
                future.completeExceptionally(throwable);
                future.cancel(true);
            }
        };
    }

    static final class InterruptibleCompletableFuture extends CompletableFuture {

        private volatile FutureTask backingTask;

        private void completedBy(FutureTask task) {
            backingTask = task;
        }

        @Override
        public boolean cancel(boolean mayInterruptIfRunning) {
            var task = backingTask;
            if (task != null) {
                task.cancel(mayInterruptIfRunning);
            }
            return super.cancel(mayInterruptIfRunning);
        }
    }
}