Spring multi thread transaction tools

1.import MultiThreadTransactionManager.java
2.DI PlatformTransactionManager
3.new MultiThreadTransactionManager<>(transactionManager)
4.addTask()
5.execute()
6.getResult()

    
    @Autowired
    TestMapper2 testMapper;
    
    @Autowired
    private PlatformTransactionManager transactionManager;

    private void test() {
        // new manager
        MultiThreadTransactionManager<String> manager = new MultiThreadTransactionManager<>(transactionManager);
        manager.setTimeout(10);

        // add task
        manager.addTask((t) -> {
            testMapper.insertTest(10, t);
            return "task-1";
        });
        manager.addTask((t) -> {
            testMapper.insertTest2(20, t);
            return "task-2";
        });
        
        Instant start = Instant.now();
        // execute
        boolean success = manager.execute();
        manager.getResult().forEach(r -> {
            System.out.println(r);
        });;
        System.out.println("use time : " + Duration.between(start, Instant.now()).getSeconds() + ", success=" + success);
    }

//-------------------------------

public class MultiThreadTransactionManager<R> {
    private PlatformTransactionManager transactionManager;
    private ExecutorService executorService;
    private int threadCount;
    private int timeout = 20;
    private Meta meta;
    
    List<Function<String, R>> tasks = new ArrayList<>();
    Map<Integer, R> result;
    
    public MultiThreadTransactionManager(PlatformTransactionManager transactionManager) {
        this.transactionManager = transactionManager;
    }

    public MultiThreadTransactionManager(PlatformTransactionManager transactionManager, int timeout, ExecutorService executorService) {
        this.transactionManager = transactionManager;
        this.timeout = timeout;
        this.executorService = executorService;
    }
    
    public void addTask(Function<String, R> task) {
        tasks.add(task);
    }

    private void init() {
        threadCount = tasks.size();
        result = new ConcurrentHashMap<>(threadCount);
        
        Assert.isTrue(threadCount > 0, "threadCount must great than 1");
        
        meta = new Meta();
        meta.resultList = new LinkedBlockingDeque<Boolean>(threadCount);
        meta.subMonitor = new CountDownLatch(threadCount);
        meta.mainMonitor = new CountDownLatch(1);
        meta.rollback = new Rollback(false);
        if (executorService == null) {
            executorService = new ForkJoinPool();
        }
    }

    /**
     * For sub thread to update db records
     * 
     * @param consumer
     * @return
     */
    private void startSubThread(Function<String, R> function, int index) {
        CountDownLatch subMonitor = meta.getSubMonitor();
        CountDownLatch mainMonitor = meta.getMainMonitor();
        Rollback rollback = meta.getRollback();
        BlockingDeque<Boolean> resultList = meta.getResultList();

        // start sub thread
        executorService.submit(() -> {
            String threadFlag = "sub thread id = " + Thread.currentThread().getId() + " name="
                    + Thread.currentThread().getName();

            DefaultTransactionDefinition definition = new DefaultTransactionDefinition();
            definition.setName(UUID.randomUUID().toString());
            definition.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRED);
            TransactionStatus txStatus = transactionManager.getTransaction(definition);
            R r = null;
            try {
                // call mapper to update DB
                r = function.apply(threadFlag);

                resultList.add(true);
                // save result
                this.result.put(index, r);
            } catch (Exception e) {
                e.printStackTrace();
                resultList.add(false);
            } catch (Error e) {
                e.printStackTrace();
                resultList.add(false);
            } finally {
                // current sub thread done
                subMonitor.countDown();
                try {
                    // sub thread await until all others finish update db
                    mainMonitor.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }

            if (rollback.isNeedRollBack()) {
                shouldRollBack(txStatus);
            } else {
                submit(txStatus);
            }
        });
    }

    public Boolean execute() {
        init();
        
        CountDownLatch subMonitor = meta.getSubMonitor();
        CountDownLatch mainMonitor = meta.getMainMonitor();
        Rollback rollback = meta.getRollback();
        BlockingDeque<Boolean> resultList = meta.getResultList();

        // call all sub thread and start
        for (int i = 0, len = tasks.size(); i < len; i++) {
            startSubThread(tasks.get(i), i);
        }

        try {
            // current thread await until all sub thread finish update db or time out
            boolean await = subMonitor.await(timeout, TimeUnit.SECONDS);
            if (!await) {
                rollback.needRollBack();
            } else {
                // judge if need rollback
                for (int i = 0, len = resultList.size(); i < len; i++) {
                    Boolean result = resultList.take();
                    if (!result) {
                        rollback.needRollBack();
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            rollback.needRollBack();
        } catch (Error e) {
            e.printStackTrace();
            rollback.needRollBack();
        } finally {
            // release sub thread to do commit/rollback
            mainMonitor.countDown();
            executorService.shutdown();
        }

        return !rollback.isNeedRollBack();
    }
    
    public Map<Integer, R> getResultMap() {
        return result;
    }
    
    public List<R> getResult() {
        List<R> list = new ArrayList<R>();
        for (int i = 0, len = result.size(); i < len; i++) {
            list.add(result.get(i));
        }
        return list;
    }

    private void shouldRollBack(TransactionStatus status) {
        transactionManager.rollback(status);
    }

    private void submit(TransactionStatus status) {
        transactionManager.commit(status);
    }

    public ExecutorService getExecutorService() {
        return executorService;
    }

    public void setExecutorService(ExecutorService executorService) {
        this.executorService = executorService;
    }

    public int getTimeout() {
        return timeout;
    }

    public void setTimeout(int timeout) {
        this.timeout = timeout;
    }

    class Rollback {
        private Boolean rollBackFlag = false;

        public Rollback(Boolean rollBackFlag) {
            this.rollBackFlag = rollBackFlag;
        }

        public Boolean isNeedRollBack() {
            return rollBackFlag;
        }

        public void needRollBack() {
            this.rollBackFlag = true;
        }

        public void notNeddRollBack() {
            this.rollBackFlag = true;
        }
    }

    class Meta {
        private BlockingDeque<Boolean> resultList;
        private CountDownLatch subMonitor;
        private CountDownLatch mainMonitor;
        private Rollback rollback;

        public BlockingDeque<Boolean> getResultList() {
            return resultList;
        }

        public void setResultList(BlockingDeque<Boolean> resultList) {
            this.resultList = resultList;
        }

        public CountDownLatch getSubMonitor() {
            return subMonitor;
        }

        public void setSubMonitor(CountDownLatch subMonitor) {
            this.subMonitor = subMonitor;
        }

        public CountDownLatch getMainMonitor() {
            return mainMonitor;
        }

        public void setMainMonitor(CountDownLatch mainMonitor) {
            this.mainMonitor = mainMonitor;
        }

        public Rollback getRollback() {
            return rollback;
        }

        public void setRollback(Rollback rollback) {
            this.rollback = rollback;
        }
    }
}