task-72: 文件作业数据库原子 claim，避免重复派发同一 job

2026-08-30 20:01:18 +08:00
parent c07e26a400
commit baa43d80ba
3 changed files with 282 additions and 2 deletions
@@ -86,6 +86,72 @@ public class TaskFileJobService {
                .last("limit " + Math.max(1, Math.min(limit, 100))));
    }

+    /**
+     * 原子 claim：先查候选（PENDING/FAILED 且未达重试上限），再逐个以条件
+     * UPDATE（状态+retryCount 匹配 → RUNNING）翻转；只有翻转成功的行才返回。
+     * 多个 worker/实例并发取数时，同一 job 只会被一个消费者 claim 到，
+     * 从源头避免重复派发同一 job。
+     */
+    public List<TaskFileJobEntity> claimRunnableJobs(int limit) {
+        List<TaskFileJobEntity> candidates = listRunnableJobs(limit);
+        return claimCandidates(candidates);
+    }
+
+    /** 带 owner 过滤的原子 claim：owner 任务优先，其余补足；同 {@link #claimRunnableJobs} 的原子性。 */
+    public List<TaskFileJobEntity> claimRunnableJobsForOwner(int limit, String owner) {
+        String normalizedOwner = owner == null ? "" : owner.trim();
+        if (normalizedOwner.isBlank()) {
+            return claimRunnableJobs(limit);
+        }
+        String ownerMarker = ":owner:" + normalizedOwner;
+        int safeLimit = Math.max(1, Math.min(limit, 100));
+
+        List<TaskFileJobEntity> ownerJobs = taskFileJobMapper.selectList(new LambdaQueryWrapper<TaskFileJobEntity>()
+                .in(TaskFileJobEntity::getStatus, List.of("PENDING", "FAILED"))
+                .lt(TaskFileJobEntity::getRetryCount, MAX_RETRY_COUNT)
+                .in(TaskFileJobEntity::getModuleType, List.of("APPEARANCE_PATENT", "SIMILAR_ASIN", "PUBLISH", "SHOP_DATA_CRAWL"))
+                .like(TaskFileJobEntity::getScopeKey, ownerMarker)
+                .orderByAsc(TaskFileJobEntity::getUpdatedAt)
+                .last("limit " + safeLimit));
+
+        List<TaskFileJobEntity> candidates = new ArrayList<>(ownerJobs);
+        if (ownerJobs.size() < safeLimit) {
+            LambdaQueryWrapper<TaskFileJobEntity> genericWrapper = new LambdaQueryWrapper<TaskFileJobEntity>()
+                    .in(TaskFileJobEntity::getStatus, List.of("PENDING", "FAILED"))
+                    .lt(TaskFileJobEntity::getRetryCount, MAX_RETRY_COUNT)
+                    .and(wrapper -> wrapper
+                            .notIn(TaskFileJobEntity::getModuleType, List.of("APPEARANCE_PATENT", "SIMILAR_ASIN", "PUBLISH", "SHOP_DATA_CRAWL"))
+                            .or()
+                            .isNull(TaskFileJobEntity::getScopeKey)
+                            .or()
+                            .notLike(TaskFileJobEntity::getScopeKey, ":owner:"))
+                    .orderByAsc(TaskFileJobEntity::getUpdatedAt)
+                    .last("limit " + (safeLimit - ownerJobs.size()));
+            candidates.addAll(taskFileJobMapper.selectList(genericWrapper));
+        }
+        return claimCandidates(candidates);
+    }
+
+    private List<TaskFileJobEntity> claimCandidates(List<TaskFileJobEntity> candidates) {
+        if (candidates == null || candidates.isEmpty()) {
+            return List.of();
+        }
+        List<TaskFileJobEntity> claimed = new ArrayList<>();
+        for (TaskFileJobEntity candidate : candidates) {
+            if (candidate == null || candidate.getId() == null) {
+                continue;
+            }
+            if (!markRunning(candidate.getId())) {
+                continue;
+            }
+            TaskFileJobEntity claim = taskFileJobMapper.selectById(candidate.getId());
+            if (claim != null && "RUNNING".equals(claim.getStatus())) {
+                claimed.add(claim);
+            }
+        }
+        return claimed;
+    }
+
    public List<TaskFileJobEntity> listRunnableJobsForOwner(int limit, String owner) {
        String normalizedOwner = owner == null ? "" : owner.trim();
        if (normalizedOwner.isBlank()) {
@@ -85,11 +85,13 @@ public class TaskResultFileJobWorker {
        if (!localWorkerEnabled) {
            return;
        }
-        List<TaskFileJobEntity> jobs = taskFileJobService.listRunnableJobsForOwner(batchSize, currentInstanceId());
+        // 原子 claim：每个候选以条件 UPDATE 翻转为 RUNNING，只处理 claim 成功的行，
+        // 并发 worker/实例不会重复派发同一 job。
+        List<TaskFileJobEntity> jobs = taskFileJobService.claimRunnableJobsForOwner(batchSize, currentInstanceId());
        if (jobs == null || jobs.isEmpty()) {
            return;
        }
-        log.info("[task-file-job] scheduled worker picked jobs count={}", jobs.size());
+        log.info("[task-file-job] scheduled worker claimed jobs count={}", jobs.size());
        for (TaskFileJobEntity job : jobs) {
            boolean dispatched = taskFileJobLocalDispatcher.dispatch(
                    job.getId(),
@@ -0,0 +1,212 @@
+package com.nanri.aiimage.modules.task.service;
+
+import com.baomidou.mybatisplus.core.MybatisConfiguration;
+import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
+import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
+import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
+import com.nanri.aiimage.modules.task.mapper.TaskFileJobMapper;
+import com.nanri.aiimage.modules.task.model.entity.TaskFileJobEntity;
+import org.apache.ibatis.builder.MapperBuilderAssistant;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.context.ApplicationEventPublisher;
+
+import java.time.LocalDateTime;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Task 72：为文件作业实现数据库原子 claim，避免重复派发同一 job。
+ * claimRunnableJobs / claimRunnableJobsForOwner 先查询候选，再逐个以
+ * 条件 UPDATE（PENDING/FAILED + retryCount 未达上限 → RUNNING）原子翻转；
+ * 只有翻转成功的行才返回，多个 worker/实例并发取数不会重复拿到同一 job。
+ */
+@ExtendWith(MockitoExtension.class)
+class TaskFileJobClaimTest {
+
+    @Mock private TaskFileJobMapper taskFileJobMapper;
+    @Mock private ApplicationEventPublisher applicationEventPublisher;
+
+    private TaskFileJobService service;
+
+    @BeforeAll
+    static void initializeTableInfo() {
+        TableInfoHelper.initTableInfo(
+                new MapperBuilderAssistant(new MybatisConfiguration(), ""),
+                TaskFileJobEntity.class);
+    }
+
+    @BeforeEach
+    void setUp() {
+        service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+    }
+
+    private static TaskFileJobEntity runnableJob(Long id) {
+        TaskFileJobEntity job = new TaskFileJobEntity();
+        job.setId(id);
+        job.setTaskId(1000L + id);
+        job.setResultId(2000L + id);
+        job.setModuleType("SIMILAR_ASIN");
+        job.setScopeKey("task:1:owner:instance-a");
+        job.setJobType("ASSEMBLE_RESULT");
+        job.setStatus("PENDING");
+        job.setRetryCount(0);
+        job.setUpdatedAt(LocalDateTime.now());
+        return job;
+    }
+
+    private static TaskFileJobEntity runningJob(Long id) {
+        TaskFileJobEntity job = runnableJob(id);
+        job.setStatus("RUNNING");
+        return job;
+    }
+
+    private void stubCandidates(List<TaskFileJobEntity> candidates) {
+        when(taskFileJobMapper.selectList(any(LambdaQueryWrapper.class))).thenReturn(candidates);
+    }
+
+    private void stubClaimResults(int... results) {
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class)))
+                .thenReturn(results[0], java.util.Arrays.stream(results, 1, results.length).boxed().toArray(Integer[]::new));
+    }
+
+    private void stubClaimedRows(List<TaskFileJobEntity> claimed) {
+        when(taskFileJobMapper.selectById(any())).thenAnswer(invocation -> {
+            Long id = invocation.getArgument(0);
+            return claimed.stream().filter(job -> job.getId().equals(id)).findFirst().orElse(null);
+        });
+    }
+
+    @Test
+    void test_task_072_job_normal_default_path() {
+        // 默认路径：候选全部翻转成功，返回对应 RUNNING 行。
+        stubCandidates(List.of(runnableJob(1L), runnableJob(2L)));
+        stubClaimResults(1, 1);
+        stubClaimedRows(List.of(runningJob(1L), runningJob(2L)));
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(10);
+
+        assertEquals(2, claimed.size(), "两个候选全部 claim 成功");
+        assertTrue(claimed.stream().allMatch(job -> "RUNNING".equals(job.getStatus())));
+        verify(taskFileJobMapper, times(2)).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_normal_multiple_items() {
+        // 批量场景：5 个候选，其中 1 个被另一实例抢先（update 返回 0），
+        // 只返回 4 个；每个候选都尝试一次原子翻转，顺序稳定。
+        List<TaskFileJobEntity> candidates = new ArrayList<>();
+        for (long i = 1; i <= 5; i++) {
+            candidates.add(runnableJob(i));
+        }
+        stubCandidates(candidates);
+        stubClaimResults(1, 0, 1, 1, 1);
+        stubClaimedRows(List.of(runningJob(1L), runningJob(3L), runningJob(4L), runningJob(5L)));
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(10);
+
+        assertEquals(4, claimed.size(), "被抢的 job 不返回");
+        assertEquals(List.of(1L, 3L, 4L, 5L),
+                claimed.stream().map(TaskFileJobEntity::getId).toList(), "顺序与候选一致");
+        verify(taskFileJobMapper, times(5)).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_normal_repeated_operation_is_idempotent() {
+        // 幂等：已 claim（RUNNING）的 job 不再出现在候选里，重复 claim 返回空，
+        // 不产生重复翻转或重复派发。
+        stubCandidates(List.of());
+
+        List<TaskFileJobEntity> first = service.claimRunnableJobs(10);
+        List<TaskFileJobEntity> second = service.claimRunnableJobs(10);
+
+        assertTrue(first.isEmpty());
+        assertTrue(second.isEmpty());
+        verify(taskFileJobMapper, never()).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_boundary_empty_input() {
+        // 空输入：无候选时不发起任何原子翻转，返回空列表。
+        stubCandidates(List.of());
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(10);
+
+        assertTrue(claimed.isEmpty());
+        verify(taskFileJobMapper, never()).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_boundary_single_item() {
+        // 单元素：单候选单翻转单返回，不依赖批量路径。
+        stubCandidates(List.of(runnableJob(9L)));
+        stubClaimResults(1);
+        stubClaimedRows(List.of(runningJob(9L)));
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(10);
+
+        assertEquals(List.of(9L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+        verify(taskFileJobMapper, times(1)).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_boundary_limit_and_overflow() {
+        // 上限/超限：limit 非法（0/负/超上限）时按现有 clamp 回退，不发生无界取数。
+        stubCandidates(List.of(runnableJob(1L)));
+        stubClaimResults(1);
+        stubClaimedRows(List.of(runningJob(1L)));
+
+        service.claimRunnableJobs(0);
+        service.claimRunnableJobs(200);
+
+        ArgumentCaptor<LambdaQueryWrapper<TaskFileJobEntity>> captor =
+                ArgumentCaptor.forClass(LambdaQueryWrapper.class);
+        verify(taskFileJobMapper, times(2)).selectList(captor.capture());
+        assertTrue(captor.getAllValues().get(0).getCustomSqlSegment().contains("limit 1"), "limit=0 回退到 1");
+        assertTrue(captor.getAllValues().get(1).getCustomSqlSegment().contains("limit 100"), "limit=200 钳制到 100");
+    }
+
+    @Test
+    void test_task_072_job_invalid_input_rejected() {
+        // 非法参数：空白 owner 回退到通用 claim 路径（仅一次通用查询），行为确定。
+        stubCandidates(List.of());
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobsForOwner(10, "   ");
+
+        assertTrue(claimed.isEmpty());
+        verify(taskFileJobMapper, times(1)).selectList(any(LambdaQueryWrapper.class));
+        verify(taskFileJobMapper, never()).update(isNull(), any(LambdaUpdateWrapper.class));
+    }
+
+    @Test
+    void test_task_072_job_dependency_failure_releases_resources() {
+        // 依赖失败：候选查询抛异常时传播且不发起任何翻转（无半态 claim）；
+        // 恢复后再次 claim 正常返回，未产生重复或丢失。
+        when(taskFileJobMapper.selectList(any(LambdaQueryWrapper.class)))
+                .thenThrow(new RuntimeException("db down"));
+
+        assertThrows(RuntimeException.class, () -> service.claimRunnableJobs(10));
+        verify(taskFileJobMapper, never()).update(isNull(), any(LambdaUpdateWrapper.class));
+
+        stubCandidates(List.of(runnableJob(7L)));
+        stubClaimResults(1);
+        stubClaimedRows(List.of(runningJob(7L)));
+        List<TaskFileJobEntity> recovered = service.claimRunnableJobs(10);
+        assertEquals(List.of(7L), recovered.stream().map(TaskFileJobEntity::getId).toList(), "恢复后 claim 成功");
+    }
+}