控制线程池消费速度

控制线程池消费速度

场景

在使用线程池的时候，如果生产者过快，消费者过慢，有可能将 消费者线程池打满，线程池会执行拒绝策略

目的

目的就是当线程池满时，可以让生产者阻塞

解决 方案

当线程池执行拒绝策略时，将任务重新提交给线程池的阻塞队列 ，这样，就会阻塞生产者提交

实现

/**
 * 控制线程池消费速度
 * 如果生产者过快，消费者过慢，有可能将 消费者线程池打满，线程池会执行拒绝策略
 * 目的就是当线程池满时，可以让生产者阻塞
 *
 * 解决方案：
 *  当线程池执行拒绝策略时，将任务重新提交给线程池的阻塞队列 ，这样，就会阻塞生产者提交
 */
@Component
@Slf4j
public class QueueService {

    private BlockingQueue<String> queue = new ArrayBlockingQueue<>(1000);
    AtomicBoolean locked = new AtomicBoolean(false);
    ThreadPoolExecutor executor = new ThreadPoolExecutor(5, 100, 60, TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(1000), new SimplePrefixThreadFactory("queue"), new RejectedExecutionHandler() {
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            log.info("{} 线程池已满", Thread.currentThread().getName());
            try {
                executor.getQueue().put(r);
            } catch (InterruptedException e) {
            }
            log.info("{} 线程池可用，放入任务", Thread.currentThread().getName());
        }
    });

    public QueueService() {
        new Thread(() -> {
            if (locked.compareAndSet(false, true)) {
                start();
            }
        }).start();
    }

    private void start() {
        while (locked.get()) {
            try {
                String poll = queue.poll(10, TimeUnit.SECONDS);
                if (poll == null) {
                    continue;
                }
                executor.execute(() -> {
                    log.info("{} handle 任务 {}", Thread.currentThread().getName(),poll);
                    try {
                        TimeUnit.SECONDS.sleep(1);
                    } catch (InterruptedException e) {
                    }
                });
            } catch (InterruptedException e) {
            }
        }

    }

    public void put(String value) {
        try {
            queue.put(value);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    public void shutdown() {
        locked.compareAndSet(true, false);
    }
}

结果:

queue54 handle 任务 b28ca9e3-ece8-4e7b-94fd-9f0e75b7a82a
Thread-1 线程池可用，放入任务
Thread-1 线程池已满
queue65 handle 任务 f32c380a-410b-4372-a132-97e631e5a4ae
Thread-1 线程池可用，放入任务
Thread-1 线程池已满
queue53 handle 任务 5b974d7b-e9fd-45db-b628-0460b703475a
Thread-1 线程池可用，放入任务
Thread-1 线程池已满

可以发现，线程消费速度已经控制 ，而且 线程池发挥了最大作用