From 2748c80095409f1d6eba4f995ece024fb3ed3083 Mon Sep 17 00:00:00 2001 From: lan Date: Mon, 25 May 2026 14:08:06 +0800 Subject: [PATCH] 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 --- internal/cache/conversation_cache.go | 58 ++++++++++++++++++++++++++-- internal/cache/seq_buffer.go | 29 +++++++++++--- internal/repository/message_repo.go | 36 +++++++++++------ internal/service/chat_service.go | 6 +-- internal/service/group_service.go | 4 +- internal/service/message_service.go | 2 +- 6 files changed, 109 insertions(+), 26 deletions(-) diff --git a/internal/cache/conversation_cache.go b/internal/cache/conversation_cache.go index 6cee2fc..cd60906 100644 --- 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 { diff --git a/internal/cache/seq_buffer.go b/internal/cache/seq_buffer.go index d2e6ce4..11037bc 100644 --- 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 区间 diff --git a/internal/repository/message_repo.go b/internal/repository/message_repo.go index f18f719..3ec4fc9 100644 --- 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 } diff --git a/internal/service/chat_service.go b/internal/service/chat_service.go index 047ba7f..f20cc02 100644 --- 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) } diff --git a/internal/service/group_service.go b/internal/service/group_service.go index 5140832..32c4237 100644 --- 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), diff --git a/internal/service/message_service.go b/internal/service/message_service.go index 963654c..d645666 100644 --- 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) }