自定义Java ThreadPool暂停提交并取消排队任务时异常排查

需求与问题描述 #

我需要实现一个自定义Java ThreadPool，具备以下能力：

• 支持暂停功能，暂停时删除当前排队的任务；
• 暂停期间可接收任务，恢复后执行新排队的任务。

当前实现已满足上述需求，但单元测试时，ThreadPool被暂停后会有一个或多个任务抛出CancellationException异常。

ThreadPool实现代码 #

import java.util.ArrayList;
import java.util.Collections;
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;

public class QueriesThreadPoolExecutor extends ThreadPoolExecutor {
    private final List<Future<?>> futuresList = Collections.synchronizedList(new ArrayList<>());
    private boolean isPaused = false;
    private Lock pauseLock = new ReentrantLock();
    private Condition unpaused = pauseLock.newCondition();

    public QueriesThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    @Override
    public void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
//        System.out.println("beforeExecute " + t.getName());
        pauseLock.lock();
        try {
            while (isPaused) {
                unpaused.await();
            }
        } catch (InterruptedException e) {
            t.interrupt();
        } finally {
            pauseLock.unlock();
        }
    }

    @Override
    public void afterExecute(Runnable r, Throwable t) {
        super.afterExecute(r, t);
//        System.out.println("afterExecute " + Thread.currentThread().getName());
    }

    @Override
    public void terminated() {
        super.terminated();
        System.out.println("Thread pool terminated");
    }

    public List<Runnable> shutdownNow() {
        List<Runnable> tasks = super.shutdownNow();

        System.out.println("Shutting down thread pool, active tasks remaining: " + tasks.size());
        return tasks;
    }

    public <T> Future<T> submit(Callable<T> task) {
        System.out.println("submit task");
        Future<T> res = super.submit(task);
        futuresList.add(res);

        return res;
    }

    public void pause() {
        cancelRunningTasks();

        System.out.println("pause");
        pauseLock.lock();
        try {
            isPaused = true;
        } finally {
            pauseLock.unlock();
        }
    }

    public void resume() {
        System.out.println("resume");
        pauseLock.lock();
        try {
            isPaused = false;
            unpaused.signal();
        } finally {
            pauseLock.unlock();
        }
    }

    public void cancelRunningTasks() {
        synchronized (futuresList) {
            for (Future<?> future : futuresList) {
                if (!future.isDone()) {
                    future.cancel(true);
                    System.out.println("cancelRunningTasks");
                }
            }
        }
    }
}

单元测试代码 #

import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

class QueriesThreadPoolExecutorTests {
    private static QueriesThreadPoolExecutor queriesThreadPoolExecutor;
    private AtomicInteger longCounter = new AtomicInteger(0);
    private AtomicInteger shortCounter = new AtomicInteger(0);

    @BeforeAll
    static void setup() {
        queriesThreadPoolExecutor = new QueriesThreadPoolExecutor(5, 5, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>());
    }

    @AfterAll
    static void tearDown() {
        queriesThreadPoolExecutor.shutdownNow();
    }

    @Test
    void testCancelResumeTasks() {
        Callable<String> longRunningTask = () -> {
            Thread.sleep(20_000);
            return "long running task" + longCounter.incrementAndGet();
        };

        Callable<String> shortRunningTask = () -> {
            Thread.sleep(1_000);
            return "short running task" + shortCounter.incrementAndGet();
        };
        List<Future<String>> futureList = new ArrayList<>();

        submitTasks(futureList, shortRunningTask, longRunningTask);

        ScheduledExecutorService cancelScheduledExecutor = Executors.newScheduledThreadPool(2);
        cancelScheduledExecutor.schedule(() -> {
            queriesThreadPoolExecutor.pause();
        }, 2, TimeUnit.SECONDS);

        printResults(futureList);

        submitTasks(futureList, shortRunningTask, longRunningTask);

        queriesThreadPoolExecutor.resume();

        printResults(futureList);
    }

    private void submitTasks(List<Future<String>> futureList, Callable<String> shortRunningTask, Callable<String> longRunningTask) {
        for (int i = 0; i < 5; i++) {
            futureList.add(queriesThreadPoolExecutor.submit(shortRunningTask));
        }
        for (int i = 0; i < 5; i++) {
            futureList.add(queriesThreadPoolExecutor.submit(longRunningTask));
        }
    }

    private void printResults(List<Future<String>> futureList) {
        System.out.println("***********Printing results*****************");
        for (Future<String> future : futureList) {
            try {
                String status = future.isCancelled() ? "cancelled" : "done";
                String result = future.isCancelled() ? "cancelled" : future.get();
                System.out.println("task status - " + status + " - result " + result);
            } catch (InterruptedException | ExecutionException e) {
                e.printStackTrace();
            } catch (CancellationException e) {
                System.out.println("Exception - Task was cancelled");
            }
        }
    }
}

异常输出 #

task status - done - result short running task1
cancelRunningTasks
cancelRunningTasks
cancelRunningTasks
cancelRunningTasks
cancelRunningTasks
Exception - Task was cancelled  <--------------------WHY?
task status - cancelled - result cancelled
task status - cancelled - result cancelled
task status - cancelled - result cancelled
task status - cancelled - result cancelled
submit task
pause
submit task

异常产生的核心原因 #

1. 任务取消时机与future.get()的竞态
   调用pause()时，cancelRunningTasks()会遍历所有未完成的Future并执行cancel(true)。但取消操作和后续的printResults()存在竞态：代码中先判断future.isCancelled()，若返回false则调用future.get()，但就在判断和get()的间隙，任务可能被取消，此时get()会抛出CancellationException。

2. futuresList未清理已完成任务
   维护的futuresList只添加新任务的Future，从未移除已完成或取消的任务。这会导致每次cancelRunningTasks()都要遍历历史任务，且printResults()时部分已取消任务的状态同步不及时，出现判断不准确的情况。

修复方案 #

1. 调整printResults()逻辑
   无需提前判断isCancelled()，直接调用get()并捕获异常，这是更可靠的方式：

   private void printResults(List<Future<String>> futureList) {
       System.out.println("***********Printing results*****************");
       for (Future<String> future : futureList) {
           try {
               String result = future.get();
               System.out.println("task status - done - result " + result);
           } catch (InterruptedException | ExecutionException e) {
               e.printStackTrace();
           } catch (CancellationException e) {
               System.out.println("task status - cancelled - result cancelled");
           }
       }
   }
   

2. 维护futuresList清洁度
   在afterExecute()中移除已完成任务的Future，避免无效遍历和状态混乱：

   @Override
   public void afterExecute(Runnable r, Throwable t) {
       super.afterExecute(r, t);
       Iterator<Future<?>> iterator = futuresList.iterator();
       while (iterator.hasNext()) {
           Future<?> future = iterator.next();
           if (future.isDone()) {
               iterator.remove();
           }
       }
   }
   

3. 优化cancelRunningTasks()锁使用
   futuresList已是Collections.synchronizedList，无需额外加synchronized锁：

   public void cancelRunningTasks() {
       for (Future<?> future : futuresList) {
           if (!future.isDone()) {
               future.cancel(true);
               System.out.println("cancelRunningTasks");
           }
       }
   }
   

内容的提问来源于stack exchange，提问作者RBA