57 lines
1.1 KiB
Go
57 lines
1.1 KiB
Go
|
package service
|
||
|
|
|
||
|
|
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
|
||
|
|
`
|
||
|
|
|
||
|
|
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
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseSemaphore(ctx context.Context, key string) error {
|
||
|
|
_, err := g.Redis().Do(ctx, "EVAL", releaseLua, 1, key)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|