fix(cache): enhance sequence generation atomicity and concurrency safety

Implement Lua scripts for atomic sequence incrementing and cold-start initialization in Redis to prevent race conditions. Improve the sequence buffer manager with CAS-like logic to prevent interval leakage during concurrent allocations. Update the message repository to use `SELECT ... FOR UPDATE` and conditional updates for `last_seq` to ensure database consistency during high concurrency. - Add `nextSeqLua` and `initSeqLua` for atomic Redis operations - Implement thread-safe buffer management in `seq_buffer.go` - Add row-level locking in `message_repo.go` - Update service layers to support group-aware sequence retrieval
2026-05-25 14:08:06 +08:00
parent 3fc8b38184
commit 2748c80095
6 changed files with 109 additions and 26 deletions
--- a/internal/cache/conversation_cache.go
+++ b/internal/cache/conversation_cache.go
@@ -13,6 +13,24 @@ import (
 	"go.uber.org/zap"
 )

+// nextSeqLua 原子递增 seq：key 存在时 INCR 返回新值，不存在时返回 0 触发冷启动
+var nextSeqLua = redis.NewScript(`
+local seqKey = KEYS[1]
+if redis.call('EXISTS', seqKey) == 1 then
+  return redis.call('INCR', seqKey)
+else
+  return 0
+end
+`)
+
+// initSeqLua 冷启动原子初始化：SETNX 防止覆盖 + INCR 递增，返回第一个可用 seq
+var initSeqLua = redis.NewScript(`
+local seqKey = KEYS[1]
+local initVal = tonumber(ARGV[1])
+redis.call('SETNX', seqKey, initVal)
+return redis.call('INCR', seqKey)
+`)
+
 // CachedConversation 带缓存元数据的会话
 type CachedConversation struct {
 	Data      *model.Conversation // 实际数据
@@ -700,17 +718,48 @@ func (c *ConversationCache) ComputeAllUnreadCount(ctx context.Context, userID st
 func (c *ConversationCache) GetNextSeq(ctx context.Context, convID string) (int64, error) {
 	// 优先使用 seq 预分配
 	if c.seqBufferMgr != nil && c.seqBufferMgr.IsEnabled() {
-		// 判断是否群聊（通过会话类型判断，这里先按私聊处理，由上层传入 isGroup）
-		// 由于接口只传 convID，这里默认私聊逻辑，群聊的 isGroup 由调用方决定
 		return c.seqBufferMgr.GetNextSeq(ctx, convID, false)
 	}

-	// 降级到旧 INCR 逻辑
+	// 降级到旧逻辑：使用 Lua 脚本原子执行 INCR + 首次初始化对齐，避免竞态
 	seqKey := MessageSeqKey(convID)

+	// 尝试用 Lua 脚本原子递增（key 存在时 INCR，不存在时返回 0 触发冷启动）
+	if lc, ok := c.cache.(*LayeredCache); ok && lc.IsEnabled() {
+		rdb := lc.GetRedisClient()
+		if rdb != nil {
+			result, err := nextSeqLua.Run(ctx, rdb.GetClient(), []string{seqKey}).Int64()
+			if err == nil {
+				if result > 0 {
+					// key 已存在，INCR 成功
+					_ = c.cache.Expire(ctx, seqKey, c.settings.SeqTTL)
+					return result, nil
+				}
+				// result == 0: 冷启动，key 不存在，需从 DB 初始化
+				if c.repo != nil {
+					conv, dbErr := c.repo.GetConversationByID(convID)
+					if dbErr != nil {
+						return 0, fmt.Errorf("failed to load conversation for seq init: %w", dbErr)
+					}
+					initVal := conv.LastSeq // SETNX 使用 last_seq，INCR 返回 last_seq + 1
+					initResult, initErr := initSeqLua.Run(ctx, rdb.GetClient(), []string{seqKey}, initVal).Int64()
+					if initErr != nil {
+						// 初始化失败，降级到非原子路径
+						zap.L().Warn("initSeqLua failed, falling back", zap.Error(initErr))
+					} else {
+						_ = c.cache.Expire(ctx, seqKey, c.settings.SeqTTL)
+						return initResult, nil
+					}
+				}
+			}
+			// Lua 脚本失败，降级到非原子路径
+		}
+	}
+
+	// 非原子降级路径（Redis 不可用或 Lua 不支持）
 	newSeq, err := c.cache.Incr(ctx, seqKey)
 	if err != nil {
-		// Redis 不可用时降级到 DB，拿到值后回写 Redis 以便下次命中缓存
+		// Redis 不可用时降级到 DB
 		if c.msgRepo != nil {
 			if dbSeq, dbErr := c.msgRepo.GetNextSeq(convID); dbErr == nil {
 				c.cache.Set(seqKey, dbSeq, c.settings.SeqTTL)
@@ -721,6 +770,7 @@ func (c *ConversationCache) GetNextSeq(ctx context.Context, convID string) (int6
 	}

 	// 首次调用时 Redis key 不存在，INCR 从 0 返回 1，需要从 DB 补齐
+	// 注意：此路径存在竞态窗口，优先使用上面的 Lua 原子路径
 	if newSeq == 1 && c.repo != nil {
 		conv, err := c.repo.GetConversationByID(convID)
 		if err != nil {
--- a/internal/cache/seq_buffer.go
+++ b/internal/cache/seq_buffer.go
@@ -85,6 +85,8 @@ end
 -- 缓冲区耗尽，扩展 LAST
 local newLast = last + bufSize
 redis.call('HSET', bufKey, 'LAST', newLast)
+-- 推进 CURR 防止下一个调用者获得相同的 start seq
+redis.call('HSET', bufKey, 'CURR', curr + 1)
 -- 同步更新 msg_seq（兼容旧 INCR key，确保 READ 的值不回退）
 local oldSeq = tonumber(redis.call('GET', seqKey))
 if oldSeq and oldSeq > newLast then
@@ -178,12 +180,29 @@ func (m *SeqBufferManager) GetNextSeq(ctx context.Context, convID string, isGrou
 		return 0, fmt.Errorf("allocate from redis: %w", err)
 	}

-	// 3. 更新本地缓冲区
+	// 3. 更新本地缓冲区（CAS 防止并发分配时丢失区间）
 	newBuf := &localSeqBuffer{curr: startSeq, last: endSeq}
-	m.buffers.Store(convID, newBuf)
-
-	// startSeq 已经被 Redis 分配，返回它
-	return startSeq, nil
+	for {
+		existing, loaded := m.buffers.LoadOrStore(convID, newBuf)
+		if !loaded {
+			// 首次插入，直接使用
+			return startSeq, nil
+		}
+		// 已有缓冲区，检查是否被其他 goroutine 刷新
+		oldBuf := existing.(*localSeqBuffer)
+		oldBuf.mu.Lock()
+		if oldBuf.curr < oldBuf.last {
+			// 旧缓冲区仍有余量，直接使用（丢弃本次分配，避免区间泄漏）
+			seq := oldBuf.curr + 1
+			oldBuf.curr = seq
+			oldBuf.mu.Unlock()
+			return seq, nil
+		}
+		// 旧缓冲区耗尽，替换为新分配的区间
+		m.buffers.Store(convID, newBuf)
+		oldBuf.mu.Unlock()
+		return startSeq, nil
+	}
 }

 // allocateFromRedis 向 Redis 预分配 seq 区间
--- a/internal/repository/message_repo.go
+++ b/internal/repository/message_repo.go
@@ -414,14 +414,21 @@ func (r *messageRepository) UpdateConversationLastSeq(conversationID string, seq
 		}).Error
 }

-// GetNextSeq 获取会话的下一个seq值
+// GetNextSeq 获取会话的下一个seq值（使用 SELECT ... FOR UPDATE 保证原子性）
 func (r *messageRepository) GetNextSeq(conversationID string) (int64, error) {
-	var conv model.Conversation
-	err := r.db.Select("last_seq").Where("id = ?", conversationID).First(&conv).Error
-	if err != nil {
-		return 0, err
-	}
-	return conv.LastSeq + 1, nil
+	var nextSeq int64
+	err := r.db.Transaction(func(tx *gorm.DB) error {
+		var conv model.Conversation
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ?", conversationID).First(&conv).Error; err != nil {
+			return err
+		}
+		nextSeq = conv.LastSeq + 1
+		// 立即写回，防止并发读取同一个 last_seq 值
+		return tx.Model(&model.Conversation{}).
+			Where("id = ?", conversationID).
+			Update("last_seq", nextSeq).Error
+	})
+	return nextSeq, err
 }

 // CreateMessageWithSeq 创建消息并更新seq（事务操作）
@@ -433,13 +440,20 @@ func (r *messageRepository) CreateMessageWithSeq(msg *model.Message) error {
 			return err
 		}

-		// 更新会话的last_seq
-		if err := tx.Model(&model.Conversation{}).
-			Where("id = ?", msg.ConversationID).
+		// 更新会话last_seq：仅当新 seq 更大时才更新，防止并发写入导致 last_seq 回退
+		result := tx.Model(&model.Conversation{}).
+			Where("id = ? AND last_seq < ?", msg.ConversationID, msg.Seq).
 			Updates(map[string]any{
 				"last_seq":      msg.Seq,
 				"last_msg_time": gorm.Expr("CURRENT_TIMESTAMP"),
-			}).Error; err != nil {
+			})
+		if result.Error != nil {
+			return result.Error
+		}
+		// 当 last_seq >= msg.Seq 时（并发写入已更新），仍需刷新 last_msg_time
+		if err := tx.Model(&model.Conversation{}).
+			Where("id = ? AND last_seq >= ?", msg.ConversationID, msg.Seq).
+			Update("last_msg_time", gorm.Expr("CURRENT_TIMESTAMP")).Error; err != nil {
 			return err
 		}

--- a/internal/service/chat_service.go
+++ b/internal/service/chat_service.go
@@ -481,7 +481,7 @@ func (s *chatServiceImpl) SendMessage(ctx context.Context, senderID string, conv

 	// 从 Redis 获取下一个 seq
 	if s.conversationCache != nil {
-		seq, err := s.conversationCache.GetNextSeq(ctx, conversationID)
+		seq, err := s.conversationCache.GetNextSeqWithGroup(ctx, conversationID, conv.Type == model.ConversationTypeGroup)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get next seq for conversation %s: %w", conversationID, err)
 		}
@@ -1042,7 +1042,7 @@ func (s *chatServiceImpl) IsUserOnline(userID string) bool {
 // 适用于群聊等由调用方自行负责推送的场景
 func (s *chatServiceImpl) SaveMessage(ctx context.Context, senderID string, conversationID string, segments model.MessageSegments, replyToID *string) (*model.Message, error) {
 	// 验证会话是否存在
-	_, err := s.getConversation(ctx, conversationID)
+	conv, err := s.getConversation(ctx, conversationID)
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, errors.New("会话不存在，请重新创建会话")
@@ -1075,7 +1075,7 @@ func (s *chatServiceImpl) SaveMessage(ctx context.Context, senderID string, conv

 	// 从 Redis 获取下一个 seq（SaveMessage 方法）
 	if s.conversationCache != nil {
-		seq, err := s.conversationCache.GetNextSeq(ctx, conversationID)
+		seq, err := s.conversationCache.GetNextSeqWithGroup(ctx, conversationID, conv.Type == model.ConversationTypeGroup)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get next seq for conversation %s: %w", conversationID, err)
 		}
--- a/internal/service/group_service.go
+++ b/internal/service/group_service.go
@@ -473,7 +473,7 @@ func (s *groupService) broadcastMemberJoinNotice(groupID string, targetUserID st
 			}
 			// 从 Redis 获取下一个 seq
 			if s.conversationCache != nil {
-				seq, seqErr := s.conversationCache.GetNextSeq(context.Background(), conv.ID)
+				seq, seqErr := s.conversationCache.GetNextSeqWithGroup(context.Background(), conv.ID, true)
 				if seqErr != nil {
 					zap.L().Warn("get next seq failed, skipping system message",
 						zap.String("convID", conv.ID),
@@ -1398,7 +1398,7 @@ func (s *groupService) MuteMember(userID string, groupID string, targetUserID st

 			// 从 Redis 获取下一个 seq
 			if s.conversationCache != nil {
-				seq, seqErr := s.conversationCache.GetNextSeq(context.Background(), conv.ID)
+				seq, seqErr := s.conversationCache.GetNextSeqWithGroup(context.Background(), conv.ID, true)
 				if seqErr != nil {
 					zap.L().Warn("get next seq failed, skipping system message",
 						zap.String("convID", conv.ID),
--- a/internal/service/message_service.go
+++ b/internal/service/message_service.go
@@ -90,7 +90,7 @@ func (s *MessageService) SendMessage(ctx context.Context, senderID, receiverID s

 	// 从 Redis 获取下一个 seq
 	if s.conversationCache != nil {
-		seq, seqErr := s.conversationCache.GetNextSeq(context.Background(), conv.ID)
+		seq, seqErr := s.conversationCache.GetNextSeqWithGroup(context.Background(), conv.ID, false)
 		if seqErr != nil {
 			return nil, fmt.Errorf("failed to get next seq for conversation %s: %w", conv.ID, seqErr)
 		}