diff --git a/backend-java/src/main/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlTaskService.java b/backend-java/src/main/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlTaskService.java index 9c17d889..021ffe2e 100644 --- a/backend-java/src/main/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlTaskService.java +++ b/backend-java/src/main/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlTaskService.java @@ -1280,8 +1280,12 @@ public class ShopDataCrawlTaskService { chunk.setPayloadHash(payloadHash); chunk.setCreatedAt(now); chunk.setUpdatedAt(now); + int receivedChunkCount; + boolean insertedThisCall = false; try { taskChunkMapper.insert(chunk); + insertedThisCall = true; + receivedChunkCount = nextScopeChunkCount(scope, chunkTotal); } catch (DuplicateKeyException ex) { transientPayloadStorageService.deletePayloadIfPresent(storedPayload); TaskChunkEntity winner = findResultChunk(taskId, scopeHash, chunkIndex); @@ -1289,13 +1293,27 @@ public class ShopDataCrawlTaskService { throw new BusinessException("店铺数据抓取结果分片并发写入失败,请重试"); } validateExistingChunk(winner, chunkTotal, payloadHash); + receivedChunkCount = scopeChunkCount(scope); } catch (RuntimeException ex) { transientPayloadStorageService.deletePayloadIfPresent(storedPayload); throw ex; } - int receivedChunkCount = countResultChunks(taskId, scopeHash); - persistResultScope(taskId, scopeKey, scopeHash, chunkTotal, receivedChunkCount); + try { + persistResultScope(taskId, scopeKey, scopeHash, chunkTotal, receivedChunkCount); + } catch (RuntimeException ex) { + // 分片行与 scope 计数器是一致性单元:状态写入失败则回滚本次插入的分片行与 payload, + // 客户端重试会走全新插入路径,计数器不会因重试重复计数或永久欠计。 + if (insertedThisCall) { + taskChunkMapper.delete(new LambdaQueryWrapper() + .eq(TaskChunkEntity::getTaskId, taskId) + .eq(TaskChunkEntity::getModuleType, MODULE_TYPE) + .eq(TaskChunkEntity::getScopeHash, scopeHash) + .eq(TaskChunkEntity::getChunkIndex, chunkIndex)); + transientPayloadStorageService.deletePayloadIfPresent(storedPayload); + } + throw ex; + } log.info("[shop-data-crawl] result chunk received taskId={} shop={} chunk={}/{} received={}", taskId, shopKey, chunkIndex, chunkTotal, receivedChunkCount); return new ResultChunkReceipt(scopeHash, chunkTotal, receivedChunkCount >= chunkTotal); @@ -1343,12 +1361,23 @@ public class ShopDataCrawlTaskService { } } - private int countResultChunks(Long taskId, String scopeHash) { - Long count = taskChunkMapper.selectCount(new LambdaQueryWrapper() - .eq(TaskChunkEntity::getTaskId, taskId) - .eq(TaskChunkEntity::getModuleType, MODULE_TYPE) - .eq(TaskChunkEntity::getScopeHash, scopeHash)); - return count == null ? 0 : count.intValue(); + /** + * 以 scope 计数器替代 chunk 表全量 COUNT(*):新分片插入成功后 +1 并钳制到 chunk_total。 + * 任务级分布式锁串行化同一任务的接收,读改写在单锁内完成,无需再统计 chunk 行数。 + */ + private int nextScopeChunkCount(TaskScopeStateEntity scope, int chunkTotal) { + int previous = scope == null ? 0 : Math.max(0, scope.getReceivedChunkCount() == null ? 0 : scope.getReceivedChunkCount()); + return Math.min(chunkTotal, previous + 1); + } + + /** 重复提交(唯一键冲突,分片已计过数)时计数器保持不变,钳制到 chunk_total。 */ + private int scopeChunkCount(TaskScopeStateEntity scope) { + if (scope == null || scope.getReceivedChunkCount() == null) { + return 0; + } + int count = Math.max(0, scope.getReceivedChunkCount()); + Integer chunkTotal = scope.getChunkTotal(); + return chunkTotal != null && chunkTotal > 0 ? Math.min(chunkTotal, count) : count; } private void persistResultScope(Long taskId, diff --git a/backend-java/src/test/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlScopeCounterTest.java b/backend-java/src/test/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlScopeCounterTest.java new file mode 100644 index 00000000..0da7ca55 --- /dev/null +++ b/backend-java/src/test/java/com/nanri/aiimage/modules/shopdatacrawl/service/ShopDataCrawlScopeCounterTest.java @@ -0,0 +1,608 @@ +package com.nanri.aiimage.modules.shopdatacrawl.service; + +import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper; +import com.baomidou.mybatisplus.core.metadata.TableInfoHelper; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.nanri.aiimage.common.exception.BusinessException; +import com.nanri.aiimage.config.InstanceMetadata; +import com.nanri.aiimage.config.TaskPressureProperties; +import com.nanri.aiimage.modules.file.service.oss.OssStorageService; +import com.nanri.aiimage.modules.shopdatacrawl.model.dto.ShopDataCrawlCountryResultDto; +import com.nanri.aiimage.modules.shopdatacrawl.model.dto.ShopDataCrawlRowDto; +import com.nanri.aiimage.modules.shopdatacrawl.model.dto.ShopDataCrawlShopPayloadDto; +import com.nanri.aiimage.modules.shopdatacrawl.model.dto.ShopDataCrawlSubmitResultRequest; +import com.nanri.aiimage.modules.task.mapper.FileResultMapper; +import com.nanri.aiimage.modules.task.mapper.FileTaskMapper; +import com.nanri.aiimage.modules.task.mapper.TaskChunkMapper; +import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper; +import com.nanri.aiimage.modules.task.model.entity.FileResultEntity; +import com.nanri.aiimage.modules.task.model.entity.FileTaskEntity; +import com.nanri.aiimage.modules.task.model.entity.TaskChunkEntity; +import com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity; +import com.nanri.aiimage.modules.task.service.TaskDistributedLockService; +import com.nanri.aiimage.modules.task.service.TaskFileJobService; +import com.nanri.aiimage.modules.task.service.TaskProgressSnapshotService; +import com.nanri.aiimage.modules.task.service.TaskResultItemService; +import com.nanri.aiimage.modules.task.service.TransientPayloadStorageService; +import com.nanri.aiimage.modules.ziniao.service.ZiniaoShopSwitchService; +import org.apache.ibatis.builder.MapperBuilderAssistant; +import org.apache.ibatis.session.Configuration; +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.Mock; +import org.mockito.Spy; +import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.dao.DuplicateKeyException; + +import java.time.LocalDateTime; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +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.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.nullable; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Task 28:以 scope 计数器替代每个 chunk 的 COUNT(*) 完整统计。 + * 每个分片接收不再对 chunk 表执行 selectCount 全量统计,而是读取 + * biz_task_scope_state 的 received_chunk_count 计数器:新插入分片 +1、 + * 重复提交(唯一键冲突,chunk 已计过数)保持不变,并钳制到 chunk_total。 + * 任务级分布式锁串行化同一任务的接收,计数器读写安全。 + */ +@ExtendWith(MockitoExtension.class) +class ShopDataCrawlScopeCounterTest { + + private static final String MODULE_TYPE = "SHOP_DATA_CRAWL"; + private static final String SHOP_NAME = "Demo Shop"; + + @BeforeAll + static void initializeMybatisMetadata() { + MapperBuilderAssistant assistant = new MapperBuilderAssistant(new Configuration(), ""); + TableInfoHelper.initTableInfo(assistant, FileTaskEntity.class); + TableInfoHelper.initTableInfo(assistant, FileResultEntity.class); + TableInfoHelper.initTableInfo(assistant, TaskChunkEntity.class); + TableInfoHelper.initTableInfo(assistant, TaskScopeStateEntity.class); + } + + @Mock private FileTaskMapper fileTaskMapper; + @Mock private FileResultMapper fileResultMapper; + @Mock private ShopDataCrawlResolveService shopDataCrawlResolveService; + @Mock private ShopDataCrawlExcelAssemblyService excelAssemblyService; + @Mock private ShopDataCrawlTaskCacheService taskCacheService; + @Mock private OssStorageService ossStorageService; + @Mock private ZiniaoShopSwitchService ziniaoShopSwitchService; + @Mock private TaskPressureProperties taskPressureProperties; + @Mock private TaskFileJobService taskFileJobService; + @Mock private TaskResultItemService taskResultItemService; + @Mock private TaskProgressSnapshotService taskProgressSnapshotService; + @Mock private TaskDistributedLockService taskDistributedLockService; + @Mock private TaskChunkMapper taskChunkMapper; + @Mock private TaskScopeStateMapper taskScopeStateMapper; + @Mock private TransientPayloadStorageService transientPayloadStorageService; + @Mock private InstanceMetadata instanceMetadata; + @Mock private ShopDataCrawlDailyFileService dailyFileService; + @Spy private ObjectMapper objectMapper = new ObjectMapper().findAndRegisterModules(); + + private ShopDataCrawlTaskService service; + + private final List storedChunks = new ArrayList<>(); + private final List storedScopes = new ArrayList<>(); + private final Map rustfsPayloads = new LinkedHashMap<>(); + private FileTaskEntity task; + private FileResultEntity result; + private int nextPayloadId; + private boolean storeFails; + private boolean insertFails; + private boolean scopeUpdateFails; + + @BeforeEach + void configureStorage() { + service = new ShopDataCrawlTaskService( + fileTaskMapper, + fileResultMapper, + shopDataCrawlResolveService, + excelAssemblyService, + taskCacheService, + ossStorageService, + ziniaoShopSwitchService, + objectMapper, + taskPressureProperties, + taskFileJobService, + taskResultItemService, + taskProgressSnapshotService, + taskDistributedLockService, + taskChunkMapper, + taskScopeStateMapper, + transientPayloadStorageService, + instanceMetadata, + dailyFileService, + null); + + storedChunks.clear(); + storedScopes.clear(); + rustfsPayloads.clear(); + nextPayloadId = 0; + storeFails = false; + insertFails = false; + scopeUpdateFails = false; + + lenient().when(instanceMetadata.getInstanceId()).thenReturn("instance-a"); + lenient().when(taskDistributedLockService.acquire(eq(MODULE_TYPE), anyLong())) + .thenReturn(mock(TaskDistributedLockService.LockHandle.class)); + lenient().when(taskCacheService.getTaskCacheBatch(any())).thenReturn(Map.of()); + lenient().when(taskCacheService.getAllShopMergedPayload(anyLong())).thenReturn(Map.of()); + lenient().when(taskResultItemService.listResultSnapshots(anyLong(), eq(MODULE_TYPE), any())).thenReturn(List.of()); + lenient().when(taskFileJobService.findAssembleJob(anyLong(), eq(MODULE_TYPE), anyLong())).thenReturn(null); + lenient().when(excelAssemblyService.countRows(any())).thenReturn(1); + + lenient().when(fileTaskMapper.selectById(anyLong())).thenAnswer(invocation -> { + Long taskId = invocation.getArgument(0); + return task != null && Objects.equals(taskId, task.getId()) ? task : null; + }); + lenient().when(fileTaskMapper.updateById(any(FileTaskEntity.class))).thenReturn(1); + lenient().when(fileResultMapper.selectById(anyLong())).thenAnswer(invocation -> { + Long resultId = invocation.getArgument(0); + return result != null && Objects.equals(resultId, result.getId()) ? result : null; + }); + lenient().when(fileResultMapper.selectList(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + return result != null && Objects.equals(taskId, result.getTaskId()) ? List.of(result) : List.of(); + }); + lenient().when(fileResultMapper.updateById(any(FileResultEntity.class))).thenReturn(1); + + lenient().when(ziniaoShopSwitchService.normalizeShopName(nullable(String.class))).thenAnswer(invocation -> { + String value = invocation.getArgument(0); + return value == null ? "" : value.trim(); + }); + + configureTransientPayloadStorage(); + configureChunkMapper(); + configureScopeMapper(); + } + + @Test + void test_task_028_scope_counter_normal_default_path() { + // 正常输入:顺序接收 1/2、2/2,scope 计数器 1→2,齐集完成任务; + // 全程不执行 chunk 表 selectCount 全量统计。 + givenRunningTask(1101L, 2101L); + + service.submitResult(task.getId(), request(chunk(1, 2, "DE", row("2026-07-25", "B001")))); + + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "首个分片计数器为 1"); + assertEquals(0, storedScopes.get(0).getCompleted()); + assertEquals(-1, result.getSuccess()); + + service.submitResult(task.getId(), request(chunk(2, 2, "UK", row("2026-07-26", "B002")))); + + assertEquals(2, storedScopes.get(0).getReceivedChunkCount(), "齐集后计数器为 2"); + assertEquals(1, storedScopes.get(0).getCompleted()); + assertEquals(1, result.getSuccess()); + assertTrue(task.getResultJson().indexOf("B001") < task.getResultJson().indexOf("B002")); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_normal_multiple_items() { + // 批量场景:乱序分片 2/3、3/3、1/3,计数器按到达分片计数 1→2→3, + // 结果按 chunk_index 顺序合并,不丢失。 + givenRunningTask(1102L, 2102L); + + service.submitResult(task.getId(), request(chunk(2, 3, "DE", row("2026-07-25", "B002")))); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "乱序到达按已收分片计数"); + + service.submitResult(task.getId(), request(chunk(3, 3, "UK", row("2026-07-26", "B003")))); + assertEquals(2, storedScopes.get(0).getReceivedChunkCount()); + + service.submitResult(task.getId(), request(chunk(1, 3, "FR", row("2026-07-27", "B001")))); + + assertEquals(3, storedScopes.get(0).getReceivedChunkCount()); + assertEquals(1, storedScopes.get(0).getCompleted()); + assertEquals(1, result.getSuccess()); + String json = task.getResultJson(); + assertTrue(json.indexOf("B001") < json.indexOf("B002"), "合并顺序按 chunk_index"); + assertTrue(json.indexOf("B002") < json.indexOf("B003")); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_normal_repeated_operation_is_idempotent() { + // 重复执行:同一分片重试命中唯一键,已计过数的分片不再 +1, + // 计数器保持不变,无重复 scope 状态。 + givenRunningTask(1103L, 2103L); + ShopDataCrawlSubmitResultRequest request = request(chunk(1, 2, "DE", row("2026-07-25", "B001"))); + + service.submitResult(task.getId(), request); + service.submitResult(task.getId(), request); + + assertEquals(1, storedScopes.size(), "scope 状态只有一行"); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "重试不重复计数"); + assertEquals(0, storedScopes.get(0).getCompleted()); + assertEquals(1, storedChunks.size()); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_boundary_empty_input() { + // 空输入:无可处理数据的分片在落库前拒绝,不产生 payload/分片行/scope 状态, + // 也不触发任何计数统计。 + givenRunningTask(1104L, 2104L); + + ShopDataCrawlShopPayloadDto emptyItems = legacyChunk(false, "DE", null); + emptyItems.setChunkIndex(1); + emptyItems.setChunkTotal(1); + emptyItems.setCountryResults(List.of(countryWithItems("DE", List.of()))); + + BusinessException error = assertThrows(BusinessException.class, () -> + service.submitResult(task.getId(), request(emptyItems))); + assertTrue(error.getMessage().contains("内容为空")); + + assertEquals(0, storedChunks.size()); + assertEquals(0, rustfsPayloads.size()); + assertEquals(0, storedScopes.size()); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_boundary_single_item() { + // 单元素:单分片 1/1 到达即齐集,计数器直接达到 chunk_total 并完成。 + givenRunningTask(1105L, 2105L); + + service.submitResult(task.getId(), request(chunk(1, 1, "DE", row("2026-07-25", "B001")))); + + assertEquals(1, storedScopes.get(0).getReceivedChunkCount()); + assertEquals(1, storedScopes.get(0).getCompleted()); + assertEquals(1, result.getSuccess()); + assertTrue(task.getResultJson().contains("B001")); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_boundary_limit_and_overflow() { + // 上限/超限:已收分片在齐集前重放 → 唯一键冲突吸收、计数器不重复推进; + // 随后 5/5 全量到达,计数器达到 chunk_total 并完成,不溢出。 + givenRunningTask(1106L, 2106L); + + service.submitResult(task.getId(), request(chunk(1, 5, "UK", row("2026-07-25", "B001")))); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount()); + + service.submitResult(task.getId(), request(chunk(1, 5, "UK", row("2026-07-25", "B001")))); + + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "重放不重复计数"); + assertEquals(1, storedChunks.size(), "重放不产生新分片行"); + assertEquals(1, rustfsPayloads.size(), "重放重存的 payload 被清理"); + + String[] countries = {"UK", "DE", "FR", "ES", "IT"}; + String[] dates = {"2026-07-25", "2026-07-26", "2026-07-27", "2026-07-28", "2026-07-29"}; + for (int i = 2; i <= 5; i++) { + service.submitResult(task.getId(), request(chunk(i, 5, countries[i - 1], row(dates[i - 1], "B00" + i)))); + } + + assertEquals(5, storedScopes.get(0).getReceivedChunkCount(), "计数器与 chunk_total 一致"); + assertEquals(1, storedScopes.get(0).getCompleted()); + assertEquals(1, result.getSuccess()); + assertEquals(5, storedChunks.size()); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_invalid_input_rejected() { + // 非法参数:非正 chunk_index/chunk_total、跨分片改 chunk_total → + // 明确异常,计数器与 scope 状态不变,不触发计数统计。 + givenRunningTask(1107L, 2107L); + + BusinessException zeroIndex = assertThrows(BusinessException.class, () -> + service.submitResult(task.getId(), request(chunk(0, 1, "DE", row("2026-07-25", "B001"))))); + assertTrue(zeroIndex.getMessage().contains("chunk_index")); + + BusinessException zeroTotal = assertThrows(BusinessException.class, () -> + service.submitResult(task.getId(), request(chunk(1, 0, "DE", row("2026-07-25", "B001"))))); + assertTrue(zeroTotal.getMessage().contains("chunk_total")); + + service.submitResult(task.getId(), request(chunk(1, 2, "DE", row("2026-07-25", "B001")))); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount()); + + BusinessException changedTotal = assertThrows(BusinessException.class, () -> + service.submitResult(task.getId(), request(chunk(2, 3, "UK", row("2026-07-26", "B002"))))); + assertTrue(changedTotal.getMessage().contains("chunk_total")); + + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "非法输入不改变计数器"); + assertEquals(1, storedChunks.size()); + verify(taskChunkMapper, never()).selectCount(any()); + } + + @Test + void test_task_028_scope_counter_dependency_failure_releases_resources() { + // 依赖失败:分片插入异常 → payload 清理、无计数器更新,恢复后重试成功; + // scope 状态写入失败 → 已插入的分片行与 payload 被回滚补偿,恢复后重试为全新插入、计数器收敛。 + givenRunningTask(1108L, 2108L); + + insertFails = true; + assertThrows(RuntimeException.class, () -> + service.submitResult(task.getId(), request(chunk(1, 2, "DE", row("2026-07-25", "B001"))))); + assertEquals(0, storedChunks.size()); + assertEquals(0, rustfsPayloads.size(), "插入失败清理已存 payload"); + assertEquals(0, storedScopes.size(), "插入失败不更新计数器"); + insertFails = false; + + service.submitResult(task.getId(), request(chunk(1, 2, "DE", row("2026-07-25", "B001")))); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "故障恢复后重试成功"); + + scopeUpdateFails = true; + assertThrows(RuntimeException.class, () -> + service.submitResult(task.getId(), request(chunk(2, 2, "UK", row("2026-07-26", "B002"))))); + assertEquals(1, storedChunks.size(), "状态写入失败回滚已插入的分片行"); + assertEquals(1, rustfsPayloads.size(), "状态写入失败清理本次 payload"); + assertEquals(1, storedScopes.get(0).getReceivedChunkCount(), "状态写入失败计数器未推进"); + scopeUpdateFails = false; + + service.submitResult(task.getId(), request(chunk(2, 2, "UK", row("2026-07-26", "B002")))); + + assertEquals(1, result.getSuccess(), "状态写入恢复后重试成功"); + assertEquals(2, storedScopes.get(0).getReceivedChunkCount(), "重试为全新插入,计数器收敛到 chunk_total"); + assertEquals(1, storedScopes.get(0).getCompleted()); + verify(taskChunkMapper, never()).selectCount(any()); + } + + private void configureTransientPayloadStorage() { + lenient().when(transientPayloadStorageService.isSharedWriteEnabled()).thenReturn(true); + lenient().when(transientPayloadStorageService.storeChunkPayloadVersioned( + anyString(), anyLong(), anyString(), any(), anyString())).thenAnswer(invocation -> { + if (storeFails) { + throw new RuntimeException("rustfs store down"); + } + String pointer = "rustfs:payload-" + (++nextPayloadId); + rustfsPayloads.put(pointer, invocation.getArgument(4)); + return pointer; + }); + lenient().when(transientPayloadStorageService.extractPointer(any())).thenAnswer(invocation -> { + String value = invocation.getArgument(0); + if (value == null) { + return null; + } + return value.startsWith("rustfs:") || value.startsWith("local:") || value.startsWith("oss:") + ? value : null; + }); + lenient().when(transientPayloadStorageService.resolvePayload(any(), any())).thenAnswer(invocation -> { + String pointer = invocation.getArgument(0); + String payload = rustfsPayloads.get(pointer); + if (payload == null) { + throw new IllegalStateException("missing test RustFS payload: " + pointer); + } + return payload; + }); + lenient().doAnswer(invocation -> { + rustfsPayloads.remove(invocation.getArgument(0)); + return null; + }).when(transientPayloadStorageService).deletePayloadIfPresent(any()); + } + + private void configureChunkMapper() { + lenient().when(taskChunkMapper.insert(any(TaskChunkEntity.class))).thenAnswer(invocation -> { + if (insertFails) { + throw new RuntimeException("db down"); + } + TaskChunkEntity chunk = invocation.getArgument(0); + boolean duplicate = storedChunks.stream().anyMatch(existing -> + Objects.equals(existing.getTaskId(), chunk.getTaskId()) + && Objects.equals(existing.getScopeHash(), chunk.getScopeHash()) + && Objects.equals(existing.getChunkIndex(), chunk.getChunkIndex())); + if (duplicate) { + throw new DuplicateKeyException("duplicate chunk key: " + chunk.getScopeHash() + "/" + chunk.getChunkIndex()); + } + chunk.setId((long) storedChunks.size() + 1L); + storedChunks.add(chunk); + return 1; + }); + lenient().when(taskChunkMapper.selectOne(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + String scopeHash = queryScopeHash(query); + Integer chunkIndex = queryInteger(query); + return storedChunks.stream() + .filter(chunk -> matchesChunk(chunk, taskId, scopeHash, chunkIndex)) + .findFirst() + .orElse(null); + }); + lenient().when(taskChunkMapper.selectCount(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + String scopeHash = queryScopeHash(query); + return storedChunks.stream() + .filter(chunk -> matchesChunk(chunk, taskId, scopeHash, null)) + .count(); + }); + lenient().when(taskChunkMapper.selectList(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + String scopeHash = queryScopeHash(query); + return storedChunks.stream() + .filter(chunk -> matchesChunk(chunk, taskId, scopeHash, null)) + .sorted(Comparator.comparing(TaskChunkEntity::getChunkIndex)) + .toList(); + }); + lenient().when(taskChunkMapper.delete(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + String scopeHash = queryScopeHash(query); + Integer chunkIndex = queryInteger(query); + int before = storedChunks.size(); + storedChunks.removeIf(chunk -> matchesChunk(chunk, taskId, scopeHash, chunkIndex)); + return before - storedChunks.size(); + }); + } + + private void configureScopeMapper() { + lenient().when(taskScopeStateMapper.insert(any(TaskScopeStateEntity.class))).thenAnswer(invocation -> { + TaskScopeStateEntity scope = invocation.getArgument(0); + scope.setId((long) storedScopes.size() + 1L); + storedScopes.add(scope); + return 1; + }); + lenient().when(taskScopeStateMapper.selectOne(any())).thenAnswer(invocation -> { + LambdaQueryWrapper query = invocation.getArgument(0); + Long taskId = queryLong(query); + String scopeHash = queryScopeHash(query); + return storedScopes.stream() + .filter(scope -> Objects.equals(taskId, scope.getTaskId())) + .filter(scope -> scopeHash == null || Objects.equals(scopeHash, scope.getScopeHash())) + .findFirst() + .map(ShopDataCrawlScopeCounterTest::copyScope) + .orElse(null); + }); + lenient().when(taskScopeStateMapper.updateById(any(TaskScopeStateEntity.class))).thenAnswer(invocation -> { + if (scopeUpdateFails) { + throw new RuntimeException("scope update down"); + } + TaskScopeStateEntity updated = invocation.getArgument(0); + for (int i = 0; i < storedScopes.size(); i++) { + if (Objects.equals(storedScopes.get(i).getId(), updated.getId())) { + storedScopes.set(i, copyScope(updated)); + } + } + return 1; + }); + } + + private static TaskScopeStateEntity copyScope(TaskScopeStateEntity source) { + TaskScopeStateEntity copy = new TaskScopeStateEntity(); + copy.setId(source.getId()); + copy.setTaskId(source.getTaskId()); + copy.setModuleType(source.getModuleType()); + copy.setScopeKey(source.getScopeKey()); + copy.setScopeHash(source.getScopeHash()); + copy.setChunkTotal(source.getChunkTotal()); + copy.setReceivedChunkCount(source.getReceivedChunkCount()); + copy.setCompleted(source.getCompleted()); + copy.setLastChunkAt(source.getLastChunkAt()); + copy.setLastError(source.getLastError()); + copy.setCreatedAt(source.getCreatedAt()); + copy.setUpdatedAt(source.getUpdatedAt()); + return copy; + } + + private boolean matchesChunk(TaskChunkEntity chunk, Long taskId, String scopeHash, Integer chunkIndex) { + return (taskId == null || Objects.equals(taskId, chunk.getTaskId())) + && (scopeHash == null || Objects.equals(scopeHash, chunk.getScopeHash())) + && (chunkIndex == null || Objects.equals(chunkIndex, chunk.getChunkIndex())); + } + + private Long queryLong(LambdaQueryWrapper query) { + query.getSqlSegment(); + return query.getParamNameValuePairs().values().stream() + .filter(Long.class::isInstance) + .map(Long.class::cast) + .findFirst() + .orElse(null); + } + + private Integer queryInteger(LambdaQueryWrapper query) { + query.getSqlSegment(); + return query.getParamNameValuePairs().values().stream() + .filter(Integer.class::isInstance) + .map(Integer.class::cast) + .findFirst() + .orElse(null); + } + + private String queryScopeHash(LambdaQueryWrapper query) { + query.getSqlSegment(); + return query.getParamNameValuePairs().values().stream() + .filter(String.class::isInstance) + .map(String.class::cast) + .filter(value -> value.length() == 64) + .findFirst() + .orElse(null); + } + + private void givenRunningTask(long taskId, long resultId) { + task = new FileTaskEntity(); + task.setId(taskId); + task.setUserId(7L); + task.setModuleType(MODULE_TYPE); + task.setStatus("RUNNING"); + task.setRequestJson("{\"ownerInstanceId\":\"instance-a\"}"); + task.setResultJson("[]"); + task.setCreatedAt(LocalDateTime.now()); + task.setUpdatedAt(LocalDateTime.now()); + + result = new FileResultEntity(); + result.setId(resultId); + result.setTaskId(taskId); + result.setUserId(7L); + result.setModuleType(MODULE_TYPE); + result.setSourceFilename(SHOP_NAME); + result.setSourceFileUrl("shop-1"); + result.setSuccess(-1); + result.setCreatedAt(LocalDateTime.now()); + } + + private ShopDataCrawlSubmitResultRequest request(ShopDataCrawlShopPayloadDto payload) { + ShopDataCrawlSubmitResultRequest request = new ShopDataCrawlSubmitResultRequest(); + request.setShops(List.of(payload)); + return request; + } + + private ShopDataCrawlShopPayloadDto chunk(int chunkIndex, + int chunkTotal, + String country, + ShopDataCrawlRowDto row) { + ShopDataCrawlShopPayloadDto payload = legacyChunk(false, country, row); + payload.setChunkIndex(chunkIndex); + payload.setChunkTotal(chunkTotal); + return payload; + } + + private ShopDataCrawlShopPayloadDto legacyChunk(boolean shopDone, + String country, + ShopDataCrawlRowDto row) { + ShopDataCrawlShopPayloadDto payload = new ShopDataCrawlShopPayloadDto(); + payload.setShopName(SHOP_NAME); + if (row != null) { + payload.setCountryResults(List.of(country(country, row))); + } + payload.setShopDone(shopDone); + return payload; + } + + private ShopDataCrawlCountryResultDto country(String country, ShopDataCrawlRowDto row) { + return countryWithItems(country, List.of(row)); + } + + private ShopDataCrawlCountryResultDto countryWithItems(String country, List items) { + ShopDataCrawlCountryResultDto result = new ShopDataCrawlCountryResultDto(); + result.setCountry(country); + result.setItems(items); + return result; + } + + private ShopDataCrawlRowDto row(String date, String asin) { + ShopDataCrawlRowDto row = new ShopDataCrawlRowDto(); + row.setDate(date); + row.setAsin(asin); + row.setCommodityImage("https://m.media-amazon.com/images/I/" + asin + ".jpg"); + row.setInventorySales("10"); + row.setSalesRank("20"); + row.setPageViews("30"); + row.setUnitsSold("40"); + row.setPrice("50"); + row.setRecommendedOffer("60"); + return row; + } +}