使用信号量控制线程池任务提交速率的一个坑

​0 背景

不久前，标哥（此乃一位技术非常厉害的我的同事）曾让我看过一段很有意思的代码。这段代码的目的是通过信号量控制向线程池中提交任务，避免在线程池中大量积压任务影响程序执行效率。

我简化了一下，大概的逻辑如下：

// 实际工作的线程数，假设是4
static final int THREAD_CNT = 4;
public static void doWork() throws InterruptedException {
    // 使用了核心线程数和最大线程数都为4，且使用SynchronousQueue作为等待队列的线程池
    ThreadPoolExecutor executor = new ThreadPoolExecutor(THREAD_CNT, THREAD_CNT, 0l, TimeUnit.MINUTES, new SynchronousQueue<>());
    // 使用令牌数等于线程数，都是4的信号量控制任务提交速率
    Semaphore semaphore = new Semaphore(THREAD_CNT);
    while (true) {
        semaphore.acquire();
        executor.execute(() -> {
            // 这里面是业务逻辑，假设是打印时间戳
            System.out.println(System.currentTimeMillis());
            semaphore.release();
        });
    }
}

在提交任务的主线程中调用Semaphore.acquire()，在具体的执行线程执行结束时调用Semaphore.release()方法。

这段代码跑起来后，立刻就报错了：

这是为啥呢？

1 SynchronousQueue

SynchronousQueue有一对方法：

• offer()：向队列中插入元素，如果有线程等待获取则插入成功，否则插入失败；

• take()：从队列中获取元素，如果有则立刻返回，否则等待其他线程向队列中插入数据；

public static void main(String[] args) throws InterruptedException {
    SynchronousQueue<String> queue = new SynchronousQueue<>();
    // 向队列中插入数据，返回失败了
    System.out.println("[Producer]: insert when no consumer: " + queue.offer("Hello World"));

    // 先注册一个线程，在插入数据，这时返回成功了
    new Thread(() -> {
        try {
            System.out.println("[CONSUMER]: " + queue.take());
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }).start();
    Thread.sleep(10); // 等待子线程执行到queue.take()
    System.out.println("[Producer]: " + queue.offer("Hello World"));
}

返回结果：

[Producer]: insert when no consumer: false
[Producer]: true
[CONSUMER]: Hello World

2 ThreadPoolExecutor

2.1 子线程执行逻辑

final void runWorker(Worker w) {
    ...
    try {
        while (task != null || (task = getTask()) != null) {
            ...
            try {
                beforeExecute(wt, task);
                try {
                    task.run();
                    afterExecute(task, null);
                } catch (Throwable ex) {
                    afterExecute(task, ex);
                    throw ex;
                }
            } finally {
                task = null;
                w.completedTasks++;
                w.unlock();
            }
        }
        completedAbruptly = false;
    } finally {
        processWorkerExit(w, completedAbruptly);
    }
}

• task.run()方法执行的是Runnable中的方法，也就是下面这段：

// 这里面是业务逻辑，假设是打印时间戳
System.out.println(System.currentTimeMillis());
semaphore.release();

• getTask()方法是从workQueue队列中获取元素

由于是多线程执行，所以任务线程在执行完semaphore.release()，主线程便开始向workQueue中添加数据了。此时，子线程很有可能还未执行到getTask()的位置。

2.2 主线程提交逻辑

public void execute(Runnable command) {
    ...
    int c = ctl.get();
    if (workerCountOf(c) < corePoolSize) {
        if (addWorker(command, true))
            return;
        c = ctl.get();
    }
    if (isRunning(c) && workQueue.offer(command)) {
        int recheck = ctl.get();
        if (! isRunning(recheck) && remove(command))
            reject(command);
        else if (workerCountOf(recheck) == 0)
            addWorker(null, false);
    }
    else if (!addWorker(command, false))
        reject(command);
}

• 主线程提交走到isRunning(c) && workQueue.offer(command)，调用semaphore.release()的线程若还没有执行到getTask()方法的话，这时候执行workQueue.offer(command)添加任务会失败，workQueue中还是空的。

• 失败后会执行else if (!addWorker(command, false))，由于核心线程数和最大线程数都是4，系统认为当前所有线程都忙，添加任务失败，然后就抛异常了。

3 解决方案

根本原因是因为任务线程调用semaphore.release()后还需要执行一些额外的指令才能正常接收下一个任务。而这段期间很有可能添加任务的主线程已经完成任务的添加了。因此，解决方案大概会有两种：

方案一：多增加一些任务线程，保证主线程提交时有多余线程接收新提交任务。

static final int THREAD_CNT = 4;
public static void doWork() throws InterruptedException {
    ThreadPoolExecutor executor = new ThreadPoolExecutor(THREAD_CNT + 1, THREAD_CNT + 1, 0l, TimeUnit.MINUTES, new SynchronousQueue<>());
    Semaphore semaphore = new Semaphore(THREAD_CNT);
    while (true) {
        semaphore.acquire();
        executor.execute(() -> {
            System.out.println(System.currentTimeMillis());
            semaphore.release();
        });
    }
}

比如上面代码，将线程数改为5，信号量还保持4，这样基本能保证程序正常运行，但极端情况下还是会有可能导致任务提交失败的。如果要保证程序100%正确，则线程数应该是信号量的2倍，很浪费资源。

方案二：在提交任务之前的位置，增加一个小延迟，保证任务线程执行到getTask()后才会有任务提交。

下面这段代码在提交任务的地方增加了10ms的延迟，这样基本能保证子线程可以在提交任务前执行到getTask()的位置。极端情况下这种方案仍可能会出问题，但出问题的概率已经非常小了。

import java.util.concurrent.Semaphore;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class FlowControlExecutorUtils {
    public static void main(String[] args) throws InterruptedException {
        doWork();
    }

    static final int THREAD_CNT = 4;
    public static void doWork() throws InterruptedException {
        ThreadPoolExecutor executor = new ThreadPoolExecutor(THREAD_CNT, THREAD_CNT, 0l, TimeUnit.MINUTES, new SynchronousQueue<>());
        Semaphore semaphore = new Semaphore(THREAD_CNT);
        while (true) {
            semaphore.acquire();
            Thread.sleep(10);
            executor.execute(() -> {
                System.out.println(System.currentTimeMillis());
                semaphore.release();
            });
        }
    }
}