Spring Boot中长时间运行API请求的优雅取消策略

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.Map;

public class TaskRegistry {
    private static final Map<String, Future<?>> activeTasks = new ConcurrentHashMap<>();

    public static void addTask(String taskId, Future<?> future) {
        activeTasks.put(taskId, future);
    }

    public static Future<?> getTask(String taskId) {
        return activeTasks.get(taskId);
    }

    public static void removeTask(String taskId) {
        activeTasks.remove(taskId);
    }
}

import org.springframework.scheduling.annotation.Async;
import org.springframework.scheduling.annotation.AsyncResult;
import org.springframework.stereotype.Service;

import java.util.concurrent.Future;

@Service
public class QueryService {

    @Async("taskExecutor") // 指定使用哪个线程池
    public Future<Void> runLongRunningQuery(String taskId, int timeToRun) {
        System.out.println("Task " + taskId + " started on thread: " + Thread.currentThread().getName());
        try {
            for (int i = 0; i < timeToRun; i++) {
                // 模拟耗时操作
                Thread.sleep(1000);
                System.out.println("Task " + taskId + " processing... " + (i + 1) + "/" + timeToRun);

                // 检查是否被中断（取消）
                if (Thread.currentThread().isInterrupted()) {
                    System.out.println("Task " + taskId + " was interrupted. Stopping execution.");
                    break;
                }
            }
            System.out.println("Task " + taskId + " completed.");
        } catch (InterruptedException e) {
            System.out.println("Task " + taskId + " caught InterruptedException. Stopping execution.");
            Thread.currentThread().interrupt(); // 重新设置中断标志
        } finally {
            TaskRegistry.removeTask(taskId); // 任务完成或取消后移除
            System.out.println("Task " + taskId + " removed from registry.");
        }
        return new AsyncResult<>(null); // 返回一个Future<Void>
    }
}

import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.UUID;
import java.util.concurrent.Future;

@RestController
@RequestMapping("/api/query")
public class QueryController {

    private final QueryService queryService;

    public QueryController(QueryService queryService) {
        this.queryService = queryService;
    }

    @PostMapping("/run/{timeToRun}")
    public ResponseEntity<String> runQuery(@PathVariable int timeToRun) {
        String taskId = UUID.randomUUID().toString();
        System.out.println("Received request to run query for " + timeToRun + " seconds. Task ID: " + taskId);
        Future<Void> future = queryService.runLongRunningQuery(taskId, timeToRun);
        TaskRegistry.addTask(taskId, future);
        return ResponseEntity.ok("Query started. Task ID: " + taskId);
    }

    @PostMapping("/cancel/{taskId}")
    public ResponseEntity<String> cancelQuery(@PathVariable String taskId) {
        Future<?> future = TaskRegistry.getTask(taskId);
        if (future != null) {
            boolean cancelled = future.cancel(true); // 尝试中断任务
            if (cancelled) {
                // TaskRegistry.removeTask(taskId); // 实际的移除会在任务的finally块中执行
                return ResponseEntity.ok("Task " + taskId + " cancellation requested.");
            } else {
                return ResponseEntity.status(409).body("Task " + taskId + " could not be cancelled (might be already completed or in an uncancelable state).");
            }
        } else {
            return ResponseEntity.status(404).body("Task " + taskId + " not found or already completed/cancelled.");
        }
    }
}