Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/repository/batch_image_queue.go

@@ -0,0 +1,354 @@
+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)