CountDownLatch+ForkJoinPool的使用

SetfUn

于 2024-01-17 14:06:18 发布

阅读量420

点赞数 9

文章标签： windows linux 运维

本文链接：https://blog.csdn.net/superhaox/article/details/135649019

版权

文章描述了一个Java方法，用于处理任务调度，包括抠图任务（使用ForkJoinPool并发处理）和商品图片生成任务。方法首先查询任务详情，然后根据不同任务类型（如抠图、生成图或阴影）调用对应执行器并更新任务状态。

摘要由CSDN通过智能技术生成

    public void startTask(String taskId) {
        DoucaiTaskBatchDO doucaiTaskBatchDO = taskMapper.findOne(DoucaiTaskBatchDOMetaData.taskId.is(taskId).and(DoucaiTaskBatchDOMetaData.isDelete.is(0)));
        String taskType = doucaiTaskBatchDO.getType();
        CountDownLatch countDownLatch = new CountDownLatch(1);
        if(TaskTypeEnum.CUT.getCode().equals(taskType)) {
            // 抠图
            taskGenExecutor.execute(() -> {
                sendKoutuTask(doucaiTaskBatchDO, countDownLatch);
            });
        }else {
            lubanAITBThreadPoolExecutor.execute(() -> {
                try {
                    countDownLatch.await();
                    log.info("开始批量生图, taskId: {}", taskId);

                    log.info("查询任务的商品列表, taskId: {}", taskId);
                    // 查询商品列表
                    List<DoucaiProductTaskDO> doucaiProductTaskList = productTaskMapper.find(DoucaiProductTaskDOMetaData.taskId.is(taskId));

                    log.info("查询任务的商品详情列表, taskId: {}", taskId);
//            DoucaiTaskBatchDO doucaiTaskBatchDO = taskMapper.findOne(DoucaiTaskBatchDOMetaData.taskId.is(taskId).and(DoucaiTaskBatchDOMetaData.isDelete.is(0)));

            // 获取商品详情列表
            List<ProductBatchDO> productBatchDOList = doucaiProductTaskList.stream()
                .map(it -> {
                    ProductBatchDO productBatchDO = productMapper.findOne(ProductBatchDOMetaData.itemId.is(it.getItemId()),
                        ProductBatchDOMetaData.platformId.is(doucaiTaskBatchDO.getPlatformId()),
                        ProductBatchDOMetaData.isDelete.is(0));
                    if(StringUtils.isNotEmpty(doucaiTaskBatchDO.getCategoryId())){
                        productBatchDO.setCategoryId(doucaiTaskBatchDO.getCategoryId());
                    }
                    return productBatchDO;
                })
                .filter(Objects::nonNull).collect(Collectors.toList());


            log.info("发送生图任务, taskId: {}, productSize: {}, itemIds: {}", taskId, productBatchDOList.size(), productBatchDOList.stream().map(ProductBatchDO::getItemId).collect(Collectors.toList()));
            // 发起任务
//            String taskType = doucaiTaskBatchDO.getType();
                    doucaiTaskBatchDO.setStatus("FINISHED");
                    doucaiTaskBatchDO.setKoutuStatus("DONE");
                    doucaiTaskBatchDO.setFinishTime(new Date());
                    taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);

                    if (taskType.contains(TaskTypeEnum.GEN.getCode())) {
                        //临时修改
                        sendImg2ProductTask(taskId, productBatchDOList);

                    }

                    if (taskType.contains(TaskTypeEnum.SHADOW.getCode())) {
                        sendImg2LightTask(taskId, productBatchDOList);
                    }
                }catch (Exception e){
                    throw new RuntimeException(e);
                }

            });
        }

    }

代码2

    public void sendKoutuTask(DoucaiTaskBatchDO doucaiTaskBatchDO, CountDownLatch countDownLatch) {
        log.info("sendKoutuTask:【开始抠图任务】任务id:{}", doucaiTaskBatchDO.getTaskId());
        if (TaskBatch.KoutuStatus.DONE.name().equalsIgnoreCase(doucaiTaskBatchDO.getKoutuStatus())) {
            countDownLatch.countDown();
            //继续下一个任务
            startNewTask();

            return;
        }
        String taskId = doucaiTaskBatchDO.getTaskId();
        LoginUser user = ThreadContext.getUser();

        //获取待抠图商品
        List<DoucaiProductTaskDO> doucaiProductTaskList = productTaskMapper.findUnFinishList(taskId, TaskTypeEnum.CUT.getCode());
        if (CollectionUtils.isEmpty(doucaiProductTaskList)) {
            //更新跑批任务抠图状态
            doucaiTaskBatchDO.setGmtModified(new Date());
            doucaiTaskBatchDO.setKoutuStatus(TaskBatch.KoutuStatus.DONE.name());
            taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);

            countDownLatch.countDown();

            //继续下一个任务
            startNewTask();

            return;
        }
        //更新抠图状态
        if (!TaskBatch.KoutuStatus.DOING.name().equalsIgnoreCase(doucaiTaskBatchDO.getKoutuStatus())) {
            doucaiTaskBatchDO.setKoutuStatus(TaskBatch.KoutuStatus.DOING.name());
            doucaiTaskBatchDO.setGmtModified(new Date());

            taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);
        }

        // 抠图要限速
        ForkJoinPool forkJoinPool = new ForkJoinPool(6);
        forkJoinPool.submit(() -> {
            try {
                doucaiProductTaskList.parallelStream().forEach(
                        productTask -> {
                            String itemId = productTask.getItemId();
                            List<ProductBatchDO> productList = productMapper.find(ProductBatchDOMetaData.itemId.is(itemId).and(ProductBatchDOMetaData.isDelete.is(0))
                                    .and(ProductBatchDOMetaData.platformId.is(doucaiTaskBatchDO.getPlatformId())));
                            if (CollectionUtils.isNotEmpty(productList)) {
                                ProductBatchDO product = productList.get(0);

                                String kouTuOssKey = kouTuService.kouTuByPiXian(product.getOriginPic(), CommonUtils.nanoId());
                                DoucaiTaskBatchResultDO resultDO = new DoucaiTaskBatchResultDO();
                                resultDO.setGmtCreate(new Date());
                                resultDO.setGmtModified(new Date());
                                resultDO.setItemId(itemId);
                                resultDO.setTaskId(taskId);
                                String itemTaskId = taskId + "_" + itemId;
                                resultDO.setItemTaskId(itemTaskId);
                                resultDO.setType(TaskTypeEnum.CUT.getCode());
                                resultDO.setCreator(doucaiTaskBatchDO.getCreator());
                                resultDO.setSelected(Boolean.FALSE);
                                resultDO.setHighQuality(Boolean.FALSE);
                                resultDO.setDelete(Boolean.FALSE);
                                if (StringUtils.isNotEmpty(kouTuOssKey)) {
                                    resultDO.setResult(kouTuOssKey);
                                    resultDO.setStatus("SUCCESS");
                                    resultMapper.insert(resultDO);
                                    // 更新doucai_task_batch表
                                    doucaiTaskBatchDO.setStatus("FINISHED");
                                    taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);
                                    //更新aigc表
                                    List<ProductAigcDO> productAigcList = douCaiProductAigcMapper.find(
                                            ProductAigcDOMetaData.batchId.is(taskId),
                                            ProductAigcDOMetaData.itemId.is(itemId),
                                            ProductAigcDOMetaData.taskId.is(itemTaskId)
                                    );
                                    if (CollectionUtils.isNotEmpty(productAigcList)) {
                                        productAigcList.stream().forEach(
                                                productAigcDO -> {
                                                    productAigcDO.setKouTuCrop(kouTuOssKey);
                                                    douCaiProductAigcMapper.updateByPrimaryKeySelective(productAigcDO);
                                                }
                                        );
                                    }
                                } else {
                                    resultDO.setStatus("FAIL");
                                    Integer tbPrice = QuotaOpTypeEnum.kouTu.getTbPrice();
                                    resultDO.setReturnQuotaNum(tbPrice);
                                    resultMapper.insert(resultDO);
                                    // 更新doucai_task_batch表
                                    doucaiTaskBatchDO.setStatus(TaskBatchResultStatusEnum.FAIL.name());
                                    taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);
                                    if (TaskTypeEnum.CUT.getCode().equals(doucaiTaskBatchDO.getType())) {
                                        quotaService.returnQuotaTB(user.getId(), tbPrice, QuotaOpTypeEnum.kouTu.getCode(), itemTaskId);
                                    }
                                }
                            }
                        }
                );
            } catch (Exception e) {
                throw new MsBizException(e);
            } finally {
                //更新跑批任务抠图状态
                doucaiTaskBatchDO.setGmtModified(new Date());
                doucaiTaskBatchDO.setKoutuStatus(TaskBatch.KoutuStatus.DONE.name());
                taskMapper.updateByPrimaryKeySelective(doucaiTaskBatchDO);

                countDownLatch.countDown();

                //继续下一个任务
                startNewTask();
            }
        });
    }