97 lines
2.4 KiB
Go
97 lines
2.4 KiB
Go
package cache
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"strings"
|
||
"time"
|
||
|
||
"hyapp-admin-server/internal/config"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
type Redis struct {
|
||
client *redis.Client
|
||
prefix string
|
||
}
|
||
|
||
type UnlockFunc func(context.Context) error
|
||
|
||
func NewRedis(ctx context.Context, cfg config.RedisConfig) (*Redis, error) {
|
||
if !cfg.Enabled {
|
||
return nil, nil
|
||
}
|
||
client := redis.NewClient(&redis.Options{
|
||
Addr: cfg.Addr,
|
||
Password: cfg.Password,
|
||
DB: cfg.DB,
|
||
})
|
||
if err := client.Ping(ctx).Err(); err != nil {
|
||
_ = client.Close()
|
||
return nil, err
|
||
}
|
||
return &Redis{client: client, prefix: cfg.KeyPrefix}, nil
|
||
}
|
||
|
||
func (r *Redis) Close() error {
|
||
if r == nil || r.client == nil {
|
||
return nil
|
||
}
|
||
return r.client.Close()
|
||
}
|
||
|
||
func (r *Redis) Ping(ctx context.Context) error {
|
||
if r == nil || r.client == nil {
|
||
return nil
|
||
}
|
||
return r.client.Ping(ctx).Err()
|
||
}
|
||
|
||
func (r *Redis) TryLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, UnlockFunc, error) {
|
||
if r == nil || r.client == nil {
|
||
return true, func(context.Context) error { return nil }, nil
|
||
}
|
||
if strings.TrimSpace(owner) == "" {
|
||
return false, nil, errors.New("lock owner is required")
|
||
}
|
||
redisKey := r.key(key)
|
||
ok, err := r.client.SetNX(ctx, redisKey, owner, ttl).Result()
|
||
if err != nil || !ok {
|
||
return ok, nil, err
|
||
}
|
||
return true, func(unlockCtx context.Context) error {
|
||
const script = `
|
||
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
return redis.call("DEL", KEYS[1])
|
||
end
|
||
return 0`
|
||
return r.client.Eval(unlockCtx, script, []string{redisKey}, owner).Err()
|
||
}, nil
|
||
}
|
||
|
||
// RefreshLock 只延长仍由当前 owner 持有的锁；任务 worker 续租 DB lease 时同步续租 Redis，
|
||
// 防止大导出超过初始 TTL 后被另一实例重复执行并覆盖产物。
|
||
func (r *Redis) RefreshLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, error) {
|
||
if r == nil || r.client == nil {
|
||
return true, nil
|
||
}
|
||
if strings.TrimSpace(owner) == "" || ttl <= 0 {
|
||
return false, errors.New("lock owner and positive ttl are required")
|
||
}
|
||
const script = `
|
||
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
||
return redis.call("PEXPIRE", KEYS[1], ARGV[2])
|
||
end
|
||
return 0`
|
||
result, err := r.client.Eval(ctx, script, []string{r.key(key)}, owner, ttl.Milliseconds()).Int64()
|
||
return result == 1, err
|
||
}
|
||
|
||
func (r *Redis) key(key string) string {
|
||
if r.prefix == "" {
|
||
return key
|
||
}
|
||
return r.prefix + strings.TrimPrefix(key, r.prefix)
|
||
}
|