1

common/util.go

@@ -1,11 +1,24 @@
 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)
@@ -20,3 +33,76 @@ 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):
+		}
+	}
+}