如何扩展AbstractExecutorService保证线程标签多样性，规避同数据源死锁

1. 核心改造点梳理 #

你的需求本质是实现按标签的串行执行调度：同标签任务同一时间最多一个运行，不同标签任务可以并行执行，这个需求完全可以基于JDK自带的线程池体系扩展实现，不需要完全重写执行逻辑。

2. 第一步：给任务打标签 #

首先定义带标签的任务封装类，继承FutureTask，用来在任务全生命周期携带标签信息：

public class LabeledFutureTask<T> extends FutureTask<T> {
    // 存储任务对应的数据源标签
    private final String label;

    // 适配Runnable类型任务
    public LabeledFutureTask(Runnable runnable, T result, String label) {
        super(runnable, result);
        this.label = label;
    }

    // 适配Callable类型任务
    public LabeledFutureTask(Callable<T> callable, String label) {
        super(callable);
        this.label = label;
    }

    public String getLabel() {
        return label;
    }
}

然后重写AbstractExecutorService的newTaskFor方法，同时重载线程池的submit方法，支持传入标签参数：

public class LabelAwareThreadPoolExecutor extends ThreadPoolExecutor {
    // 记录当前正在运行的任务标签
    private final Set<String> runningLabels = ConcurrentHashMap.newKeySet();

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

    // 重载submit方法，支持传入标签
    public <T> Future<T> submit(Runnable task, String label, T result) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<T> ftask = newTaskFor(task, result, label);
        execute(ftask);
        return ftask;
    }

    public <T> Future<T> submit(Callable<T> task, String label) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<T> ftask = newTaskFor(task, label);
        execute(ftask);
        return ftask;
    }

    // 重写newTaskFor生成带标签的任务
    protected <T> RunnableFuture<T> newTaskFor(Runnable runnable, T value, String label) {
        return new LabeledFutureTask<>(runnable, value, label);
    }

    protected <T> RunnableFuture<T> newTaskFor(Callable<T> callable, String label) {
        return new LabeledFutureTask<>(callable, label);
    }

3. 第二步：实现标签调度控制 #

重写线程池的beforeExecute和afterExecute方法，配合自定义的任务队列实现调度逻辑：

方案A：轻量实现（适合任务量不大的场景） #

自定义阻塞队列，重写take方法，每次从队列中取出第一个标签不在运行中的任务即可：

public class LabelAwareBlockingQueue extends LinkedBlockingQueue<Runnable> {
    private final Set<String> runningLabels;

    public LabelAwareBlockingQueue(Set<String> runningLabels) {
        this.runningLabels = runningLabels;
    }

    @Override
    public Runnable take() throws InterruptedException {
        lock.lock();
        try {
            while (true) {
                // 遍历队列找第一个可运行的任务
                Iterator<Runnable> iterator = iterator();
                while (iterator.hasNext()) {
                    Runnable task = iterator.next();
                    if (task instanceof LabeledFutureTask<?>) {
                        String label = ((LabeledFutureTask<?>) task).getLabel();
                        if (!runningLabels.contains(label)) {
                            iterator.remove();
                            return task;
                        }
                    } else {
                        // 无标签的任务直接执行
                        iterator.remove();
                        return task;
                    }
                }
                // 没有符合条件的任务，等待新任务提交
                notEmpty.await();
            }
        } finally {
            lock.unlock();
        }
    }

    // 同理重写poll方法适配非阻塞取任务逻辑
}

然后在beforeExecute里把任务标签加入运行集合，afterExecute里移除：

@Override
protected void beforeExecute(Thread t, Runnable r) {
    super.beforeExecute(t, r);
    if (r instanceof LabeledFutureTask<?>) {
        String label = ((LabeledFutureTask<?>) r).getLabel();
        runningLabels.add(label);
    }
}

@Override
protected void afterExecute(Runnable r, Throwable t) {
    super.afterExecute(r, t);
    if (r instanceof LabeledFutureTask<?>) {
        String label = ((LabeledFutureTask<?>) r).getLabel();
        runningLabels.remove(label);
    }
}

这个方案天然满足你要求的运行线程数上限规则：最多同时运行的任务数就是当前不同标签的数量，不会超过线程池配置的最大池大小。

方案B：高性能实现（适合任务量极大的场景） #

如果你的任务提交量很大，遍历队列的性能损耗较高，可以给每个标签维护单独的子队列，调度时直接从空闲标签的子队列取任务即可，时间复杂度可以降到O(1)。

4. 注意事项 #

• 标签需要全局唯一标识数据源，建议提前对标签做标准化处理（比如统一小写、去前后空格），避免同一个数据源被识别为不同标签
• 若同一个标签的任务堆积过多，可以额外配置单标签最大任务队列长度，避免内存溢出
• 如果需要支持任务优先级，可以在标签队列的基础上增加优先级判断逻辑

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