108 lines
3.2 KiB
Go
108 lines
3.2 KiB
Go
package queue
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"math"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
)
|
||
|
||
// ===== 严格 queue_limit：Redis 原子闸门 =====
|
||
//
|
||
// 背景：原来的 queue_limit 通过“Count + Insert”做近似控制，分布式并发创建时会短暂超限。
|
||
// 目标：以 Redis Lua 脚本实现原子校验 + 入队占位，做到严格不超限。
|
||
//
|
||
// 计数口径与原逻辑保持一致：只统计 state=0/1（排队中/执行中）。
|
||
// - CreateTask 成功入库后占用 1 个 slot
|
||
// - 任务成功/失败（state->2/3）释放 slot
|
||
// - 失败任务重试（state 3->0）需要再次占用 slot，若占位失败则暂不重试（留在 state=3，下次 cleaner 再尝试）
|
||
//
|
||
// 说明：为避免极端情况下“占位泄漏”导致永久占满，采用 ZSET + 过期时间的方式自动回收。
|
||
// 只要任务实际生命周期远小于 gateTTLSeconds，就可保持严格。
|
||
|
||
const (
|
||
queueGateKeyPrefix = "asynch:qgate:" // asynch:qgate:{modelName}
|
||
)
|
||
|
||
// Lua：清理过期 slot，然后按 limit 做原子判定并占位
|
||
var queueGateAcquireLua = `
|
||
local key = KEYS[1]
|
||
local now = tonumber(ARGV[1])
|
||
local limit = tonumber(ARGV[2])
|
||
local expireAt = tonumber(ARGV[3])
|
||
local member = ARGV[4]
|
||
local keyTTL = tonumber(ARGV[5])
|
||
|
||
-- 先清理过期的占位
|
||
redis.call("ZREMRANGEBYSCORE", key, "-inf", now)
|
||
|
||
local current = tonumber(redis.call("ZCARD", key) or "0")
|
||
if current >= limit then
|
||
return 0
|
||
end
|
||
redis.call("ZADD", key, expireAt, member)
|
||
redis.call("EXPIRE", key, keyTTL)
|
||
return 1
|
||
`
|
||
|
||
// Lua：释放 slot（幂等）
|
||
var queueGateReleaseLua = `
|
||
local key = KEYS[1]
|
||
local member = ARGV[1]
|
||
redis.call("ZREM", key, member)
|
||
return 1
|
||
`
|
||
|
||
func queueGateKey(modelName string) string {
|
||
return fmt.Sprintf("%s%s", queueGateKeyPrefix, modelName)
|
||
}
|
||
|
||
// calcGateTTLSeconds 计算闸门占位的“自动回收 TTL”
|
||
// 取 expectedSeconds 的倍数并做上下限，避免任务异常导致永久占位。
|
||
func calcGateTTLSeconds(expectedSeconds int) int {
|
||
// 默认至少 1 小时；最多 24 小时
|
||
minTTL := 3600
|
||
maxTTL := 24 * 3600
|
||
if expectedSeconds <= 0 {
|
||
return minTTL
|
||
}
|
||
ttl := int(math.Ceil(float64(expectedSeconds) * 10)) // 预计耗时 * 10 做兜底
|
||
if ttl < minTTL {
|
||
ttl = minTTL
|
||
}
|
||
if ttl > maxTTL {
|
||
ttl = maxTTL
|
||
}
|
||
return ttl
|
||
}
|
||
|
||
// AcquireQueueSlot 严格入队：原子占位（成功返回 true）
|
||
func AcquireQueueSlot(ctx context.Context, modelName, taskId string, limit int, expectedSeconds int) (bool, error) {
|
||
if limit <= 0 {
|
||
return true, nil
|
||
}
|
||
key := queueGateKey(modelName)
|
||
now := time.Now().Unix()
|
||
ttl := calcGateTTLSeconds(expectedSeconds)
|
||
expireAt := now + int64(ttl)
|
||
// keyTTL 要略大于 member TTL，避免 key 先过期导致计数丢失
|
||
keyTTL := ttl + 60
|
||
r, err := g.Redis().Do(ctx, "EVAL", queueGateAcquireLua, 1, key, now, limit, expireAt, taskId, keyTTL)
|
||
if err != nil {
|
||
return false, fmt.Errorf("queue gate acquire failed: %w", err)
|
||
}
|
||
return gconv.Int(r) == 1, nil
|
||
}
|
||
|
||
// ReleaseQueueSlot 释放占位（幂等）
|
||
func ReleaseQueueSlot(ctx context.Context, modelName, taskId string) {
|
||
if taskId == "" || modelName == "" {
|
||
return
|
||
}
|
||
key := queueGateKey(modelName)
|
||
_, _ = g.Redis().Do(ctx, "EVAL", queueGateReleaseLua, 1, key, taskId)
|
||
}
|