355 lines
10 KiB
Go
355 lines
10 KiB
Go
|
package repository
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"crypto/rand"
|
||
|
|
"encoding/hex"
|
||
|
|
"errors"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
|
|
"github.com/redis/go-redis/v9"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
defaultBatchImageReadyKey = "batch_image:queue:ready"
|
||
|
|
defaultBatchImageDelayedKey = "batch_image:queue:delayed"
|
||
|
|
defaultBatchImageActiveKey = "batch_image:queue:active"
|
||
|
|
defaultBatchImageInflightPrefix = "batch_image:queue:inflight:"
|
||
|
|
defaultBatchImageLockPrefix = "batch_image:queue:lock:"
|
||
|
|
defaultBatchImageInflightTTL = 7 * 24 * time.Hour
|
||
|
|
defaultBatchImageJobLockTTL = 5 * time.Minute
|
||
|
|
|
||
|
|
// batchImageReservePollInterval 是原子 Reserve 脚本空轮询的间隔。
|
||
|
|
// 用轮询替代 BRPop 是为了保证 "弹出 + 写 active" 的原子性。
|
||
|
|
batchImageReservePollInterval = time.Second
|
||
|
|
)
|
||
|
|
|
||
|
|
var batchImageMoveDueDelayedScript = redis.NewScript(`
|
||
|
|
local jobs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, ARGV[2])
|
||
|
|
for _, job in ipairs(jobs) do
|
||
|
|
redis.call("ZREM", KEYS[1], job)
|
||
|
|
redis.call("LPUSH", KEYS[2], job)
|
||
|
|
end
|
||
|
|
return #jobs
|
||
|
|
`)
|
||
|
|
|
||
|
|
var batchImageRecoverStaleActiveScript = redis.NewScript(`
|
||
|
|
local jobs = redis.call("ZRANGEBYSCORE", KEYS[1], "-inf", ARGV[1], "LIMIT", 0, ARGV[2])
|
||
|
|
for _, job in ipairs(jobs) do
|
||
|
|
redis.call("ZREM", KEYS[1], job)
|
||
|
|
redis.call("LPUSH", KEYS[2], job)
|
||
|
|
end
|
||
|
|
return #jobs
|
||
|
|
`)
|
||
|
|
|
||
|
|
var batchImageReleaseLockScript = redis.NewScript(`
|
||
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call("DEL", KEYS[1])
|
||
|
|
end
|
||
|
|
return 0
|
||
|
|
`)
|
||
|
|
|
||
|
|
var batchImageRefreshLockScript = redis.NewScript(`
|
||
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
|
|
return redis.call("PEXPIRE", KEYS[1], ARGV[2])
|
||
|
|
end
|
||
|
|
return 0
|
||
|
|
`)
|
||
|
|
|
||
|
|
// batchImageReserveScript 原子地从 ready 弹出并写入 active zset。
|
||
|
|
// BRPop + ZAdd 两步方案在两步之间进程崩溃时 job 会脱离所有队列结构，
|
||
|
|
// 且 inflight 去重键（默认 7 天）会挡住所有重新入队。
|
||
|
|
var batchImageReserveScript = redis.NewScript(`
|
||
|
|
local job = redis.call("RPOP", KEYS[1])
|
||
|
|
if not job then
|
||
|
|
return nil
|
||
|
|
end
|
||
|
|
redis.call("ZADD", KEYS[2], ARGV[1], job)
|
||
|
|
return job
|
||
|
|
`)
|
||
|
|
|
||
|
|
// batchImageEnqueueScript 原子地设置 inflight 去重键并推入 ready。
|
||
|
|
// SetNX + LPush 两步方案在两步之间进程崩溃时，inflight 键（默认 7 天）
|
||
|
|
// 会挡住所有后续入队，而 job 从未进入 ready。
|
||
|
|
var batchImageEnqueueScript = redis.NewScript(`
|
||
|
|
if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then
|
||
|
|
redis.call("LPUSH", KEYS[2], ARGV[1])
|
||
|
|
return 1
|
||
|
|
end
|
||
|
|
return 0
|
||
|
|
`)
|
||
|
|
|
||
|
|
type batchImageQueue struct {
|
||
|
|
rdb *redis.Client
|
||
|
|
readyKey string
|
||
|
|
delayedKey string
|
||
|
|
activeKey string
|
||
|
|
inflightPrefix string
|
||
|
|
lockPrefix string
|
||
|
|
inflightTTL time.Duration
|
||
|
|
lockTTL time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewBatchImageQueue(rdb *redis.Client, cfg *config.Config) service.BatchImageQueue {
|
||
|
|
return newBatchImageQueueWithOptions(rdb, batchImageQueueOptionsFromConfig(cfg))
|
||
|
|
}
|
||
|
|
|
||
|
|
type batchImageQueueOptions struct {
|
||
|
|
ReadyKey string
|
||
|
|
DelayedKey string
|
||
|
|
ActiveKey string
|
||
|
|
InflightPrefix string
|
||
|
|
LockPrefix string
|
||
|
|
InflightTTL time.Duration
|
||
|
|
LockTTL time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func newBatchImageQueueWithOptions(rdb *redis.Client, opts batchImageQueueOptions) *batchImageQueue {
|
||
|
|
opts = normalizeBatchImageQueueOptions(opts)
|
||
|
|
return &batchImageQueue{
|
||
|
|
rdb: rdb,
|
||
|
|
readyKey: opts.ReadyKey,
|
||
|
|
delayedKey: opts.DelayedKey,
|
||
|
|
activeKey: opts.ActiveKey,
|
||
|
|
inflightPrefix: opts.InflightPrefix,
|
||
|
|
lockPrefix: opts.LockPrefix,
|
||
|
|
inflightTTL: opts.InflightTTL,
|
||
|
|
lockTTL: opts.LockTTL,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func batchImageQueueOptionsFromConfig(cfg *config.Config) batchImageQueueOptions {
|
||
|
|
if cfg == nil {
|
||
|
|
return batchImageQueueOptions{}
|
||
|
|
}
|
||
|
|
return batchImageQueueOptions{
|
||
|
|
ReadyKey: cfg.BatchImage.QueueReadyKey,
|
||
|
|
DelayedKey: cfg.BatchImage.QueueDelayedKey,
|
||
|
|
ActiveKey: cfg.BatchImage.QueueActiveKey,
|
||
|
|
InflightPrefix: cfg.BatchImage.InflightKeyPrefix,
|
||
|
|
LockPrefix: cfg.BatchImage.LockKeyPrefix,
|
||
|
|
InflightTTL: time.Duration(cfg.BatchImage.InflightTTLSeconds) * time.Second,
|
||
|
|
LockTTL: time.Duration(cfg.BatchImage.JobLockTTLSeconds) * time.Second,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func normalizeBatchImageQueueOptions(opts batchImageQueueOptions) batchImageQueueOptions {
|
||
|
|
if opts.ReadyKey == "" {
|
||
|
|
opts.ReadyKey = defaultBatchImageReadyKey
|
||
|
|
}
|
||
|
|
if opts.DelayedKey == "" {
|
||
|
|
opts.DelayedKey = defaultBatchImageDelayedKey
|
||
|
|
}
|
||
|
|
if opts.ActiveKey == "" {
|
||
|
|
opts.ActiveKey = defaultBatchImageActiveKey
|
||
|
|
}
|
||
|
|
if opts.InflightPrefix == "" {
|
||
|
|
opts.InflightPrefix = defaultBatchImageInflightPrefix
|
||
|
|
}
|
||
|
|
if opts.LockPrefix == "" {
|
||
|
|
opts.LockPrefix = defaultBatchImageLockPrefix
|
||
|
|
}
|
||
|
|
if opts.InflightTTL <= 0 {
|
||
|
|
opts.InflightTTL = defaultBatchImageInflightTTL
|
||
|
|
}
|
||
|
|
if opts.LockTTL <= 0 {
|
||
|
|
opts.LockTTL = defaultBatchImageJobLockTTL
|
||
|
|
}
|
||
|
|
return opts
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) Enqueue(ctx context.Context, batchID string) error {
|
||
|
|
if !service.IsValidBatchImageID(batchID) {
|
||
|
|
return service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
|
||
|
|
applied, err := batchImageEnqueueScript.Run(ctx, q.rdb,
|
||
|
|
[]string{q.inflightKey(batchID), q.readyKey},
|
||
|
|
batchID, q.inflightTTL.Milliseconds(),
|
||
|
|
).Int()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if applied == 0 {
|
||
|
|
return service.ErrBatchImageAlreadyQueued
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) Reserve(ctx context.Context, blockTimeout time.Duration) (service.ReservedBatchImageJob, error) {
|
||
|
|
deadline := time.Now().Add(blockTimeout)
|
||
|
|
for {
|
||
|
|
batchID, err := q.reserveOnce(ctx)
|
||
|
|
if err == nil {
|
||
|
|
return service.ReservedBatchImageJob{BatchID: batchID}, nil
|
||
|
|
}
|
||
|
|
if !errors.Is(err, service.ErrBatchImageQueueEmpty) {
|
||
|
|
return service.ReservedBatchImageJob{}, err
|
||
|
|
}
|
||
|
|
remaining := time.Until(deadline)
|
||
|
|
if remaining <= 0 {
|
||
|
|
return service.ReservedBatchImageJob{}, service.ErrBatchImageQueueEmpty
|
||
|
|
}
|
||
|
|
wait := batchImageReservePollInterval
|
||
|
|
if remaining < wait {
|
||
|
|
wait = remaining
|
||
|
|
}
|
||
|
|
timer := time.NewTimer(wait)
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
timer.Stop()
|
||
|
|
return service.ReservedBatchImageJob{}, ctx.Err()
|
||
|
|
case <-timer.C:
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) reserveOnce(ctx context.Context) (string, error) {
|
||
|
|
raw, err := batchImageReserveScript.Run(ctx, q.rdb, []string{q.readyKey, q.activeKey}, time.Now().UnixMilli()).Result()
|
||
|
|
if errors.Is(err, redis.Nil) {
|
||
|
|
return "", service.ErrBatchImageQueueEmpty
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
batchID, ok := raw.(string)
|
||
|
|
if !ok || !service.IsValidBatchImageID(batchID) {
|
||
|
|
// 非法 payload 已被脚本写入 active，必须移除，否则 stale 恢复会把它
|
||
|
|
// 无限重投回 ready。
|
||
|
|
if ok && batchID != "" {
|
||
|
|
_ = q.rdb.ZRem(ctx, q.activeKey, batchID).Err()
|
||
|
|
}
|
||
|
|
return "", service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
return batchID, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) RequeueAfter(ctx context.Context, batchID string, delay time.Duration) error {
|
||
|
|
if !service.IsValidBatchImageID(batchID) {
|
||
|
|
return service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
pipe := q.rdb.TxPipeline()
|
||
|
|
pipe.ZRem(ctx, q.activeKey, batchID)
|
||
|
|
pipe.ZRem(ctx, q.delayedKey, batchID)
|
||
|
|
if delay <= 0 {
|
||
|
|
pipe.LPush(ctx, q.readyKey, batchID)
|
||
|
|
} else {
|
||
|
|
pipe.ZAdd(ctx, q.delayedKey, redis.Z{
|
||
|
|
Score: float64(time.Now().Add(delay).UnixMilli()),
|
||
|
|
Member: batchID,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
_, err := pipe.Exec(ctx)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) Ack(ctx context.Context, batchID string) error {
|
||
|
|
if !service.IsValidBatchImageID(batchID) {
|
||
|
|
return service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
pipe := q.rdb.TxPipeline()
|
||
|
|
pipe.ZRem(ctx, q.activeKey, batchID)
|
||
|
|
pipe.ZRem(ctx, q.delayedKey, batchID)
|
||
|
|
pipe.Del(ctx, q.inflightKey(batchID))
|
||
|
|
_, err := pipe.Exec(ctx)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) Heartbeat(ctx context.Context, batchID string) error {
|
||
|
|
if !service.IsValidBatchImageID(batchID) {
|
||
|
|
return service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
// XX：只刷新已存在的 active 成员。无条件 ZAdd 会在 Ack/Requeue 之后的
|
||
|
|
// 竞态心跳里把幽灵成员塞回 active zset。
|
||
|
|
return q.rdb.ZAddXX(ctx, q.activeKey, redis.Z{
|
||
|
|
Score: float64(time.Now().UnixMilli()),
|
||
|
|
Member: batchID,
|
||
|
|
}).Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) MoveDueDelayedToReady(ctx context.Context, limit int) (int, error) {
|
||
|
|
if limit <= 0 {
|
||
|
|
limit = 100
|
||
|
|
}
|
||
|
|
return batchImageMoveDueDelayedScript.Run(ctx, q.rdb, []string{q.delayedKey, q.readyKey}, time.Now().UnixMilli(), limit).Int()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) RecoverStaleActive(ctx context.Context, staleAfter time.Duration, limit int) (int, error) {
|
||
|
|
if staleAfter <= 0 {
|
||
|
|
return 0, service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
if limit <= 0 {
|
||
|
|
limit = 100
|
||
|
|
}
|
||
|
|
cutoff := time.Now().Add(-staleAfter).UnixMilli()
|
||
|
|
return batchImageRecoverStaleActiveScript.Run(ctx, q.rdb, []string{q.activeKey, q.readyKey}, cutoff, limit).Int()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) TryAcquireJobLock(ctx context.Context, batchID string, ttl time.Duration) (service.BatchImageJobLock, bool, error) {
|
||
|
|
if !service.IsValidBatchImageID(batchID) {
|
||
|
|
return nil, false, service.ErrInvalidBatchImageQueuePayload
|
||
|
|
}
|
||
|
|
if ttl <= 0 {
|
||
|
|
ttl = q.lockTTL
|
||
|
|
}
|
||
|
|
token, err := newBatchImageLockToken()
|
||
|
|
if err != nil {
|
||
|
|
return nil, false, err
|
||
|
|
}
|
||
|
|
key := q.lockKey(batchID)
|
||
|
|
ok, err := q.rdb.SetNX(ctx, key, token, ttl).Result()
|
||
|
|
if err != nil {
|
||
|
|
return nil, false, err
|
||
|
|
}
|
||
|
|
if !ok {
|
||
|
|
return nil, false, nil
|
||
|
|
}
|
||
|
|
return &batchImageRedisJobLock{rdb: q.rdb, key: key, token: token}, true, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) inflightKey(batchID string) string {
|
||
|
|
return q.inflightPrefix + batchID
|
||
|
|
}
|
||
|
|
|
||
|
|
func (q *batchImageQueue) lockKey(batchID string) string {
|
||
|
|
return q.lockPrefix + batchID
|
||
|
|
}
|
||
|
|
|
||
|
|
type batchImageRedisJobLock struct {
|
||
|
|
rdb *redis.Client
|
||
|
|
key string
|
||
|
|
token string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *batchImageRedisJobLock) Release(ctx context.Context) error {
|
||
|
|
if l == nil || l.rdb == nil || l.key == "" || l.token == "" {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return batchImageReleaseLockScript.Run(ctx, l.rdb, []string{l.key}, l.token).Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
// Refresh 在仍持有锁（token 匹配）时续期 TTL，供长处理任务的心跳调用。
|
||
|
|
func (l *batchImageRedisJobLock) Refresh(ctx context.Context, ttl time.Duration) error {
|
||
|
|
if l == nil || l.rdb == nil || l.key == "" || l.token == "" {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
if ttl <= 0 {
|
||
|
|
ttl = defaultBatchImageJobLockTTL
|
||
|
|
}
|
||
|
|
return batchImageRefreshLockScript.Run(ctx, l.rdb, []string{l.key}, l.token, ttl.Milliseconds()).Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
var _ service.BatchImageJobLockRefresher = (*batchImageRedisJobLock)(nil)
|
||
|
|
|
||
|
|
func newBatchImageLockToken() (string, error) {
|
||
|
|
var b [16]byte
|
||
|
|
if _, err := rand.Read(b[:]); err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
return hex.EncodeToString(b[:]), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
var _ service.BatchImageQueue = (*batchImageQueue)(nil)
|