Skip to content

Commit 98d49ba

Browse files
ltrinthewindltrinthewind
committed
fix: prevent races between document chunking and deletion
1 parent 146850f commit 98d49ba

17 files changed

Lines changed: 807 additions & 44 deletions

File tree

‎bootstrap/pom.xml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,18 @@
104104
<scope>test</scope>
105105
</dependency>
106106

107+
<dependency>
108+
<groupId>org.testcontainers</groupId>
109+
<artifactId>junit-jupiter</artifactId>
110+
<scope>test</scope>
111+
</dependency>
112+
113+
<dependency>
114+
<groupId>org.testcontainers</groupId>
115+
<artifactId>postgresql</artifactId>
116+
<scope>test</scope>
117+
</dependency>
118+
107119
<dependency>
108120
<groupId>org.springframework.boot</groupId>
109121
<artifactId>spring-boot-starter-validation</artifactId>

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/core/ingest/DefaultIngestionKernel.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import java.util.HashMap;
3737
import java.util.List;
3838
import java.util.Map;
39+
import java.util.function.BiConsumer;
3940

4041
/**
4142
* 摄取内核默认实现：固定五步骨架，全文唯一一条摄取执行序列
@@ -66,7 +67,8 @@ public class DefaultIngestionKernel implements IngestionKernel {
6667
public IngestionOutcome run(DocumentRef doc,
6768
byte[] bytes,
6869
IngestionSpec spec,
69-
VectorTarget target) {
70+
VectorTarget target,
71+
BiConsumer<String, Integer> beforeIndexCommit) {
7072
if (bytes == null || bytes.length == 0) {
7173
throw new ClientException("文件内容为空：docId=" + doc.docId());
7274
}
@@ -103,7 +105,8 @@ public IngestionOutcome run(DocumentRef doc,
103105

104106
// ⑤ index：扇出到全部落点，事务边界在写入器内
105107
long indexStart = System.currentTimeMillis();
106-
chunkIndexWriter.replaceDocument(target, doc, embedded);
108+
chunkIndexWriter.replaceDocument(target, doc, embedded,
109+
() -> beforeIndexCommit.accept(mimeType, embedded.size()));
107110
long indexMillis = System.currentTimeMillis() - indexStart;
108111

109112
return new IngestionOutcome(mimeType, parser.getParserType(), blocks.size(), chunks,

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/core/ingest/IngestionKernel.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
package com.nageoffer.ai.ragent.core.ingest;
1919

20+
import java.util.function.BiConsumer;
21+
2022
/**
2123
* 摄取内核：固定五步骨架，调用方不可跳过、不可换序、不可替换
2224
* <pre>
@@ -37,10 +39,12 @@ public interface IngestionKernel {
3739
* @param bytes 文件字节
3840
* @param spec 文档级配置：解析档位 + 分块预算
3941
* @param target 向量落点：逻辑分区 + 嵌入模型 + 维度
42+
* @param beforeIndexCommit 全部落点写入成功后、索引事务提交前执行的回调，入参为 MIME 和块数
4043
* @return 摄取结果
4144
*/
4245
IngestionOutcome run(DocumentRef doc,
4346
byte[] bytes,
4447
IngestionSpec spec,
45-
VectorTarget target);
48+
VectorTarget target,
49+
BiConsumer<String, Integer> beforeIndexCommit);
4650
}

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/core/ingest/sink/ChunkIndexWriter.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,12 @@ public class ChunkIndexWriter {
4343
/**
4444
* 整体替换该文档的块：全部落点在同一个事务里
4545
*/
46-
public void replaceDocument(VectorTarget target, DocumentRef doc, List<EmbeddedChunk> chunks) {
47-
transactionOperations.executeWithoutResult(status ->
48-
sinks.forEach(sink -> sink.replaceDocument(target, doc, chunks)));
46+
public void replaceDocument(VectorTarget target, DocumentRef doc, List<EmbeddedChunk> chunks,
47+
Runnable beforeCommit) {
48+
transactionOperations.executeWithoutResult(status -> {
49+
sinks.forEach(sink -> sink.replaceDocument(target, doc, chunks));
50+
beforeCommit.run();
51+
});
4952
log.info("块索引写入完成 docId={} 分区={} 块数={} 落点数={}",
5053
doc.docId(), target.partition(), chunks.size(), sinks.size());
5154
}

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/dao/entity/KnowledgeDocumentDO.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ public class KnowledgeDocumentDO {
136136
* - running：向量化中
137137
* - failed：向量化失败
138138
* - success：向量化完成
139+
* - deleting：删除中
139140
*/
140141
private String status;
141142

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/dao/mapper/KnowledgeDocumentMapper.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,46 @@
1919

2020
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
2121
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeDocumentDO;
22+
import org.apache.ibatis.annotations.Param;
23+
import org.apache.ibatis.annotations.Select;
24+
import org.apache.ibatis.annotations.Update;
25+
26+
import java.util.List;
2227

2328
public interface KnowledgeDocumentMapper extends BaseMapper<KnowledgeDocumentDO> {
29+
30+
/**
31+
* 仅当文档未删除且当前状态命中允许集合时，原子转换状态。
32+
*/
33+
@Update("<script>"
34+
+ "UPDATE t_knowledge_document "
35+
+ "SET status = #{toStatus}, updated_by = #{updatedBy}, update_time = NOW() "
36+
+ "WHERE id = #{docId} AND deleted = 0 AND status IN "
37+
+ "<foreach collection='fromStatuses' item='source' open='(' close=')' separator=','>"
38+
+ "#{source}"
39+
+ "</foreach>"
40+
+ "</script>")
41+
int casStatus(@Param("docId") String docId,
42+
@Param("fromStatuses") List<String> fromStatuses,
43+
@Param("toStatus") String toStatus,
44+
@Param("updatedBy") String updatedBy);
45+
46+
/**
47+
* 分块成功收尾：状态和块数在同一条 CAS SQL 中写回。
48+
*/
49+
@Update("UPDATE t_knowledge_document "
50+
+ "SET status = #{successStatus}, chunk_count = #{chunkCount}, "
51+
+ "updated_by = #{updatedBy}, update_time = NOW() "
52+
+ "WHERE id = #{docId} AND deleted = 0 AND status = #{runningStatus}")
53+
int markSuccessIfRunning(@Param("docId") String docId,
54+
@Param("chunkCount") int chunkCount,
55+
@Param("updatedBy") String updatedBy,
56+
@Param("runningStatus") String runningStatus,
57+
@Param("successStatus") String successStatus);
58+
59+
/**
60+
* 写块前在当前事务内锁定文档行。
61+
*/
62+
@Select("SELECT status FROM t_knowledge_document WHERE id = #{docId} FOR UPDATE")
63+
String selectStatusForUpdate(@Param("docId") String docId);
2464
}

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/enums/DocumentStatus.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,12 @@ public enum DocumentStatus {
4747
/**
4848
* 文档处理成功
4949
*/
50-
SUCCESS("success");
50+
SUCCESS("success"),
51+
52+
/**
53+
* 文档删除中
54+
*/
55+
DELETING("deleting");
5156

5257
/**
5358
* 状态码

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/schedule/DocumentStatusHelper.java‎

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,16 @@
3636
public class DocumentStatusHelper {
3737

3838
private static final String SYSTEM_USER = "system";
39+
private static final List<String> CHUNKABLE_STATUSES = List.of(
40+
DocumentStatus.PENDING.getCode(),
41+
DocumentStatus.FAILED.getCode(),
42+
DocumentStatus.SUCCESS.getCode()
43+
);
44+
private static final List<String> DELETABLE_STATUSES = List.of(
45+
DocumentStatus.PENDING.getCode(),
46+
DocumentStatus.FAILED.getCode(),
47+
DocumentStatus.SUCCESS.getCode()
48+
);
3949

4050
private final KnowledgeDocumentMapper documentMapper;
4151

@@ -49,18 +59,49 @@ public boolean tryMarkRunning(String docId) {
4959
.eq(KnowledgeDocumentDO::getId, docId)
5060
.eq(KnowledgeDocumentDO::getDeleted, 0)
5161
.eq(KnowledgeDocumentDO::getEnabled, 1)
52-
.ne(KnowledgeDocumentDO::getStatus, DocumentStatus.RUNNING.getCode())
62+
.in(KnowledgeDocumentDO::getStatus, CHUNKABLE_STATUSES)
63+
) > 0;
64+
}
65+
66+
public boolean tryStartChunk(String docId, String updatedBy) {
67+
return documentMapper.casStatus(
68+
docId,
69+
CHUNKABLE_STATUSES,
70+
DocumentStatus.RUNNING.getCode(),
71+
updatedBy
72+
) > 0;
73+
}
74+
75+
public boolean tryMarkDeleting(String docId, String updatedBy) {
76+
return documentMapper.casStatus(
77+
docId,
78+
DELETABLE_STATUSES,
79+
DocumentStatus.DELETING.getCode(),
80+
updatedBy
81+
) > 0;
82+
}
83+
84+
public boolean tryMarkSuccess(String docId, int chunkCount, String updatedBy) {
85+
return documentMapper.markSuccessIfRunning(
86+
docId,
87+
chunkCount,
88+
updatedBy,
89+
DocumentStatus.RUNNING.getCode(),
90+
DocumentStatus.SUCCESS.getCode()
91+
) > 0;
92+
}
93+
94+
public boolean tryMarkFailed(String docId, String updatedBy) {
95+
return documentMapper.casStatus(
96+
docId,
97+
List.of(DocumentStatus.RUNNING.getCode()),
98+
DocumentStatus.FAILED.getCode(),
99+
updatedBy
53100
) > 0;
54101
}
55102

56103
public void markFailedIfRunning(String docId) {
57-
documentMapper.update(
58-
Wrappers.lambdaUpdate(KnowledgeDocumentDO.class)
59-
.set(KnowledgeDocumentDO::getStatus, DocumentStatus.FAILED.getCode())
60-
.set(KnowledgeDocumentDO::getUpdatedBy, SYSTEM_USER)
61-
.eq(KnowledgeDocumentDO::getId, docId)
62-
.eq(KnowledgeDocumentDO::getStatus, DocumentStatus.RUNNING.getCode())
63-
);
104+
tryMarkFailed(docId, SYSTEM_USER);
64105
}
65106

66107
public void applyRefreshedFileMetadata(String docId, StoredFileDTO stored) {

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/service/impl/KnowledgeDocumentServiceImpl.java‎

Lines changed: 32 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
import com.nageoffer.ai.ragent.knowledge.handler.RemoteFileFetcher;
7272
import com.nageoffer.ai.ragent.knowledge.mq.event.KnowledgeDocumentChunkEvent;
7373
import com.nageoffer.ai.ragent.knowledge.schedule.CronScheduleHelper;
74+
import com.nageoffer.ai.ragent.knowledge.schedule.DocumentStatusHelper;
7475
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeChunkService;
7576
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeDocumentScheduleService;
7677
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeDocumentService;
@@ -127,6 +128,7 @@ public class KnowledgeDocumentServiceImpl implements KnowledgeDocumentService {
127128
private final RemoteFileFetcher remoteFileFetcher;
128129
private final VectorTargetResolver vectorTargetResolver;
129130
private final BizChangeLogContext bizChangeLogContext;
131+
private final DocumentStatusHelper documentStatusHelper;
130132

131133
@Value("knowledge-document-chunk_topic${unique-name:}")
132134
private String chunkTopic;
@@ -210,19 +212,10 @@ public void startChunk(String docId) {
210212
"文档分块",
211213
event,
212214
arg -> {
213-
// Wrapper 更新不触发 updateTime 自动填充, 显式刷新, 使卡死恢复以分块开始时刻为基准
214-
int updated = documentMapper.update(
215-
new LambdaUpdateWrapper<KnowledgeDocumentDO>()
216-
.set(KnowledgeDocumentDO::getStatus, DocumentStatus.RUNNING.getCode())
217-
.set(KnowledgeDocumentDO::getUpdatedBy, event.getOperator())
218-
.set(KnowledgeDocumentDO::getUpdateTime, new Date())
219-
.eq(KnowledgeDocumentDO::getId, docId)
220-
.ne(KnowledgeDocumentDO::getStatus, DocumentStatus.RUNNING.getCode())
221-
);
222-
if (updated == 0) {
215+
if (!documentStatusHelper.tryStartChunk(docId, event.getOperator())) {
223216
KnowledgeDocumentDO documentDO = documentMapper.selectById(docId);
224217
Assert.notNull(documentDO, () -> new ClientException("文档不存在"));
225-
throw new ClientException("文档分块操作正在进行中，请稍后再试");
218+
throw new ClientException("文档状态已变更，无法开始分块");
226219
}
227220
KnowledgeDocumentDO documentDO = documentMapper.selectById(docId);
228221
event.setKbId(documentDO.getKbId());
@@ -239,6 +232,10 @@ public void executeChunk(String docId) {
239232
log.warn("文档不存在，跳过分块任务, docId={}", docId);
240233
return;
241234
}
235+
if (!DocumentStatus.RUNNING.getCode().equals(documentDO.getStatus())) {
236+
log.warn("文档已非运行状态，跳过分块任务, docId={}, status={}", docId, documentDO.getStatus());
237+
return;
238+
}
242239

243240
runChunkTask(documentDO);
244241
}
@@ -282,16 +279,23 @@ private void runChunkTask(KnowledgeDocumentDO documentDO) {
282279
throw new ClientException("管道模式重构中，暂不可用，请改用直接分块：docId=" + docId);
283280
}
284281

285-
IngestionOutcome outcome = ingestionKernel.run(doc, readFileBytes(documentDO), spec, target);
282+
IngestionOutcome outcome = ingestionKernel.run(
283+
doc,
284+
readFileBytes(documentDO),
285+
spec,
286+
target,
287+
// 所有落点成功后、释放行锁前回填 MIME、块数和 SUCCESS，避免超时恢复抢占提交窗口，
288+
// 导致 chunk 和向量已经写入成功，但文档状态被改为 FAILED。
289+
(detectedMimeType, savedCount) -> {
290+
refreshMimeType(docId, detectedMimeType);
291+
markChunkSucceeded(docId, savedCount);
292+
}
293+
);
286294
extractDuration = outcome.timings().parseMillis();
287295
chunkDuration = outcome.timings().chunkMillis();
288296
embedDuration = outcome.timings().embedMillis();
289297
persistDuration = outcome.timings().indexMillis();
290298
int savedCount = outcome.chunkCount();
291-
// 回填字节探测出的真实 MIME；展示用的 file_type 仍由扩展名决定，两者互不导出
292-
refreshMimeType(docId, outcome.mimeType());
293-
294-
markChunkSucceeded(docId, savedCount);
295299
long totalDuration = System.currentTimeMillis() - totalStartTime;
296300
updateChunkLog(chunkLog.getId(), DocumentStatus.SUCCESS.getCode(), savedCount,
297301
extractDuration, chunkDuration, embedDuration, persistDuration, totalDuration, null);
@@ -309,12 +313,9 @@ private DocumentRef documentRef(KnowledgeDocumentDO documentDO) {
309313
}
310314

311315
private void markChunkSucceeded(String docId, int chunkCount) {
312-
documentMapper.updateById(KnowledgeDocumentDO.builder()
313-
.id(docId)
314-
.chunkCount(chunkCount)
315-
.status(DocumentStatus.SUCCESS.getCode())
316-
.updatedBy(UserContext.getUsername())
317-
.build());
316+
if (!documentStatusHelper.tryMarkSuccess(docId, chunkCount, UserContext.getUsername())) {
317+
throw new IllegalStateException("文档运行权已失效，无法写回分块成功状态: docId=" + docId);
318+
}
318319
}
319320

320321
private void refreshMimeType(String docId, String mimeType) {
@@ -404,11 +405,9 @@ public void chunkDocument(KnowledgeDocumentDO documentDO) {
404405

405406
private void markChunkFailed(String docId) {
406407
transactionOperations.executeWithoutResult(status -> {
407-
KnowledgeDocumentDO update = new KnowledgeDocumentDO();
408-
update.setId(docId);
409-
update.setStatus(DocumentStatus.FAILED.getCode());
410-
update.setUpdatedBy(UserContext.getUsername());
411-
documentMapper.updateById(update);
408+
if (!documentStatusHelper.tryMarkFailed(docId, UserContext.getUsername())) {
409+
log.warn("文档运行权已失效，跳过分块失败状态写回: docId={}", docId);
410+
}
412411
});
413412
}
414413

@@ -429,11 +428,16 @@ public void delete(String docId) {
429428
bizChangeLogContext.putName(documentDO.getDocName());
430429
KnowledgeDocumentDO before = BeanUtil.copyProperties(documentDO, KnowledgeDocumentDO.class);
431430

432-
// 禁止在文档分块运行时删除
431+
// 常见路径快速失败，避免分块写入持有文档行锁时，删除请求等待写入事务结束后才失败。
433432
if (DocumentStatus.RUNNING.getCode().equals(documentDO.getStatus())) {
434433
throw new ClientException("文档正在分块中，无法删除");
435434
}
436435

436+
// 原子抢占删除权；分块或其它删除流程已抢占时立即失败。
437+
if (!documentStatusHelper.tryMarkDeleting(docId, UserContext.getUsername())) {
438+
throw new ClientException("文档正在分块或已被处理，无法删除");
439+
}
440+
437441
scheduleService.deleteByDocId(docId);
438442
chunkLogMapper.delete(Wrappers.lambdaQuery(KnowledgeDocumentChunkLogDO.class)
439443
.eq(KnowledgeDocumentChunkLogDO::getDocId, docId));

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/sink/RelationalChunkSink.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import com.nageoffer.ai.ragent.infra.token.TokenCounterService;
2828
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeChunkDO;
2929
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeChunkMapper;
30+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeDocumentMapper;
31+
import com.nageoffer.ai.ragent.knowledge.enums.DocumentStatus;
3032
import lombok.RequiredArgsConstructor;
3133
import lombok.extern.slf4j.Slf4j;
3234
import org.springframework.core.Ordered;
@@ -51,9 +53,19 @@ public class RelationalChunkSink implements ChunkSink {
5153

5254
private final KnowledgeChunkMapper chunkMapper;
5355
private final TokenCounterService tokenCounterService;
56+
private final KnowledgeDocumentMapper documentMapper;
5457

5558
@Override
5659
public void replaceDocument(VectorTarget target, DocumentRef doc, List<EmbeddedChunk> chunks) {
60+
// 分块可能执行超过恢复阈值：恢复任务将 RUNNING 改为 FAILED 后，用户可以取得删除权；
61+
// 此时原分块任务仍可能在执行，并在完成 Embedding 后重新写入数据。
62+
// 这里通过持有文档行锁，使删除状态 CAS 等待到分块写入事务结束，避免删除过程中被旧任务重新写入数据。
63+
String status = documentMapper.selectStatusForUpdate(doc.docId());
64+
if (!DocumentStatus.RUNNING.getCode().equals(status)) {
65+
throw new IllegalStateException(
66+
"文档状态已变更，放弃分块写入: docId=" + doc.docId() + ", status=" + status
67+
);
68+
}
5769
deleteDocument(target, doc);
5870
if (chunks.isEmpty()) {
5971
return;

0 commit comments

Comments
 (0)