JUC系列(2)——自定义线程池

上一篇刚说完 JDK 中的线程池，这次尝试着自己来定义一个线程池

参考 ThreadPoolExecutor 的实现可以知道，要自定义一个线程池，总少不了以下一些基本元素：

• 核心线程数 CoreSize

• 任务队列 TaskQueue

• 拒绝策略 RejectPolicy

• 任务超时时间 Timeout

• 超时时间单位 Timeunit

下面分步骤实现.

一. 自定义任务队列

java

class BlockingQueue<T> {

    private Deque<T> queue;

    private ReentrantLock lock;

    private Condition fullWaitSignal;

    private Condition emptyWaitSignal;

    private int capacity;

    public BlockingQueue(int capacity) {
        queue = new ArrayDeque<>();
        lock = new ReentrantLock();
        fullWaitSignal = lock.newCondition();
        emptyWaitSignal = lock.newCondition();
        this.capacity = capacity;
    }

    public T poll(long timeout, TimeUnit timeUnit) {
        lock.lock();
        try {
            long nanos = timeUnit.toNanos(timeout);
            while (queue.isEmpty()) {
                if (nanos <= 0) {
                    return null;
                }
                try {
                    nanos = emptyWaitSignal.awaitNanos(nanos);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            T t = queue.removeFirst();
            System.out.printf("task [%s] removed from queue\n", t);
            fullWaitSignal.signal();
            return t;
        } finally {
            lock.unlock();
        }

    }

    public T take() {
        return poll(0, TimeUnit.NANOSECONDS);
    }

    public void put(T t) {
        lock.lock();
        try {
            while (queue.size() == capacity) {
                try {
                    fullWaitSignal.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            System.out.printf("[put]-add task [%s] to queue\n", t);
            queue.addLast(t);
            emptyWaitSignal.signal();

        } finally {
            lock.unlock();
        }
    }

    public boolean offer(T t, long timeout, TimeUnit timeUnit) {
        lock.lock();
        try {
            long nanos = timeUnit.toNanos(timeout);
            while (queue.size() == capacity) {
                // if time's up
                if (nanos <= 0)
                    return false;

                try {
                    nanos = fullWaitSignal.awaitNanos(nanos);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            System.out.printf("[offer]-task [%s] add into queue\n", t);
            queue.addLast(t);
            emptyWaitSignal.signal();
            return true;

        } finally {
            lock.unlock();
        }
    }

    public boolean tryPut(T t, RejectedPolicy policy) {
        lock.lock();
        try {
            if (queue.size() == capacity) {
                policy.reject(t, this);
                return false;
            } else {
                System.out.printf("[try put]-task [%s] add into queue\n", t);
                queue.addLast(t);
                emptyWaitSignal.signal();
                return true;
            }

        } finally {
            lock.unlock();
        }

    }

    public int size() {
        lock.lock();
        try {
            return queue.size();
        } finally {
            lock.unlock();
        }
    }
}

二. 自定义拒绝策略接口

java

@FunctionalInterface
interface RejectedPolicy<T> {

    void reject(T task, BlockingQueue<T> queue);
}

三. 自定义线程池

java

public class CustomizedThreadPool {

    private int coresize;

    private BlockingQueue<Runnable> taskQueue;

    private Set<Worker> workers;

    private long timeout;

    private TimeUnit timeUnit;

    private RejectedPolicy<Runnable> rejectedPolicy;

    public CustomizedThreadPool(int coresize, int queueSize, long timeout, TimeUnit timeUnit, RejectedPolicy<Runnable> rejectedPolicy) {
        this.coresize = coresize;
        this.taskQueue = new BlockingQueue<>(queueSize);
        workers = new HashSet<>();
        this.timeout = timeout;
        this.timeUnit = timeUnit;
        this.rejectedPolicy = rejectedPolicy;
    }


    public void execute(Runnable runnable) {
        synchronized (workers) {
            if (workers.size() < coresize) {
                Worker worker = new Worker(runnable);
                workers.add(worker);
                System.out.printf("worker [%s] created, task [%s]\n", worker, runnable);
                worker.start();
            } else {
                // try to put the task into task queue
                taskQueue.tryPut(runnable, rejectedPolicy);
            }
        }
    }
    
    class Worker extends Thread {

        private Runnable task;

        public Worker(Runnable task) {
            this.task = task;
        }

        @Override
        public void run() {

            while (task != null || (task = taskQueue.poll(timeout, timeUnit)) != null) {
                try {
                    System.out.printf("running [%s]\n", task);
                    task.run();
                } finally {
                   task = null;
                }
                synchronized (workers) {
                    System.out.printf("remove worker [%s]", this);
                    workers.remove(this);
                    System.out.printf("worker size is:%d\n", workers.size());
                }

            }

        }
    }
}

四. 编写测试用例

java

public static void main(String[] args) {

    CustomizedThreadPool threadPool = new CustomizedThreadPool(4, 5, 1L, TimeUnit.SECONDS,
                                                               (task, queue) -> {
                                                                   task.run();
                                                               });
    for (int i=0; i < 20; i++) {
        int idx = i;
        threadPool.execute(() -> {
            try {
                // simulate task execution process
                TimeUnit.MILLISECONDS.sleep(2000L);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            System.out.printf("%d\n", idx);
        });
    }
}