task-72: 文件作业数据库原子 claim，避免重复派发同一 job
This commit is contained in:
+212
@@ -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 成功");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user