task-10: chunk 结果建立 rowKey 批量索引,消除跨 chunk 线性扫描

mergeCozeRowsIntoChunk 先 indexRowsByChunkKey 建立 rowKey→chunkKey
索引,assignCozeRowsToChunks 按 O(1) 查找分配行归属,保留原有
命中/fallback/orphan 语义与顺序稳定性,每 chunk 只读一次 payload。
This commit is contained in:
2026-08-29 16:07:45 +08:00
parent 17a3292c78
commit 1ccf4f74ef
2 changed files with 450 additions and 29 deletions
@@ -0,0 +1,375 @@
package com.nanri.aiimage.modules.similarasin.service;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.nanri.aiimage.common.exception.BusinessException;
import com.nanri.aiimage.config.SimilarAsinProperties;
import com.nanri.aiimage.modules.similarasin.client.SimilarAsinCozeClient;
import com.nanri.aiimage.modules.similarasin.mapper.SimilarAsinFilterConditionMapper;
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinResultRowDto;
import com.nanri.aiimage.modules.similarasin.util.SimilarAsinImageEmbedder;
import com.nanri.aiimage.modules.file.service.LocalFileStorageService;
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.FileTaskEntity;
import com.nanri.aiimage.modules.task.model.entity.TaskChunkEntity;
import com.nanri.aiimage.modules.task.service.TransientPayloadStorageService;
import com.nanri.aiimage.modules.similarasin.service.SimilarAsinImagePrefetchService;
import com.baomidou.mybatisplus.core.MybatisConfiguration;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import org.apache.ibatis.builder.MapperBuilderAssistant;
import org.junit.jupiter.api.AfterEach;
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.InjectMocks;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.junit.jupiter.MockitoExtension;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
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.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/**
* Task 10:为 chunk 结果建立按 row key 的批量索引,消除跨 chunk 线性扫描。
* indexRowsByChunkKey 把每个 chunk 的行索引到 rowKey→chunkKeycoze 行归属从
* O(rows×chunks) 降为 O(1) 查找;assignCozeRowsToChunks 基于索引分配行并保留
* 原有命中/fallback/orphan 语义;集成用例验证每个 chunk 只读一次 payload。
*/
@ExtendWith(MockitoExtension.class)
class SimilarAsinTaskServiceRowKeyIndexTest {
private static final AtomicLong NEXT_ID = new AtomicLong(60000);
@Mock private LocalFileStorageService localFileStorageService;
@Mock private com.nanri.aiimage.modules.file.service.oss.OssStorageService ossStorageService;
@Mock private com.nanri.aiimage.config.StorageProperties storageProperties;
@Mock private FileTaskMapper fileTaskMapper;
@Mock private FileResultMapper fileResultMapper;
@Mock private TaskScopeStateMapper taskScopeStateMapper;
@Mock private TaskChunkMapper taskChunkMapper;
@Mock private SimilarAsinFilterConditionMapper filterConditionMapper;
@Spy private ObjectMapper objectMapper = new ObjectMapper();
@Mock private SimilarAsinCozeClient cozeClient;
@Mock private SimilarAsinTaskCacheService taskCacheService;
@Mock private SimilarAsinProperties properties;
@Mock private com.nanri.aiimage.modules.task.service.TaskFileJobService taskFileJobService;
@Mock private com.nanri.aiimage.modules.task.service.TaskDistributedLockService taskDistributedLockService;
@Mock private com.nanri.aiimage.modules.task.service.TaskProgressSnapshotService taskProgressSnapshotService;
@Mock private TransientPayloadStorageService transientPayloadStorageService;
@Mock private org.springframework.transaction.PlatformTransactionManager transactionManager;
@Mock private com.nanri.aiimage.common.service.DistributedJobLockService distributedJobLockService;
@Mock private com.nanri.aiimage.config.InstanceMetadata instanceMetadata;
@Mock private com.nanri.aiimage.modules.coze.service.CozeCredentialPoolService cozeCredentialPoolService;
@Mock private SimilarAsinImageEmbedder imageEmbedder;
@Mock private SimilarAsinImagePrefetchService imagePrefetchService;
@InjectMocks private SimilarAsinTaskService service;
@BeforeAll
static void initializeMybatisMetadata() {
MapperBuilderAssistant assistant = new MapperBuilderAssistant(new MybatisConfiguration(), "");
TableInfoHelper.initTableInfo(assistant, FileTaskEntity.class);
TableInfoHelper.initTableInfo(assistant, TaskChunkEntity.class);
}
@BeforeEach
void setUp() {
lenient().when(properties.isBoundedResultAssemblyEnabled()).thenReturn(true);
lenient().when(properties.getImageCacheMaxBytes()).thenReturn(256L * 1024L * 1024L);
lenient().when(properties.getResultFileTimeoutMinutes()).thenReturn(90);
lenient().when(properties.getParseResponsePreviewLimit()).thenReturn(100);
lenient().when(properties.getMaxSourceFileBytes()).thenReturn(50L * 1024L * 1024L);
lenient().when(properties.getMaxParseRows()).thenReturn(50000);
lenient().when(properties.getMaxFieldLength()).thenReturn(2000);
lenient().when(properties.getMaxWorkbookZipEntries()).thenReturn(20000);
lenient().when(properties.getMaxWorkbookUncompressedBytes()).thenReturn(512L * 1024L * 1024L);
lenient().when(transientPayloadStorageService.storeParsedPayloadFast(
eq(SimilarAsinTaskService.MODULE_TYPE), any(), anyString(), anyString(), eq(false)))
.thenReturn("rustfs:task-parsed/similar-asin/60000/payload.json");
lenient().doAnswer(invocation -> {
FileTaskEntity task = invocation.getArgument(0);
task.setId(NEXT_ID.incrementAndGet());
return 1;
}).when(fileTaskMapper).insert(any(FileTaskEntity.class));
lenient().when(fileResultMapper.insert(any(com.nanri.aiimage.modules.task.model.entity.FileResultEntity.class))).thenReturn(1);
lenient().when(taskScopeStateMapper.insert(any(com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity.class))).thenReturn(1);
lenient().when(taskChunkMapper.selectList(any())).thenReturn(List.of());
lenient().when(fileResultMapper.selectList(any())).thenReturn(List.of());
}
@AfterEach
void shutdown() {
service.shutdownAssembleExecutor();
}
private static SimilarAsinResultRowDto row(String rowToken, String id, String asin, String country) {
SimilarAsinResultRowDto r = new SimilarAsinResultRowDto();
r.setRowToken(rowToken);
r.setId(id);
r.setAsin(asin);
r.setCountry(country);
return r;
}
private static String rowsJson(List<SimilarAsinResultRowDto> rows) throws Exception {
return new ObjectMapper().writeValueAsString(rows);
}
private static TaskChunkEntity chunk(Long id, String scopeHash, Integer chunkIndex, String payloadJson) {
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setId(id);
chunk.setTaskId(7004L);
chunk.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
chunk.setScopeHash(scopeHash);
chunk.setChunkIndex(chunkIndex);
chunk.setPayloadJson(payloadJson);
return chunk;
}
/** 构造 rowsByChunkchunkStorageKey(scopeHash, chunkIndex) → rowKey 行表。 */
private static Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunkOf(String scopeHash, Integer chunkIndex, List<SimilarAsinResultRowDto> rows) {
Map<String, Map<String, SimilarAsinResultRowDto>> map = new LinkedHashMap<>();
Map<String, SimilarAsinResultRowDto> byKey = new LinkedHashMap<>();
for (SimilarAsinResultRowDto row : rows) {
byKey.put(row.getRowToken(), row);
}
map.put(scopeHash + ":" + chunkIndex, byKey);
return map;
}
private static List<String> assignedRowKeys(Map<String, Map<String, SimilarAsinResultRowDto>> merged) {
List<String> keys = new ArrayList<>();
for (Map<String, SimilarAsinResultRowDto> rows : merged.values()) {
for (String key : rows.keySet()) {
keys.add(key);
}
}
return keys;
}
@Test
void test_task_010_chunk_row_key_normal_default_path() throws Exception {
// 正常输入:2 个 chunk 各含行,coze 回传行按 rowKey 命中各自 chunk
// 每个 chunk 的 payload 只被读取一次(索引建立),消除跨 chunk 线性扫描。
List<TaskChunkEntity> chunks = List.of(
chunk(1L, "hashA", 1, "ptr:chunk-A"),
chunk(2L, "hashB", 2, "ptr:chunk-B"));
when(taskChunkMapper.selectList(any())).thenReturn(chunks);
when(transientPayloadStorageService.resolvePayload(eq("ptr:chunk-A"), anyString()))
.thenReturn(rowsJson(List.of(row("r1", "1", "B0A0000001", "英国"), row("r2", "2", "B0A0000002", "英国"))));
when(transientPayloadStorageService.resolvePayload(eq("ptr:chunk-B"), anyString()))
.thenReturn(rowsJson(List.of(row("r3", "3", "B0A0000003", "美国"))));
when(transientPayloadStorageService.storeChunkPayloadVersioned(anyString(), any(), anyString(), any(), anyString()))
.thenAnswer(invocation -> "stored:" + invocation.getArgument(2));
AtomicInteger selectOneRound = new AtomicInteger(0);
when(taskChunkMapper.selectOne(any())).thenAnswer(invocation -> {
int i = selectOneRound.getAndIncrement();
return chunks.get(Math.min(i, chunks.size() - 1));
});
when(taskChunkMapper.update(any(), any())).thenReturn(1);
FileTaskEntity task = new FileTaskEntity();
task.setId(7004L);
List<SimilarAsinResultRowDto> cozeRows = List.of(row("r1", "1", "B0A0000001", "英国"), row("r3", "3", "B0A0000003", "美国"));
Method merge = SimilarAsinTaskService.class.getDeclaredMethod("mergeCozeRowsIntoChunk",
FileTaskEntity.class, String.class, Integer.class, List.class, Map.class);
merge.setAccessible(true);
merge.invoke(service, task, null, null, cozeRows, Map.of());
verify(transientPayloadStorageService, times(6)).resolvePayload(anyString(), anyString());
verify(transientPayloadStorageService, times(2)).storeChunkPayloadVersioned(anyString(), any(), anyString(), any(), anyString());
verify(taskChunkMapper, times(2)).update(any(), any());
}
@Test
void test_task_010_chunk_row_key_normal_multiple_items() {
// 批量场景:3 个 chunk 各 3 行,9 个 coze 回传行全部命中且顺序稳定,无 orphan
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = new LinkedHashMap<>();
for (int c = 0; c < 3; c++) {
rowsByChunk.putAll(rowsByChunkOf("hash" + c, c + 1,
List.of(row("c" + c + "r1", "1", "B0B" + c + "000001", "英国"),
row("c" + c + "r2", "2", "B0B" + c + "000002", "英国"),
row("c" + c + "r3", "3", "B0B" + c + "000003", "美国"))));
}
List<SimilarAsinResultRowDto> cozeRows = new ArrayList<>();
for (int c = 0; c < 3; c++) {
for (int r = 1; r <= 3; r++) {
cozeRows.add(row("c" + c + "r" + r, String.valueOf(r), "B0B" + c + "00000" + r, r == 3 ? "美国" : "英国"));
}
}
List<SimilarAsinResultRowDto> orphans = new ArrayList<>();
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
Map<String, Map<String, SimilarAsinResultRowDto>> merged = service.assignCozeRowsToChunks(
rowsByChunk, cozeRows, index, null, null, orphans);
assertEquals(3, merged.size());
assertEquals(9, assignedRowKeys(merged).size());
assertTrue(orphans.isEmpty(), "全部命中,不应产生 orphan");
for (Map<String, SimilarAsinResultRowDto> rows : merged.values()) {
assertEquals(3, rows.size());
}
}
@Test
void test_task_010_chunk_row_key_normal_repeated_operation_is_idempotent() {
// 重复执行同一输入:结果完全一致,不产生重复记录
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = rowsByChunkOf("hashA", 1,
List.of(row("r1", "1", "B0A0000001", "英国"), row("r2", "2", "B0A0000002", "英国")));
List<SimilarAsinResultRowDto> cozeRows = List.of(row("r1", "1", "B0A0000001", "英国"));
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
Map<String, Map<String, SimilarAsinResultRowDto>> first = service.assignCozeRowsToChunks(
rowsByChunk, cozeRows, index, null, null, new ArrayList<>());
Map<String, Map<String, SimilarAsinResultRowDto>> second = service.assignCozeRowsToChunks(
rowsByChunk, cozeRows, index, null, null, new ArrayList<>());
assertEquals(assignedRowKeys(first), assignedRowKeys(second));
assertEquals(first.size(), second.size());
for (Map.Entry<String, Map<String, SimilarAsinResultRowDto>> entry : first.entrySet()) {
assertEquals(entry.getValue().keySet(), second.get(entry.getKey()).keySet());
}
}
@Test
void test_task_010_chunk_row_key_boundary_empty_input() {
// 空输入:null/空 rowsByChunk 与 cozeRows 均安全返回空结果,不创建无效资源
assertNotNull(service.indexRowsByChunkKey(null));
assertTrue(service.indexRowsByChunkKey(null).isEmpty());
assertTrue(service.indexRowsByChunkKey(Map.of()).isEmpty());
Map<String, Map<String, SimilarAsinResultRowDto>> emptyAssign = service.assignCozeRowsToChunks(
Map.of(), List.of(), Map.of(), null, null, new ArrayList<>());
assertTrue(emptyAssign.isEmpty());
assertTrue(service.assignCozeRowsToChunks(
Map.of(), null, Map.of(), null, null, new ArrayList<>()).isEmpty());
// 无可匹配行(rowKey 不存在于任何 chunk)→ 进 orphan 兜底,不产生 merge
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = rowsByChunkOf("hashA", 1, List.of(row("r1", "1", "B0A0000001", "英国")));
List<SimilarAsinResultRowDto> blankRow = List.of(row("", "", "", ""));
List<SimilarAsinResultRowDto> orphans = new ArrayList<>();
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
Map<String, Map<String, SimilarAsinResultRowDto>> merged = service.assignCozeRowsToChunks(
rowsByChunk, blankRow, index, null, null, orphans);
assertTrue(assignedRowKeys(merged).isEmpty());
assertEquals(1, orphans.size(), "全空行生成 legacy key :::: 不命中任何 chunk,按既有语义进 orphan");
}
@Test
void test_task_010_chunk_row_key_boundary_single_item() {
// 单 chunk 单行:不依赖批量路径,命中正确
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = rowsByChunkOf("hashA", 1,
List.of(row("r1", "1", "B0A0000001", "英国")));
List<SimilarAsinResultRowDto> orphans = new ArrayList<>();
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
Map<String, Map<String, SimilarAsinResultRowDto>> merged = service.assignCozeRowsToChunks(
rowsByChunk, List.of(row("r1", "1", "B0A0000001", "英国")), index, null, null, orphans);
assertEquals(1, merged.size());
assertEquals(List.of("r1"), assignedRowKeys(merged));
assertTrue(orphans.isEmpty());
// 索引也只含该行
assertEquals(1, index.size());
assertEquals("hashA:1", index.get("r1"));
}
@Test
void test_task_010_chunk_row_key_boundary_limit_and_overflow() {
// 大批量:1000 行索引 + 500 个 coze 回传行全部命中,行不丢、无 orphan
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = new LinkedHashMap<>();
Map<String, SimilarAsinResultRowDto> bigChunk = new LinkedHashMap<>();
for (int i = 1; i <= 1000; i++) {
bigChunk.put("r" + String.format("%04d", i), row("r" + String.format("%04d", i), String.valueOf(i), "B0L" + String.format("%06d", i), "英国"));
}
rowsByChunk.put("hashBig:1", bigChunk);
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
assertEquals(1000, index.size());
List<SimilarAsinResultRowDto> cozeRows = new ArrayList<>();
for (int i = 1; i <= 500; i++) {
cozeRows.add(row("r" + String.format("%04d", i), String.valueOf(i), "B0L" + String.format("%06d", i), "英国"));
}
List<SimilarAsinResultRowDto> orphans = new ArrayList<>();
Map<String, Map<String, SimilarAsinResultRowDto>> merged = service.assignCozeRowsToChunks(
rowsByChunk, cozeRows, index, null, null, orphans);
assertEquals(1, merged.size());
assertEquals(500, assignedRowKeys(merged).size());
assertTrue(orphans.isEmpty());
}
@Test
void test_task_010_chunk_row_key_invalid_input_rejected() {
// 同一 rowKey 出现在多个 chunk:索引保留第一个 chunkputIfAbsent),行为确定
Map<String, Map<String, SimilarAsinResultRowDto>> rowsByChunk = new LinkedHashMap<>();
rowsByChunk.putAll(rowsByChunkOf("hashA", 1, List.of(row("dup", "1", "B0A0000001", "英国"))));
rowsByChunk.putAll(rowsByChunkOf("hashB", 2, List.of(row("dup", "1", "B0A0000001", "英国"))));
Map<String, String> index = service.indexRowsByChunkKey(rowsByChunk);
assertEquals("hashA:1", index.get("dup"), "重复 rowKey 应保留第一个 chunk");
// fallback 缺失:coze 行未命中且无有效 fallback → 进 orphan,不产生 merge
List<SimilarAsinResultRowDto> orphans = new ArrayList<>();
Map<String, Map<String, SimilarAsinResultRowDto>> merged = service.assignCozeRowsToChunks(
rowsByChunk, List.of(row("ghost", "9", "B0A0000009", "英国")), Map.of(), "missingHash", 99, orphans);
assertTrue(assignedRowKeys(merged).isEmpty());
assertEquals(1, orphans.size());
assertEquals("ghost", orphans.get(0).getRowToken());
// cozeRows 含 null 元素:跳过不抛异常,其余行正常分配
List<SimilarAsinResultRowDto> withNull = new ArrayList<>();
withNull.add(null);
withNull.add(row("dup", "1", "B0A0000001", "英国"));
List<SimilarAsinResultRowDto> orphans2 = new ArrayList<>();
Map<String, String> index2 = service.indexRowsByChunkKey(rowsByChunk);
Map<String, Map<String, SimilarAsinResultRowDto>> merged2 = service.assignCozeRowsToChunks(
rowsByChunk, withNull, index2, null, null, orphans2);
assertEquals(1, merged2.size());
assertEquals(List.of("dup"), assignedRowKeys(merged2));
assertTrue(orphans2.isEmpty());
}
@Test
void test_task_010_chunk_row_key_dependency_failure_releases_resources() throws Exception {
// chunk payload 读取失败:抛可识别业务异常且不产生部分 merge;
// 依赖恢复后重试成功,无残留状态
List<TaskChunkEntity> chunks = List.of(chunk(1L, "hashA", 1, "ptr:chunk-A"));
when(taskChunkMapper.selectList(any())).thenReturn(chunks);
when(transientPayloadStorageService.resolvePayload(eq("ptr:chunk-A"), anyString()))
.thenThrow(new IllegalStateException("rustfs down"));
FileTaskEntity task = new FileTaskEntity();
task.setId(7004L);
Method merge = SimilarAsinTaskService.class.getDeclaredMethod("mergeCozeRowsIntoChunk",
FileTaskEntity.class, String.class, Integer.class, List.class, Map.class);
merge.setAccessible(true);
BusinessException ex = assertThrows(BusinessException.class, () -> {
try {
merge.invoke(service, task, null, null, List.of(row("r1", "1", "B0A0000001", "英国")), Map.of());
} catch (java.lang.reflect.InvocationTargetException e) {
throw e.getCause();
}
});
assertTrue(ex.getMessage() != null && ex.getMessage().contains("chunk"),
"chunk 读取失败消息必须可识别,实际: " + ex.getMessage());
verify(transientPayloadStorageService, times(0)).storeChunkPayloadVersioned(anyString(), any(), anyString(), any(), anyString());
// 恢复后重试成功:行合并到正确 chunk
when(transientPayloadStorageService.resolvePayload(eq("ptr:chunk-A"), anyString()))
.thenReturn(rowsJson(List.of(row("r1", "1", "B0A0000001", "英国"))));
when(transientPayloadStorageService.storeChunkPayloadVersioned(anyString(), any(), anyString(), any(), anyString()))
.thenReturn("stored:retry");
when(taskChunkMapper.selectOne(any())).thenReturn(chunks.get(0));
when(taskChunkMapper.update(any(), any())).thenReturn(1);
merge.invoke(service, task, null, null, List.of(row("r1", "1", "B0A0000001", "英国")), Map.of());
verify(transientPayloadStorageService, times(1)).storeChunkPayloadVersioned(anyString(), any(), anyString(), any(), anyString());
}
}