126 lines
3.7 KiB
Go
126 lines
3.7 KiB
Go
package common
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/database/gredis"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/gcache"
|
||
)
|
||
|
||
// ErrLockHeld 锁被他人持有（重试耗尽仍拿不到时返回）
|
||
var ErrLockHeld = errors.New("lock held")
|
||
|
||
// WithLock 互斥临界区唯一入口（泛型）：业务返回值经 T 原样透出。
|
||
// 内部按 config.yml 自动选择锁实现：配置了 redis 节点 → redis 锁（跨实例互斥，
|
||
// SET NX EX + token 对比删除防误删他人锁）；未配置 → gcache 内存锁（单实例互斥）。
|
||
// 拿不到锁最多重试 retries 次、每次间隔 retryInterval（retries=0 立即失败；
|
||
// ctx 取消/超时同样终止）；中间件故障不重试直接返回。defer 自动释放：无论 fn
|
||
// 成功、失败还是 panic。expire 必须 > 0（进程崩溃兜底不死锁），fn 耗时须在 expire 前完成，
|
||
// fn 内禁止长耗时 IO（LLM/DB 调用）；锁粒度按业务唯一键尽量小。
|
||
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, errors.New("lock expire must be positive")
|
||
}
|
||
lock, err := newLock(ctx, key, expire)
|
||
if err != nil {
|
||
return zero, err
|
||
}
|
||
var ok bool
|
||
for attempt := 0; ; attempt++ {
|
||
ok, err = lock.TryAcquire(ctx)
|
||
if err != nil {
|
||
return zero, err
|
||
}
|
||
if ok {
|
||
break
|
||
}
|
||
if attempt >= retries {
|
||
return zero, ErrLockHeld
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
return zero, ctx.Err()
|
||
case <-time.After(retryInterval):
|
||
}
|
||
}
|
||
defer lock.Release(ctx)
|
||
return fn()
|
||
}
|
||
|
||
type lock interface {
|
||
TryAcquire(ctx context.Context) (bool, error)
|
||
Release(ctx context.Context)
|
||
}
|
||
|
||
func newLock(ctx context.Context, key string, expire time.Duration) (lock, error) {
|
||
token := fmt.Sprintf("%d-%d", time.Now().UnixNano(), time.Now().UnixMilli()%1e9)
|
||
if g.Cfg().MustGet(ctx, "redis.default.address", "").String() != "" {
|
||
return &redisLock{key: "lock:" + key, token: token, expire: expire, client: g.Redis()}, nil
|
||
}
|
||
return &memoryLock{key: "lock:" + key, token: token, expire: expire}, nil
|
||
}
|
||
|
||
type redisLock struct {
|
||
key string
|
||
token string
|
||
expire time.Duration
|
||
client *gredis.Redis
|
||
}
|
||
|
||
func (l *redisLock) TryAcquire(ctx context.Context) (bool, error) {
|
||
// SetNX 无 TTL 参数，SET NX 与 TTL 分两步；当前项目未配置 redis 节点此路径不可达，
|
||
// 若崩溃于两步之间仅残留无 TTL 锁（token 归属明确，可手工清除），可接受
|
||
ok, err := l.client.SetNX(ctx, l.key, l.token)
|
||
if err != nil || !ok {
|
||
return ok, err
|
||
}
|
||
if _, err := l.client.PExpire(ctx, l.key, l.expire.Milliseconds()); err != nil {
|
||
if _, derr := l.client.Del(ctx, l.key); derr != nil {
|
||
g.Log().Warningf(ctx, "清理未设 TTL 的锁失败: %v", derr)
|
||
}
|
||
return false, err
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
func (l *redisLock) Release(ctx context.Context) {
|
||
v, err := l.client.Get(ctx, l.key)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "释放锁失败（读取）: %v", err)
|
||
return
|
||
}
|
||
if v.String() == l.token {
|
||
if _, err := l.client.Del(ctx, l.key); err != nil {
|
||
g.Log().Warningf(ctx, "释放锁失败（删除）: %v", err)
|
||
}
|
||
}
|
||
}
|
||
|
||
type memoryLock struct {
|
||
key string
|
||
token string
|
||
expire time.Duration
|
||
}
|
||
|
||
func (l *memoryLock) TryAcquire(ctx context.Context) (bool, error) {
|
||
return gcache.SetIfNotExist(ctx, l.key, l.token, l.expire)
|
||
}
|
||
|
||
func (l *memoryLock) Release(ctx context.Context) {
|
||
v, err := gcache.Get(ctx, l.key)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "释放锁失败（读取）: %v", err)
|
||
return
|
||
}
|
||
if v.String() == l.token {
|
||
if _, err := gcache.Remove(ctx, l.key); err != nil {
|
||
g.Log().Warningf(ctx, "释放锁失败（删除）: %v", err)
|
||
}
|
||
}
|
||
}
|