From f069f98b353e19c68201f87c1b30b33554dc488b Mon Sep 17 00:00:00 2001
From: zizi <zizi2000mail@gmail.com>
Date: Wed, 20 May 2026 14:27:18 +0800
Subject: [PATCH] feat: add concurrency guard service

Add Redis-backed user and channel concurrency slots plus user RPM checks with fail-open behavior.
---
 service/concurrency.go      | 77 +++++++++++++++++++++++++++++++++++++
 service/concurrency_test.go | 59 ++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+)
 create mode 100644 service/concurrency.go
 create mode 100644 service/concurrency_test.go

diff --git a/service/concurrency.go b/service/concurrency.go
new file mode 100644
index 00000000..8c1b9c51
--- /dev/null
+++ b/service/concurrency.go
@@ -0,0 +1,77 @@
+package service
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+)
+
+const (
+	userConcurrencyKeyPrefix    = "user:concurrency:"
+	channelConcurrencyKeyPrefix = "channel:concurrency:"
+	userRPMKeyPrefix            = "user:rpm:"
+
+	concurrencyTTL = 5 * time.Minute
+	rpmTTL         = 2 * time.Minute
+)
+
+func AcquireUserConcurrency(userId, limit int) (bool, func()) {
+	return acquireConcurrency(userConcurrencyKeyPrefix+strconv.Itoa(userId), limit)
+}
+
+func AcquireChannelConcurrency(channelId, limit int) (bool, func()) {
+	return acquireConcurrency(channelConcurrencyKeyPrefix+strconv.Itoa(channelId), limit)
+}
+
+func CheckUserRPM(userId, limit int) bool {
+	if limit <= 0 || !redisAvailable() {
+		return true
+	}
+
+	key := userRPMKeyPrefix + strconv.Itoa(userId) + ":" + time.Now().Format("200601021504")
+	ctx := context.Background()
+
+	count, err := common.RDB.Incr(ctx, key).Result()
+	if err != nil {
+		return true
+	}
+	if err = common.RDB.Expire(ctx, key, rpmTTL).Err(); err != nil {
+		return true
+	}
+	return count <= int64(limit)
+}
+
+func acquireConcurrency(key string, limit int) (bool, func()) {
+	release := func() {}
+	if limit <= 0 || !redisAvailable() {
+		return true, release
+	}
+
+	ctx := context.Background()
+	count, err := common.RDB.Incr(ctx, key).Result()
+	if err != nil {
+		return true, release
+	}
+
+	release = func() {
+		if redisAvailable() {
+			_ = common.RDB.Decr(context.Background(), key).Err()
+		}
+	}
+
+	if err = common.RDB.Expire(ctx, key, concurrencyTTL).Err(); err != nil {
+		return true, release
+	}
+	if count > int64(limit) {
+		release()
+		return false, func() {}
+	}
+
+	return true, release
+}
+
+func redisAvailable() bool {
+	return common.RedisEnabled && common.RDB != nil
+}
diff --git a/service/concurrency_test.go b/service/concurrency_test.go
new file mode 100644
index 00000000..cc4f1011
--- /dev/null
+++ b/service/concurrency_test.go
@@ -0,0 +1,59 @@
+package service
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAcquireConcurrencyAllowsWhenLimitDisabled(t *testing.T) {
+	previousRedisEnabled := common.RedisEnabled
+	previousRDB := common.RDB
+	common.RedisEnabled = true
+	common.RDB = nil
+	t.Cleanup(func() {
+		common.RedisEnabled = previousRedisEnabled
+		common.RDB = previousRDB
+	})
+
+	allowed, release := AcquireUserConcurrency(1, 0)
+	require.True(t, allowed)
+	require.NotNil(t, release)
+	release()
+
+	allowed, release = AcquireChannelConcurrency(1, -1)
+	require.True(t, allowed)
+	require.NotNil(t, release)
+	release()
+
+	require.True(t, CheckUserRPM(1, 0))
+}
+
+func TestAcquireConcurrencyAllowsWhenRedisUnavailable(t *testing.T) {
+	for _, redisEnabled := range []bool{false, true} {
+		t.Run("redis_enabled_"+strconv.FormatBool(redisEnabled), func(t *testing.T) {
+			previousRedisEnabled := common.RedisEnabled
+			previousRDB := common.RDB
+			common.RedisEnabled = redisEnabled
+			common.RDB = nil
+			t.Cleanup(func() {
+				common.RedisEnabled = previousRedisEnabled
+				common.RDB = previousRDB
+			})
+
+			allowed, release := AcquireUserConcurrency(1, 1)
+			require.True(t, allowed)
+			require.NotNil(t, release)
+			release()
+
+			allowed, release = AcquireChannelConcurrency(1, 1)
+			require.True(t, allowed)
+			require.NotNil(t, release)
+			release()
+
+			require.True(t, CheckUserRPM(1, 1))
+		})
+	}
+}