缓解io压力

internal/task/polling_base.go

@@ -16,28 +16,32 @@ import (
 // acquireConcurrencyScript 原子获取并发信号量的 Lua 脚本
 // INCR + EXPIRE 合并为单个服务端操作，消除二者之间的崩溃窗口：
 // 若 Worker 在 INCR 后、EXPIRE 前崩溃，key 将永久留在 Redis 导致计数器卡死。
-// KEYS[1]: 当前并发计数 key
-// ARGV[1]: 最大并发数；ARGV[2]: key TTL（秒）
-// 返回 -1 表示超额拒绝，>0 表示成功获取后的计数值
+// KEYS[1]: 全部轮询计数；KEYS[2]: 分类轮询计数
+// ARGV[1]: 总量上限；ARGV[2]: 分类上限；ARGV[3]: key TTL（秒）
+// 返回 -1 表示任一上限超额，>0 表示成功获取后的总计数值
 var acquireConcurrencyScript = redis.NewScript(`
-    local current = redis.call('INCR', KEYS[1])
-    if tonumber(current) > tonumber(ARGV[1]) then
+    local total = redis.call('INCR', KEYS[1])
+    local kind = redis.call('INCR', KEYS[2])
+    if tonumber(total) > tonumber(ARGV[1]) or tonumber(kind) > tonumber(ARGV[2]) then
         redis.call('DECR', KEYS[1])
+        redis.call('DECR', KEYS[2])
         return -1
     end
-    redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]))
-    return current
+    redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3]))
+    redis.call('EXPIRE', KEYS[2], tonumber(ARGV[3]))
+    return total
 `)
 
 var releaseConcurrencyScript = redis.NewScript(`
-    local current = tonumber(redis.call('GET', KEYS[1]) or '0') or 0
-    if current <= 0 then
-        if redis.call('EXISTS', KEYS[1]) == 1 then
-            redis.call('SET', KEYS[1], 0, 'KEEPTTL')
+    for _, key in ipairs(KEYS) do
+        local current = tonumber(redis.call('GET', key) or '0') or 0
+        if current <= 0 then
+            if redis.call('EXISTS', key) == 1 then redis.call('SET', key, 0, 'KEEPTTL') end
+        else
+            redis.call('DECR', key)
         end
-        return 0
     end
-    return redis.call('DECR', KEYS[1])
+    return 0
 `)
 
 const pollingFallbackOperationTimeout = 5 * time.Second
@@ -52,13 +56,14 @@ func pollingFallbackContext() (context.Context, context.CancelFunc) {
 // PollingBase 轮询共享基类
 // 封装并发控制、卡缓存、重入队、配置间隔查询等公共方法，所有 Handler 共享
 type PollingBase struct {
-	redis        *redis.Client
-	queueMgr     *polling.PollingQueueManager
-	configMgr    *polling.PollingConfigManager
-	iotCardStore *postgres.IotCardStore
-	logger       *zap.Logger
-	verboseLog   bool
-	trafficLock  *cardtrafficlock.Lock
+	redis               *redis.Client
+	queueMgr            *polling.PollingQueueManager
+	configMgr           *polling.PollingConfigManager
+	iotCardStore        *postgres.IotCardStore
+	logger              *zap.Logger
+	verboseLog          bool
+	totalMaxConcurrency int
+	trafficLock         *cardtrafficlock.Lock
 }
 
 // NewPollingBase 创建轮询共享基类
@@ -69,15 +74,17 @@ func NewPollingBase(
 	iotCardStore *postgres.IotCardStore,
 	logger *zap.Logger,
 	verboseLog bool,
+	totalMaxConcurrency int,
 ) *PollingBase {
 	return &PollingBase{
-		redis:        redisClient,
-		queueMgr:     queueMgr,
-		configMgr:    configMgr,
-		iotCardStore: iotCardStore,
-		logger:       logger,
-		verboseLog:   verboseLog,
-		trafficLock:  cardtrafficlock.New(redisClient),
+		redis:               redisClient,
+		queueMgr:            queueMgr,
+		configMgr:           configMgr,
+		iotCardStore:        iotCardStore,
+		logger:              logger,
+		verboseLog:          verboseLog,
+		totalMaxConcurrency: totalMaxConcurrency,
+		trafficLock:         cardtrafficlock.New(redisClient),
 	}
 }
 
@@ -88,22 +95,33 @@ func (b *PollingBase) acquireConcurrency(ctx context.Context, taskType string) b
 	shortType := shortTaskType(taskType)
 	configKey := constants.RedisPollingConcurrencyConfigKey(shortType)
 	currentKey := constants.RedisPollingConcurrencyCurrentKey(taskType)
+	totalKey := constants.RedisPollingConcurrencyTotalCurrentKey()
 
 	maxConcurrency, err := b.redis.Get(ctx, configKey).Int()
-	if err != nil {
+	if err != nil || maxConcurrency < 1 || maxConcurrency > constants.PollingMaxConcurrencyLimit {
 		maxConcurrency = constants.PollingDefaultMaxConcurrency
 	}
+	totalMax := b.totalMaxConcurrency
+	if totalMax < 1 || totalMax > constants.PollingMaxConcurrencyLimit {
+		totalMax = constants.PollingDefaultTotalMaxConcurrency
+	}
 
 	result, err := acquireConcurrencyScript.Run(
-		ctx, b.redis, []string{currentKey},
-		maxConcurrency, constants.PollingConcurrencyKeyTTL,
+		ctx, b.redis, []string{totalKey, currentKey},
+		totalMax, maxConcurrency, constants.PollingConcurrencyKeyTTL,
 	).Int64()
 	if err != nil {
 		b.logger.Warn("获取并发计数失败，放行任务", zap.Error(err))
 		return true
 	}
 
-	return result != -1
+	if result == -1 {
+		b.logger.Info("轮询因并发令牌不足延后", zap.String("task_type", taskType),
+			zap.Int("max_concurrency", maxConcurrency), zap.Int("total_max_concurrency", totalMax),
+			zap.String("metric", "polling.deferred.concurrency_limit"))
+		return false
+	}
+	return true
 }
 
 // releaseConcurrency 释放并发信号量
@@ -112,7 +130,7 @@ func (b *PollingBase) releaseConcurrency(_ context.Context, taskType string) {
 	defer cancel()
 
 	currentKey := constants.RedisPollingConcurrencyCurrentKey(taskType)
-	if err := releaseConcurrencyScript.Run(ctx, b.redis, []string{currentKey}).Err(); err != nil {
+	if err := releaseConcurrencyScript.Run(ctx, b.redis, []string{constants.RedisPollingConcurrencyTotalCurrentKey(), currentKey}).Err(); err != nil {
 		b.logger.Warn("释放并发计数失败", zap.String("task_type", taskType), zap.Error(err))
 	}
 }