net.dongliu.commons.concurrent.RichExecutorService Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of commons Show documentation
There is a newer version: 6.7.0
package net.dongliu.commons.concurrent;

import net.dongliu.commons.PreChecks;

import javax.annotation.Nonnull;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import static java.util.stream.Collectors.*;

/**
 * Executor service wrapper can wait all task finished
 *
 * @author Liu Dong
 */
public class RichExecutorService implements ExecutorService {
    private final ExecutorService executorService;

    private final Lock lock = new ReentrantLock();
    private final Condition condition = lock.newCondition();
    private int count = 0;

    public RichExecutorService(ExecutorService executorService) {
        this.executorService = executorService;
    }

    /**
     * Task num not finished
     */
    public int unFinishedTasks() {
        lock.lock();
        try {
            return count;
        } finally {
            lock.unlock();
        }

    }

    /**
     * No task wait to be run
     */
    public boolean idle() {
        lock.lock();
        try {
            return count == 0;
        } finally {
            lock.unlock();
        }
    }

    /**
     * Wait till no task to be run
     */
    public void waitTillIdle() throws InterruptedException {
        lock.lock();
        try {
            while (count > 0) {
                condition.await();
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * Wait till no task to be run
     */
    public void waitTillIdle(long time, TimeUnit timeUnit) throws InterruptedException, TimeoutException {
        long nanos = timeUnit.toNanos(time);
        lock.lock();
        try {
            while (count > 0) {
                nanos = condition.awaitNanos(nanos);
                if (nanos <= 0) {
                    throw new TimeoutException();
                }
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void shutdown() {
        executorService.shutdown();
    }

    @Nonnull
    @Override
    public List shutdownNow() {
        return executorService.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return executorService.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return executorService.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return executorService.awaitTermination(timeout, unit);
    }

    @Nonnull
    @Override
    public  Future submit(Callable task) {
        increaseCount();
        return executorService.submit(() -> callAndArrive(task));
    }

    @Nonnull
    @Override
    public  Future submit(Runnable task, T result) {
        increaseCount();
        return executorService.submit(() -> runAndArrive(task), result);
    }

    @Nonnull
    @Override
    public Future submit(Runnable task) {
        increaseCount();
        return executorService.submit(() -> runAndArrive(task));
    }

    @Nonnull
    @Override
    public  List> invokeAll(Collection> tasks) throws InterruptedException {
        PreChecks.nonNull(tasks);
        tasks = tasks.stream().map(task -> (Callable) () -> callAndArrive(task)).collect(toList());
        increaseCount(tasks.size());
        return executorService.invokeAll(tasks);
    }

    @Nonnull
    @Override
    public  List> invokeAll(Collection> tasks, long timeout, TimeUnit unit)
            throws InterruptedException {
        PreChecks.nonNull(tasks);
        tasks = tasks.stream().map(task -> (Callable) () -> callAndArrive(task)).collect(toList());
        increaseCount(tasks.size());
        return executorService.invokeAll(tasks, timeout, unit);
    }

    @Nonnull
    @Override
    public  T invokeAny(Collection> tasks) throws InterruptedException, ExecutionException {
        PreChecks.nonNull(tasks);
        tasks = tasks.stream().map(task -> (Callable) () -> callAndArrive(task)).collect(toList());
        increaseCount(tasks.size());
        return executorService.invokeAny(tasks);
    }

    @Override
    public  T invokeAny(Collection> tasks, long timeout, TimeUnit unit)
            throws InterruptedException, ExecutionException, TimeoutException {
        PreChecks.nonNull(tasks);
        tasks = tasks.stream().map(task -> (Callable) () -> callAndArrive(task)).collect(toList());
        increaseCount(tasks.size());
        return executorService.invokeAny(tasks, timeout, unit);
    }

    @Override
    public void execute(Runnable command) {
        increaseCount();
        executorService.execute(() -> runAndArrive(command));
    }

    private  T callAndArrive(Callable task) throws Exception {
        try {
            return task.call();
        } finally {
            decreaseCountAndNotify();
        }
    }

    private void runAndArrive(Runnable task) {
        try {
            task.run();
        } finally {
            decreaseCountAndNotify();
        }
    }

    private void increaseCount() {
        increaseCount(1);
    }

    private void increaseCount(int num) {
        lock.lock();
        try {
            count += num;
        } finally {
            lock.unlock();
        }
    }

    private void decreaseCountAndNotify() {
        lock.lock();
        try {
            if (--count == 0) {
                condition.signal();
            }
        } finally {
            lock.unlock();
        }
    }


}