字节跳动Java后端面试深度解析（六）：TaskManager 任务生命周期管理 — 状态变更的原子性与并发安全

1
用户提交 Prompt "画一只猫" → 创建 Task-1 → AI 处理中
2
            ↓
3
用户编辑 Prompt 为 "画一只蓝色的猫"
4
            ↓
5
需要：取消 Task-1 + 创建 Task-2（原子操作）

1
// 危险代码：非原子操作
2
public String editPrompt(String taskId, String newPrompt) {
3
    cancelTask(taskId);           // 步骤1：取消旧任务
4
    return createTask(newPrompt); // 步骤2：创建新任务
5
}

1
public enum TaskStatus {
2
    PENDING,       // 已创建，等待处理
3
    QUEUED,        // 已进入队列
4
    PROCESSING,    // AI 处理中
5
    COMPLETED,     // 已完成
6
    FAILED,        // 失败
7
    CANCELLED,     // 已取消
8
    EDITING        // 编辑中（临时状态，锁定任务防止并发修改）
9
}

1
public class TaskStateMachine {
2

3
    /**
4
     * 合法的状态转移映射
5
     * key: 当前状态, value: 允许转移到的目标状态集合
6
     */
7
    private static final Map<TaskStatus, Set<TaskStatus>> TRANSITIONS = Map.of(
8
        PENDING,    Set.of(QUEUED, CANCELLED, EDITING),
9
        QUEUED,     Set.of(PROCESSING, CANCELLED, EDITING),
10
        PROCESSING, Set.of(COMPLETED, FAILED, CANCELLED),
11
        COMPLETED,  Set.of(),          // 终态，不可转移
12
        FAILED,     Set.of(),          // 终态，不可转移
13
        CANCELLED,  Set.of(),          // 终态，不可转移
14
        EDITING,    Set.of(PENDING, CANCELLED)  // 编辑完成后回到 PENDING 或取消
15
    );
16

17
    /**
18
     * 校验状态转移是否合法
19
     */
20
    public static void validateTransition(TaskStatus from, TaskStatus to) {
21
        Set<TaskStatus> allowed = TRANSITIONS.getOrDefault(from, Set.of());
22
        if (!allowed.contains(to)) {
23
            throw new IllegalTaskStateException(
24
                String.format("Cannot transition from %s to %s", from, to));
25
        }
26
    }
27
}

1
PENDING → QUEUED → PROCESSING → COMPLETED
2
  │         │          │
3
  │         │          └→ FAILED
4
  │         │
5
  │         └→ CANCELLED
6
  │
7
  └→ EDITING → PENDING（编辑后重新排队）
8
         │
9
         └→ CANCELLED（编辑后取消）

1
@Entity
2
@Table(name = "ai_task")
3
public class AiTask {
4

5
    @Id
6
    @GeneratedValue(strategy = GenerationType.IDENTITY)
7
    private Long id;
8

9
    @Column(name = "task_id", unique = true, nullable = false)
10
    private String taskId;
11

12
    @Column(name = "user_id", nullable = false)
13
    private String userId;
14

15
    @Column(name = "prompt", nullable = false, columnDefinition = "TEXT")
16
    private String prompt;
17

18
    @Enumerated(EnumType.STRING)
19
    @Column(nullable = false)
20
    private TaskStatus status;
21

22
    @Column(name = "result_url")
23
    private String resultUrl;
24

25
    @Column(name = "error_message")
26
    private String errorMessage;
27

28
    @Column(name = "created_at")
29
    private LocalDateTime createdAt;
30

31
    @Column(name = "updated_at")
32
    private LocalDateTime updatedAt;
33

34
    @Version
35
    private int version;  // 乐观锁
36

37
    @Column(name = "edit_lock_until")
38
    private LocalDateTime editLockUntil;  // 编辑锁过期时间
39
}

1
@Repository
2
public interface AiTaskRepository extends JpaRepository<AiTask, Long> {
3

4
    Optional<AiTask> findByTaskId(String taskId);
5

6
    /**
7
     * CAS 更新状态：只有当前状态匹配时才更新
8
     * 返回影响行数，0 表示状态已被其他线程修改
9
     */
10
    @Modifying
11
    @Query("""
12
        UPDATE AiTask t
13
        SET t.status = :newStatus,
14
            t.updatedAt = :now,
15
            t.version = t.version + 1
16
        WHERE t.taskId = :taskId
17
          AND t.status = :currentStatus
18
          AND t.version = :version
19
        """)
20
    int casUpdateStatus(
21
        @Param("taskId") String taskId,
22
        @Param("currentStatus") TaskStatus currentStatus,
23
        @Param("newStatus") TaskStatus newStatus,
24
        @Param("version") int version,
25
        @Param("now") LocalDateTime now
26
    );
27
}

1
@Service
2
@Transactional
3
public class TaskManager {
4

5
    private final AiTaskRepository taskRepository;
6
    private final RedisTemplate<String, String> redisTemplate;
7
    private final TaskQueueProducer queueProducer;
8
    private final ApplicationEventPublisher eventPublisher;
9

10
    private static final Duration EDIT_LOCK_TTL = Duration.ofSeconds(30);
11
    private static final Duration TASK_LOCK_TTL = Duration.ofSeconds(10);
12

13
    /**
14
     * 创建新任务
15
     */
16
    public AiTask createTask(String userId, String prompt) {
17
        AiTask task = new AiTask();
18
        task.setTaskId(UUID.randomUUID().toString());
19
        task.setUserId(userId);
20
        task.setPrompt(prompt);
21
        task.setStatus(TaskStatus.PENDING);
22
        task.setCreatedAt(LocalDateTime.now());
23
        task.setUpdatedAt(LocalDateTime.now());
24

25
        taskRepository.save(task);
26

27
        // 投递到队列
28
        queueProducer.enqueue(task.getTaskId());
29

30
        return task;
31
    }
32

33
    /**
34
     * 编辑 Prompt 并重启任务（原子操作）
35
     *
36
     * 核心流程：
37
     * 1. 获取分布式锁（防止并发编辑）
38
     * 2. CAS 将旧任务标记为 EDITING（锁定）
39
     * 3. 尝试取消旧任务的 AI 处理
40
     * 4. 创建新任务
41
     * 5. CAS 将旧任务标记为 CANCELLED
42
     * 6. 释放锁
43
     */
44
    public AiTask editAndRestart(String userId, String oldTaskId, String newPrompt) {
45
        String lockKey = "task:edit:lock:" + oldTaskId;
46

47
        // 1. 获取分布式锁
48
        Boolean locked = redisTemplate.opsForValue()
49
            .setIfAbsent(lockKey, userId, TASK_LOCK_TTL);
50
        if (!Boolean.TRUE.equals(locked)) {
51
            throw new TaskLockedException("Task is being edited by another request");
52
        }
53

54
        try {
55
            // 2. 查询旧任务
56
            AiTask oldTask = taskRepository.findByTaskId(oldTaskId)
57
                .orElseThrow(() -> new TaskNotFoundException(oldTaskId));
58

59
            // 校验所有权
60
            if (!oldTask.getUserId().equals(userId)) {
61
                throw new AccessDeniedException("Not your task");
62
            }
63

64
            // 3. CAS 将旧任务标记为 EDITING
65
            int updated = taskRepository.casUpdateStatus(
66
                oldTaskId,
67
                oldTask.getStatus(),
68
                TaskStatus.EDITING,
69
                oldTask.getVersion(),
70
                LocalDateTime.now()
71
            );
72
            if (updated == 0) {
73
                throw new ConcurrentModificationException(
74
                    "Task status changed during edit, please retry");
75
            }
76

77
            // 4. 尝试取消旧任务的 AI 处理（尽力而为）
78
            tryCancelAiProcessing(oldTaskId);
79

80
            // 5. 创建新任务
81
            AiTask newTask = createTask(userId, newPrompt);
82

83
            // 6. CAS 将旧任务标记为 CANCELLED
84
            // 重新查询获取最新的 version
85
            AiTask editingTask = taskRepository.findByTaskId(oldTaskId).orElseThrow();
86
            taskRepository.casUpdateStatus(
87
                oldTaskId,
88
                TaskStatus.EDITING,
89
                TaskStatus.CANCELLED,
90
                editingTask.getVersion(),
91
                LocalDateTime.now()
92
            );
93

94
            // 7. 发布取消事件
95
            eventPublisher.publishEvent(new TaskCancelledEvent(
96
                oldTaskId, "User edited prompt", LocalDateTime.now()));
97

98
            return newTask;
99

100
        } finally {
101
            redisTemplate.delete(lockKey);
102
        }
103
    }
104

105
    /**
106
     * 取消任务
107
     */
108
    public void cancelTask(String userId, String taskId) {
109
        String lockKey = "task:cancel:lock:" + taskId;
110

111
        Boolean locked = redisTemplate.opsForValue()
112
            .setIfAbsent(lockKey, userId, TASK_LOCK_TTL);
113
        if (!Boolean.TRUE.equals(locked)) {
114
            throw new TaskLockedException("Task is being modified");
115
        }
116

117
        try {
118
            AiTask task = taskRepository.findByTaskId(taskId)
119
                .orElseThrow(() -> new TaskNotFoundException(taskId));
120

121
            if (!task.getUserId().equals(userId)) {
122
                throw new AccessDeniedException("Not your task");
123
            }
124

125
            // CAS 更新为 CANCELLED
126
            int updated = taskRepository.casUpdateStatus(
127
                taskId,
128
                task.getStatus(),
129
                TaskStatus.CANCELLED,
130
                task.getVersion(),
131
                LocalDateTime.now()
132
            );
133
            if (updated == 0) {
134
                throw new ConcurrentModificationException(
135
                    "Task status changed, please retry");
136
            }
137

138
            // 取消 AI 处理
139
            tryCancelAiProcessing(taskId);
140

141
            eventPublisher.publishEvent(new TaskCancelledEvent(
142
                taskId, "User cancelled", LocalDateTime.now()));
143

144
        } finally {
145
            redisTemplate.delete(lockKey);
146
        }
147
    }
148

149
    /**
150
     * 尝试取消 AI 处理（尽力而为）
151
     * 如果 AI 任务已经在处理中，可能无法取消
152
     */
153
    private void tryCancelAiProcessing(String taskId) {
154
        try {
155
            queueProducer.cancel(taskId);
156
        } catch (Exception e) {
157
            // 记录日志但不抛异常，因为任务状态已经在 DB 中标记为 CANCELLED
158
            // AI 处理完成后回调时会检查状态，发现已取消则忽略结果
159
            log.warn("Failed to cancel AI processing for task {}: {}",
160
                taskId, e.getMessage());
161
        }
162
    }
163
}

1
第一层：分布式锁（Redis）    — 防止同一任务的并发编辑
2
第二层：CAS 更新（数据库）   — 防止状态更新的竞态条件
3
第三层：乐观锁（version）    — 防止并发写覆盖
4
第四层：状态机校验           — 防止非法状态转移

1
/**
2
 * 更健壮的分布式锁实现（Redisson 风格）
3
 */
4
@Component
5
public class DistributedTaskLock {
6

7
    private final RedisTemplate<String, String> redisTemplate;
8

9
    private static final String LOCK_PREFIX = "task:lock:";
10
    private static final String UNLOCK_SCRIPT = """
11
        if redis.call('get', KEYS[1]) == ARGV[1] then
12
            return redis.call('del', KEYS[1])
13
        else
14
            return 0
15
        end
16
        """;
17

18
    /**
19
     * 获取锁（带重试）
20
     */
21
    public TaskLock acquire(String taskId, String ownerId, Duration timeout) {
22
        String lockKey = LOCK_PREFIX + taskId;
23
        long maxRetries = 3;
24
        long retryDelayMs = 100;
25

26
        for (int i = 0; i < maxRetries; i++) {
27
            Boolean acquired = redisTemplate.opsForValue()
28
                .setIfAbsent(lockKey, ownerId, timeout);
29
            if (Boolean.TRUE.equals(acquired)) {
30
                return new TaskLock(lockKey, ownerId, this);
31
            }
32

33
            // 短暂等待后重试
34
            try {
35
                Thread.sleep(retryDelayMs * (i + 1));
36
            } catch (InterruptedException e) {
37
                Thread.currentThread().interrupt();
38
                throw new LockAcquisitionException("Interrupted while acquiring lock");
39
            }
40
        }
41

42
        throw new LockAcquisitionException(
43
            "Failed to acquire lock for task: " + taskId);
44
    }
45

46
    /**
47
     * 释放锁（Lua 脚本保证原子性）
48
     */
49
    public void release(String lockKey, String ownerId) {
50
        Long result = redisTemplate.execute(
51
            new DefaultRedisScript<>(UNLOCK_SCRIPT, Long.class),
52
            List.of(lockKey),
53
            ownerId
54
        );
55
        if (result == null || result == 0) {
56
            log.warn("Lock already released or owned by another process: {}", lockKey);
57
        }
58
    }
59
}
60

61
/**
62
 * 锁对象（实现 AutoCloseable，支持 try-with-resources）
63
 */
64
public class TaskLock implements AutoCloseable {
65
    private final String lockKey;
66
    private final String ownerId;
67
    private final DistributedTaskLock lockManager;
68
    private boolean released = false;
69

70
    @Override
71
    public void close() {
72
        if (!released) {
73
            lockManager.release(lockKey, ownerId);
74
            released = true;
75
        }
76
    }
77
}

1
public AiTask editAndRestart(String userId, String oldTaskId, String newPrompt) {
2
    // try-with-resources 自动释放锁
3
    try (TaskLock lock = lockManager.acquire(oldTaskId, userId, Duration.ofSeconds(10))) {
4
        // ... 业务逻辑 ...
5
    }
6
}

1
@Component
2
public class TaskCompensationJob {
3

4
    /**
5
     * 每 5 分钟扫描 EDITING 状态超过 1 分钟的任务
6
     * 这些任务可能是编辑操作中途失败遗留的
7
     */
8
    @Scheduled(fixedRate = 300_000)
9
    @Transactional
10
    public void compensateStuckEditingTasks() {
11
        LocalDateTime cutoff = LocalDateTime.now().minusMinutes(1);
12
        List<AiTask> stuckTasks = taskRepository
13
            .findByStatusAndUpdatedAtBefore(TaskStatus.EDITING, cutoff);
14

15
        for (AiTask task : stuckTasks) {
16
            log.warn("Found stuck EDITING task: {}, reverting to PENDING",
17
                task.getTaskId());
18
            taskRepository.casUpdateStatus(
19
                task.getTaskId(),
20
                TaskStatus.EDITING,
21
                TaskStatus.PENDING,
22
                task.getVersion(),
23
                LocalDateTime.now()
24
            );
25
            // 重新投递到队列
26
            queueProducer.enqueue(task.getTaskId());
27
        }
28
    }
29

30
    /**
31
     * 每小时清理超过 24 小时的 CANCELLED 任务
32
     */
33
    @Scheduled(cron = "0 0 * * * *")
34
    @Transactional
35
    public void cleanupCancelledTasks() {
36
        LocalDateTime cutoff = LocalDateTime.now().minusHours(24);
37
        List<AiTask> expired = taskRepository
38
            .findByStatusAndUpdatedAtBefore(TaskStatus.CANCELLED, cutoff);
39

40
        if (!expired.isEmpty()) {
41
            taskRepository.deleteAll(expired);
42
            log.info("Cleaned up {} cancelled tasks", expired.size());
43
        }
44
    }
45
}

1
@EventListener
2
public void onTaskCancelled(TaskCancelledEvent event) {
3
    // 检查是否有对应的"编辑补偿"标记
4
    String pendingEdit = redisTemplate.opsForValue()
5
        .get("task:pending_edit:" + event.taskId());
6

7
    if (pendingEdit != null) {
8
        // 有待补偿的编辑操作
9
        EditContext ctx = JsonUtil.fromJson(pendingEdit, EditContext.class);
10
        try {
11
            taskManager.createTask(ctx.userId(), ctx.newPrompt());
12
            redisTemplate.delete("task:pending_edit:" + event.taskId());
13
        } catch (Exception e) {
14
            log.error("Failed to compensate edit for task {}", event.taskId(), e);
15
        }
16
    }
17
}