refactor(message): implement atomic message merging and patching to prevent race conditions

Introduce `mergeMessages` and `patchMessages` to the message store to handle updates atomically within Zustand's `set` callback. This replaces manual read-modify-write patterns in services, preventing message loss during concurrent WebSocket updates and synchronization processes. - Add `mergeMessages` for atomic merging of new messages into existing lists - Add `patchMessages` for efficient field updates (e.g., sender info, status) - Update `MessageSendService`, `MessageSyncService`, and `WSMessageHandler` to use these new atomic operations - Remove reliance on external `mergeMessagesById` in service layers to ensure state consistency
2026-05-25 01:44:23 +08:00
parent 7a17323e8d
commit f4db4eb1ed
4 changed files with 86 additions and 89 deletions
--- a/src/stores/message/services/MessageSendService.ts
+++ b/src/stores/message/services/MessageSendService.ts
@@ -7,7 +7,7 @@ import type { MessageResponse, MessageSegment, ConversationResponse } from '../.
 import { messageService } from '@/services/message';
 import { messageRepository } from '@/database';
 import type { IMessageSendService } from '../types';
-import { useMessageStore, mergeMessagesById } from '../store';
+import { useMessageStore } from '../store';

 export class MessageSendService implements IMessageSendService {
  private getCurrentUserId: () => string | null;
@@ -36,8 +36,7 @@ export class MessageSendService implements IMessageSendService {
      if (response) {
        const currentUserId = this.getCurrentUserId();
        
-        // 添加到本地消息列表
-        const existingMessages = store.getMessages(conversationId);
+        // 原子性合并新消息到内存（防止并发 WS 消息丢失）
        const newMessage: MessageResponse = {
          id: response.id,
          conversation_id: conversationId,
@@ -48,8 +47,7 @@ export class MessageSendService implements IMessageSendService {
          status: 'normal',
        };

-        const updatedMessages = mergeMessagesById(existingMessages, [newMessage]);
-        store.setMessages(conversationId, updatedMessages);
+        store.mergeMessages(conversationId, [newMessage]);

        // 更新会话最后消息
        this.updateConversationWithNewMessage(conversationId, newMessage);
--- a/src/stores/message/services/MessageSyncService.ts
+++ b/src/stores/message/services/MessageSyncService.ts
@@ -14,7 +14,7 @@ import {
  CONVERSATION_LIST_PAGE_SIZE,
 } from '../sources';
 import type { IMessageSyncService, MessageManagerConversationListDeps, IUserCacheService } from '../types';
-import { useMessageStore, normalizeConversationId, mergeMessagesById } from '../store';
+import { useMessageStore, normalizeConversationId } from '../store';
 import { ReadReceiptManager } from './ReadReceiptManager';

 export class MessageSyncService implements IMessageSyncService {
@@ -181,7 +181,7 @@ export class MessageSyncService implements IMessageSyncService {
   */
  async fetchMessages(conversationId: string, afterSeq?: number): Promise<void> {
    const store = useMessageStore.getState();
-    
+
    // 防止重复加载
    if (store.isLoadingMessages(conversationId)) {
      return;
@@ -190,15 +190,10 @@ export class MessageSyncService implements IMessageSyncService {
    store.setLoadingMessages(conversationId, true);

    try {
-      const existingMessagesAtStart = store.getMessages(conversationId);
-      const hasInMemoryMessages = existingMessagesAtStart.length > 0;
-      let baselineMaxSeq = hasInMemoryMessages
-        ? existingMessagesAtStart.reduce((max, m) => Math.max(max, m.seq || 0), 0)
-        : 0;
-
      if (!afterSeq) {
-        // 先从本地数据库加载
-        if (!hasInMemoryMessages) {
+        // 先从本地数据库加载（使用 merge 避免 WS 消息丢失）
+        const existingMessagesAtStart = store.getMessages(conversationId);
+        if (existingMessagesAtStart.length === 0) {
          try {
            const localMessages = await messageRepository.getByConversation(conversationId, 20);

@@ -213,9 +208,8 @@ export class MessageSyncService implements IMessageSyncService {
                created_at: m.createdAt,
              }));

-              store.setMessages(conversationId, formattedMessages);
-            } else {
-              store.setMessages(conversationId, []);
+              // 原子性合并本地缓存消息，防止并发 WS 消息丢失
+              store.mergeMessages(conversationId, formattedMessages);
            }
          } catch (error) {
            console.warn('[MessageSyncService] 读取本地消息失败:', error);
@@ -224,7 +218,7 @@ export class MessageSyncService implements IMessageSyncService {

        // 更新 baseline：合并本地DB后重新计算
        const currentMessages = store.getMessages(conversationId);
-        baselineMaxSeq = currentMessages.reduce((max, m) => Math.max(max, m.seq || 0), 0);
+        const baselineMaxSeq = currentMessages.reduce((max, m) => Math.max(max, m.seq || 0), 0);

        // 策略：有本地数据时仅增量同步，无数据时拉快照
        if (baselineMaxSeq > 0) {
@@ -234,9 +228,7 @@ export class MessageSyncService implements IMessageSyncService {
            const newMessages = incrementalResp?.messages || [];

            if (newMessages.length > 0) {
-              const existingMsgs = store.getMessages(conversationId);
-              const mergedMessages = mergeMessagesById(existingMsgs, newMessages);
-              store.setMessages(conversationId, mergedMessages);
+              store.mergeMessages(conversationId, newMessages);

              messageRepository.saveMessagesBatch(
                MessageMapper.toCachedMessages(newMessages, conversationId)
@@ -254,7 +246,8 @@ export class MessageSyncService implements IMessageSyncService {
            const snapshotMessages = snapshotResp?.messages || [];

            if (snapshotMessages.length > 0) {
-              store.setMessages(conversationId, snapshotMessages);
+              // 快照来自分页接口，后端返回 DESC 顺序；mergeMessages 内部按 seq ASC 排序并合并
+              store.mergeMessages(conversationId, snapshotMessages);

              messageRepository.saveMessagesBatch(
                MessageMapper.toCachedMessages(snapshotMessages, conversationId)
@@ -272,10 +265,7 @@ export class MessageSyncService implements IMessageSyncService {

        if (response?.messages && response.messages.length > 0) {
          const newMessages = response.messages;
-          const existingMessages = store.getMessages(conversationId);
-          const mergedMessages = mergeMessagesById(existingMessages, newMessages);
-
-          store.setMessages(conversationId, mergedMessages);
+          store.mergeMessages(conversationId, newMessages);

            messageRepository.saveMessagesBatch(
              MessageMapper.toCachedMessages(newMessages, conversationId)
@@ -288,13 +278,22 @@ export class MessageSyncService implements IMessageSyncService {
      console.error('[MessageSyncService] 获取消息失败:', error);
    } finally {
      // 异步填充用户信息（不阻塞消息显示）
+      // 使用 merge 而非覆盖，防止异步期间新到的 WS 消息丢失
      const currentMessages = store.getMessages(conversationId);
      if (currentMessages.length > 0) {
        this.userCacheService.enrichMessagesWithSenderInfo(
          conversationId,
          currentMessages,
          (convId, enrichedMessages) => {
-            store.setMessages(convId, enrichedMessages);
+            const patches = new Map<string, Partial<MessageResponse>>();
+            for (const m of enrichedMessages) {
+              if (m.sender && m.sender_id) {
+                patches.set(String(m.id), { sender: m.sender });
+              }
+            }
+            if (patches.size > 0) {
+              useMessageStore.getState().patchMessages(convId, patches);
+            }
          }
        );
      }
@@ -323,9 +322,7 @@ export class MessageSyncService implements IMessageSyncService {
          created_at: m.createdAt,
        }));

-        const existingMessages = store.getMessages(conversationId);
-        const mergedMessages = this.mergeOlderMessages(existingMessages, formattedMessages);
-        store.setMessages(conversationId, mergedMessages);
+        store.mergeMessages(conversationId, formattedMessages);

        return formattedMessages;
      }
@@ -340,9 +337,7 @@ export class MessageSyncService implements IMessageSyncService {
          MessageMapper.toCachedMessages(serverMessages, conversationId)
        );

-        const existingMessages = store.getMessages(conversationId);
-        const mergedMessages = this.mergeOlderMessages(existingMessages, serverMessages);
-        store.setMessages(conversationId, mergedMessages);
+        store.mergeMessages(conversationId, serverMessages);

        return serverMessages;
      }
@@ -612,19 +607,4 @@ export class MessageSyncService implements IMessageSyncService {
    return true;
  }

-  private mergeOlderMessages(existing: MessageResponse[], incoming: MessageResponse[]): MessageResponse[] {
-    if (incoming.length === 0) return existing;
-
-    const incomingAsc = [...incoming].sort((a, b) => a.seq - b.seq);
-    if (existing.length === 0) return incomingAsc;
-
-    const existingMinSeq = existing[0]?.seq ?? Number.MAX_SAFE_INTEGER;
-    const incomingMaxSeq = incomingAsc[incomingAsc.length - 1]?.seq ?? Number.MIN_SAFE_INTEGER;
-
-    if (incomingMaxSeq < existingMinSeq) {
-      return [...incomingAsc, ...existing];
-    }
-
-    return mergeMessagesById(existing, incomingAsc);
-  }
-}
+  }
--- a/src/stores/message/services/WSMessageHandler.ts
+++ b/src/stores/message/services/WSMessageHandler.ts
@@ -25,7 +25,7 @@ import type {
  HandleNewMessageOptions,
 } from '../types';
 import { MAX_FLUSH_ITERATIONS } from '../constants';
-import { useMessageStore, normalizeConversationId, mergeMessagesById } from '../store';
+import { useMessageStore, normalizeConversationId } from '../store';

 export class WSMessageHandler implements IWSMessageHandler {
  private deduplication: IMessageDeduplication;
@@ -326,13 +326,9 @@ export class WSMessageHandler implements IWSMessageHandler {
    if (message.type === 'group_message' && sender_id && sender_id !== currentUserId && sender_id !== '10000') {
      this.userCacheService.getSenderInfo(sender_id).then(user => {
        if (user) {
-          const currentMessages = store.getMessages(normalizedConversationId);
-          if (currentMessages) {
-            const updatedMessages = currentMessages.map(m =>
-              m.id === id ? { ...m, sender: user } : m
-            );
-            store.setMessages(normalizedConversationId, updatedMessages);
-          }
+          const patches = new Map<string, Partial<MessageResponse>>();
+          patches.set(String(id), { sender: user });
+          useMessageStore.getState().patchMessages(normalizedConversationId, patches);
        }
      }).catch(error => {
        console.error('[WSMessageHandler] 获取发送者信息失败:', { userId: sender_id, error });
@@ -350,13 +346,11 @@ export class WSMessageHandler implements IWSMessageHandler {
      status: 'normal',
    };

-    // 立即更新内存中的消息列表
-    const existingMessages = store.getMessages(normalizedConversationId);
-    const messageExists = existingMessages.some(m => String(m.id) === String(id));
+    // 立即更新内存中的消息列表 — 使用原子性 mergeMessages 防止并发写入丢失
+    const messageExists = useMessageStore.getState().getMessages(normalizedConversationId).some(m => String(m.id) === String(id));

    if (!messageExists) {
-      const updatedMessages = mergeMessagesById(existingMessages, [newMessage]);
-      store.setMessages(normalizedConversationId, updatedMessages);
+      useMessageStore.getState().mergeMessages(normalizedConversationId, [newMessage]);
    }

    // 更新会话信息
@@ -496,13 +490,12 @@ export class WSMessageHandler implements IWSMessageHandler {
      }
    }

-    // 补一条系统消息到当前会话消息流
+    // 补一条系统消息到当前会话消息流（原子性 merge）
    if (message_id && typeof seq === 'number' && seq > 0) {
      const conversationId = this.findConversationIdByGroupId(groupIdStr);
      if (conversationId) {
-        const existingMessages = store.getMessages(conversationId);
-        const exists = existingMessages.some(m => String(m.id) === String(message_id));
-        if (!exists) {
+        const messageExists = store.getMessages(conversationId).some(m => String(m.id) === String(message_id));
+        if (!messageExists) {
          const noticeText = this.buildGroupNoticeText(notice_type, data);
          const systemNoticeMessage: MessageResponse = {
            id: String(message_id),
@@ -517,8 +510,7 @@ export class WSMessageHandler implements IWSMessageHandler {
            created_at: (() => { const d = new Date(timestamp || Date.now()); return Number.isNaN(d.getTime()) ? new Date().toISOString() : d.toISOString(); })(),
          };

-          const updatedMessages = [...existingMessages, systemNoticeMessage].sort((a, b) => a.seq - b.seq);
-          store.setMessages(conversationId, updatedMessages);
+          useMessageStore.getState().mergeMessages(conversationId, [systemNoticeMessage]);
        }
      }
    }
@@ -557,26 +549,9 @@ export class WSMessageHandler implements IWSMessageHandler {
  }

  private markMessageAsRecalled(conversationId: string, messageId: string): void {
-    const store = useMessageStore.getState();
-    const messages = store.getMessages(conversationId);
-    if (!messages) return;
-
-    let changed = false;
-    const updatedMessages: MessageResponse[] = messages.map((m): MessageResponse => {
-      if (String(m.id) !== String(messageId) || m.status === 'recalled') {
-        return m;
-      }
-      changed = true;
-      return {
-        ...m,
-        status: 'recalled' as MessageResponse['status'],
-        segments: [],
-      };
-    });
-
-    if (!changed) return;
-
-    store.setMessages(conversationId, updatedMessages);
+    const patches = new Map<string, Partial<MessageResponse>>();
+    patches.set(String(messageId), { status: 'recalled' as MessageResponse['status'], segments: [] });
+    useMessageStore.getState().patchMessages(conversationId, patches);
  }

  private syncConversationLastMessageOnRecall(conversationId: string, messageId: string): void {
--- a/src/stores/message/store.ts
+++ b/src/stores/message/store.ts
@@ -84,6 +84,8 @@ export interface MessageActions {
  updateConversation: (conversation: ConversationResponse) => void;
  removeConversation: (conversationId: string) => void;
  setMessages: (conversationId: string, messages: MessageResponse[]) => void;
+  mergeMessages: (conversationId: string, incoming: MessageResponse[]) => void;
+  patchMessages: (conversationId: string, patches: Map<string, Partial<MessageResponse>>) => void;
  setUnreadCount: (total: number, system: number) => void;
  setLastSystemMessageAt: (time: string | null) => void;
  setSSEConnected: (connected: boolean) => void;
@@ -331,6 +333,48 @@ export const useMessageStore = create<MessageStore>((set, get) => ({
    });
  },

+  /**
+   * 原子性合并消息到 messagesMap
+   * 在 zustand set 回调内完成 read-merge-write，消除并发竞态
+   */
+  mergeMessages: (conversationId: string, incoming: MessageResponse[]) => {
+    if (incoming.length === 0) return;
+    const normalizedId = normalizeConversationId(conversationId);
+    set(state => {
+      const existing = state.messagesMap.get(normalizedId) || [];
+      const merged = mergeMessagesById(existing, incoming);
+      const newMessagesMap = new Map(state.messagesMap);
+      newMessagesMap.set(normalizedId, merged);
+      return { messagesMap: newMessagesMap };
+    });
+  },
+
+  /**
+   * 原子性 patch：按消息 ID 更新指定字段（如 sender、status）
+   * 不影响其他消息，不丢失并发写入
+   */
+  patchMessages: (conversationId: string, patches: Map<string, Partial<MessageResponse>>) => {
+    if (patches.size === 0) return;
+    const normalizedId = normalizeConversationId(conversationId);
+    set(state => {
+      const existing = state.messagesMap.get(normalizedId);
+      if (!existing) return state;
+      let changed = false;
+      const updated = existing.map(m => {
+        const patch = patches.get(String(m.id));
+        if (patch) {
+          changed = true;
+          return { ...m, ...patch };
+        }
+        return m;
+      });
+      if (!changed) return state;
+      const newMessagesMap = new Map(state.messagesMap);
+      newMessagesMap.set(normalizedId, updated);
+      return { messagesMap: newMessagesMap };
+    });
+  },
+
  setUnreadCount: (total: number, system: number) => {
    set({ totalUnreadCount: total, systemUnreadCount: system });
  },