diff --git a/backend-java/src/main/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskService.java b/backend-java/src/main/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskService.java index 044d51a1..3781cc73 100644 --- a/backend-java/src/main/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskService.java +++ b/backend-java/src/main/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskService.java @@ -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 entries = new ArrayList<>(); for (FileTaskEntity task : runningTasks) { Map 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 refreshTaskIds = entries.stream() + .filter(entry -> !entry.finalizeThrew) + .map(entry -> entry.task.getId()) + .toList(); + Map 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() .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()); diff --git a/backend-java/src/main/java/com/nanri/aiimage/modules/task/service/TaskFileJobService.java b/backend-java/src/main/java/com/nanri/aiimage/modules/task/service/TaskFileJobService.java index 490e0bce..86d18129 100644 --- a/backend-java/src/main/java/com/nanri/aiimage/modules/task/service/TaskFileJobService.java +++ b/backend-java/src/main/java/com/nanri/aiimage/modules/task/service/TaskFileJobService.java @@ -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 claimed = new ArrayList<>(); + List 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 claimedById = refreshJobsByIds(claimedIds); + List 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 exhaustedJobs = new ArrayList<>(); + List eventJobIds = new ArrayList<>(); + List exhaustedJobIds = new ArrayList<>(); + Map originalExhaustedById = new LinkedHashMap<>(); + Map 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 refreshedById = refreshJobsByIds(eventJobIds); + for (Long jobId : eventJobIds) { + publishDispatchEvent(refreshedById.get(jobId)); + } + } + if (!exhaustedJobIds.isEmpty()) { + List refreshIds = exhaustedJobIds.stream() + .filter(id -> !originalExhaustedById.containsKey(id)) + .toList(); + Map 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 refreshJobsByIds(List jobIds) { + return taskFileJobMapper.selectBatchIds(jobIds).stream() + .collect(Collectors.toMap(TaskFileJobEntity::getId, job -> job, (a, b) -> a)); + } + public record StuckJobResetResult(int resetCount, List exhaustedJobs) { } diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskServiceTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskServiceTest.java new file mode 100644 index 00000000..838c3b76 --- /dev/null +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/deletebrand/service/DeleteBrandStaleTaskServiceTest.java @@ -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 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 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; + } +} diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobClaimTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobClaimTest.java index 39e407d8..ac751b44 100644 --- a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobClaimTest.java +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobClaimTest.java @@ -86,9 +86,10 @@ class TaskFileJobClaimTest { } private void stubClaimedRows(List 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 ids = invocation.getArgument(0); + return claimed.stream().filter(job -> ids.contains(job.getId())).toList(); }); } diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobN1BatchFixTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobN1BatchFixTest.java new file mode 100644 index 00000000..6a8554d7 --- /dev/null +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobN1BatchFixTest.java @@ -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 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 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 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 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 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 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 events = ArgumentCaptor.forClass(Object.class); + verify(applicationEventPublisher, times(2)).publishEvent(events.capture()); + List 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 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; + } +} diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobOwnerColumnTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobOwnerColumnTest.java index 4c703b62..11e91f35 100644 --- a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobOwnerColumnTest.java +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobOwnerColumnTest.java @@ -92,9 +92,10 @@ class TaskFileJobOwnerColumnTest { } private void stubClaimedRows(List 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 ids = invocation.getArgument(0); + return claimed.stream().filter(job -> ids.contains(job.getId())).toList(); }); } diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobServiceTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobServiceTest.java index 8b068194..eff68a26 100644 --- a/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobServiceTest.java +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/task/service/TaskFileJobServiceTest.java @@ -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);