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

backend/internal/repository/batch_image_download_limiter.go

@@ -0,0 +1,102 @@
+package repository
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	defaultBatchImageDownloadActivePrefix = "batch_image:download:active:"
+	defaultBatchImageDownloadActiveTTL    = 10 * time.Minute
+	defaultBatchImageDownloadConcurrency  = 2
+)
+
+var batchImageDownloadAcquireScript = redis.NewScript(`
+local current = tonumber(redis.call("GET", KEYS[1]) or "0")
+local max = tonumber(ARGV[1])
+if current >= max then
+  return 0
+end
+redis.call("INCR", KEYS[1])
+redis.call("EXPIRE", KEYS[1], ARGV[2])
+return 1
+`)
+
+var batchImageDownloadReleaseScript = redis.NewScript(`
+local current = tonumber(redis.call("GET", KEYS[1]) or "0")
+if current <= 1 then
+  redis.call("DEL", KEYS[1])
+  return 0
+end
+return redis.call("DECR", KEYS[1])
+`)
+
+type batchImageDownloadLimiter struct {
+	rdb          *redis.Client
+	activePrefix string
+	maxActive    int
+	ttl          time.Duration
+}
+
+func NewBatchImageDownloadLimiter(rdb *redis.Client, cfg *config.Config) service.BatchImageDownloadLimiter {
+	maxActive := defaultBatchImageDownloadConcurrency
+	ttl := defaultBatchImageDownloadActiveTTL
+	if cfg != nil {
+		if cfg.BatchImage.MaxDownloadConcurrencyPerUser > 0 {
+			maxActive = cfg.BatchImage.MaxDownloadConcurrencyPerUser
+		}
+		if cfg.BatchImage.MaxDownloadDurationSeconds > 0 {
+			ttl = time.Duration(cfg.BatchImage.MaxDownloadDurationSeconds) * time.Second
+		}
+	}
+	return &batchImageDownloadLimiter{
+		rdb:          rdb,
+		activePrefix: defaultBatchImageDownloadActivePrefix,
+		maxActive:    maxActive,
+		ttl:          ttl,
+	}
+}
+
+func (l *batchImageDownloadLimiter) Acquire(ctx context.Context, userID string, kind string) (service.BatchImageDownloadPermit, error) {
+	if l == nil || l.rdb == nil {
+		return nil, service.ErrBatchImageDownloadLimited
+	}
+	key := l.activeKey(userID)
+	ok, err := batchImageDownloadAcquireScript.Run(ctx, l.rdb, []string{key}, l.maxActive, int(l.ttl.Seconds())).Int()
+	if err != nil {
+		return nil, err
+	}
+	if ok != 1 {
+		return nil, service.ErrBatchImageDownloadLimited
+	}
+	return &batchImageDownloadPermit{rdb: l.rdb, key: key}, nil
+}
+
+func (l *batchImageDownloadLimiter) activeKey(userID string) string {
+	return l.activePrefix + userID
+}
+
+type batchImageDownloadPermit struct {
+	rdb  *redis.Client
+	key  string
+	once sync.Once
+	err  error
+}
+
+func (p *batchImageDownloadPermit) Release(ctx context.Context) error {
+	if p == nil || p.rdb == nil || p.key == "" {
+		return nil
+	}
+	p.once.Do(func() {
+		_, p.err = batchImageDownloadReleaseScript.Run(ctx, p.rdb, []string{p.key}).Result()
+	})
+	return p.err
+}
+
+var _ service.BatchImageDownloadLimiter = (*batchImageDownloadLimiter)(nil)
+var _ service.BatchImageDownloadPermit = (*batchImageDownloadPermit)(nil)