172 lines
5.9 KiB
Go
172 lines
5.9 KiB
Go
// Package utils 提供统一的 Redis 计数原语（与 redislock 的「互斥锁」互补，本包是「计数」一族）。
|
||
//
|
||
// 只保留 model-gateway 实际业务场景用到、且全仓有调用方的两个原语：
|
||
// - Semaphore：INCR 信号量（model-gateway/service/queue/semaphore.go），并发额度计数
|
||
// —— asyncWorker 按模型并发上限抢信号量，超限放回队列（task/worker.go:51）；
|
||
// - SafeDecr：安全回减（model-gateway/model_call_service.go decrSlotLua），计数防负
|
||
// —— 并发名额释放、超限回滚时原子夹在 >=0。
|
||
//
|
||
// （ZSET 槽位幂等闸门 queue_gate.go、固定窗口限流 rate_limiter.go 全仓无调用方，已删除；后续用到再加。）
|
||
//
|
||
// 与 redislock 一致，本包不写 Lua 脚本：
|
||
// - 多步原子操作（判定 + 写）经 go-redis WATCH/MULTI/EXEC 事务组合原生命令，
|
||
// 通过 gogf 官方 escape hatch GetAdapter().Client() 取底层客户端执行。
|
||
//
|
||
// 组合模式（如 model-gateway reserveSlot 的 锁+计数）：
|
||
// redislock.WithLock + 本包 SafeDecr。
|
||
//
|
||
// 兼容红线：只用 gf v2.9.5 已有 API（Set/Get/Incr/Decr、GetAdapter().Client()）。
|
||
package utils
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
goredis "github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
// errSkip 内部信号：WATCH 事务里判定不通过，不执行写。
|
||
// 调用方据 errors.Is(err, errSkip) 视为「本次操作失败但非异常」。
|
||
var errSkip = errors.New("rediscount: 条件不满足，跳过写入")
|
||
|
||
// rawClient 取底层 go-redis 客户端（gogf 官方 escape hatch）。
|
||
// 各服务 main.go 空导入 contrib/nosql/redis/v2 后 Adapter 为 go-redis 实现，
|
||
// Client() 返回 redis.UniversalClient（单节点 / 哨兵 / 集群均实现 Watch）。
|
||
func rawClient() (goredis.UniversalClient, error) {
|
||
universal, ok := g.Redis().GetAdapter().Client().(goredis.UniversalClient)
|
||
if !ok {
|
||
return nil, errors.New("redis 底层客户端非 UniversalClient，无法执行 WATCH 原子事务")
|
||
}
|
||
return universal, nil
|
||
}
|
||
|
||
// tx 在 WATCH/MULTI/EXEC 事务里执行 fn（先读后写，无 Lua）：
|
||
// fn 里先发读命令判定，命中条件后经 tx.TxPipelined 排队写命令，EXEC 原子提交。
|
||
// 判定不通过时 fn 返回 errSkip（不写）；WATCH 冲突（TxFailedErr，读与写之间被并发改动）自动重试。
|
||
func tx(ctx context.Context, key string, fn func(tx *goredis.Tx) error) error {
|
||
c, err := rawClient()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
for {
|
||
err := c.Watch(ctx, func(tx *goredis.Tx) error { return fn(tx) }, key)
|
||
if errors.Is(err, goredis.TxFailedErr) {
|
||
continue // 冲突，重读重试
|
||
}
|
||
return err
|
||
}
|
||
}
|
||
|
||
// pipedInt 取出事务写队列中第 n 条命令（INCR/DECR）的执行结果。
|
||
func pipedInt(cmds []goredis.Cmder, n int) (int64, error) {
|
||
return cmds[n].(*goredis.IntCmd).Result()
|
||
}
|
||
|
||
// SemaphoreAcquire 获取并发额度：WATCH key → GET 判满 → MULTI/INCR + 首设 EXPIRE/EXEC。
|
||
// 原子判定未满才 INCR；首次（计数从 0/不存在起）顺手设 TTL 自动回收防泄漏；超限不写、返回 false。
|
||
// 对应 model-gateway asyncWorker 按模型并发上限抢信号量。max<=0 表示不限制；ttlSeconds<=0 时按 3600 兜底。
|
||
func SemaphoreAcquire(ctx context.Context, key string, max int, ttlSeconds int64) (bool, error) {
|
||
if max <= 0 {
|
||
return true, nil
|
||
}
|
||
if ttlSeconds <= 0 {
|
||
ttlSeconds = 3600
|
||
}
|
||
err := tx(ctx, key, func(tx *goredis.Tx) error {
|
||
current, err := tx.Get(ctx, key).Int64()
|
||
if errors.Is(err, goredis.Nil) {
|
||
current = 0
|
||
} else if err != nil {
|
||
return err
|
||
}
|
||
if current >= int64(max) {
|
||
return errSkip // 已满
|
||
}
|
||
_, err = tx.TxPipelined(ctx, func(pipe goredis.Pipeliner) error {
|
||
pipe.Incr(ctx, key)
|
||
if current == 0 {
|
||
pipe.Expire(ctx, key, time.Duration(ttlSeconds)*time.Second)
|
||
}
|
||
return nil
|
||
})
|
||
return err
|
||
})
|
||
if errors.Is(err, errSkip) {
|
||
return false, nil
|
||
}
|
||
if err != nil {
|
||
return false, fmt.Errorf("获取并发额度失败: %w", err)
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
// SemaphoreRelease 释放并发额度（幂等，计数归零自动删除 key）。
|
||
// WATCH key → GET → MULTI：计数>1 仅 DECR；==1 时 DECR+DEL；<=0 直接 DEL（清理残留，不造负）。
|
||
// 全部原子，避免「DECR 后发现归零、中间被并发 INCR 抢占、DEL 误删新额度」的丢失更新。
|
||
func SemaphoreRelease(ctx context.Context, key string) error {
|
||
err := tx(ctx, key, func(tx *goredis.Tx) error {
|
||
current, err := tx.Get(ctx, key).Int64()
|
||
if errors.Is(err, goredis.Nil) {
|
||
current = 0
|
||
} else if err != nil {
|
||
return err
|
||
}
|
||
_, err = tx.TxPipelined(ctx, func(pipe goredis.Pipeliner) error {
|
||
switch {
|
||
case current <= 0:
|
||
pipe.Del(ctx, key)
|
||
case current == 1:
|
||
pipe.Decr(ctx, key)
|
||
pipe.Del(ctx, key)
|
||
default:
|
||
pipe.Decr(ctx, key)
|
||
}
|
||
return nil
|
||
})
|
||
return err
|
||
})
|
||
if err != nil {
|
||
return fmt.Errorf("释放并发额度失败: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// SafeDecr 安全回减：WATCH key → GET，仅当 key 存在且值>0 才 MULTI/DECR/EXEC，
|
||
// 杜绝 Redis DECR 对缺失 Key 按 0-1 处理把计数造负。
|
||
// 对应 model-gateway decrSlot：并发名额释放 / 超限回滚时原子夹在 >=0。
|
||
// 返回回减后的计数值；未回减（key 不存在或已为 0）返回 -1。
|
||
func SafeDecr(ctx context.Context, key string) (int64, error) {
|
||
var result int64
|
||
err := tx(ctx, key, func(tx *goredis.Tx) error {
|
||
current, err := tx.Get(ctx, key).Int64()
|
||
if errors.Is(err, goredis.Nil) {
|
||
return errSkip // 不存在，不写
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if current <= 0 {
|
||
return errSkip // 已为 0，不造负
|
||
}
|
||
cmds, err := tx.TxPipelined(ctx, func(pipe goredis.Pipeliner) error {
|
||
pipe.Decr(ctx, key)
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
result, err = pipedInt(cmds, 0)
|
||
return err
|
||
})
|
||
if errors.Is(err, errSkip) {
|
||
return -1, nil
|
||
}
|
||
if err != nil {
|
||
return -1, err
|
||
}
|
||
return result, nil
|
||
}
|