| | |
| | | import io.jsonwebtoken.lang.Collections; |
| | | import io.swagger.annotations.Api; |
| | | import io.swagger.annotations.ApiOperation; |
| | | import reactor.core.publisher.Flux; |
| | | import reactor.core.publisher.Mono; |
| | | |
| | | import org.springframework.web.bind.annotation.* ; |
| | |
| | | import com.linghu.model.dto.TaskResultResponse.QuestionResult; |
| | | import com.linghu.model.dto.TaskResultResponse.UserResult; |
| | | import reactor.core.scheduler.Schedulers; |
| | | |
| | | import java.util.stream.Collectors; |
| | | import java.util.stream.IntStream; |
| | | import java.util.stream.Stream; |
| | | |
| | | @RestController |
| | | @RequestMapping("/collect") |
| | |
| | | private UserService userService; |
| | | @Autowired |
| | | private OrderService orderService; |
| | | |
| | | |
| | | /* @PostMapping("/search") |
| | | private static final Queue<SearchTaskRequest> taskQueue = new LinkedList<>(); |
| | | private static boolean isProcessing = false; |
| | | @PostMapping("/search") |
| | | @ApiOperation(value = "开始采集") |
| | | public Mono<SearchTaskResponse> createSearchTask( |
| | | public Mono<ResponseResult<?>> createSearchTask( |
| | | @RequestBody SearchTaskRequest searchTaskRequest, |
| | | HttpServletRequest request) throws JsonProcessingException { |
| | | |
| | | return webClient.post() |
| | | .uri(baseUrl + "/api/v1/search") |
| | | .contentType(MediaType.APPLICATION_JSON) |
| | | .bodyValue(searchTaskRequest) |
| | | .retrieve() |
| | | .onStatus(HttpStatus::is4xxClientError, response -> response.bodyToMono(String.class) |
| | | .flatMap(errorBody -> Mono.error(new RuntimeException(errorBody)))) |
| | | .bodyToMono(new ParameterizedTypeReference<SearchTaskResponse>() { |
| | | }) |
| | | .flatMap(responseResult -> { |
| | | // 提取任务ID |
| | | SearchTaskResponse taskResponse = responseResult; |
| | | if (taskResponse != null && taskResponse.getTask_id() != null) { |
| | | // 保存任务ID到关键词 |
| | | LambdaUpdateWrapper<Keyword> updateWrapper = new LambdaUpdateWrapper<>(); |
| | | updateWrapper.eq(Keyword::getKeyword_id, searchTaskRequest.getKeyword_id()); |
| | | updateWrapper.set(Keyword::getStatus,"Submitted"); |
| | | updateWrapper.set(Keyword::getTask_id, taskResponse.getTask_id()); |
| | | keywordService.update(updateWrapper); |
| | | // 首先检查服务器资源 |
| | | return getServerResource() |
| | | .flatMap(resourceResponse -> { |
| | | double cpuUsage = parseUsage(resourceResponse.getCpu_usage_percent()); |
| | | double memoryUsage = parseUsage(resourceResponse.getMemory_usage_percent()); |
| | | |
| | | // 可选:更新响应中的其他信息 |
| | | if (cpuUsage >= 90.0 || memoryUsage >= 90.0) { |
| | | |
| | | |
| | | String errorMsg = String.format("服务器资源不足,请稍后再试"); |
| | | |
| | | log.warn(errorMsg); |
| | | return Mono.just(ResponseResult.error(503, errorMsg)); |
| | | } |
| | | return Mono.just(taskResponse); |
| | | |
| | | // 将新的任务请求加入队列 |
| | | taskQueue.add(searchTaskRequest); |
| | | |
| | | // 如果当前没有任务在处理中,则启动任务队列的处理 |
| | | if (!isProcessing) { |
| | | processNextTaskInQueue(); |
| | | } |
| | | |
| | | // 返回响应,通知用户任务已开始 |
| | | return Mono.just(ResponseResult.success("任务已加入队列,正在处理...")); |
| | | }) |
| | | .onErrorResume(e -> { |
| | | // return Mono.just(ResponseResult.error("调用失败: " + e.getMessage())); |
| | | SearchTaskResponse task = new SearchTaskResponse(); |
| | | task.setMessage("调用失败: " + e.getMessage()); |
| | | return Mono.just(task); |
| | | log.error("检查服务器资源失败: {}", e.getMessage(), e); |
| | | return Mono.just(ResponseResult.error("检查服务器资源失败: " + e.getMessage())); |
| | | }); |
| | | }*/ |
| | | } |
| | | |
| | | // public SearchTaskController(WebClient.Builder webClientBuilder, KeywordService keywordService) { |
| | | // this.webClient = webClientBuilder.build(); |
| | | // this.keywordService = keywordService; |
| | | // } |
| | | private void processNextTaskInQueue() { |
| | | // 设置为正在处理 |
| | | isProcessing = true; |
| | | |
| | | /* @PostMapping("/search") |
| | | @ApiOperation(value = "开始采集") |
| | | public Mono<SearchTaskResponse> createSearchTask( |
| | | @RequestBody SearchTaskRequest searchTaskRequest, |
| | | HttpServletRequest request) throws JsonProcessingException { |
| | | // 从队列中取出下一个任务 |
| | | SearchTaskRequest nextTaskRequest = taskQueue.poll(); |
| | | |
| | | if (nextTaskRequest != null) { |
| | | // 处理任务 |
| | | executeBatchTask(nextTaskRequest) |
| | | .doFinally(signal -> { |
| | | // 完成后,继续处理下一个任务 |
| | | isProcessing = false; |
| | | if (!taskQueue.isEmpty()) { |
| | | processNextTaskInQueue(); // 继续处理队列中的下一个任务 |
| | | } |
| | | }) |
| | | .subscribe(); |
| | | } |
| | | } |
| | | |
| | | private Mono<ResponseResult<String>> executeBatchTask(SearchTaskRequest searchTaskRequest) { |
| | | log.info("开始处理任务:{}", searchTaskRequest); |
| | | log.info("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"); |
| | | Integer keywordId = searchTaskRequest.getKeyword_id(); |
| | | int maxConcurrentUsers = searchTaskRequest.getConfig() != null ? |
| | | searchTaskRequest.getConfig().getMax_concurrent_users() : 3; |
| | | List<List<UserDto>> userBatches = splitUsersIntoBatches(searchTaskRequest.getUsers(), maxConcurrentUsers); |
| | | List<List<UserDto>> userBatches = splitUsersIntoBatches(searchTaskRequest.getUsers(), maxConcurrentUsers, keywordId); |
| | | |
| | | return processBatchesSequentially(userBatches, searchTaskRequest) |
| | | .onErrorResume(e -> { |
| | | SearchTaskResponse task = new SearchTaskResponse(); |
| | | task.setMessage("调用失败: " + e.getMessage()); |
| | | return Mono.just(task); |
| | | Queue<List<UserDto>> batchQueue = new LinkedList<>(userBatches); // 用队列存储批次 |
| | | |
| | | return Mono.just(ResponseResult.success("第一个批次已开始")) |
| | | .doOnTerminate(() -> { |
| | | // 启动后台任务,继续处理批次 |
| | | executeBatchTask(batchQueue, searchTaskRequest, keywordId) |
| | | .subscribe(); // 使用subscribe()启动后台任务 |
| | | }); |
| | | } |
| | | |
| | | private List<List<UserDto>> splitUsersIntoBatches(List<UserDto> users, int batchSize) { |
| | | List<List<UserDto>> batches = new ArrayList<>(); |
| | | for (int i = 0; i < users.size(); i += batchSize) { |
| | | batches.add(users.subList(i, Math.min(i + batchSize, users.size()))); |
| | | } |
| | | return batches; |
| | | private Mono<ResponseResult<?>> executeBatchTask(Queue<List<UserDto>> batchQueue, SearchTaskRequest searchTaskRequest, Integer keywordId) { |
| | | // 如果队列为空,说明所有批次已经完成 |
| | | if (batchQueue.isEmpty()) { |
| | | return Mono.just(ResponseResult.success("所有批次已完成")); |
| | | } |
| | | |
| | | private Mono<SearchTaskResponse> processBatchesSequentially(List<List<UserDto>> userBatches, SearchTaskRequest originalRequest) { |
| | | Mono<SearchTaskResponse> resultMono = Mono.empty(); |
| | | for (List<UserDto> batch : userBatches) { |
| | | List<UserDto> currentBatch = batchQueue.poll(); // 从队列中获取当前批次 |
| | | SearchTaskRequest batchRequest = new SearchTaskRequest(); |
| | | batchRequest.setUsers(batch); |
| | | batchRequest.setQuestions(originalRequest.getQuestions()); |
| | | batchRequest.setConfig(originalRequest.getConfig()); |
| | | batchRequest.setSave_to_database(originalRequest.getSave_to_database()); |
| | | batchRequest.setWebhook_url(originalRequest.getWebhook_url()); |
| | | batchRequest.setKeyword_id(originalRequest.getKeyword_id()); |
| | | batchRequest.setUsers(currentBatch); |
| | | batchRequest.setQuestions(searchTaskRequest.getQuestions()); |
| | | batchRequest.setConfig(searchTaskRequest.getConfig()); |
| | | batchRequest.setSave_to_database(searchTaskRequest.getSave_to_database()); |
| | | batchRequest.setWebhook_url(searchTaskRequest.getWebhook_url()); |
| | | batchRequest.setKeyword_id(keywordId); |
| | | |
| | | resultMono = resultMono.then(createSingleBatchTask(batchRequest)); |
| | | } |
| | | return resultMono; |
| | | } |
| | | |
| | | private Mono<SearchTaskResponse> createSingleBatchTask(SearchTaskRequest batchRequest) { |
| | | return webClient.post() |
| | | .uri(baseUrl + "/api/v1/search") |
| | | .contentType(MediaType.APPLICATION_JSON) |
| | | .bodyValue(batchRequest) |
| | | .retrieve() |
| | | .onStatus(HttpStatus::is4xxClientError, response -> response.bodyToMono(String.class) |
| | | .flatMap(errorBody -> Mono.error(new RuntimeException(errorBody)))) |
| | | .bodyToMono(new ParameterizedTypeReference<SearchTaskResponse>() { |
| | | }) |
| | | .flatMap(responseResult -> { |
| | | SearchTaskResponse taskResponse = responseResult; |
| | | return createSingleBatchTask(batchRequest) |
| | | .flatMap(taskResponse -> { |
| | | if (taskResponse != null && taskResponse.getTask_id() != null) { |
| | | LambdaUpdateWrapper<Keyword> updateWrapper = new LambdaUpdateWrapper<>(); |
| | | updateWrapper.eq(Keyword::getKeyword_id, batchRequest.getKeyword_id()); |
| | | updateWrapper.set(Keyword::getStatus, "Submitted"); |
| | | updateWrapper.set(Keyword::getTask_id, taskResponse.getTask_id()); |
| | | keywordService.update(updateWrapper); |
| | | // 保存任务关联到数据库 |
| | | return saveKeywordTasks(keywordId, taskResponse) |
| | | .then(waitForTaskCompletion(taskResponse.getTask_id(), batchQueue, searchTaskRequest, keywordId)); |
| | | } else { |
| | | return Mono.just(ResponseResult.error("创建批次任务失败")); |
| | | } |
| | | return waitForTaskCompletion(taskResponse.getTask_id()) |
| | | .then(Mono.just(taskResponse)); |
| | | }); |
| | | } |
| | | |
| | | private Mono<Void> waitForTaskCompletion(String taskId) { |
| | | return Flux.interval(Duration.ofSeconds(5)) // 每5秒执行一次 |
| | | .flatMap(tick -> webClient.get() |
| | | private Mono<Void> saveKeywordTasks(Integer keywordId, SearchTaskResponse taskResponse) { |
| | | if (taskResponse == null || taskResponse.getTask_id() == null) { |
| | | return Mono.error(new RuntimeException("任务响应无效或任务ID为空")); |
| | | } |
| | | |
| | | KeywordTask keywordTask = new KeywordTask(); |
| | | keywordTask.setKeyword_id(keywordId); |
| | | keywordTask.setTask_id(taskResponse.getTask_id()); |
| | | keywordTask.setStatus("pending"); |
| | | |
| | | // 将 MyBatis-Plus 的同步方法包装为 Mono<Void> |
| | | return Mono.fromRunnable(() -> { |
| | | boolean success = keywordTaskService.saveOrUpdate(keywordTask); |
| | | if (!success) { |
| | | throw new RuntimeException("保存关键词任务关联失败"); |
| | | } |
| | | }) |
| | | .doFinally(signalType -> log.info("成功保存关键词任务关联: Task ID {}", taskResponse.getTask_id())) |
| | | .then(); |
| | | } |
| | | private Mono<ResponseResult<?>> waitForTaskCompletion(String taskId, Queue<List<UserDto>> batchQueue, SearchTaskRequest searchTaskRequest, Integer keywordId) { |
| | | // 查询任务状态 |
| | | return getTaskStatus(taskId) |
| | | .flatMap(statusResponse -> { |
| | | // 如果任务状态是"submitted"或"running",继续轮询 |
| | | if ("submitted".equalsIgnoreCase(statusResponse.getStatus()) || "running".equalsIgnoreCase(statusResponse.getStatus())) { |
| | | return Mono.delay(Duration.ofSeconds(5)) // 延迟 5 秒后再次查询 |
| | | .flatMap(aLong -> waitForTaskCompletion(taskId, batchQueue, searchTaskRequest, keywordId)); // 递归调用继续等待 |
| | | } else { |
| | | // 如果状态为其他状态,则继续处理下一个批次 |
| | | return executeBatchTask(batchQueue, searchTaskRequest, keywordId); |
| | | } |
| | | }) |
| | | .onErrorResume(e -> { |
| | | // 处理查询任务状态时的错误 |
| | | return Mono.just(ResponseResult.error("查询任务状态失败: " + e.getMessage())); |
| | | }); |
| | | } |
| | | @ApiOperation(value = "查询任务状态") |
| | | @GetMapping("/status") |
| | | public Mono<TaskStatusResponse> getTaskStatus(String taskId) { |
| | | return webClient.get() |
| | | .uri(baseUrl + "/api/v1/tasks/" + taskId) |
| | | .accept(MediaType.APPLICATION_JSON) |
| | | .retrieve() |
| | | .onStatus(HttpStatus::isError, response -> response.bodyToMono(TaskStatusResponse.class) |
| | | .flatMap(errorBody -> Mono.error(new RuntimeException(errorBody.getDetail())))) |
| | | .bodyToMono(TaskStatusResponse.class) |
| | | ) |
| | | .filter(response -> "completed".equals(response.getStatus())) |
| | | .next() // 找到第一个完成的响应后结束流 |
| | | .then(); // 转换为Mono<Void> |
| | | }*/ |
| | | .onErrorResume(e -> { |
| | | // 处理错误,创建一个自定义的错误响应对象 |
| | | TaskStatusResponse errorResponse = new TaskStatusResponse(); |
| | | errorResponse.setStatus("ERROR"); |
| | | errorResponse.setMessage(e.getMessage()); |
| | | errorResponse.setDetail(e.getMessage()); |
| | | return Mono.just(errorResponse); |
| | | }); |
| | | } |
| | | |
| | | |
| | | |
| | | // 添加一个辅助方法来安全地将字符串转换为double |
| | | private double parseUsage(String usageStr) { |
| | | try { |
| | |
| | | return 0.0; |
| | | } |
| | | } |
| | | @PostMapping("/search") |
| | | /*@PostMapping("/search") |
| | | @ApiOperation(value = "开始采集") |
| | | public Mono<ResponseResult<?>> createSearchTask( |
| | | @RequestBody SearchTaskRequest searchTaskRequest, |
| | |
| | | log.error("检查服务器资源失败: {}", e.getMessage(), e); |
| | | return Mono.just(ResponseResult.error("检查服务器资源失败: " + e.getMessage())); |
| | | }); |
| | | } |
| | | }*/ |
| | | |
| | | private Mono<Void> saveKeywordTasks(Integer keywordId, List<SearchTaskResponse> taskResponses) { |
| | | List<KeywordTask> keywordTasks = taskResponses.stream() |
| | | .filter(response -> response.getTask_id() != null) |
| | | .map(response -> { |
| | | KeywordTask keywordTask = new KeywordTask(); |
| | | keywordTask.setKeyword_id(keywordId); |
| | | keywordTask.setTask_id(response.getTask_id()); |
| | | keywordTask.setStatus("pending"); |
| | | return keywordTask; |
| | | }) |
| | | .collect(Collectors.toList()); |
| | | |
| | | // 将 MyBatis-Plus 的同步方法包装为 Mono<Void> |
| | | return Mono.fromRunnable(() -> { |
| | | boolean success = keywordTaskService.saveOrUpdateBatch(keywordTasks); |
| | | if (!success) { |
| | | // throw new RuntimeException("保存关键词任务关联失败"); |
| | | // 添加异常处理 |
| | | Mono.error( new RuntimeException("保存关键词任务关联失败")); |
| | | } |
| | | }) |
| | | .doFinally(signalType -> log.info("成功保存 {} 个关键词任务关联", keywordTasks.size())) |
| | | .then(); |
| | | } |
| | | // private Mono<Void> saveKeywordTasks(Integer keywordId, List<SearchTaskResponse> taskResponses) { |
| | | // List<KeywordTask> keywordTasks = taskResponses.stream() |
| | | // .filter(response -> response.getTask_id() != null) |
| | | // .map(response -> { |
| | | // KeywordTask keywordTask = new KeywordTask(); |
| | | // keywordTask.setKeyword_id(keywordId); |
| | | // keywordTask.setTask_id(response.getTask_id()); |
| | | // keywordTask.setStatus("pending"); |
| | | // return keywordTask; |
| | | // }) |
| | | // .collect(Collectors.toList()); |
| | | // |
| | | // // 将 MyBatis-Plus 的同步方法包装为 Mono<Void> |
| | | // return Mono.fromRunnable(() -> { |
| | | // boolean success = keywordTaskService.saveOrUpdateBatch(keywordTasks); |
| | | // if (!success) { |
| | | //// throw new RuntimeException("保存关键词任务关联失败"); |
| | | // // 添加异常处理 |
| | | // Mono.error( new RuntimeException("保存关键词任务关联失败")); |
| | | // } |
| | | // }) |
| | | // .doFinally(signalType -> log.info("成功保存 {} 个关键词任务关联", keywordTasks.size())) |
| | | // .then(); |
| | | // } |
| | | |
| | | private List<List<UserDto>> splitUsersIntoBatches(List<UserDto> users, int batchSize,Integer keywordId) { |
| | | |
| | |
| | | } |
| | | |
| | | // 移除原来的waitForTaskCompletion方法,不再需要同步等待 |
| | | @ApiOperation(value = "查询任务状态") |
| | | @GetMapping("/status") |
| | | public Mono<TaskStatusResponse> getTaskStatus(String taskId) { |
| | | return webClient.get() |
| | | .uri(baseUrl + "/api/v1/tasks/" + taskId) |
| | | .accept(MediaType.APPLICATION_JSON) |
| | | .retrieve() |
| | | .onStatus(HttpStatus::isError, response -> response.bodyToMono(TaskStatusResponse.class) |
| | | .flatMap(errorBody -> Mono.error(new RuntimeException(errorBody.getDetail())))) |
| | | .bodyToMono(TaskStatusResponse.class) |
| | | .flatMap(result -> { |
| | | TaskStatusResponse taskStatusResponse = result; |
| | | if (taskStatusResponse != null && taskStatusResponse.getStatus() != null) { |
| | | List<Question> updateQuestions = taskStatusResponse.getQuestions_status().stream() |
| | | .map(qs -> { |
| | | Question question = new Question(); |
| | | question.setQuestion_id(qs.getQuestion_id()); |
| | | question.setStatus(qs.getStatus()); |
| | | return question; |
| | | }).collect(Collectors.toList()); |
| | | |
| | | // 包装成响应式操作 |
| | | return Mono.fromCallable(() -> { |
| | | questionService.updateBatchById(updateQuestions); |
| | | return result; |
| | | }); |
| | | |
| | | } |
| | | return Mono.just(result); |
| | | }) |
| | | .onErrorResume(e -> { |
| | | // 创建一个自定义的错误响应对象 |
| | | TaskStatusResponse errorResponse = new TaskStatusResponse(); |
| | | errorResponse.setStatus("ERROR"); |
| | | errorResponse.setMessage(e.getMessage()); |
| | | errorResponse.setDetail(e.getMessage()); |
| | | |
| | | return Mono.just(errorResponse); |
| | | }); |
| | | } |
| | | // @ApiOperation(value = "查询任务状态") |
| | | // @GetMapping("/status") |
| | | // public Mono<TaskStatusResponse> getTaskStatus(String taskId) { |
| | | // return webClient.get() |
| | | // .uri(baseUrl + "/api/v1/tasks/" + taskId) |
| | | // .accept(MediaType.APPLICATION_JSON) |
| | | // .retrieve() |
| | | // .onStatus(HttpStatus::isError, response -> response.bodyToMono(TaskStatusResponse.class) |
| | | // .flatMap(errorBody -> Mono.error(new RuntimeException(errorBody.getDetail())))) |
| | | // .bodyToMono(TaskStatusResponse.class) |
| | | // .flatMap(result -> { |
| | | // TaskStatusResponse taskStatusResponse = result; |
| | | // if (taskStatusResponse != null && taskStatusResponse.getStatus() != null) { |
| | | // List<Question> updateQuestions = taskStatusResponse.getQuestions_status().stream() |
| | | // .map(qs -> { |
| | | // Question question = new Question(); |
| | | // question.setQuestion_id(qs.getQuestion_id()); |
| | | // question.setStatus(qs.getStatus()); |
| | | // return question; |
| | | // }).collect(Collectors.toList()); |
| | | // |
| | | // // 包装成响应式操作 |
| | | // return Mono.fromCallable(() -> { |
| | | // questionService.updateBatchById(updateQuestions); |
| | | // return result; |
| | | // }); |
| | | // |
| | | // } |
| | | // return Mono.just(result); |
| | | // }) |
| | | // .onErrorResume(e -> { |
| | | // // 创建一个自定义的错误响应对象 |
| | | // TaskStatusResponse errorResponse = new TaskStatusResponse(); |
| | | // errorResponse.setStatus("ERROR"); |
| | | // errorResponse.setMessage(e.getMessage()); |
| | | // errorResponse.setDetail(e.getMessage()); |
| | | // |
| | | // return Mono.just(errorResponse); |
| | | // }); |
| | | // } |
| | | |
| | | @PostMapping("/cancel/{taskId}") |
| | | @ApiOperation(value = "取消任务") |