使用ForkJoin批量插入造大量数据

事出有因

很多时候为了测试一些接口或者大量数据迁移等测试方面相关的时候，就会需要大量的数据，而这块数据本身的质量其实是不关心的。
当然为了这个需求直接写Java代码看起来是有点成本较高的，不过谁让咱只会写Java呢。

模拟开搞

首先，插入数据需要两部分，一个是要执行的sql, 一个是数据源，这个看自己情况决定使用哪个，
我这里为了简单，直接使用JdbcTemplate， 可以按照实际情况自行决定。

1. 新建一个fork-join的拆分后执行的子任务

   /**
    * 
   任务数据源
    *
    * @author Snowball
    * @version 1.0
    * @date 2021/04/25 15:55
    */
   public class ActionDatasource {
   
       /**
        * sql操作对象
        */
       private final JdbcTemplate jdbcTemplate;
   
       /**
        * 存储所有要执行的sql的集合
        */
       private final List<String> executeSql;
   
       public ActionDatasource(JdbcTemplate jdbcTemplate, List<String> executeSql) {
           this.jdbcTemplate = jdbcTemplate;
           this.executeSql = executeSql;
       }
   
       /**
        * 执行对应角标的任务
        *
        * @param index
        */
       public void execute(long index) {
           if (index < executeSql.size()) {
               jdbcTemplate.execute(executeSql.get((int) index));
           }
       }
   }
   

2. 定义任务拆分逻辑

   /**
    * 批量插入数据任务拆分计算类
    *
    * @author dongfang.ding
    * @date 2021/4/25 16:54
    **/
   public class BatchInsertAction extends RecursiveAction {
   
       /**
        * 任务拆分开始标志
        */
       private final long start;
   
       /**
        * 任务拆分结束标志
        */
       private final long end;
   
       /**
        * 具体要做的任务类
        */
       private final ActionDatasource task;
   
       /**
        * 任务拆分临界值，小于这个值，任务不再继续拆分
        */
       private final long thurshold;
   
       public BatchInsertAction(ActionDatasource task, long start, long end, long thurshold) {
           this.task = task;
           this.start = start;
           this.end = end;
           this.thurshold = thurshold;
       }
   
       @Override
       protected void compute() {
           if ((end - start) < thurshold) {
               // 达到临界点开始执行任务，具体情况具体代码；这里由于是到集合中取要执行的sql, 所以会将当前角标传过去
               for (long i = start; i <= end; i++) {
                   task.execute(i);
               }
           } else {
               long middle = (end + start) / 2;
   
               BatchInsertAction left = new BatchInsertAction(task, start, middle, thurshold);
               left.fork();
   
               BatchInsertAction right = new BatchInsertAction(task, middle + 1, end, thurshold);
               right.fork();
           }
       }
   }
   

3. 一个实际例子

   1. 需求插入1000W条数据
   2. 控制每1000条sql合成一条批量sql，以提高插入效率
   3. 最终1000W条会被分割成10000条批量语句。

   /**
    * 
   description
    *
    * @author Snowball
    * @version 1.0
    * @date 2021/04/25 17:00
    */
   @Component
   public class BatchInsertActionExample {
   
       @Autowired
       private JdbcTemplate jdbcTemplate;
   
       /**
        * 要插入的sql定义语句， 到values即可
        * 如
        * insert into t_user(id, name) values
        */
       private String dmlSql;
   
       /**
        * 要执行的每条插入语句sql，不需要处理各种分隔符
        * 如
        * (1, 'ddf')
        */
       private String valueSq;
   
       public void execute() {
           String sql = "INSERT INTO `t_user`(id, name) VALUES ";
           String varSql = "('%s', '%s')";
   
           StringBuilder sbl = new StringBuilder();
   
           // 预期要插入的数据行数
           int expectRecordSize = 1000000;
           // 循环结束角标规则
           int loopEndIndex = 2 * expectRecordSize;
           // 批量sql的数量，会将指定数量的sql合并成一条sql执行
           int batchSize = 1000;
           List<String> executeSql = new ArrayList<>(expectRecordSize);
   
           for (int i = expectRecordSize + 1; i <= loopEndIndex; i++) {
               if (sbl.length() == 0) {
                   sbl.append(sql);
               }
               sbl.append(String.format(varSql, i, "ddf" + i));
               if (i == loopEndIndex || i % batchSize == 0) {
                   sbl.append(";");
                   executeSql.add(sbl.toString());
                   sbl.delete(0, sbl.length());
               } else {
                   sbl.append(",");
               }
           }
   
   
           ForkJoinPool forkJoinPool = new ForkJoinPool(Runtime.getRuntime().availableProcessors());
           final ActionDatasource task = new ActionDatasource(jdbcTemplate, executeSql);
           final BatchInsertAction action = new BatchInsertAction(task,0L, 10000, 1000);
           forkJoinPool.invoke(action);
       }
   }