package shortrecommend import ( "bytes" "context" cryptorand "crypto/rand" "crypto/sha256" "encoding/hex" "errors" "fmt" "hash/fnv" "math/rand" "sort" "strconv" "strings" "sync/atomic" "time" "91porn-server/common/redis" "91porn-server/models/v/vidmod" ) const ( CurrentKey = "recommend:short:current" BuildLockKey = "recommend:short:build-lock" defaultKeyTTL = 72 * time.Hour minKeyTTLHours = 24 maxKeyTTLHours = 30 * 24 buildingKeyTTL = time.Hour takeIdempotencyTTL = 30 * time.Second reservationTTL = 15 * time.Second maxTakeRequestIDLength = 128 healthTTLDriftTolerance = 2 * time.Second queuePushBatchSize = 1000 queuePipelineBatchCount = 16 publishCleanupTimeout = 5 * time.Second ) var keyTTLSeconds atomic.Int64 func init() { keyTTLSeconds.Store(int64(defaultKeyTTL / time.Second)) } // SetKeyTTLHours 配置正式队列、元数据和用户偏移的统一绝对过期时间。 func SetKeyTTLHours(hours int) error { if hours < minKeyTTLHours || hours > maxKeyTTLHours { return fmt.Errorf( "short recommend key TTL hours must be between %d and %d", minKeyTTLHours, maxKeyTTLHours, ) } keyTTLSeconds.Store(int64(time.Duration(hours) * time.Hour / time.Second)) return nil } // KeyTTL 返回当前并发安全的短推荐键过期时间配置。 func KeyTTL() time.Duration { seconds := keyTTLSeconds.Load() if seconds <= 0 { return defaultKeyTTL } return time.Duration(seconds) * time.Second } type QueueMeta struct { Version string Length int GeneratedAt time.Time StartOffset int } func QueueKey(version string) string { return "recommend:short:queue:" + version } func OffsetKey(version string) string { return "recommend:short:offset:" + version } func MetaKey(version string) string { return "recommend:short:meta:" + version } func LockKey(_ string) string { return BuildLockKey } func VersionAt(t time.Time) string { return t.In(time.FixedZone("CST", 8*60*60)).Format("20060102") } // VersionWithRevision 为同一天的每次实际重建生成独立队列版本。 // revision 必须由调用方提供非空且唯一的构建标识。 func VersionWithRevision(t time.Time, revision string) string { return VersionAt(t) + "-" + revision } func versionMatchesDate(version, date string) bool { return version == date || (len(version) > len(date)+1 && strings.HasPrefix(version, date+"-")) } type AssembleStats struct { TotalCount int HighCount int NewCount int BlockCount int } // Assemble 按每块17条高分+3条新视频组装,视频全局去重,尾块保留实际数量。 func Assemble(input []vidmod.RecommendCandidate, generatedAt time.Time) []string { ids, _ := AssembleWithStats(input, generatedAt) return ids } // AssembleWithStats 在组装队列的同时返回高分池、新视频池和块数指标。 func AssembleWithStats( input []vidmod.RecommendCandidate, generatedAt time.Time, ) ([]string, AssembleStats) { cutoff := generatedAt.Add(-24 * time.Hour) news := make([]*vidmod.RecommendCandidate, 0) high := make([]*vidmod.RecommendCandidate, 0) for i := range input { c := &input[i] if c.ReviewAt.After(generatedAt) { continue } if !c.ReviewAt.Before(cutoff) && !c.ReviewAt.After(generatedAt) { news = append(news, c) } else { high = append(high, c) } } sort.Slice(news, func(i, j int) bool { if news[i].ReviewAt.Equal(news[j].ReviewAt) { return bytes.Compare(news[i].ID[:], news[j].ID[:]) > 0 } return news[i].ReviewAt.After(news[j].ReviewAt) }) sort.Slice(high, func(i, j int) bool { if high[i].RecommendScore != high[j].RecommendScore { return high[i].RecommendScore > high[j].RecommendScore } if !high[i].ReviewAt.Equal(high[j].ReviewAt) { return high[i].ReviewAt.After(high[j].ReviewAt) } return bytes.Compare(high[i].ID[:], high[j].ID[:]) > 0 }) out := make([]string, 0, len(input)) hi, ni, block := 0, 0, 0 for hi < len(high) || ni < len(news) { blockStart := len(out) highCount := 17 newCount := min(3, len(news)-ni) highCount += 3 - newCount highCount = min(highCount, len(high)-hi) for n := 0; n < highCount; n++ { out = append(out, high[hi].ID.Hex()) hi++ } for n := 0; n < newCount; n++ { out = append(out, news[ni].ID.Hex()) ni++ } if len(out) == blockStart { // 高分耗尽后继续消费剩余新视频,仍不复制凑数 for ni < len(news) && len(out)-blockStart < 3 { out = append(out, news[ni].ID.Hex()) ni++ } } seed := deterministicSeed(VersionAt(generatedAt), block) part := out[blockStart:] rand.New(rand.NewSource(seed)).Shuffle(len(part), func(i, j int) { part[i], part[j] = part[j], part[i] }) block++ } return out, AssembleStats{ TotalCount: len(out), HighCount: len(high), NewCount: len(news), BlockCount: block, } } func deterministicSeed(version string, block int) int64 { h := fnv.New64a() _, _ = h.Write([]byte(version + ":" + strconv.Itoa(block))) return int64(h.Sum64()) } func min(a, b int) int { if a < b { return a } return b } type publishClient interface { DelContext(ctx context.Context, keys ...string) (int64, error) EvalContext( ctx context.Context, script string, keys []string, args ...interface{}, ) (interface{}, error) RPushPipelineContext( ctx context.Context, key string, values []string, batchSize int, expirations ...time.Duration, ) (int64, error) } // Publish 完整写入新版本后再切换current,构建失败不会破坏旧版本。 func Publish(ctx context.Context, client *redis.Client, meta QueueMeta, ids []string, lockKey, lockToken string) error { return publish(ctx, client, meta, ids, lockKey, lockToken) } func publish( ctx context.Context, client publishClient, meta QueueMeta, ids []string, lockKey, lockToken string, ) error { if len(ids) == 0 { return fmt.Errorf("short recommend queue is empty") } if lockKey == "" || lockToken == "" { return fmt.Errorf("short recommend build lock is required") } if !versionMatchesDate(meta.Version, VersionAt(meta.GeneratedAt)) { return fmt.Errorf("short recommend version does not match generatedAt") } if meta.Length != len(ids) { return fmt.Errorf("short recommend metadata length does not match queue length") } queueTTL := KeyTTL() queueKey, metaKey := QueueKey(meta.Version), MetaKey(meta.Version) suffix := ":" + strconv.FormatInt(time.Now().UnixNano(), 10) tmpQueueKey, tmpMetaKey := queueKey+":building"+suffix, metaKey+":building"+suffix const firstPushScript = ` for i = 1, #ARGV - 1 do redis.call('RPUSH', KEYS[1], ARGV[i]) end redis.call('EXPIRE', KEYS[1], ARGV[#ARGV]) return 1 ` if err := ctx.Err(); err != nil { return err } firstEnd := min(queuePushBatchSize, len(ids)) firstValues := make([]interface{}, firstEnd) for i := 0; i < firstEnd; i++ { firstValues[i] = ids[i] } firstArgs := append(firstValues, int64(buildingKeyTTL/time.Second)) if _, err := client.EvalContext(ctx, firstPushScript, []string{tmpQueueKey}, firstArgs...); err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } pipelineWindowSize := queuePushBatchSize * queuePipelineBatchCount for start := firstEnd; start < len(ids); start += pipelineWindowSize { if err := ctx.Err(); err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } end := min(start+pipelineWindowSize, len(ids)) length, err := client.RPushPipelineContext( ctx, tmpQueueKey, ids[start:end], queuePushBatchSize, buildingKeyTTL, ) if err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } if length != int64(end) { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return fmt.Errorf( "short recommend queue length mismatch after pipeline: got %d, want %d", length, end, ) } } if err := ctx.Err(); err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } const writeMetaScript = ` redis.call('HSET', KEYS[1], 'length', ARGV[1], 'generatedAt', ARGV[2], 'startOffset', ARGV[3]) redis.call('EXPIRE', KEYS[1], ARGV[4]) return 1 ` if _, err := client.EvalContext(ctx, writeMetaScript, []string{tmpMetaKey}, len(ids), meta.GeneratedAt.UTC().Format(time.RFC3339), meta.StartOffset, int64(buildingKeyTTL/time.Second)); err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } if err := ctx.Err(); err != nil { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) return err } const swapScript = ` if redis.call('GET', KEYS[7]) ~= ARGV[3] then return 0 end local redisTime = redis.call('TIME') if tonumber(redisTime[1]) >= tonumber(ARGV[4]) then return -1 end if redis.call('LLEN', KEYS[1]) ~= tonumber(ARGV[5]) then return -2 end if tonumber(redis.call('HGET', KEYS[2], 'length') or '-1') ~= tonumber(ARGV[5]) then return -3 end local oldVersion = redis.call('GET', KEYS[6]) redis.call('UNLINK', KEYS[5]) if oldVersion and oldVersion ~= ARGV[1] then redis.call('UNLINK', ARGV[6] .. oldVersion) end redis.call('RENAME', KEYS[1], KEYS[3]) redis.call('RENAME', KEYS[2], KEYS[4]) local expiresAtMs = tonumber(redisTime[1]) * 1000 + math.floor(tonumber(redisTime[2]) / 1000) + tonumber(ARGV[2]) * 1000 redis.call('HSET', KEYS[4], 'expiresAtMs', expiresAtMs) redis.call('PEXPIREAT', KEYS[3], expiresAtMs) redis.call('PEXPIREAT', KEYS[4], expiresAtMs) redis.call('SET', KEYS[6], ARGV[1]) return 1 ` swapped, err := client.EvalContext(ctx, swapScript, []string{ tmpQueueKey, tmpMetaKey, queueKey, metaKey, OffsetKey(meta.Version), CurrentKey, lockKey, }, meta.Version, int64(queueTTL/time.Second), lockToken, versionValidUntil(meta).Unix(), len(ids), "recommend:short:offset:") if err != nil || fmt.Sprint(swapped) != "1" { cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey) if err != nil { return err } switch fmt.Sprint(swapped) { case "0": return fmt.Errorf("short recommend build lock lost before queue swap") case "-1": return fmt.Errorf("short recommend version expired before queue swap") case "-2": return fmt.Errorf("short recommend queue length mismatch before queue swap") case "-3": return fmt.Errorf("short recommend metadata length mismatch before queue swap") default: return fmt.Errorf("short recommend queue swap rejected: %v", swapped) } } return nil } func cleanupPublishKeys(client publishClient, keys ...string) { cleanupCtx, cancel := context.WithTimeout(context.Background(), publishCleanupTimeout) defer cancel() _, _ = client.DelContext(cleanupCtx, keys...) } func versionValidUntil(meta QueueMeta) time.Time { cst := time.FixedZone("CST", 8*60*60) generatedAt := meta.GeneratedAt.In(cst) return time.Date( generatedAt.Year(), generatedAt.Month(), generatedAt.Day()+1, 0, 0, 0, 0, cst, ) } var ( // ErrVersionChanged 表示调用方固定的队列版本已被每日切版替换。 ErrVersionChanged = errors.New("short recommend queue version changed") // ErrQueueUnhealthy 表示current指向的队列、元数据或TTL不完整。 ErrQueueUnhealthy = errors.New("short recommend queue is unhealthy") // ErrReservationBusy 表示同一用户已有另一个请求正在读取当前偏移。 ErrReservationBusy = errors.New("short recommend offset reservation is busy") // ErrReservationExpired 表示预留在提交前已过期,偏移未推进。 ErrReservationExpired = errors.New("short recommend offset reservation expired") // ErrReservationConflict 表示预留内容或当前偏移与提交参数不一致。 ErrReservationConflict = errors.New("short recommend offset reservation conflict") ) const takeScriptSource = ` local function readBatch(queueKey, version, length, offset, wanted) local result = {version, tostring(length), tostring(offset)} local firstEnd = math.min(offset + wanted - 1, length - 1) local first = redis.call('LRANGE', queueKey, offset, firstEnd) for _, item in ipairs(first) do table.insert(result, item) end local remaining = wanted - #first if remaining > 0 then local wrapped = redis.call('LRANGE', queueKey, 0, remaining - 1) for _, item in ipairs(wrapped) do table.insert(result, item) end end return result end local version = redis.call('GET', KEYS[1]) if not version then return {} end if ARGV[6] ~= '' and version ~= ARGV[6] then return {version, '-1', '0'} end local queueKey = ARGV[1] .. version local offsetKey = ARGV[2] .. version local metaKey = ARGV[3] .. version local offsetKeyTTL = redis.call('PTTL', offsetKey) local length = redis.call('LLEN', queueKey) local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1') local queueTTL = redis.call('PTTL', queueKey) local metaTTL = redis.call('PTTL', metaKey) if length == 0 or metaLength ~= length or queueTTL <= 0 or metaTTL <= 0 then return {version, '-2', '0'} end local idempotencyKey = '' if ARGV[7] ~= '' then idempotencyKey = ARGV[7] .. ':' .. version local cached = redis.call('LRANGE', idempotencyKey, 0, -1) if #cached == 4 then local cachedVersion = cached[1] local cachedLength = tonumber(cached[2]) local cachedOffset = tonumber(cached[3]) local cachedWanted = tonumber(cached[4]) if cachedVersion == version and cachedLength == length and cachedOffset and cachedWanted then return readBatch( queueKey, cachedVersion, cachedLength, cachedOffset, cachedWanted ) end end if #cached > 0 then redis.call('DEL', idempotencyKey) end end local wanted = tonumber(ARGV[5]) if wanted > length then wanted = length end local storedOffset = redis.call('HGET', offsetKey, ARGV[4]) local offset if storedOffset then offset = tonumber(storedOffset) or 0 else offset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0') end offset = offset % length local result = readBatch(queueKey, version, length, offset, wanted) redis.call('HSET', offsetKey, ARGV[4], (offset + wanted) % length) local expiresAtMs = tonumber(redis.call('HGET', metaKey, 'expiresAtMs') or '0') if expiresAtMs <= 0 then local redisTime = redis.call('TIME') expiresAtMs = tonumber(redisTime[1]) * 1000 + math.floor(tonumber(redisTime[2]) / 1000) + math.min(queueTTL, metaTTL) end if offsetKeyTTL <= 0 then redis.call('PEXPIREAT', offsetKey, expiresAtMs) end if idempotencyKey ~= '' then redis.call( 'RPUSH', idempotencyKey, version, tostring(length), tostring(offset), tostring(wanted) ) redis.call('EXPIRE', idempotencyKey, tonumber(ARGV[8])) end return result ` var takeRedisScript = redis.NewScript(takeScriptSource) type TakeResult struct { Version string Length int Offset int IDs []string } type scriptClient interface { RunScriptContext( ctx context.Context, script *redis.Script, keys []string, args ...interface{}, ) (interface{}, error) } // Take 原子读取并推进单个用户的下一偏移,保留原调用方式。 func Take(client *redis.Client, uid uint64, size int) (TakeResult, error) { return TakeVersionContext(context.Background(), client, uid, size, "") } // TakeVersion 原子读取指定版本;切版后返回 ErrVersionChanged 且不推进新队列。 func TakeVersion( client *redis.Client, uid uint64, size int, expectedVersion string, ) (TakeResult, error) { return TakeVersionContext(context.Background(), client, uid, size, expectedVersion) } // TakeVersionContext 与 TakeVersion 相同,并将请求取消和超时传递到Redis。 func TakeVersionContext( ctx context.Context, client *redis.Client, uid uint64, size int, expectedVersion string, ) (TakeResult, error) { return takeContext(ctx, client, uid, size, expectedVersion, "") } // TakeIdempotentContext 为一次HTTP请求的单个补位批次提供短时幂等保护。 // 同一uid/requestID/batchIndex重试会返回原批次,不会再次推进偏移; // expectedVersion可在首批留空,后续批次应固定为首批返回的Version。 func TakeIdempotentContext( ctx context.Context, client *redis.Client, uid uint64, size int, expectedVersion, requestID string, batchIndex int, ) (TakeResult, error) { return takeIdempotentContext( ctx, client, uid, size, expectedVersion, requestID, batchIndex, ) } // TakeCurrentIdempotentContext 用于首批读取:Lua原子读取current,并按该版本隔离幂等缓存, // 因而不需要额外执行Health/GET;切版后不会复用上一版本缓存。 func TakeCurrentIdempotentContext( ctx context.Context, client *redis.Client, uid uint64, size int, requestID string, batchIndex int, ) (TakeResult, error) { return takeIdempotentContext(ctx, client, uid, size, "", requestID, batchIndex) } func takeIdempotentContext( ctx context.Context, client scriptClient, uid uint64, size int, expectedVersion, requestID string, batchIndex int, ) (TakeResult, error) { requestID = strings.TrimSpace(requestID) if requestID == "" { return TakeResult{}, fmt.Errorf("short recommend request ID is required") } if len(requestID) > maxTakeRequestIDLength { return TakeResult{}, fmt.Errorf( "short recommend request ID exceeds %d bytes", maxTakeRequestIDLength, ) } if batchIndex < 0 { return TakeResult{}, fmt.Errorf("short recommend batch index must not be negative") } requestHash := sha256.Sum256([]byte(requestID)) idempotencyKey := fmt.Sprintf( "recommend:short:take-cache:%d:%x:%d", uid, requestHash[:16], batchIndex, ) return takeContext(ctx, client, uid, size, expectedVersion, idempotencyKey) } func takeContext( ctx context.Context, client scriptClient, uid uint64, size int, expectedVersion, idempotencyKey string, ) (TakeResult, error) { if size <= 0 { return TakeResult{}, nil } if ctx == nil { return TakeResult{}, fmt.Errorf("short recommend take context must not be nil") } if client == nil { return TakeResult{}, fmt.Errorf("short recommend take Redis client must not be nil") } if err := ctx.Err(); err != nil { return TakeResult{}, err } raw, err := client.RunScriptContext(ctx, takeRedisScript, []string{CurrentKey}, "recommend:short:queue:", "recommend:short:offset:", "recommend:short:meta:", strconv.FormatUint(uid, 10), size, expectedVersion, idempotencyKey, int64(takeIdempotencyTTL/time.Second)) if err != nil { return TakeResult{}, err } items, ok := raw.([]interface{}) if !ok || len(items) < 3 { return TakeResult{}, nil } result := TakeResult{Version: fmt.Sprint(items[0])} if result.Version == "" { return TakeResult{}, fmt.Errorf("short recommend take returned empty version") } result.Length, err = strconv.Atoi(fmt.Sprint(items[1])) if err != nil { return TakeResult{}, fmt.Errorf("short recommend take returned invalid length: %w", err) } result.Offset, err = strconv.Atoi(fmt.Sprint(items[2])) if err != nil { return TakeResult{}, fmt.Errorf("short recommend take returned invalid offset: %w", err) } switch result.Length { case -1: result.Length = 0 return result, ErrVersionChanged case -2: result.Length = 0 return result, ErrQueueUnhealthy } if result.Length <= 0 || result.Offset < 0 || result.Offset >= result.Length { return TakeResult{}, fmt.Errorf( "short recommend take returned invalid bounds: length=%d offset=%d", result.Length, result.Offset, ) } for _, item := range items[3:] { result.IDs = append(result.IDs, fmt.Sprint(item)) } return result, nil } // Reservation 是一次请求对用户偏移的临时预留。Reserve 只读取队列, // Commit 成功后才推进偏移;预留过期或 Abort 均不会改变用户进度。 type Reservation struct { Version string Length int Offset int Reserved int IDs []string LeaseToken string ReceiptID string AlreadyCommitted bool } const reserveScriptSource = ` local function readBatch(queueKey, version, status, length, offset, wanted, token) local result = { version, status, tostring(length), tostring(offset), tostring(wanted), token } local firstEnd = math.min(offset + wanted - 1, length - 1) local first = redis.call('LRANGE', queueKey, offset, firstEnd) for _, item in ipairs(first) do table.insert(result, item) end local remaining = wanted - #first if remaining > 0 then local wrapped = redis.call('LRANGE', queueKey, 0, remaining - 1) for _, item in ipairs(wrapped) do table.insert(result, item) end end return result end local receiptKey = ARGV[6] .. ARGV[4] .. ':' .. ARGV[10] local receipt = redis.call('LRANGE', receiptKey, 0, -1) if #receipt == 4 then local receiptVersion = receipt[1] local receiptLength = tonumber(receipt[2]) local receiptOffset = tonumber(receipt[3]) local receiptConsumed = tonumber(receipt[4]) local receiptQueueKey = ARGV[1] .. receiptVersion if receiptLength and receiptOffset and receiptConsumed and receiptLength > 0 and receiptConsumed > 0 and receiptConsumed <= receiptLength and redis.call('LLEN', receiptQueueKey) == receiptLength and redis.call('PTTL', receiptQueueKey) > 0 then return readBatch( receiptQueueKey, receiptVersion, 'COMMITTED', receiptLength, receiptOffset, receiptConsumed, ARGV[9] ) end redis.call('DEL', receiptKey) elseif #receipt > 0 then redis.call('DEL', receiptKey) end local version = redis.call('GET', KEYS[1]) if not version then return {} end if ARGV[8] ~= '' and version ~= ARGV[8] then return {version, 'VERSION_CHANGED', '0', '0', '0', ''} end local queueKey = ARGV[1] .. version local offsetKey = ARGV[2] .. version local metaKey = ARGV[3] .. version local length = redis.call('LLEN', queueKey) local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1') local queueTTL = redis.call('PTTL', queueKey) local metaTTL = redis.call('PTTL', metaKey) if length == 0 or metaLength ~= length or queueTTL <= 0 or metaTTL <= 0 then return {version, 'UNHEALTHY', '0', '0', '0', ''} end local leaseKey = ARGV[5] .. version .. ':' .. ARGV[4] local lease = redis.call( 'HMGET', leaseKey, 'token', 'receiptID', 'length', 'offset', 'reserved' ) if lease[1] then if lease[1] ~= ARGV[9] or lease[2] ~= ARGV[10] then return {version, 'BUSY', tostring(length), '0', '0', ''} end local leaseLength = tonumber(lease[3]) local leaseOffset = tonumber(lease[4]) local leaseReserved = tonumber(lease[5]) if leaseLength == length and leaseOffset and leaseReserved and leaseOffset >= 0 and leaseOffset < length and leaseReserved > 0 and leaseReserved <= length then redis.call('PEXPIRE', leaseKey, tonumber(ARGV[11])) return readBatch( queueKey, version, 'RESERVED', length, leaseOffset, leaseReserved, ARGV[9] ) end redis.call('DEL', leaseKey) end local wanted = tonumber(ARGV[7]) if wanted > length then wanted = length end local storedOffset = redis.call('HGET', offsetKey, ARGV[4]) local offset if storedOffset then offset = tonumber(storedOffset) or 0 else offset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0') end offset = offset % length redis.call('HSET', leaseKey, 'token', ARGV[9], 'receiptID', ARGV[10], 'length', tostring(length), 'offset', tostring(offset), 'reserved', tostring(wanted)) redis.call('PEXPIRE', leaseKey, tonumber(ARGV[11])) return readBatch( queueKey, version, 'RESERVED', length, offset, wanted, ARGV[9] ) ` const commitReservationScriptSource = ` local receiptKey = ARGV[14] .. ARGV[5] .. ':' .. ARGV[8] local receipt = redis.call('LRANGE', receiptKey, 0, -1) if #receipt == 4 then if receipt[1] == ARGV[6] and tonumber(receipt[2]) == tonumber(ARGV[9]) and tonumber(receipt[3]) == tonumber(ARGV[10]) and tonumber(receipt[4]) == tonumber(ARGV[12]) then return 2 end return -5 elseif #receipt > 0 then return -5 end local version = redis.call('GET', KEYS[1]) if version ~= ARGV[6] then return -1 end local queueKey = ARGV[1] .. version local offsetKey = ARGV[2] .. version local metaKey = ARGV[3] .. version local length = redis.call('LLEN', queueKey) local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1') if length <= 0 or metaLength ~= length or length ~= tonumber(ARGV[9]) or redis.call('PTTL', queueKey) <= 0 or redis.call('PTTL', metaKey) <= 0 then return -2 end local leaseKey = ARGV[4] .. version .. ':' .. ARGV[5] local lease = redis.call( 'HMGET', leaseKey, 'token', 'receiptID', 'length', 'offset', 'reserved' ) if not lease[1] then return -3 end if lease[1] ~= ARGV[7] or lease[2] ~= ARGV[8] or tonumber(lease[3]) ~= tonumber(ARGV[9]) or tonumber(lease[4]) ~= tonumber(ARGV[10]) or tonumber(lease[5]) ~= tonumber(ARGV[11]) then return -4 end local consumed = tonumber(ARGV[12]) if not consumed or consumed <= 0 or consumed > tonumber(lease[5]) then return -4 end local storedOffset = redis.call('HGET', offsetKey, ARGV[5]) local currentOffset if storedOffset then currentOffset = tonumber(storedOffset) or 0 else currentOffset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0') end currentOffset = currentOffset % length if currentOffset ~= tonumber(lease[4]) then return -4 end redis.call('HSET', offsetKey, ARGV[5], (currentOffset + consumed) % length) local offsetKeyTTL = redis.call('PTTL', offsetKey) local expiresAtMs = tonumber(redis.call('HGET', metaKey, 'expiresAtMs') or '0') if expiresAtMs <= 0 then local redisTime = redis.call('TIME') expiresAtMs = tonumber(redisTime[1]) * 1000 + math.floor(tonumber(redisTime[2]) / 1000) + math.min(redis.call('PTTL', queueKey), redis.call('PTTL', metaKey)) end if offsetKeyTTL <= 0 then redis.call('PEXPIREAT', offsetKey, expiresAtMs) end redis.call('DEL', leaseKey) redis.call( 'RPUSH', receiptKey, version, tostring(length), tostring(currentOffset), tostring(consumed) ) redis.call('EXPIRE', receiptKey, tonumber(ARGV[13])) return 1 ` const abortReservationScriptSource = ` local leaseKey = ARGV[1] .. ARGV[2] .. ':' .. ARGV[3] local lease = redis.call('HMGET', leaseKey, 'token', 'receiptID') if lease[1] ~= ARGV[4] or lease[2] ~= ARGV[5] then return 0 end return redis.call('DEL', leaseKey) ` var ( reserveRedisScript = redis.NewScript(reserveScriptSource) commitReservationRedisScript = redis.NewScript(commitReservationScriptSource) abortReservationRedisScript = redis.NewScript(abortReservationScriptSource) ) // ReserveCurrentContext 原子预留当前版本的一段队列,但不推进用户偏移。 func ReserveCurrentContext( ctx context.Context, client *redis.Client, uid uint64, size int, requestID string, ) (Reservation, error) { if client == nil { return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil") } return reserveContext(ctx, client, uid, size, "", requestID) } // ReserveVersionContext 与 ReserveCurrentContext 相同,但要求current仍为指定版本。 func ReserveVersionContext( ctx context.Context, client *redis.Client, uid uint64, size int, expectedVersion, requestID string, ) (Reservation, error) { if client == nil { return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil") } return reserveContext(ctx, client, uid, size, expectedVersion, requestID) } func reserveContext( ctx context.Context, client scriptClient, uid uint64, size int, expectedVersion, requestID string, ) (Reservation, error) { if size <= 0 { return Reservation{}, nil } if ctx == nil { return Reservation{}, fmt.Errorf("short recommend reserve context must not be nil") } if client == nil { return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil") } if uid == 0 { return Reservation{}, fmt.Errorf("short recommend reserve uid must be positive") } if err := ctx.Err(); err != nil { return Reservation{}, err } requestID = strings.TrimSpace(requestID) if len(requestID) > maxTakeRequestIDLength { return Reservation{}, fmt.Errorf( "short recommend request ID exceeds %d bytes", maxTakeRequestIDLength, ) } leaseToken, err := newReservationToken() if err != nil { return Reservation{}, err } receiptID := leaseToken if requestID != "" { sum := sha256.Sum256([]byte(strconv.FormatUint(uid, 10) + ":" + requestID)) receiptID = hex.EncodeToString(sum[:16]) } keys := []string{CurrentKey} args := []interface{}{ "recommend:short:queue:", "recommend:short:offset:", "recommend:short:meta:", strconv.FormatUint(uid, 10), "recommend:short:reservation:", "recommend:short:commit-receipt:", size, expectedVersion, leaseToken, receiptID, int64(reservationTTL / time.Millisecond), } raw, err := client.RunScriptContext(ctx, reserveRedisScript, keys, args...) if err != nil && ctx.Err() == nil { // 预留脚本可能已经执行成功但响应在网络中丢失。使用完全相同的 // token/receipt 重试可命中原租约,不会创建第二个预留或推进偏移。 raw, err = client.RunScriptContext(ctx, reserveRedisScript, keys, args...) } if err != nil { return Reservation{}, err } return parseReservation(raw, receiptID) } func newReservationToken() (string, error) { value := make([]byte, 16) if _, err := cryptorand.Read(value); err != nil { return "", fmt.Errorf("generate short recommend reservation token: %w", err) } return hex.EncodeToString(value), nil } func parseReservation(raw interface{}, receiptID string) (Reservation, error) { items, ok := raw.([]interface{}) if !ok || len(items) == 0 { return Reservation{}, nil } if len(items) < 6 { return Reservation{}, fmt.Errorf("short recommend reserve returned malformed result") } result := Reservation{ Version: fmt.Sprint(items[0]), LeaseToken: fmt.Sprint(items[5]), ReceiptID: receiptID, } status := fmt.Sprint(items[1]) var err error result.Length, err = strconv.Atoi(fmt.Sprint(items[2])) if err != nil { return Reservation{}, fmt.Errorf("short recommend reserve returned invalid length: %w", err) } result.Offset, err = strconv.Atoi(fmt.Sprint(items[3])) if err != nil { return Reservation{}, fmt.Errorf("short recommend reserve returned invalid offset: %w", err) } result.Reserved, err = strconv.Atoi(fmt.Sprint(items[4])) if err != nil { return Reservation{}, fmt.Errorf("short recommend reserve returned invalid size: %w", err) } switch status { case "VERSION_CHANGED": return result, ErrVersionChanged case "UNHEALTHY": return result, ErrQueueUnhealthy case "BUSY": return result, ErrReservationBusy case "COMMITTED": result.AlreadyCommitted = true case "RESERVED": default: return Reservation{}, fmt.Errorf("short recommend reserve returned unknown status %q", status) } if result.Version == "" || result.Length <= 0 || result.Offset < 0 || result.Offset >= result.Length || result.Reserved <= 0 || result.Reserved > result.Length || len(items)-6 != result.Reserved { return Reservation{}, fmt.Errorf( "short recommend reserve returned invalid bounds: version=%q length=%d offset=%d reserved=%d ids=%d", result.Version, result.Length, result.Offset, result.Reserved, len(items)-6, ) } for _, item := range items[6:] { result.IDs = append(result.IDs, fmt.Sprint(item)) } return result, nil } // CommitReservationContext 原子校验预留并推进实际扫描的条数。 func CommitReservationContext( ctx context.Context, client *redis.Client, uid uint64, reservation Reservation, consumed int, ) error { if client == nil { return fmt.Errorf("short recommend commit Redis client must not be nil") } return commitReservationContext(ctx, client, uid, reservation, consumed) } func commitReservationContext( ctx context.Context, client scriptClient, uid uint64, reservation Reservation, consumed int, ) error { if ctx == nil { return fmt.Errorf("short recommend commit context must not be nil") } if client == nil { return fmt.Errorf("short recommend commit Redis client must not be nil") } if err := validateReservationCommit(uid, reservation, consumed); err != nil { return err } if err := ctx.Err(); err != nil { return err } raw, err := client.RunScriptContext( ctx, commitReservationRedisScript, []string{CurrentKey}, "recommend:short:queue:", "recommend:short:offset:", "recommend:short:meta:", "recommend:short:reservation:", strconv.FormatUint(uid, 10), reservation.Version, reservation.LeaseToken, reservation.ReceiptID, reservation.Length, reservation.Offset, reservation.Reserved, consumed, int64(takeIdempotencyTTL/time.Second), "recommend:short:commit-receipt:", ) if err != nil { return err } switch fmt.Sprint(raw) { case "1", "2": return nil case "-1": return ErrVersionChanged case "-2": return ErrQueueUnhealthy case "-3": return ErrReservationExpired case "-4", "-5": return ErrReservationConflict default: return fmt.Errorf("short recommend commit returned unexpected result %v", raw) } } func validateReservationCommit(uid uint64, reservation Reservation, consumed int) error { if uid == 0 { return fmt.Errorf("short recommend commit uid must be positive") } if reservation.Version == "" || reservation.Length <= 0 || reservation.Offset < 0 || reservation.Offset >= reservation.Length || reservation.Reserved <= 0 || reservation.Reserved > reservation.Length || reservation.LeaseToken == "" || reservation.ReceiptID == "" { return fmt.Errorf("short recommend commit reservation is invalid") } if consumed <= 0 || consumed > reservation.Reserved { return fmt.Errorf("short recommend consumed size is out of reservation bounds") } return nil } // AbortReservationContext 仅删除token仍匹配的未提交预留,不改变用户偏移。 func AbortReservationContext( ctx context.Context, client *redis.Client, uid uint64, reservation Reservation, ) error { if client == nil { return fmt.Errorf("short recommend abort Redis client must not be nil") } return abortReservationContext(ctx, client, uid, reservation) } func abortReservationContext( ctx context.Context, client scriptClient, uid uint64, reservation Reservation, ) error { if ctx == nil { return fmt.Errorf("short recommend abort context must not be nil") } if client == nil { return fmt.Errorf("short recommend abort Redis client must not be nil") } if uid == 0 || reservation.Version == "" || reservation.LeaseToken == "" || reservation.ReceiptID == "" { return nil } if err := ctx.Err(); err != nil { return err } _, err := client.RunScriptContext( ctx, abortReservationRedisScript, nil, "recommend:short:reservation:", reservation.Version, strconv.FormatUint(uid, 10), reservation.LeaseToken, reservation.ReceiptID, ) return err } type QueueHealth struct { Healthy bool Status string Version string QueueLength int MetadataLength int QueueTTL time.Duration MetadataTTL time.Duration } const healthScriptSource = ` local version = redis.call('GET', KEYS[1]) if not version then return {'NO_CURRENT', '', '0', '-1', '-2', '-2'} end if ARGV[1] ~= '' then local expectedDate = ARGV[1] local matchesDate = version == expectedDate or ( string.len(version) > string.len(expectedDate) + 1 and string.sub(version, 1, string.len(expectedDate) + 1) == expectedDate .. '-' ) if not matchesDate then return {'VERSION_MISMATCH', version, '0', '-1', '-2', '-2'} end end local queueKey = ARGV[2] .. version local metaKey = ARGV[3] .. version local queueLength = redis.call('LLEN', queueKey) local metadataLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1') local queueTTL = redis.call('PTTL', queueKey) local metadataTTL = redis.call('PTTL', metaKey) local status = 'OK' if queueLength <= 0 then status = 'QUEUE_EMPTY' elseif metadataLength < 0 then status = 'META_MISSING' elseif metadataLength ~= queueLength then status = 'LENGTH_MISMATCH' elseif queueTTL <= 0 or metadataTTL <= 0 then status = 'TTL_MISSING' elseif math.abs(queueTTL - metadataTTL) > tonumber(ARGV[4]) then status = 'TTL_MISMATCH' end return { status, version, tostring(queueLength), tostring(metadataLength), tostring(queueTTL), tostring(metadataTTL) } ` var healthRedisScript = redis.NewScript(healthScriptSource) // Health 原子检查current、队列长度、metadata.length和两者TTL。 func Health(client *redis.Client, expectedVersion string) (QueueHealth, error) { return HealthContext(context.Background(), client, expectedVersion) } // HealthContext 与 Health 相同,并将取消和超时传递到Redis。 func HealthContext( ctx context.Context, client *redis.Client, expectedVersion string, ) (QueueHealth, error) { return healthContext(ctx, client, expectedVersion) } func healthContext( ctx context.Context, client scriptClient, expectedVersion string, ) (QueueHealth, error) { if ctx == nil { return QueueHealth{}, fmt.Errorf("short recommend health context must not be nil") } if client == nil { return QueueHealth{}, fmt.Errorf("short recommend health Redis client must not be nil") } if err := ctx.Err(); err != nil { return QueueHealth{}, err } raw, err := client.RunScriptContext(ctx, healthRedisScript, []string{CurrentKey}, expectedVersion, "recommend:short:queue:", "recommend:short:meta:", int64(healthTTLDriftTolerance/time.Millisecond)) if err != nil { return QueueHealth{}, err } items, ok := raw.([]interface{}) if !ok || len(items) != 6 { return QueueHealth{}, fmt.Errorf("short recommend health returned malformed result") } health := QueueHealth{ Status: fmt.Sprint(items[0]), Version: fmt.Sprint(items[1]), } health.QueueLength, err = strconv.Atoi(fmt.Sprint(items[2])) if err != nil { return QueueHealth{}, fmt.Errorf("short recommend health returned invalid queue length: %w", err) } health.MetadataLength, err = strconv.Atoi(fmt.Sprint(items[3])) if err != nil { return QueueHealth{}, fmt.Errorf("short recommend health returned invalid metadata length: %w", err) } queueTTLMillis, err := strconv.ParseInt(fmt.Sprint(items[4]), 10, 64) if err != nil { return QueueHealth{}, fmt.Errorf("short recommend health returned invalid queue TTL: %w", err) } metadataTTLMillis, err := strconv.ParseInt(fmt.Sprint(items[5]), 10, 64) if err != nil { return QueueHealth{}, fmt.Errorf("short recommend health returned invalid metadata TTL: %w", err) } health.QueueTTL = time.Duration(queueTTLMillis) * time.Millisecond health.MetadataTTL = time.Duration(metadataTTLMillis) * time.Millisecond health.Healthy = health.Status == "OK" return health, nil }