58 lines
1.2 KiB
Go
58 lines
1.2 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|
)
|
|
|
|
var acquireLua = `
|
|
local current = tonumber(redis.call("GET", KEYS[1]) or "0")
|
|
local max = tonumber(ARGV[1])
|
|
local ttl = tonumber(ARGV[2])
|
|
if current >= max then
|
|
return 0
|
|
end
|
|
current = redis.call("INCR", KEYS[1])
|
|
if current == 1 then
|
|
redis.call("EXPIRE", KEYS[1], ttl)
|
|
end
|
|
if current > max then
|
|
redis.call("DECR", KEYS[1])
|
|
return 0
|
|
end
|
|
return 1
|
|
`
|
|
|
|
var releaseLua = `
|
|
local current = tonumber(redis.call("DECR", KEYS[1]) or "0")
|
|
if current <= 0 then
|
|
redis.call("DEL", KEYS[1])
|
|
end
|
|
return 1
|
|
`
|
|
|
|
// AcquireSemaphore 获取并发令牌
|
|
func AcquireSemaphore(ctx context.Context, key string, max int, ttlSeconds int64) (bool, error) {
|
|
if max <= 0 {
|
|
// 不限制
|
|
return true, nil
|
|
}
|
|
if ttlSeconds <= 0 {
|
|
ttlSeconds = 3600
|
|
}
|
|
r, err := g.Redis().Do(ctx, "EVAL", acquireLua, 1, key, max, ttlSeconds)
|
|
if err != nil {
|
|
return false, fmt.Errorf("获取并发令牌失败: %w", err)
|
|
}
|
|
return gconv.Int(r) == 1, nil
|
|
}
|
|
|
|
// ReleaseSemaphore 释放并发令牌
|
|
func ReleaseSemaphore(ctx context.Context, key string) error {
|
|
_, err := g.Redis().Do(ctx, "EVAL", releaseLua, 1, key)
|
|
return err
|
|
}
|