lan
2748c80095
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
322 lines
10 KiB
Go
322 lines
10 KiB
Go
package cache
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
"go.uber.org/zap"
|
||
|
||
redisPkg "with_you/internal/pkg/redis"
|
||
)
|
||
|
||
// SeqBufferConfig seq 预分配配置
|
||
type SeqBufferConfig struct {
|
||
Enabled bool `mapstructure:"enabled"`
|
||
PrivateBufferSize int `mapstructure:"private_buffer_size"` // 私聊预分配步长
|
||
GroupBufferSize int `mapstructure:"group_buffer_size"` // 群聊预分配步长
|
||
LockTimeoutMs int `mapstructure:"lock_timeout_ms"` // 分布式锁超时(ms)
|
||
MaxRetries int `mapstructure:"max_retries"` // 冷启动重试次数
|
||
}
|
||
|
||
// localSeqBuffer 本地 seq 缓冲区
|
||
type localSeqBuffer struct {
|
||
mu sync.Mutex
|
||
curr int64 // 当前已分配到的 seq
|
||
last int64 // 本地缓冲区的上界
|
||
}
|
||
|
||
// SeqBufferManager seq 预分配管理器
|
||
type SeqBufferManager struct {
|
||
rdb *redis.Client // 原始 Redis 客户端（用于执行 Lua 脚本）
|
||
config SeqBufferConfig
|
||
repo ConversationRepository
|
||
msgRepo MessageRepository
|
||
buffers sync.Map // convID -> *localSeqBuffer
|
||
}
|
||
|
||
// mallocLua 预分配 Lua 跟踪脚本
|
||
// KEYS[1] = seq_buf:{convID} (Hash: CURR/LAST/LOCK/TIME)
|
||
// KEYS[2] = msg_seq:{convID} (String, 兼容旧 INCR key)
|
||
// ARGV[1] = bufferSize
|
||
// ARGV[2] = lockTimeoutMs (毫秒)
|
||
// ARGV[3] = current timestamp (毫秒)
|
||
//
|
||
// 返回:
|
||
// status=1, startSeq, endSeq — 成功分配
|
||
// status=2 — 冷启动（key 不存在）
|
||
// status=0 — 获取锁失败
|
||
var mallocLua = redis.NewScript(`
|
||
local bufKey = KEYS[1]
|
||
local seqKey = KEYS[2]
|
||
local bufSize = tonumber(ARGV[1])
|
||
local lockMs = tonumber(ARGV[2])
|
||
local nowMs = tonumber(ARGV[3])
|
||
|
||
-- 尝试获取分布式锁
|
||
local lockVal = redis.call('HGET', bufKey, 'LOCK')
|
||
if lockVal then
|
||
local lockTs = tonumber(lockVal)
|
||
if lockTs > 0 and (nowMs - lockTs) < lockMs then
|
||
return {0, 0, 0}
|
||
end
|
||
end
|
||
redis.call('HSET', bufKey, 'LOCK', nowMs)
|
||
redis.call('HSET', bufKey, 'TIME', nowMs)
|
||
|
||
local curr = tonumber(redis.call('HGET', bufKey, 'CURR'))
|
||
local last = tonumber(redis.call('HGET', bufKey, 'LAST'))
|
||
|
||
if curr == nil or last == nil then
|
||
-- 冷启动：key 不存在
|
||
redis.call('HDEL', bufKey, 'LOCK')
|
||
return {2, 0, 0}
|
||
end
|
||
|
||
if curr < last then
|
||
-- 本地缓冲区仍有余量，直接分配一个 seq
|
||
local newCurr = redis.call('HINCRBY', bufKey, 'CURR', 1)
|
||
redis.call('HDEL', bufKey, 'LOCK')
|
||
return {1, newCurr, last}
|
||
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
|
||
newLast = oldSeq + bufSize
|
||
redis.call('HSET', bufKey, 'LAST', newLast)
|
||
end
|
||
redis.call('SET', seqKey, newLast)
|
||
redis.call('HDEL', bufKey, 'LOCK')
|
||
return {1, curr + 1, newLast}
|
||
`)
|
||
|
||
// initLua 冷启动初始化 Lua 跟踪脚本
|
||
// KEYS[1] = seq_buf:{convID}
|
||
// KEYS[2] = msg_seq:{convID}
|
||
// ARGV[1] = startSeq (从 DB 获取的 conv.LastSeq)
|
||
// ARGV[2] = bufferSize
|
||
// ARGV[3] = lockTimeoutMs
|
||
// ARGV[4] = current timestamp (毫秒)
|
||
var initLua = redis.NewScript(`
|
||
local bufKey = KEYS[1]
|
||
local seqKey = KEYS[2]
|
||
local startSeq = tonumber(ARGV[1])
|
||
local bufSize = tonumber(ARGV[2])
|
||
local lockMs = tonumber(ARGV[3])
|
||
local nowMs = tonumber(ARGV[4])
|
||
|
||
-- 创建 seq_buf Hash
|
||
redis.call('HSET', bufKey, 'CURR', startSeq)
|
||
redis.call('HSET', bufKey, 'LAST', startSeq + bufSize)
|
||
redis.call('HSET', bufKey, 'LOCK', 0)
|
||
redis.call('HSET', bufKey, 'TIME', nowMs)
|
||
|
||
-- 同步更新 msg_seq（兼容旧 INCR key）
|
||
local oldSeq = tonumber(redis.call('GET', seqKey))
|
||
if oldSeq and oldSeq > (startSeq + bufSize) then
|
||
redis.call('HSET', bufKey, 'LAST', oldSeq + bufSize)
|
||
redis.call('SET', seqKey, oldSeq + bufSize)
|
||
else
|
||
redis.call('SET', seqKey, startSeq + bufSize)
|
||
end
|
||
|
||
return {1, startSeq + 1, startSeq + bufSize}
|
||
`)
|
||
|
||
// NewSeqBufferManager 创建 seq 预分配管理器
|
||
func NewSeqBufferManager(rdb *redis.Client, cfg SeqBufferConfig, repo ConversationRepository, msgRepo MessageRepository) *SeqBufferManager {
|
||
if cfg.PrivateBufferSize <= 0 {
|
||
cfg.PrivateBufferSize = 50
|
||
}
|
||
if cfg.GroupBufferSize <= 0 {
|
||
cfg.GroupBufferSize = 200
|
||
}
|
||
if cfg.LockTimeoutMs <= 0 {
|
||
cfg.LockTimeoutMs = 5000
|
||
}
|
||
if cfg.MaxRetries <= 0 {
|
||
cfg.MaxRetries = 3
|
||
}
|
||
return &SeqBufferManager{
|
||
rdb: rdb,
|
||
config: cfg,
|
||
repo: repo,
|
||
msgRepo: msgRepo,
|
||
}
|
||
}
|
||
|
||
// GetNextSeq 获取下一个 seq（从本地缓冲区分配，耗尽时向 Redis 预分配新区间）
|
||
func (m *SeqBufferManager) GetNextSeq(ctx context.Context, convID string, isGroup bool) (int64, error) {
|
||
bufSize := m.config.PrivateBufferSize
|
||
if isGroup {
|
||
bufSize = m.config.GroupBufferSize
|
||
}
|
||
|
||
// 1. 检查本地缓冲区
|
||
bufVal, ok := m.buffers.Load(convID)
|
||
if ok {
|
||
buf := bufVal.(*localSeqBuffer)
|
||
buf.mu.Lock()
|
||
if buf.curr < buf.last {
|
||
seq := buf.curr + 1
|
||
buf.curr = seq
|
||
buf.mu.Unlock()
|
||
return seq, nil
|
||
}
|
||
buf.mu.Unlock()
|
||
}
|
||
|
||
// 2. 本地缓冲区耗尽或不存在，向 Redis 预分配
|
||
startSeq, endSeq, err := m.allocateFromRedis(ctx, convID, bufSize)
|
||
if err != nil {
|
||
return 0, fmt.Errorf("allocate from redis: %w", err)
|
||
}
|
||
|
||
// 3. 更新本地缓冲区（CAS 防止并发分配时丢失区间）
|
||
newBuf := &localSeqBuffer{curr: startSeq, last: endSeq}
|
||
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 区间
|
||
func (m *SeqBufferManager) allocateFromRedis(ctx context.Context, convID string, bufferSize int) (startSeq, endSeq int64, err error) {
|
||
bufKey := SeqBufferKey(convID)
|
||
seqKey := MessageSeqKey(convID)
|
||
nowMs := time.Now().UnixMilli()
|
||
|
||
for i := 0; i < m.config.MaxRetries; i++ {
|
||
result, err := mallocLua.Run(ctx, m.rdb, []string{bufKey, seqKey}, bufferSize, m.config.LockTimeoutMs, nowMs).Int64Slice()
|
||
if err != nil {
|
||
// Redis 不可用，降级到 DB
|
||
zap.L().Warn("seq malloc lua failed, falling back to DB", zap.String("convID", convID), zap.Error(err))
|
||
if m.msgRepo != nil {
|
||
if dbSeq, dbErr := m.msgRepo.GetNextSeq(convID); dbErr == nil {
|
||
return dbSeq, dbSeq, nil
|
||
}
|
||
}
|
||
return 0, 0, fmt.Errorf("lua script error: %w", err)
|
||
}
|
||
|
||
status := result[0]
|
||
switch status {
|
||
case 1:
|
||
// 成功分配
|
||
return result[1], result[2], nil
|
||
case 2:
|
||
// 冷启动，需要从 DB 初始化
|
||
startSeq, endSeq, initErr := m.initFromDB(ctx, convID, bufferSize)
|
||
if initErr != nil {
|
||
return 0, 0, fmt.Errorf("init from db: %w", initErr)
|
||
}
|
||
return startSeq, endSeq, nil
|
||
case 0:
|
||
// 获取锁失败，短暂等待后重试
|
||
time.Sleep(time.Duration(50) * time.Millisecond)
|
||
continue
|
||
}
|
||
}
|
||
|
||
// 所有重试失败，降级到 DB
|
||
zap.L().Warn("seq malloc retries exhausted, falling back to DB", zap.String("convID", convID))
|
||
if m.msgRepo != nil {
|
||
if dbSeq, dbErr := m.msgRepo.GetNextSeq(convID); dbErr == nil {
|
||
return dbSeq, dbSeq, nil
|
||
}
|
||
}
|
||
return 0, 0, fmt.Errorf("malloc retries exhausted")
|
||
}
|
||
|
||
// initFromDB 冷启动时从 DB 加载 LastSeq 并初始化 Redis seq_buf Hash
|
||
func (m *SeqBufferManager) initFromDB(ctx context.Context, convID string, bufferSize int) (startSeq, endSeq int64, err error) {
|
||
if m.repo == nil {
|
||
return 0, 0, fmt.Errorf("conversation repository not configured")
|
||
}
|
||
|
||
conv, err := m.repo.GetConversationByID(convID)
|
||
if err != nil {
|
||
return 0, 0, fmt.Errorf("load conversation: %w", err)
|
||
}
|
||
|
||
dbLastSeq := conv.LastSeq
|
||
if dbLastSeq <= 0 {
|
||
dbLastSeq = 0
|
||
}
|
||
|
||
bufKey := SeqBufferKey(convID)
|
||
seqKey := MessageSeqKey(convID)
|
||
nowMs := time.Now().UnixMilli()
|
||
|
||
result, err := initLua.Run(ctx, m.rdb, []string{bufKey, seqKey}, dbLastSeq, bufferSize, m.config.LockTimeoutMs, nowMs).Int64Slice()
|
||
if err != nil {
|
||
return 0, 0, fmt.Errorf("init lua script error: %w", err)
|
||
}
|
||
|
||
zap.L().Info("seq buffer cold-start initialized",
|
||
zap.String("convID", convID),
|
||
zap.Int64("dbLastSeq", dbLastSeq),
|
||
zap.Int64("startSeq", result[1]),
|
||
zap.Int64("endSeq", result[2]),
|
||
)
|
||
|
||
return result[1], result[2], nil
|
||
}
|
||
|
||
// InvalidateBuffer 清除本地缓冲区（会话 seq 被外部重置时调用）
|
||
func (m *SeqBufferManager) InvalidateBuffer(convID string) {
|
||
m.buffers.Delete(convID)
|
||
}
|
||
|
||
// IsEnabled 是否启用预分配
|
||
func (m *SeqBufferManager) IsEnabled() bool {
|
||
return m.config.Enabled && m.rdb != nil
|
||
}
|
||
|
||
// NewSeqBufferManagerFromCache 从 ConversationCache 的 Cache 接口创建 SeqBufferManager
|
||
// 当 wire 中不方便注入 redis.Client 时，可通过 Cache 接口获取底层 Redis 连接
|
||
func NewSeqBufferManagerFromCache(cacheBackend Cache, cfg SeqBufferConfig, repo ConversationRepository, msgRepo MessageRepository) *SeqBufferManager {
|
||
redisCache, ok := cacheBackend.(*RedisCache)
|
||
if !ok {
|
||
zap.L().Warn("SeqBuffer: cache is not RedisCache, seq pre-allocation disabled")
|
||
return &SeqBufferManager{config: cfg, repo: repo, msgRepo: msgRepo}
|
||
}
|
||
|
||
rdb := redisCache.GetRedisClient().GetClient()
|
||
return NewSeqBufferManager(rdb, cfg, repo, msgRepo)
|
||
}
|
||
|
||
// NewSeqBufferManagerFromRedisPkg 从 pkg/redis.Client 创建 SeqBufferManager
|
||
func NewSeqBufferManagerFromRedisPkg(redisPkgClient *redisPkg.Client, cfg SeqBufferConfig, repo ConversationRepository, msgRepo MessageRepository) *SeqBufferManager {
|
||
if redisPkgClient == nil {
|
||
zap.L().Warn("SeqBuffer: redis client is nil, seq pre-allocation disabled")
|
||
return &SeqBufferManager{config: cfg, repo: repo, msgRepo: msgRepo}
|
||
}
|
||
rdb := redisPkgClient.GetClient()
|
||
return NewSeqBufferManager(rdb, cfg, repo, msgRepo)
|
||
} |