Browse Source

fix: 附件链路多用户并发竞态修复(接口级 check-then-act 清理)

多用户并发访问同一接口时的 5 处竞态:
1. 草稿附件绑定改 CAS 条件更新(session_id 仍为 0 才允许绑):
   并发抢绑只有一方成功,失败方重读归属后拒绝 —— 修掉同一附件
   被两个会话同时引用的 check-then-act 漏洞
2. purgeOrphans 改 DB 条件删除先行(删行条件 = 仍未关联且超期)、
   文件/向量后置清理:与发送并发时不会删掉刚被消息引用的附件
3. AttachmentDocIndexer 并发计数 int[]++ 丢更新 → AtomicInteger
4. 并发首建向量表 PG 报错兜底(CREATE TABLE IF NOT EXISTS 并发撞
   pg_class 唯一索引是已知行为):失败后确认表存在即视为成功
5. session.messageCount 读改写丢更新(既有)→ 原子自增 setSql

新增并发单测 2 例(CAS 抢绑互斥:胜者放行/败者 400),18/18 全绿
cc 1 week ago
parent
commit
49d453ae6b

+ 14 - 5
ai-server/src/main/java/com/zsjz/ai/module/agent/attachment/AttachmentDocIndexer.java

@@ -59,7 +59,15 @@ public class AttachmentDocIndexer {
             throw new IllegalStateException("未配置嵌入模型(model 表需存在 type='embedding' 的活跃记录),无法索引大文档");
         }
         String table = vectorTableName(caseId, spec.dimensions());
-        vectorMapper.createVectorTable(table, spec.dimensions());
+        // 并发首建兜底:PG 的 CREATE TABLE IF NOT EXISTS 并发时会撞 pg_class 唯一索引报错
+        // (多用户同时上传大文档是常态),失败后确认表已存在即视为成功
+        try {
+            vectorMapper.createVectorTable(table, spec.dimensions());
+        } catch (Exception e) {
+            if (vectorMapper.tableExists(table) == 0) {
+                throw new IllegalStateException("附件向量表创建失败: " + e.getMessage());
+            }
+        }
 
         List<String> chunks = chunk(fullText, props.getChunkSize(), props.getChunkOverlap());
         boolean truncated = chunks.size() > props.getMaxChunks();
@@ -71,7 +79,8 @@ public class AttachmentDocIndexer {
 
         Semaphore gate = new Semaphore(Math.max(1, props.getIndexConcurrency()), true);
         List<Runnable> jobs = new ArrayList<>(chunks.size());
-        int[] ok = {0};
+        // 并发累加必须用原子计数(int[] ++ 在多线程下丢更新,块数会少报)
+        java.util.concurrent.atomic.AtomicInteger ok = new java.util.concurrent.atomic.AtomicInteger();
         for (int i = 0; i < chunks.size(); i++) {
             int chunkNo = i;
             String content = chunks.get(i);
@@ -92,7 +101,7 @@ public class AttachmentDocIndexer {
                     payload.put("chunkNo", chunkNo);
                     vectorMapper.insertVectorDoc(table, attachmentId + ":" + chunkNo, content,
                             Json.toStr(payload), toVectorLiteral(vec));
-                    ok[0]++;
+                    ok.incrementAndGet();
                 } catch (Exception e) {
                     log.warn("附件 chunk 向量化失败: attachmentId={}, chunk={}, err={}",
                             attachmentId, chunkNo, e.getMessage());
@@ -113,10 +122,10 @@ public class AttachmentDocIndexer {
                 break;
             }
         }
-        if (ok[0] == 0) {
+        if (ok.get() == 0) {
             throw new IllegalStateException("附件全部 chunk 向量化失败,请检查嵌入服务");
         }
-        return new IndexResult(ok[0], truncated);
+        return new IndexResult(ok.get(), truncated);
     }
 
     /**

+ 4 - 0
ai-server/src/main/java/com/zsjz/ai/module/agent/attachment/AttachmentVectorMapper.java

@@ -26,6 +26,10 @@ public interface AttachmentVectorMapper {
             + "embedding vector(${dims}))")
     void createVectorTable(@Param("tableName") String tableName, @Param("dims") int dims);
 
+    /** 表是否已存在(并发建表兜底用) */
+    @Select("SELECT count(*) FROM information_schema.tables WHERE table_name = #{tableName}")
+    int tableExists(@Param("tableName") String tableName);
+
     @Insert("INSERT INTO public.${tableName} (doc_id, content, payload, embedding) "
             + "VALUES (#{docId}, #{content}, #{payloadJson}::jsonb, #{embedding}::vector) "
             + "ON CONFLICT (doc_id) DO NOTHING")

+ 26 - 4
ai-server/src/main/java/com/zsjz/ai/module/agent/attachment/ChatAttachmentService.java

@@ -227,12 +227,24 @@ public class ChatAttachmentService {
                 throw new ServerException(400, "附件不属于当前案件,拒绝访问: " + att.getFileName());
             }
             if (att.getSessionId() == null || att.getSessionId() == 0L) {
-                // 草稿期上传的附件:发送时绑定到当前会话(此后按会话归属校验)
-                attachmentMapper.update(null, Wrappers.<ChatAttachment>lambdaUpdate()
+                // 草稿期附件绑定:CAS 条件更新(session_id 仍是 0 才允许绑)——
+                // 多个会话/多端并发抢同一草稿附件时只有一方能绑成功,
+                // 后来者拿到 0 行更新就必须重读归属,绝不能沿用旧快照放行
+                int updated = attachmentMapper.update(null, Wrappers.<ChatAttachment>lambdaUpdate()
                         .eq(ChatAttachment::getId, att.getId())
+                        .and(w -> w.isNull(ChatAttachment::getSessionId).or().eq(ChatAttachment::getSessionId, 0L))
                         .set(ChatAttachment::getSessionId, sessionId)
                         .set(ChatAttachment::getUpdateTime, LocalDateTime.now()));
-                att.setSessionId(sessionId);
+                if (updated == 0) {
+                    // 被并发请求抢绑了:重读真实归属再判一次
+                    ChatAttachment fresh = attachmentMapper.selectById(att.getId());
+                    if (fresh == null || !sessionId.equals(fresh.getSessionId())) {
+                        throw new ServerException(400, "附件已被其他消息使用,拒绝重复访问: " + att.getFileName());
+                    }
+                    att = fresh;
+                } else {
+                    att.setSessionId(sessionId);
+                }
             } else if (!att.getSessionId().equals(sessionId)) {
                 throw new ServerException(400, "附件不属于当前会话,拒绝访问: " + att.getFileName());
             }
@@ -378,10 +390,20 @@ public class ChatAttachmentService {
                 .lt(ChatAttachment::getCreateTime, deadline));
         for (ChatAttachment att : orphans) {
             try {
+                // 条件删除先行(DB 原子判):仍须「未关联消息且超期」才真正删行 ——
+                // select 与删除之间用户可能刚刚发送(linkToMessage 抢先落库),
+                // 无条件删会把正在被消息引用的附件清掉,发送方就会撞上文件缺失
+                int deleted = attachmentMapper.delete(Wrappers.<ChatAttachment>lambdaQuery()
+                        .eq(ChatAttachment::getId, att.getId())
+                        .isNull(ChatAttachment::getLinkMessageId)
+                        .lt(ChatAttachment::getCreateTime, deadline));
+                if (deleted == 0) {
+                    continue;   // 刚被关联/刚被处理,不删
+                }
+                // 行已删才动文件与向量(失败只影响磁盘残留,不影响数据正确性)
                 docIndexer.deleteAttachmentVectors(att.getId(), att.getCaseId());
                 deleteQuietly(att.getFilePath());
                 deleteQuietly(att.getSidecarPath());
-                attachmentMapper.deleteById(att.getId());
             } catch (Exception e) {
                 log.warn("孤儿附件清理失败: id={}, err={}", att.getId(), e.getMessage());
             }

+ 3 - 2
ai-server/src/main/java/com/zsjz/ai/module/agent/service/impl/AgentChatServiceImpl.java

@@ -586,13 +586,14 @@ public class AgentChatServiceImpl extends ServiceImpl<AgentChatSessionMapper, Ag
             // 附件与消息关联(孤儿清理依据)
             chatAttachmentService.linkToMessage(attachmentIds, userMessage.getId());
 
-            int current = (session != null && session.getMessageCount() != null) ? session.getMessageCount() : 0;
+            // 原子自增而不是读改写(current+1):多用户/多端并发向同一会话发消息时
+            // 读改写会丢更新,计数越用越小
             chatSessionMapper.update(
                     null,
                     new LambdaUpdateWrapper<AgentChatSession>()
                             .eq(AgentChatSession::getId, sessionId)
                             .set(AgentChatSession::getLastMessageAt, LocalDateTime.now())
-                            .set(AgentChatSession::getMessageCount, current + 1)
+                            .setSql("message_count = COALESCE(message_count, 0) + 1")
                             .set(AgentChatSession::getUpdateAt, LocalDateTime.now()));
             log.info("用户提问落库成功: sessionId={}, messageId={}", sessionId, userMessage.getId());
         } catch (Exception e) {

+ 25 - 1
ai-server/src/test/java/com/zsjz/ai/module/agent/attachment/ChatAttachmentServiceTest.java

@@ -84,9 +84,10 @@ class ChatAttachmentServiceTest {
     @DisplayName("草稿期附件(sessionId=0)发送时绑定到当前会话")
     void bindsDraftAttachmentToSession() {
         when(mapper.selectById(10L)).thenReturn(att(10L, 1L, 1L, 0L, "TABLE"));
+        when(mapper.update(any(), any())).thenReturn(1);   // CAS 绑定成功
         List<String> blocks = service.buildBlocks(List.of(10L), 1L, 1L, 55L, "问题");
         assertEquals(1, blocks.size());
-        // 绑定动作发生在加载阶段(update 换 sessionId)
+        // 绑定动作发生在加载阶段(CAS 条件更新换 sessionId)
         verify(mapper).update(any(), any());
     }
 
@@ -155,6 +156,29 @@ class ChatAttachmentServiceTest {
         assertEquals("attachment", ChatAttachmentService.sanitizeFileName(""));
     }
 
+    @Test
+    @DisplayName("并发抢绑同一草稿附件:CAS 失败方重读后必须拒绝(多用户多端互斥)")
+    void concurrentBindOnlyOneWins() {
+        ChatAttachment draft = att(11L, 1L, 1L, 0L, "TABLE");
+        when(mapper.selectById(11L)).thenReturn(draft);
+        // 模拟对方抢先绑定:本次 CAS 更新影响 0 行,重读后归属已是对方的会话
+        when(mapper.update(any(), any())).thenReturn(0);
+        ChatAttachment stolen = att(11L, 1L, 1L, 777L, "TABLE");
+        when(mapper.selectById(11L)).thenReturn(draft, stolen);
+        ServerException e = assertThrows(ServerException.class,
+                () -> service.buildBlocks(List.of(11L), 1L, 1L, 55L, "问题"));
+        assertTrue(e.getMessage().contains("已被其他消息使用"), e.getMessage());
+    }
+
+    @Test
+    @DisplayName("CAS 成功(影响 1 行)则正常绑定放行")
+    void concurrentBindWinnerProceeds() {
+        when(mapper.selectById(12L)).thenReturn(att(12L, 1L, 1L, 0L, "TABLE"));
+        when(mapper.update(any(), any())).thenReturn(1);
+        List<String> blocks = service.buildBlocks(List.of(12L), 1L, 1L, 55L, "问题");
+        assertEquals(1, blocks.size());
+    }
+
     @Test
     @DisplayName("无附件时空列表,不动库")
     void emptyIdsNoop() {