package com.diaoyun.zion.master.thread;

import java.util.concurrent.*;

/**
 * 通过Semaphore限定线程数的线程池调用
 */
public class TaskLimitSemaphore {
    //the number of threads to keep in the pool, even
    // if they are idle, unless {@code allowCoreThreadTimeOut} is set
    private int corePoolSize = 5;
    //the maximum number of threads to allow in the pool
    private int maximumPoolSize = 20;
    //when the number of threads is greater than
    //the core, this is the maximum time that excess idle threads
    //will wait for new tasks before terminating.
    private long keepAliveTime = 60L;
    //线程池
    private ExecutorService executor;
    //控制线程数
    private Semaphore semaphore;

    public TaskLimitSemaphore(int limit) {
        corePoolSize = limit;
        this.maximumPoolSize = limit + 20;
        this.semaphore = new Semaphore(limit);
        this.executor = new ThreadPoolExecutor(corePoolSize, maximumPoolSize,
                keepAliveTime, TimeUnit.SECONDS, new SynchronousQueue<Runnable>());
    }

    /**
     * @param corePoolSize the number of threads to keep in the pool, even
     *                     if they are idle, unless {@code allowCoreThreadTimeOut} is set
     * @param limit        最大线程数 数值小于等于 maximumPoolSize
     */
    public TaskLimitSemaphore(int corePoolSize, int limit) {
        corePoolSize = limit;
        this.maximumPoolSize = limit + 20;
        this.semaphore = new Semaphore(limit);
        this.executor = new ThreadPoolExecutor(corePoolSize, maximumPoolSize,
                keepAliveTime, TimeUnit.SECONDS, new SynchronousQueue<Runnable>());
    }

    /**
     * @param corePoolSize    the number of threads to keep in the pool, even
     *                        if they are idle, unless {@code allowCoreThreadTimeOut} is set
     * @param maximumPoolSize the maximum number of threads to allow in the pool
     * @param keepAliveTime   when the number of threads is greater than
     *                        the core, this is the maximum time that excess idle threads
     *                        will wait for new tasks before terminating.
     * @param limit           最大线程数 数值小于等于 maximumPoolSize
     */
    public TaskLimitSemaphore(int corePoolSize, int maximumPoolSize, long keepAliveTime, int limit) {
        corePoolSize = limit;
        this.maximumPoolSize = limit + 20;
        this.keepAliveTime = keepAliveTime;
        this.semaphore = new Semaphore(limit);
        this.executor = new ThreadPoolExecutor(corePoolSize, maximumPoolSize,
                keepAliveTime, TimeUnit.SECONDS, new SynchronousQueue<Runnable>());
    }

    public <T> Future<T> submit(final Callable<T> task) throws InterruptedException {
        semaphore.acquire();
        return executor.submit(() -> {
            try {
                return task.call();
            } finally {
                semaphore.release();
            }
        });
    }

    /**
     * 关闭线程池
     */
    public void shutdown() {
        executor.shutdown();
    }

}
