109 lines
3.9 KiB
Go
109 lines
3.9 KiB
Go
package common
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/database/gredis"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/gcache"
|
||
|
||
_ "github.com/gogf/gf/contrib/nosql/redis/v2" // 注册 go-redis 适配器（gredis.New 需要），无 redis 配置时闲置
|
||
)
|
||
|
||
// unlockLuaScript 对比 token 再删除，防止锁过期后误删他人持有的锁
|
||
const unlockLuaScript = `if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) else return 0 end`
|
||
|
||
// RandomToken 生成 n 字节随机数的十六进制字符串（2n 位）
|
||
func RandomToken(n int) string {
|
||
b := make([]byte, n)
|
||
if _, err := rand.Read(b); err != nil {
|
||
panic("crypto/rand failed: " + err.Error())
|
||
}
|
||
return hex.EncodeToString(b)
|
||
}
|
||
|
||
// TokenFingerprint 计算令牌 SHA-256 指纹前 16 位，用于 JWT 会话校验
|
||
func TokenFingerprint(token string) string {
|
||
sum := sha256.Sum256([]byte(token))
|
||
return hex.EncodeToString(sum[:8])
|
||
}
|
||
|
||
// ErrLockHeld 锁被他人持有。WithLock 据此重试；中间件故障等其它错误不重试，直接返回。
|
||
var ErrLockHeld = gerror.New("lock held by another holder")
|
||
|
||
// WithLock 是互斥锁的唯一封装（util.go 仅此一个锁方法，禁止业务代码直接操作 gcache/gredis 自己实现）：
|
||
// 获取互斥锁后执行 fn，fn 结束（成功、失败或 panic）时自动释放锁（defer）。
|
||
// 锁实现按 config.yml 自动选择：配置了 redis 节点 → redis 锁（SET NX EX，跨实例互斥）；
|
||
// 未配置 → gcache 内存锁（单实例互斥）。
|
||
//
|
||
// 拿不到锁（被占用）时最多重试 retries 次（含首次共 retries+1 次尝试），每次间隔 retryInterval；
|
||
// retries=0 表示拿不到锁立即失败。ctx 取消/超时同样终止等待。中间件故障不重试，直接返回。
|
||
//
|
||
// fn 为泛型回调，返回 (T, error)：T 由闭包返回类型推断，业务返回值原样透出给下游；
|
||
// 方法参数/数量无约束，闭包捕获即可。仅需 error 的场景返回 (nil, err)（T 推断为 any）。
|
||
//
|
||
// doc, err := common.WithLock(ctx, "task:"+id, 30*time.Second, 3, 200*time.Millisecond, func() (string, error) {
|
||
// return s.process(ctx, id, mode) // 返回值透出给下游
|
||
// })
|
||
//
|
||
// expire 必须 > 0：锁自动过期兜底（进程崩溃不死锁），fn 耗时必须在 expire 前完成，fn 内禁止长耗时 IO。
|
||
func WithLock[T any](ctx context.Context, key string, expire time.Duration, retries int, retryInterval time.Duration, fn func() (T, error)) (T, error) {
|
||
var zero T
|
||
if expire <= 0 {
|
||
return zero, gerror.Newf("with lock %s: expire must be positive", key)
|
||
}
|
||
if retries > 0 && retryInterval <= 0 {
|
||
return zero, gerror.Newf("with lock %s: retryInterval must be positive when retries > 0", key)
|
||
}
|
||
for attempt := 0; ; attempt++ {
|
||
if err := ctx.Err(); err != nil {
|
||
return zero, gerror.Wrapf(err, "with lock %s", key)
|
||
}
|
||
var unlock func()
|
||
if !g.Cfg().MustGet(ctx, "redis", nil).IsNil() {
|
||
client := g.Redis()
|
||
token := RandomToken(16)
|
||
secs := int64(expire / time.Second)
|
||
if secs < 1 {
|
||
secs = 1
|
||
}
|
||
v, err := client.Set(ctx, key, token, gredis.SetOption{TTLOption: gredis.TTLOption{EX: &secs}, NX: true})
|
||
if err != nil {
|
||
return zero, gerror.Wrap(err, "redis lock acquire failed")
|
||
}
|
||
if !v.IsEmpty() {
|
||
unlock = func() {
|
||
if _, err := client.Eval(ctx, unlockLuaScript, 1, []string{key}, []any{token}); err != nil {
|
||
fmt.Printf("redis lock unlock failed (%s): %v\n", key, err)
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
ok, err := gcache.SetIfNotExist(ctx, key, 1, expire)
|
||
if err != nil {
|
||
return zero, err
|
||
}
|
||
if ok {
|
||
unlock = func() { gcache.Remove(ctx, key) }
|
||
}
|
||
}
|
||
if unlock != nil {
|
||
defer unlock()
|
||
return fn()
|
||
}
|
||
if attempt >= retries {
|
||
return zero, ErrLockHeld
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
case <-time.After(retryInterval):
|
||
}
|
||
}
|
||
}
|