场景

在读取文件、或则需要使用多线程批量入库的时候，往往是需要我们自己来写多线程的调度完成多线程批量入库的功能；

难点：多线程的调度、数据分批的逻辑

:::info 不仅仅用于数据库插入，只要在 固定数量 的 多线程处理 的场景都适用 :::

解决的问题

  1. 多线程调度、数据分批的逻辑
  2. 提供多线程批量插入/处理
  3. 提供多线程单条插入/处理

实现思路

  1. 使用 ArrayBlockingQueue 来调节生产方和消费方速度不一致的情况：使用阻塞的 put 来达到让生产方阻塞等待
  2. 使用 CountDownLatch 来实现，等待多线程将所有生产的数据都入库/处理完成

工具类实现

依赖

  1. // lombok，使用了里面的 @Slf4j 日志工具
  2. compileOnly 'org.projectlombok:lombok:1.18.18'
  3. testCompileOnly 'org.projectlombok:lombok:1.18.18'
  4. annotationProcessor 'org.projectlombok:lombok:1.18.18'
  6. // hutool 工具类
  7. // 比如 ExceptionUtil  来自于 https://www.hutool.cn/ 工具包中，完全可以换成手动 new  异常
  8. implementation 'cn.hutool:hutool-all:5.8.3'
  1. package cn.mrcode;
  3. import cn.hutool.core.exceptions.ExceptionUtil;
  4. import cn.hutool.core.util.StrUtil;
  5. import lombok.extern.slf4j.Slf4j;
  7. import java.sql.Struct;
  8. import java.util.ArrayList;
  9. import java.util.List;
  10. import java.util.concurrent.ArrayBlockingQueue;
  11. import java.util.concurrent.CountDownLatch;
  12. import java.util.concurrent.TimeUnit;
  13. import java.util.stream.Collectors;
  14. import java.util.stream.IntStream;
  16. /**
  17.  * 多线程批处理器
  18.  * @author mrcode
  19.  * @date 2021/6/2 17:52
  20.  */
  21. @Slf4j
  22. public class BatchProcessor<T> {
  23.     /**
  24.      * 线程名称前缀，可自定义
  25.      */
  26.     private String threadNamePrefix = "reportInsert-";
  28.     // 是否已经开始处理
  29.     private boolean started;
  30.     // 用于等待线程处理结束后的收尾处理
  31.     private CountDownLatch cdl;
  32.     // 是否还会产生数据: 用于配合 queue.size() 判断线程是否该结束
  33.     private volatile boolean isProduceData = true;
  34.     // 实体数据容器队列，队列满，则限制生产方的生产速度
  35.     private ArrayBlockingQueue<T> queue;
  37.     // 消费到一条实体数据，就调用该方法给使用方，使用方可以调用存储接口存储
  38.     private StorageConsumer<T> consumer;
  39.     // 批量插入时，每次最多插入多少条
  40.     private int maxItemCount;
  41.     private List<WorkThread> workThreads;
  44.     public BatchProcessor() {
  45.         this(1000);
  46.     }
  48.     /**
  49.      * <pre>
  50.      *    capacity ：利用队列的阻塞 put，来调节生产速度和消费速度的差别
  51.      *      当生产速度明显大于插入速度时，该参数用来限制生产的速度，达到该上限时，生成方就会阻塞，知道有新的容量空闲出来
  52.      * </pre>
  53.      *
  54.      * @param capacity 队列能接收的最大容量
  55.      */
  56.     public BatchProcessor(int capacity) {
  57.         queue = new ArrayBlockingQueue<>(capacity);
  58.     }
  60.     /**
  61.      * 配置线程名称前缀
  62.      *
  63.      * @param threadNamePrefix
  64.      */
  65.     public synchronized void setThreadNamePrefix(String threadNamePrefix) {
  66.         if (started) {
  67.             throw new RuntimeException("已经开始处理，不能再线程名称前缀");
  68.         }
  69.         this.threadNamePrefix = threadNamePrefix;
  70.     }
  72.     /**
  73.      * 默认 4 个线程，每个线程每次处理一条数据
  74.      *
  75.      * @param consumer 每次达到消费条数时，消费方的消费回调逻辑
  76.      */
  77.     public void start(StorageConsumer<T> consumer) {
  78.         this.start(consumer, 4);
  79.     }
  81.     /**
  82.      * 默认每个线程每次处理 1 条数据
  83.      *
  84.      * @param consumer        每次达到消费条数时，消费方的消费回调逻辑
  85.      * @param workThreadCount 需要并行处理的线程数量，必须大于 0
  86.      */
  87.     public void start(StorageConsumer<T> consumer, int workThreadCount) {
  88.         this.start(consumer, workThreadCount, 0);
  89.     }
  91.     /**
  92.      * @param consumer        每次达到消费条数时，消费方的消费回调逻辑
  93.      * @param workThreadCount 需要并行处理的线程数量，必须大于 0
  94.      * @param maxItemCount    每次每个线程希望的消费数据条数， 0：每个线程每次消费 1 条数据，大于 0 则按照期望的条数进行消费
  95.      */
  96.     public synchronized void start(StorageConsumer<T> consumer,
  97.                                    int workThreadCount,
  98.                                    int maxItemCount) {
  99.         this.start(consumer, workThreadCount, maxItemCount, null);
  100.     }
  102.     /**
  103.      * @param consumer                 每次达到消费条数时，消费方的消费回调逻辑
  104.      *                                 由于是线程处理，所有在消费逻辑处理的时候，建议消费方一定要将逻辑都 try 一下，否则就会进入 uncaughtExceptionHandler 处理异常，并且该工作线程退出工作
  105.      * @param workThreadCount          需要并行处理的线程数量，必须大于 0
  106.      * @param maxItemCount             每次每个线程希望的消费数据条数， 0：每个线程每次消费 1 条数据，大于 0 则按照期望的条数进行消费
  107.      * @param uncaughtExceptionHandler 当抛出异常的时候，该异常如何处理，可以为 null, 如果为 null, 将使用 @Slf4j 日志打印
  108.      */
  109.     public synchronized void start(StorageConsumer<T> consumer,
  110.                                    int workThreadCount,
  111.                                    int maxItemCount,
  112.                                    Thread.UncaughtExceptionHandler uncaughtExceptionHandler) {
  113.         if (started) {
  114.             throw new RuntimeException("处理中");
  115.         }
  116.         if (workThreadCount <= 0) {
  117.             throw new IllegalArgumentException("workThreadCount 必须大于 0");
  118.         }
  119.         if (maxItemCount < 0) {
  120.             throw new IllegalArgumentException("maxItemCount 必须大于等于 0");
  121.         }
  122.         started = true;
  123.         this.consumer = consumer;
  124.         this.maxItemCount = maxItemCount;
  125.         this.cdl = new CountDownLatch(workThreadCount);
  126.         if (uncaughtExceptionHandler == null) {
  127.             uncaughtExceptionHandler = (t, e) -> {
  128.                 log.error(StrUtil.format("工作线程异常退出，threadName={}", t.getName()), e);
  129.             };
  130.         }
  131.         Thread.UncaughtExceptionHandler finalUncaughtExceptionHandler = uncaughtExceptionHandler;
  132.         workThreads = IntStream.range(0, workThreadCount)
  133.                 .mapToObj(i -> {
  134.                     final WorkThread workThread = new WorkThread(threadNamePrefix + i, maxItemCount);
  135.                     workThread.start();
  136.                     // 如果不设置异常处理器，那么当 run 方法抛出异常的时候，会被 java.lang.ThreadGroup.uncaughtException 处理
  137.                     // 然后 ThreadGroup.uncaughtException 的默认处理是使用 System.error 打印错误，和调用 e.printStackTrace(System.err);
  138.                     // 这就会导致在生产环境中使用日志框架的时候，在日志框架里面看不到打印的错误信息，看起来就像异常被吞了
  139.                     workThread.setUncaughtExceptionHandler(finalUncaughtExceptionHandler);
  140.                     return workThread;
  141.                 })
  142.                 .collect(Collectors.toList());
  143.     }
  145.     /**
  146.      * 将实体交给处理器，处理器的线程会消费该实体；
  147.      * <pre>
  148.      *  当容器队列已满时，则会阻塞，以此达到生产方暂停生产的目的；可以防止生产速度过快（消费速度过慢），导致占用过多内存
  149.      * </pre>
  150.      *
  151.      * @param entity
  152.      */
  153.     public void put(T entity) {
  154.         try {
  155.             queue.put(entity);
  156.         } catch (InterruptedException e) {
  157.             ExceptionUtil.wrapAndThrow(e);
  158.         }
  159.     }
  161.     /**
  162.      * 等待，处理器处理完成；此方法会阻塞
  163.      */
  164.     public void await() {
  165.         if (!started) {
  166.             throw new RuntimeException("还未运行");
  167.         }
  168.         try {
  169.             isProduceData = false;
  170.             cdl.await();
  171.             for (WorkThread workThread : workThreads) {
  172.                 workThread.clearEntity();
  173.             }
  174.         } catch (InterruptedException e) {
  175.             ExceptionUtil.wrapAndThrow(e);
  176.         }
  177.     }
  179.     /**
  180.      * 立即停止，只适合在生产方不生产数据时，调用
  181.      */
  182.     public void stop() {
  183.         if (!started) {
  184.             throw new RuntimeException("还未运行");
  185.         }
  186.         isProduceData = false;
  187.         queue.clear();
  188.     }
  190.     private class WorkThread extends Thread {
  191.         // 批量插入时，用于缓存实体的容器
  192.         private List<T> batchCacheContainer;
  193.         private int maxItemCount;
  195.         public WorkThread(String name, int maxItemCount) {
  196.             super(name);
  197.             this.maxItemCount = maxItemCount;
  198.             if (maxItemCount > 0) {
  199.                 batchCacheContainer = new ArrayList<>(maxItemCount);
  200.             }
  201.         }
  203.         @Override
  204.         public void run() {
  205.             try {
  206.                 doRun();
  207.             } catch (InterruptedException e) {
  208.                 log.debug("工作线程收到中断异常退出", e);
  209.             } finally {
  210.                 cdl.countDown();
  211.             }
  212.         }
  214.         private void doRun() throws InterruptedException {
  215.             while (true) {
  216.                 // 如果不产生数据了，队列也会空，则退出线程
  217.                 if (!isProduceData && queue.size() == 0) {
  218.                     break;
  219.                 }
  220.                 final T entity;
  221.                 entity = queue.poll(500, TimeUnit.MILLISECONDS);
  222.                 if (entity == null) {
  223.                     continue;
  224.                 }
  225.                 if (maxItemCount > 0) {
  226.                     batchCacheContainer.add(entity);
  227.                     if (batchCacheContainer.size() >= maxItemCount) {
  228.                         consumer.accept(null, batchCacheContainer);
  229.                         batchCacheContainer.clear();
  230.                     }
  231.                 } else {
  232.                     consumer.accept(entity, null);
  233.                 }
  234.             }
  235.         }
  237.         public void clearEntity() {
  238.             if (maxItemCount > 0 && batchCacheContainer.size() > 0) {
  239.                 consumer.accept(null, batchCacheContainer);
  240.                 batchCacheContainer.clear();
  241.             }
  242.         }
  243.     }
  245.     public interface StorageConsumer<T> {
  246.         /**
  247.          * 需要使用方存储数据时，会调用该方法
  248.          *
  249.          * @param t
  250.          * @param ts
  251.          */
  252.         void accept(T t, List<T> ts);
  253.     }
  254. }

用法测试

  1. package cn.mrcode;
  3. import org.junit.jupiter.api.Test;
  5. /**
  6.  * @author mrcode
  7.  * @date 2021/6/3 23:24
  8.  */
  9. class BatchProcessorTest {
  10.     /**
  11.      * 批量插入测试
  12.      */
  13.     @Test
  14.     public void batchInsert() {
  15.         final BatchProcessor<DemoEntity> work = new BatchProcessor<>();
  16.         work.start((t, ts) -> {
  17.                     System.out.println("插入数据库条数：" + ts.size());
  18.                 },
  19.                 4, 4);
  20.         // 模拟生产数据
  21.         try {
  22.             for (int i = 0; i < 21; i++) {
  23.                 work.put(new DemoEntity(i, i + " name"));
  24.             }
  25.             // 等待入库完成
  26.             work.await();
  27.         } catch (Exception e) {
  28.             // 如果生产过程中有异常,立即停止掉处理器，不再入库
  29.             work.stop();
  30.         }
  31.     }
  33.     /**
  34.      * 单条插入测试
  35.      */
  36.     @Test
  37.     public void insert() {
  38.         final BatchProcessor<DemoEntity> work = new BatchProcessor<>();
  39.         work.start((t, ts) -> {
  40.                     System.out.println("插入数据库：" + t);
  41.                 },
  42.                 4, 0);
  44.         // 模拟生产数据
  45.         try {
  46.             for (int i = 0; i < 5; i++) {
  47.                 work.put(new DemoEntity(i, i + " name"));
  48.             }
  49.             // 等待入库完成
  50.             work.await();
  51.         } catch (Exception e) {
  52.             // 如果生产过程中有异常,立即停止掉处理器，不再入库
  53.             work.stop();
  54.         }
  55.     }
  57.     /**
  58.      * 异常测试 - 不自定义异常处理器
  59.      */
  60.     @Test
  61.     public void exceptionTest() {
  62.         // 看看在消费逻辑中发现业务异常，会发生什么事情
  63.         final BatchProcessor<DemoEntity> work = new BatchProcessor<>();
  64.         work.start((t, ts) -> {
  65.                     if (true) {
  66.                         // 会抛出 ArithmeticException: / by zero 异常
  67.                         int a = 1 / 0;
  68.                     }
  69.                     System.out.println("插入数据库：" + t);
  70.                 },
  71.                 2, 0,
  72.                 // 异常处理器，如果为 null， BatchProcessor 工具会捕获，并使用 Slf4j error 级别打印日志
  73.                 // 如果框架不做这个处理，jdk 会使用 System.err.out 打印到控制台，所以在线上生产环境，就不会记录到日志文件中
  74.                 // 当出现问题的时候，就很难发现出现了什么问题
  75.                 null);
  77.         // 模拟生产数据
  78.         try {
  79.             for (int i = 0; i < 5; i++) {
  80.                 work.put(new DemoEntity(i, i + " name"));
  81.             }
  82.             // 等待入库完成
  83.             work.await();
  84.             System.out.println("处理完成");
  85.         } catch (Exception e) {
  86.             // 如果生产过程中有异常,立即停止掉处理器，不再入库
  87.             System.err.println("异常处理完成");
  88.             work.stop();
  89.         }
  90.     }
  92.     /**
  93.      * 异常测试 - 自定义异常处理器
  94.      */
  95.     @Test
  96.     public void exceptionHandlerTest() {
  97.         // 看看在消费逻辑中发现业务异常，会发生什么事情
  98.         final BatchProcessor<DemoEntity> work = new BatchProcessor<>();
  99.         work.start((t, ts) -> {
  100.                     if (true) {
  101.                         // 会抛出 ArithmeticException: / by zero 异常
  102.                         int a = 1 / 0;
  103.                     }
  104.                     System.out.println("插入数据库：" + t);
  105.                 },
  106.                 2, 0,
  107.                 // 自定义异常处理器
  108.                 new Thread.UncaughtExceptionHandler() {
  109.                     @Override
  110.                     public void uncaughtException(Thread t, Throwable e) {
  111.                         System.out.println("工作线程异常退出");
  112.                         e.printStackTrace();
  113.                     }
  114.                 });
  116.         // 模拟生产数据
  117.         try {
  118.             for (int i = 0; i < 5; i++) {
  119.                 work.put(new DemoEntity(i, i + " name"));
  120.             }
  121.             // 等待入库完成
  122.             work.await();
  123.             System.out.println("处理完成");
  124.         } catch (Exception e) {
  125.             // 如果生产过程中有异常,立即停止掉处理器，不再入库
  126.             System.err.println("异常处理完成");
  127.             work.stop();
  128.         }
  129.     }
  131.     /**
  132.      * 测试 实体
  133.      */
  134.     private class DemoEntity {
  135.         private int id;
  136.         private String name;
  138.         public DemoEntity(int id, String name) {
  139.             this.id = id;
  140.             this.name = name;
  141.         }
  143.         public int getId() {
  144.             return id;
  145.         }
  147.         public void setId(int id) {
  148.             this.id = id;
  149.         }
  151.         public String getName() {
  152.             return name;
  153.         }
  155.         public void setName(String name) {
  156.             this.name = name;
  157.         }
  159.         @Override
  160.         public String toString() {
  161.             return "DemoEntity{" +
  162.                     "id=" + id +
  163.                     ", name='" + name + '\'' +
  164.                     '}';
  165.         }
  166.     }
  167. }

测试输出

  1. // 批量插入输出
  2. 插入数据库条数：4
  3. 插入数据库条数：4
  4. 插入数据库条数：4
  5. 插入数据库条数：4
  6. 插入数据库条数：3
  7. 插入数据库条数：2
  9. // 单条插入输出
  11. 插入数据库：DemoEntity{id=0, name='0 name'}
  12. 插入数据库：DemoEntity{id=1, name='1 name'}
  13. 插入数据库：DemoEntity{id=3, name='3 name'}
  14. 插入数据库：DemoEntity{id=4, name='4 name'}
  15. 插入数据库：DemoEntity{id=2, name='2 name'}

异常处理器相关测试

  1. // 异常测试 - 不自定义异常处理器
  2. 处理完成
  3. 10:52:47.847 [reportInsert-0] ERROR cn.mrcode.BatchProcessor - 工作线程异常退出，threadName=reportInsert-0
  4. java.lang.ArithmeticException: / by zero
  5.     at cn.mrcode.BatchProcessorTest.lambda$exceptionTest$2(BatchProcessorTest.java:67)
  6.     at cn.mrcode.BatchProcessor$WorkThread.doRun(BatchProcessor.java:270)
  7.     at cn.mrcode.BatchProcessor$WorkThread.run(BatchProcessor.java:244)
  8. 10:52:47.847 [reportInsert-1] ERROR cn.mrcode.BatchProcessor - 工作线程异常退出，threadName=reportInsert-1
  9. java.lang.ArithmeticException: / by zero
  10.     at cn.mrcode.BatchProcessorTest.lambda$exceptionTest$2(BatchProcessorTest.java:67)
  11.     at cn.mrcode.BatchProcessor$WorkThread.doRun(BatchProcessor.java:270)
  12.     at cn.mrcode.BatchProcessor$WorkThread.run(BatchProcessor.java:244)
  16. // 异常测试 - 自定义异常处理器
  17. 工作线程异常退出
  18. 工作线程异常退出
  19. 处理完成
  20. java.lang.ArithmeticException: / by zero
  21.     at cn.mrcode.BatchProcessorTest.lambda$exceptionHandlerTest$3(BatchProcessorTest.java:102)
  22.     at cn.mrcode.BatchProcessor$WorkThread.doRun(BatchProcessor.java:270)
  23.     at cn.mrcode.BatchProcessor$WorkThread.run(BatchProcessor.java:244)
  24. java.lang.ArithmeticException: / by zero
  25.     at cn.mrcode.BatchProcessorTest.lambda$exceptionHandlerTest$3(BatchProcessorTest.java:102)
  26.     at cn.mrcode.BatchProcessor$WorkThread.doRun(BatchProcessor.java:270)
  27.     at cn.mrcode.BatchProcessor$WorkThread.run(BatchProcessor.java:244)

注意事项

作为批量入库的注意事项

  1.  /**
  2.      * 多线程批量插入测试
  3.      */
  4.     @Test
  5.     public void batchInsert() {
  6.         final BatchInsertProcessor<DemoEntity> work = new BatchInsertProcessor<>();
  7.         work.start((t, ts) -> {
  8.                     System.out.println("插入数据库条数：" + ts.size());
  11.                     在这里调用数据库批量插入，记住是  批量插入，将 ts 整个一次性插入到数据库
  12.                 },
  13.                 4, 4);
  14.         // 模拟生产数据
  15.         try {
  16.             for (int i = 0; i < 21; i++) {
  17.                 work.put(new DemoEntity(i, i + " name"));
  18.             }
  19.             // 等待入库完成
  20.             work.await();
  21.         } catch (Exception e) {
  22.             // 如果生产过程中有异常,立即停止掉处理器，不再入库
  23.             work.stop();
  24.         }
  25.     }

在 mybatis 中批量插入如下所示

  1.   <insert id="batchInsertProductCategory" parameterType="java.util.List">
  2.     INSERT INTO
  3.     tb_product_category(product_category_name,priority,shop_id)
  4.     VALUES
  5.     <foreach collection="list" item="productCategory" index="index" separator=",">
  6.       (
  7.         #{productCategory.productCategoryName},
  8.         #{productCategory.priority},
  9.         #{productCategory.shopId}
  10.       )
  11.     </foreach>
  12.   </insert>