task-122: N+1 批量化修复（claim/stuck-reset/stale-check 逐行 selectById 改 selectBatchIds 批量回读 + Map 装配，行为等价）

2026-09-02 03:13:05 +08:00
parent 29625f703c
commit 445b7780c5
7 changed files with 536 additions and 26 deletions
@@ -36,8 +36,10 @@ import java.nio.file.attribute.FileTime;
 import java.time.Duration;
 import java.time.Instant;
 import java.time.LocalDateTime;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@@ -178,6 +180,7 @@ public class DeleteBrandStaleTaskService {
                .lt(FileTaskEntity::getUpdatedAt, threshold)
                .last("limit 200"));

+        List<StaleFinalizeEntry> entries = new ArrayList<>();
        for (FileTaskEntity task : runningTasks) {
            Map<Object, Object> progress = deleteBrandTaskCacheService.getProgress(task.getId());
            long lastHeartbeatAt = 0L;
@@ -203,28 +206,52 @@ public class DeleteBrandStaleTaskService {
            if (taskLockHandle == null) {
                continue;
            }
+            boolean finalizeThrew;
            try (taskLockHandle) {
-            try {
-                deleteBrandRunService.tryFinalizeTask(task.getId(), true);
-                FileTaskEntity refreshed = fileTaskMapper.selectById(task.getId());
+                try {
+                    deleteBrandRunService.tryFinalizeTask(task.getId(), true);
+                    finalizeThrew = false;
+                } catch (Exception ex) {
+                    finalizeThrew = true;
+                    log.warn("[stale-check] delete-brand finalize threw taskId={} msg={}", task.getId(), ex.getMessage());
+                }
+            }
+            entries.add(new StaleFinalizeEntry(task, lastHeartbeatAt, completedScopeCount, hasStartedProgress, finalizeThrew));
+        }
+        if (entries.isEmpty()) {
+            return;
+        }
+
+        // 两阶段：全部 finalize 后一次 IN 批量回读，替代逐任务 selectById。
+        List<Long> refreshTaskIds = entries.stream()
+                .filter(entry -> !entry.finalizeThrew)
+                .map(entry -> entry.task.getId())
+                .toList();
+        Map<Long, FileTaskEntity> refreshedById = refreshTaskIds.isEmpty()
+                ? Map.of()
+                : fileTaskMapper.selectBatchIds(refreshTaskIds).stream()
+                        .collect(Collectors.toMap(FileTaskEntity::getId, task -> task, (a, b) -> a));
+
+        for (StaleFinalizeEntry entry : entries) {
+            FileTaskEntity task = entry.task;
+            if (!entry.finalizeThrew) {
+                FileTaskEntity refreshed = refreshedById.get(task.getId());
                if (refreshed != null && !"RUNNING".equals(refreshed.getStatus())) {
                    deleteBrandTaskCacheService.saveTaskCache(refreshed);
                    log.info("[stale-check] delete-brand finalized before timeout-fail taskId={} status={} updatedAt={} finishedAt={}",
                            refreshed.getId(), refreshed.getStatus(), refreshed.getUpdatedAt(), refreshed.getFinishedAt());
                    continue;
                }
-            } catch (Exception ex) {
-                log.warn("[stale-check] delete-brand finalize threw taskId={} msg={}", task.getId(), ex.getMessage());
            }

            log.warn("[stale-check] delete-brand failing stale task -> taskId={} updatedAt={} createdAt={} lastHeartbeatAt={} timeoutMinutes={} completedScopes={} hasStartedProgress={}",
                    task.getId(),
                    task.getUpdatedAt(),
                    task.getCreatedAt(),
-                    lastHeartbeatAt,
+                    entry.lastHeartbeatAt,
                    minutes,
-                    completedScopeCount,
-                    hasStartedProgress);
+                    entry.completedScopeCount,
+                    entry.hasStartedProgress);
            int updated = fileTaskMapper.update(null, new LambdaUpdateWrapper<FileTaskEntity>()
                    .eq(FileTaskEntity::getId, task.getId())
                    .eq(FileTaskEntity::getModuleType, MODULE_TYPE_DELETE_BRAND)
@@ -238,10 +265,13 @@ public class DeleteBrandStaleTaskService {
                deleteBrandTaskCacheService.delete(task.getId());
                log.warn("[stale-check] delete-brand failed taskId={} reason=timeout", task.getId());
            }
-            }
        }
    }

+    private record StaleFinalizeEntry(FileTaskEntity task, long lastHeartbeatAt, int completedScopeCount,
+                                      boolean hasStartedProgress, boolean finalizeThrew) {
+    }
+
    private ProductRiskStaleCheckStats failStaleProductRiskResolveTasks() {
        ProductRiskStaleCheckStats stats = new ProductRiskStaleCheckStats();
        long minutes = Math.max(1L, deleteBrandProgressProperties.getProductRiskStaleTimeoutMinutes());
@@ -17,6 +17,7 @@ import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@@ -142,15 +143,22 @@ public class TaskFileJobService {
        if (candidates == null || candidates.isEmpty()) {
            return List.of();
        }
-        List<TaskFileJobEntity> claimed = new ArrayList<>();
+        List<Long> claimedIds = new ArrayList<>();
        for (TaskFileJobEntity candidate : candidates) {
            if (candidate == null || candidate.getId() == null) {
                continue;
            }
-            if (!markRunning(candidate.getId())) {
-                continue;
+            if (markRunning(candidate.getId())) {
+                claimedIds.add(candidate.getId());
            }
-            TaskFileJobEntity claim = taskFileJobMapper.selectById(candidate.getId());
+        }
+        if (claimedIds.isEmpty()) {
+            return List.of();
+        }
+        Map<Long, TaskFileJobEntity> claimedById = refreshJobsByIds(claimedIds);
+        List<TaskFileJobEntity> claimed = new ArrayList<>();
+        for (Long id : claimedIds) {
+            TaskFileJobEntity claim = claimedById.get(id);
            if (claim != null && "RUNNING".equals(claim.getStatus())) {
                claimed.add(claim);
            }
@@ -259,10 +267,16 @@ public class TaskFileJobService {
                .last("limit " + Math.max(1, Math.min(limit, 200))));
        int reset = 0;
        List<TaskFileJobEntity> exhaustedJobs = new ArrayList<>();
+        List<Long> eventJobIds = new ArrayList<>();
+        List<Long> exhaustedJobIds = new ArrayList<>();
+        Map<Long, TaskFileJobEntity> originalExhaustedById = new LinkedHashMap<>();
+        Map<Long, TaskFileJobEntity> originalJobsById = jobs.stream()
+                .collect(Collectors.toMap(TaskFileJobEntity::getId, job -> job, (a, b) -> a));
        for (TaskFileJobEntity job : jobs) {
            if ("FAILED".equals(job.getStatus())
                    && job.getRetryCount() != null && job.getRetryCount() >= MAX_RETRY_COUNT) {
-                exhaustedJobs.add(job);
+                exhaustedJobIds.add(job.getId());
+                originalExhaustedById.put(job.getId(), job);
                continue;
            }
            if ("PENDING".equals(job.getStatus())) {
@@ -279,7 +293,7 @@ public class TaskFileJobService {
                        .set(TaskFileJobEntity::getRetryCount, nextRetryCount)
                        .set(TaskFileJobEntity::getUpdatedAt, LocalDateTime.now()));
                if (updated > 0) {
-                    publishDispatchEvent(taskFileJobMapper.selectById(job.getId()));
+                    eventJobIds.add(job.getId());
                }
                continue;
            }
@@ -300,8 +314,7 @@ public class TaskFileJobService {
                        .set(TaskFileJobEntity::getFinishedAt, now)
                        .set(TaskFileJobEntity::getTerminalCallbackAt, null));
                if (updated > 0) {
-                    TaskFileJobEntity exhausted = taskFileJobMapper.selectById(job.getId());
-                    exhaustedJobs.add(exhausted == null ? job : exhausted);
+                    exhaustedJobIds.add(job.getId());
                }
                continue;
            }
@@ -312,12 +325,40 @@ public class TaskFileJobService {
                    .set(TaskFileJobEntity::getTerminalCallbackAt, null));
            if (updated > 0) {
                reset++;
-                publishDispatchEvent(taskFileJobMapper.selectById(job.getId()));
+                eventJobIds.add(job.getId());
+            }
+        }
+        if (!eventJobIds.isEmpty()) {
+            Map<Long, TaskFileJobEntity> refreshedById = refreshJobsByIds(eventJobIds);
+            for (Long jobId : eventJobIds) {
+                publishDispatchEvent(refreshedById.get(jobId));
+            }
+        }
+        if (!exhaustedJobIds.isEmpty()) {
+            List<Long> refreshIds = exhaustedJobIds.stream()
+                    .filter(id -> !originalExhaustedById.containsKey(id))
+                    .toList();
+            Map<Long, TaskFileJobEntity> refreshedById = refreshIds.isEmpty()
+                    ? Map.of()
+                    : refreshJobsByIds(refreshIds);
+            for (Long jobId : exhaustedJobIds) {
+                TaskFileJobEntity original = originalExhaustedById.get(jobId);
+                if (original != null) {
+                    exhaustedJobs.add(original);
+                    continue;
+                }
+                TaskFileJobEntity refreshed = refreshedById.get(jobId);
+                exhaustedJobs.add(refreshed == null ? originalJobsById.get(jobId) : refreshed);
            }
        }
        return new StuckJobResetResult(reset, List.copyOf(exhaustedJobs));
    }

+    private Map<Long, TaskFileJobEntity> refreshJobsByIds(List<Long> jobIds) {
+        return taskFileJobMapper.selectBatchIds(jobIds).stream()
+                .collect(Collectors.toMap(TaskFileJobEntity::getId, job -> job, (a, b) -> a));
+    }
+
    public record StuckJobResetResult(int resetCount, List<TaskFileJobEntity> exhaustedJobs) {
    }

@@ -0,0 +1,220 @@
+package com.nanri.aiimage.modules.deletebrand.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.config.DeleteBrandProgressProperties;
+import com.nanri.aiimage.modules.task.mapper.FileTaskMapper;
+import com.nanri.aiimage.modules.task.model.entity.FileTaskEntity;
+import com.nanri.aiimage.modules.task.service.TaskDistributedLockService;
+import org.apache.ibatis.builder.MapperBuilderAssistant;
+import org.junit.jupiter.api.BeforeAll;
+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.test.util.ReflectionTestUtils;
+
+import java.time.LocalDateTime;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.mock;
+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-122 N+1 批量化修复：failStaleDeleteBrandTasks 的逐任务
+ * selectById 回读改为两阶段（全部 finalize 后一次 selectBatchIds 批量回读 + Map 装配）。
+ */
+@ExtendWith(MockitoExtension.class)
+class DeleteBrandStaleTaskServiceTest {
+
+    @Mock private FileTaskMapper fileTaskMapper;
+    @Mock private DeleteBrandTaskCacheService deleteBrandTaskCacheService;
+    @Mock private DeleteBrandTaskStorageService deleteBrandTaskStorageService;
+    @Mock private DeleteBrandRunService deleteBrandRunService;
+    @Mock private DeleteBrandProgressProperties deleteBrandProgressProperties;
+    @Mock private TaskDistributedLockService taskDistributedLockService;
+
+    @BeforeAll
+    static void initializeTableInfo() {
+        TableInfoHelper.initTableInfo(
+                new MapperBuilderAssistant(new MybatisConfiguration(), ""),
+                FileTaskEntity.class);
+    }
+
+    @Test
+    void staleFinalizedTasksAreBatchRefreshedAndCached() {
+        FileTaskEntity t1 = runningTask(101L);
+        FileTaskEntity t2 = runningTask(102L);
+        FileTaskEntity done1 = finishedTask(101L);
+        FileTaskEntity done2 = finishedTask(102L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1, t2));
+        when(fileTaskMapper.selectBatchIds(any())).thenReturn(List.of(done1, done2));
+        lockAvailable();
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(deleteBrandTaskCacheService, times(2)).saveTaskCache(any(FileTaskEntity.class));
+        verify(fileTaskMapper, never()).update(any(), any());
+        assertBatchIds(101L, 102L);
+    }
+
+    @Test
+    void staleSingleRowStillWorks() {
+        FileTaskEntity t1 = runningTask(201L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(fileTaskMapper.selectBatchIds(any())).thenReturn(List.of(finishedTask(201L)));
+        lockAvailable();
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(deleteBrandTaskCacheService).saveTaskCache(any(FileTaskEntity.class));
+        assertBatchIds(201L);
+    }
+
+    @Test
+    void staleEmptyCandidatesDoesNotRefresh() {
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of());
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(fileTaskMapper, times(1)).selectList(any());
+        verify(fileTaskMapper, never()).selectBatchIds(any());
+        verify(deleteBrandRunService, never()).tryFinalizeTask(anyLong(), anyBoolean());
+    }
+
+    @Test
+    void staleStillRunningTaskIsFailedWithCas() {
+        FileTaskEntity t1 = runningTask(301L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(fileTaskMapper.selectBatchIds(any())).thenReturn(List.of(runningTask(301L)));
+        when(fileTaskMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        lockAvailable();
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(deleteBrandTaskCacheService, never()).saveTaskCache(any(FileTaskEntity.class));
+        verify(deleteBrandTaskCacheService).delete(301L);
+        @SuppressWarnings({"rawtypes", "unchecked"})
+        ArgumentCaptor<LambdaUpdateWrapper> update = ArgumentCaptor.forClass(LambdaUpdateWrapper.class);
+        verify(fileTaskMapper).update(isNull(), update.capture());
+        assertTrue(update.getValue().getParamNameValuePairs().containsValue("FAILED"),
+                update.getValue().getSqlSet());
+    }
+
+    @Test
+    void staleFinalizeThrowsTaskIsFailedWithoutRefreshQuery() {
+        FileTaskEntity t1 = runningTask(401L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(fileTaskMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        org.mockito.Mockito.doThrow(new IllegalStateException("boom"))
+                .when(deleteBrandRunService).tryFinalizeTask(401L, true);
+        lockAvailable();
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(fileTaskMapper, times(1)).selectList(any());
+        verify(fileTaskMapper, never()).selectBatchIds(any());
+        verify(deleteBrandTaskCacheService).delete(401L);
+    }
+
+    @Test
+    void staleQueryCountReducedNoSelectById() {
+        FileTaskEntity t1 = runningTask(501L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(fileTaskMapper.selectBatchIds(any())).thenReturn(List.of(finishedTask(501L)));
+        lockAvailable();
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(fileTaskMapper, never()).selectById(any());
+        verify(fileTaskMapper, times(1)).selectList(any());
+        verify(fileTaskMapper, times(1)).selectBatchIds(any());
+    }
+
+    @Test
+    void staleTaskWithBusyLockIsSkipped() {
+        FileTaskEntity t1 = runningTask(601L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(taskDistributedLockService.acquire(any(), any(), anyLong())).thenReturn(null);
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(deleteBrandRunService, never()).tryFinalizeTask(anyLong(), anyBoolean());
+        verify(fileTaskMapper, times(1)).selectList(any());
+        verify(fileTaskMapper, never()).selectBatchIds(any());
+    }
+
+    @Test
+    void staleTaskWithRecentHeartbeatIsSkipped() {
+        FileTaskEntity t1 = runningTask(701L);
+        when(fileTaskMapper.selectList(any())).thenReturn(List.of(t1));
+        when(deleteBrandTaskCacheService.getProgress(701L))
+                .thenReturn(Map.of("last_heartbeat_at", System.currentTimeMillis()));
+        DeleteBrandStaleTaskService service = service();
+
+        failStaleDeleteBrandTasks(service);
+
+        verify(taskDistributedLockService, never()).acquire(any(), any(), anyLong());
+        verify(deleteBrandRunService, never()).tryFinalizeTask(anyLong(), anyBoolean());
+    }
+
+    private DeleteBrandStaleTaskService service() {
+        return new DeleteBrandStaleTaskService(
+                fileTaskMapper, deleteBrandTaskCacheService, deleteBrandTaskStorageService, deleteBrandRunService,
+                null, null, null, null, null, null, null, null, null, null, null, null, null, null, null,
+                deleteBrandProgressProperties, null, taskDistributedLockService, null);
+    }
+
+    private void lockAvailable() {
+        when(deleteBrandProgressProperties.getHeartbeatTimeoutMinutes()).thenReturn(15L);
+        when(taskDistributedLockService.acquire(any(), any(), anyLong()))
+                .thenReturn(mock(TaskDistributedLockService.LockHandle.class));
+    }
+
+    private static void failStaleDeleteBrandTasks(DeleteBrandStaleTaskService service) {
+        ReflectionTestUtils.invokeMethod(service, "failStaleDeleteBrandTasks");
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    private void assertBatchIds(Long... expected) {
+        ArgumentCaptor<java.util.Collection> captor = ArgumentCaptor.forClass(java.util.Collection.class);
+        verify(fileTaskMapper).selectBatchIds(captor.capture());
+        assertEquals(List.of(expected), List.copyOf(captor.getValue()));
+    }
+
+    private static FileTaskEntity runningTask(Long id) {
+        FileTaskEntity task = new FileTaskEntity();
+        task.setId(id);
+        task.setModuleType("DELETE_BRAND");
+        task.setStatus("RUNNING");
+        task.setUpdatedAt(LocalDateTime.now().minusHours(1));
+        task.setCreatedAt(LocalDateTime.now().minusDays(1));
+        return task;
+    }
+
+    private static FileTaskEntity finishedTask(Long id) {
+        FileTaskEntity task = runningTask(id);
+        task.setStatus("SUCCESS");
+        return task;
+    }
+}
@@ -86,9 +86,10 @@ class TaskFileJobClaimTest {
    }

    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);
+        when(taskFileJobMapper.selectBatchIds(any())).thenAnswer(invocation -> {
+            @SuppressWarnings("unchecked")
+            java.util.Collection<Object> ids = invocation.getArgument(0);
+            return claimed.stream().filter(job -> ids.contains(job.getId())).toList();
        });
    }

@@ -0,0 +1,217 @@
+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.dto.TaskFileJobDispatchEvent;
+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.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.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+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-122 N+1 批量化修复：claimCandidates 与 resetStuckRunningJobsDetailed
+ * 由逐行 selectById 改为 selectBatchIds 批量回读 + Map 装配（行为等价，查询次数下降）。
+ */
+@ExtendWith(MockitoExtension.class)
+class TaskFileJobN1BatchFixTest {
+
+    @Mock private TaskFileJobMapper taskFileJobMapper;
+    @Mock private ApplicationEventPublisher applicationEventPublisher;
+
+    @BeforeAll
+    static void initializeTableInfo() {
+        TableInfoHelper.initTableInfo(
+                new MapperBuilderAssistant(new MybatisConfiguration(), ""),
+                TaskFileJobEntity.class);
+    }
+
+    @Test
+    void claimCandidatesFetchesAllClaimedInOneBatchQuery() {
+        TaskFileJobEntity c1 = job(101L, "PENDING");
+        TaskFileJobEntity c2 = job(102L, "PENDING");
+        TaskFileJobEntity claimed1 = job(101L, "RUNNING");
+        TaskFileJobEntity claimed2 = job(102L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(c1, c2));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(claimed1, claimed2));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertEquals(List.of(101L, 102L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+        assertBatchIds(101L, 102L);
+    }
+
+    @Test
+    void claimCandidatesSingleRowStillWorks() {
+        TaskFileJobEntity c1 = job(201L, "PENDING");
+        TaskFileJobEntity claimed1 = job(201L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(c1));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(claimed1));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertEquals(List.of(201L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+        assertBatchIds(201L);
+    }
+
+    @Test
+    void claimCandidatesEmptyCandidatesDoesNotQuery() {
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of());
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertTrue(claimed.isEmpty());
+        verify(taskFileJobMapper, never()).update(any(), any());
+        verify(taskFileJobMapper, never()).selectBatchIds(any());
+    }
+
+    @Test
+    void claimCandidatesSkipsUnclaimedIdsFromBatch() {
+        TaskFileJobEntity c1 = job(301L, "PENDING");
+        TaskFileJobEntity c2 = job(302L, "PENDING");
+        TaskFileJobEntity claimed2 = job(302L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(c1, c2));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(0, 1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(claimed2));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertEquals(List.of(302L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+        assertBatchIds(302L);
+    }
+
+    @Test
+    void claimCandidatesNullGuardSkipsBrokenCandidates() {
+        TaskFileJobEntity broken = job(null, "PENDING");
+        TaskFileJobEntity c2 = job(402L, "PENDING");
+        TaskFileJobEntity claimed2 = job(402L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(broken, c2));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(claimed2));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertEquals(List.of(402L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+        assertBatchIds(402L);
+    }
+
+    @Test
+    void claimCandidatesQueryCountIsReducedToSingleBatch() {
+        TaskFileJobEntity c1 = job(501L, "PENDING");
+        TaskFileJobEntity c2 = job(502L, "PENDING");
+        TaskFileJobEntity c3 = job(503L, "PENDING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(c1, c2, c3));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any()))
+                .thenReturn(List.of(job(501L, "RUNNING"), job(502L, "RUNNING"), job(503L, "RUNNING")));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        service.claimRunnableJobs(20);
+
+        verify(taskFileJobMapper, times(1)).selectList(any());
+        verify(taskFileJobMapper, times(1)).selectBatchIds(any());
+        verify(taskFileJobMapper, never()).selectById(any());
+    }
+
+    @Test
+    void claimCandidatesDropsRowsThatAreNotRunningAfterRefresh() {
+        TaskFileJobEntity c1 = job(601L, "PENDING");
+        TaskFileJobEntity c2 = job(602L, "PENDING");
+        TaskFileJobEntity stale = job(601L, "PENDING");
+        TaskFileJobEntity claimed2 = job(602L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(c1, c2));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(stale, claimed2));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        List<TaskFileJobEntity> claimed = service.claimRunnableJobs(20);
+
+        assertEquals(List.of(602L), claimed.stream().map(TaskFileJobEntity::getId).toList());
+    }
+
+    @Test
+    void resetStuckBatchesRefreshAndPreservesEventOrder() {
+        TaskFileJobEntity zombie = job(701L, "PENDING");
+        TaskFileJobEntity requeue = job(702L, "RUNNING");
+        TaskFileJobEntity exhausted = job(703L, "FAILED");
+        exhausted.setRetryCount(TaskFileJobService.MAX_RETRY_COUNT);
+        TaskFileJobEntity zombieRefreshed = job(701L, "PENDING");
+        TaskFileJobEntity requeueRefreshed = job(702L, "PENDING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(zombie, requeue, exhausted));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(zombieRefreshed, requeueRefreshed));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        TaskFileJobService.StuckJobResetResult result = service.resetStuckRunningJobsDetailed(30, 20);
+
+        assertEquals(1, result.resetCount());
+        assertEquals(List.of(703L), result.exhaustedJobs().stream().map(TaskFileJobEntity::getId).toList());
+        ArgumentCaptor<Object> events = ArgumentCaptor.forClass(Object.class);
+        verify(applicationEventPublisher, times(2)).publishEvent(events.capture());
+        List<Object> published = events.getAllValues();
+        assertEquals(701L, ((TaskFileJobDispatchEvent) published.get(0)).jobId());
+        assertEquals(702L, ((TaskFileJobDispatchEvent) published.get(1)).jobId());
+        assertBatchIds(701L, 702L);
+    }
+
+    @Test
+    void resetStuckUsesSingleBatchRefreshWithoutSelectById() {
+        TaskFileJobEntity running = job(801L, "RUNNING");
+        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(running));
+        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(job(801L, "PENDING")));
+        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);
+
+        service.resetStuckRunningJobsDetailed(30, 20);
+
+        verify(taskFileJobMapper, times(1)).selectList(any());
+        verify(taskFileJobMapper, times(1)).selectBatchIds(any());
+        verify(taskFileJobMapper, never()).selectById(any());
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    private void assertBatchIds(Long... expected) {
+        ArgumentCaptor<java.util.Collection> captor = ArgumentCaptor.forClass(java.util.Collection.class);
+        verify(taskFileJobMapper).selectBatchIds(captor.capture());
+        assertEquals(List.of(expected), List.copyOf(captor.getValue()));
+    }
+
+    private static TaskFileJobEntity job(Long id, String status) {
+        TaskFileJobEntity job = new TaskFileJobEntity();
+        job.setId(id);
+        job.setTaskId(20553L);
+        job.setResultId(23110L);
+        job.setModuleType("SIMILAR_ASIN");
+        job.setJobType("ASSEMBLE_RESULT");
+        job.setStatus(status);
+        job.setRetryCount(0);
+        job.setUpdatedAt(LocalDateTime.now());
+        return job;
+    }
+}
@@ -92,9 +92,10 @@ class TaskFileJobOwnerColumnTest {
    }

    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);
+        when(taskFileJobMapper.selectBatchIds(any())).thenAnswer(invocation -> {
+            @SuppressWarnings("unchecked")
+            java.util.Collection<Object> ids = invocation.getArgument(0);
+            return claimed.stream().filter(job -> ids.contains(job.getId())).toList();
        });
    }

@@ -47,7 +47,7 @@ class TaskFileJobServiceTest {
        pending.setStatus("PENDING");
        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(running));
        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
-        when(taskFileJobMapper.selectById(101L)).thenReturn(pending);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(pending));
        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);

        TaskFileJobService.StuckJobResetResult result = service.resetStuckRunningJobsDetailed(30, 20);
@@ -69,7 +69,7 @@ class TaskFileJobServiceTest {
        failed.setErrorMessage("result file job timeout");
        when(taskFileJobMapper.selectList(any())).thenReturn(List.of(running));
        when(taskFileJobMapper.update(isNull(), any(LambdaUpdateWrapper.class))).thenReturn(1);
-        when(taskFileJobMapper.selectById(102L)).thenReturn(failed);
+        when(taskFileJobMapper.selectBatchIds(any())).thenReturn(List.of(failed));
        TaskFileJobService service = new TaskFileJobService(taskFileJobMapper, applicationEventPublisher);

        TaskFileJobService.StuckJobResetResult result = service.resetStuckRunningJobsDetailed(30, 20);