在Spring Boot中实现长耗时API请求的异步执行与取消策略

// src/main/java/com/example/async/AsyncConfig.java
package com.example.async;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.Executor;

@Configuration
@EnableAsync
public class AsyncConfig {

    @Bean(name = "taskExecutor")
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5); // 核心线程数
        executor.setMaxPoolSize(10); // 最大线程数
        executor.setQueueCapacity(25); // 队列容量
        executor.setThreadNamePrefix("AsyncTask-"); // 线程名前缀
        executor.initialize();
        return executor;
    }
}

// src/main/java/com/example/async/LongRunningService.java
package com.example.async;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.util.concurrent.CompletableFuture;

@Service
public class LongRunningService {

    private static final Logger logger = LoggerFactory.getLogger(LongRunningService.class);

    @Async("taskExecutor") // 指定使用名为"taskExecutor"的线程池
    public CompletableFuture<String> executeLongRunningTask(String taskId, int timeToRun) {
        logger.info("Task {} started with timeToRun: {}", taskId, timeToRun);
        try {
            for (int i = 0; i < timeToRun; i++) {
                // 模拟耗时操作
                Thread.sleep(1000); // 每次操作耗时1秒

                // 检查线程是否被中断
                if (Thread.currentThread().isInterrupted()) {
                    logger.warn("Task {} was interrupted after {} seconds.", taskId, i + 1);
                    return CompletableFuture.completedFuture("Task " + taskId + " cancelled.");
                }
                logger.info("Task {} processing... {}/{} seconds", taskId, i + 1, timeToRun);
            }
            logger.info("Task {} completed successfully.", taskId);
            return CompletableFuture.completedFuture("Task " + taskId + " completed.");
        } catch (InterruptedException e) {
            // 捕获InterruptedException，通常表示线程被中断
            Thread.currentThread().interrupt(); // 重新设置中断标志
            logger.warn("Task {} caught InterruptedException. It was cancelled.", taskId);
            return CompletableFuture.completedFuture("Task " + taskId + " cancelled due to interruption.");
        } catch (Exception e) {
            logger.error("Task {} encountered an error: {}", taskId, e.getMessage());
            return CompletableFuture.completedFuture("Task " + taskId + " failed: " + e.getMessage());
        }
    }
}

// src/main/java/com/example/async/TaskController.java
package com.example.async;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CompletableFuture;

@RestController
@RequestMapping("/api/tasks")
public class TaskController {

    private static final Logger logger = LoggerFactory.getLogger(TaskController.class);
    private final LongRunningService longRunningService;

    // 使用ConcurrentHashMap来安全地存储任务Future
    private final Map<String, CompletableFuture<String>> runningTasks = new ConcurrentHashMap<>();

    public TaskController(LongRunningService longRunningService) {
        this.longRunningService = longRunningService;
    }

    /**
     * 启动一个长耗时任务
     * @param timeToRun 模拟任务运行时间（秒）
     * @return 包含任务ID的响应
     */
    @PostMapping("/start/{timeToRun}")
    public ResponseEntity<String> startTask(@PathVariable int timeToRun) {
        String taskId = UUID.randomUUID().toString();
        logger.info("Received request to start task {} with timeToRun: {}", taskId, timeToRun);

        CompletableFuture<String> future = longRunningService.executeLongRunningTask(taskId, timeToRun);
        runningTasks.put(taskId, future);

        // 异步监听任务完成，并从map中移除
        future.whenComplete((result, ex) -> {
            if (ex == null) {
                logger.info("Task {} completed with result: {}", taskId, result);
            } else {
                logger.error("Task {} failed with exception: {}", taskId, ex.getMessage());
            }
            runningTasks.remove(taskId); // 任务完成后从map中移除
        });

        return ResponseEntity.ok("Task started. Task ID: " + taskId);
    }

    /**
     * 取消一个正在运行的任务
     * @param taskId 要取消的任务ID
     * @return 取消结果
     */
    @PostMapping("/cancel/{taskId}")
    public ResponseEntity<String> cancelTask(@PathVariable String taskId) {
        CompletableFuture<String> future = runningTasks.get(taskId);

        if (future == null) {
            logger.warn("Attempted to cancel non-existent or already completed task: {}", taskId);
            return ResponseEntity.badRequest().body("Task " + taskId + " not found or already completed.");
        }

        if (future.isDone()) {
            logger.info("Task {} is already done, cannot cancel.", taskId);
            runningTasks.remove(taskId); // 确保已完成的任务被移除
            return ResponseEntity.badRequest().body("Task " + taskId + " is already done, cannot cancel.");
        }

        // 尝试取消任务，true表示如果任务正在运行，会尝试中断线程
        boolean cancelled = future.cancel(true);
        if (cancelled) {
            logger.info("Task {} successfully requested for cancellation.", taskId);
            // 任务被成功取消后，也从map中移除
            runningTasks.remove(taskId);
            return ResponseEntity.ok("Task " + taskId + " cancellation requested.");
        } else {
            logger.warn("Task {} could not be cancelled. It might be completed or not interruptible.", taskId);
            return ResponseEntity.status(500).body("Failed to cancel task " + taskId + ".");
        }
    }

    /**
     * 获取所有正在运行的任务ID
     * @return 正在运行的任务ID列表
     */
    @GetMapping("/status")
    public ResponseEntity<Map<String, String>> getRunningTasks() {
        Map<String, String> statusMap = new ConcurrentHashMap<>();
        runningTasks.forEach((id, future) -> {
            statusMap.put(id, future.isDone() ? "DONE" : (future.isCancelled() ? "CANCELLED" : "RUNNING"));
        });
        return ResponseEntity.ok(statusMap);
    }
}