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

Go to download
Show more of this group Show more artifacts with this name
Show all versions of parallel-collectors Show documentation
Parallel collection processing with customizable thread pools
The newest version!
package com.pivovarit.collectors;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.FutureTask;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import java.util.function.Supplier;

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

    private static final Runnable POISON_PILL = () -> System.out.println("Why so serious?");

    private final CompletableFuture completionSignaller = new CompletableFuture<>();
    private final BlockingQueue workingQueue = new LinkedBlockingQueue<>();

    private final Executor executor;
    private final Semaphore limiter;

    private final AtomicBoolean started = new AtomicBoolean(false);

    private volatile boolean shortCircuited = false;

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

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

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

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

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

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

    void start() {
        if (!started.getAndSet(true)) {
            Thread.ofVirtual().start(() -> {
                try {
                    while (true) {
                        try {
                            if (limiter != null) {
                                limiter.acquire();
                            }
                        } catch (InterruptedException e) {
                            handle(e);
                        }
                        Runnable task;
                        if ((task = workingQueue.take()) != POISON_PILL) {
                            retry(() -> {
                                executor.execute(() -> {
                                    try {
                                        task.run();
                                    } finally {
                                        if (limiter != null) {
                                            limiter.release();
                                        }
                                    }
                                });
                            });
                        } else {
                            break;
                        }
                    }
                } catch (Throwable e) {
                    handle(e);
                }
            });
        }
    }

    void stop() {
        try {
            workingQueue.put(POISON_PILL);
        } catch (InterruptedException e) {
            completionSignaller.completeExceptionally(e);
        }
    }

    boolean isRunning() {
        return started.get();
    }

    CompletableFuture enqueue(Supplier supplier) {
        InterruptibleCompletableFuture future = new InterruptibleCompletableFuture<>();
        workingQueue.add(completionTask(supplier, future));
        completionSignaller.exceptionally(shortcircuit(future));
        return future;
    }

    private FutureTask completionTask(Supplier supplier, InterruptibleCompletableFuture future) {
        FutureTask task = new FutureTask<>(() -> {
            try {
                if (!shortCircuited) {
                    future.complete(supplier.get());
                }
            } catch (Throwable e) {
                handle(e);
            }
        }, null);
        future.completedBy(task);
        return task;
    }

    private void handle(Throwable e) {
        shortCircuited = true;
        completionSignaller.completeExceptionally(e);
    }

    private static Function shortcircuit(InterruptibleCompletableFuture future) {
        return throwable -> {
            future.completeExceptionally(throwable);
            future.cancel(true);
            return null;
        };
    }

    static final class InterruptibleCompletableFuture extends CompletableFuture {

        private volatile FutureTask backingTask;

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

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

    private static ExecutorService defaultExecutorService() {
        return Executors.newVirtualThreadPerTaskExecutor();
    }

    private static void requireValidExecutor(Executor executor) {
        if (executor instanceof ThreadPoolExecutor tpe) {
            switch (tpe.getRejectedExecutionHandler()) {
                case ThreadPoolExecutor.DiscardPolicy __ ->
                  throw new IllegalArgumentException("Executor's RejectedExecutionHandler can't discard tasks");
                case ThreadPoolExecutor.DiscardOldestPolicy __ ->
                  throw new IllegalArgumentException("Executor's RejectedExecutionHandler can't discard tasks");
                default -> {
                    // no-op
                }
            }
        }
    }

    private static void retry(Runnable runnable) {
        try {
            runnable.run();
        } catch (RejectedExecutionException e) {
            Thread.onSpinWait();
            runnable.run();
        }
    }
}