1261 lines
38 KiB
Go
1261 lines
38 KiB
Go
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
|
|
}
|