90CNServer/server/internal/platform/lock/redis.go

package lock

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/gomodule/redigo/redis"
	"longheng.io/server/internal/platform/db"
)

// ErrNilRedisExecutor 表示 Redis 锁缺少执行器，无法访问外部分布式存储。
var ErrNilRedisExecutor = errors.New("lock redis executor is required")

// #nosec G101 -- Lua 脚本里的 token 是锁令牌变量名，不是硬编码凭证。
const redisReleaseToken = `
local token = redis.call("GET", KEYS[1])
if not token then
  return "missing"
end
if token == ARGV[1] then
  return redis.call("DEL", KEYS[1])
end
return "mismatch"
`

const RedisReleaseIfTokenScript = redisReleaseToken

// RedisOptions 配置 Redis 分布式锁管理器。
type RedisOptions struct {
	Name           string
	Executor       db.RedisExecutor
	KeyPrefix      string
	Address        string
	Password       string
	DB             int
	PoolSize       int
	Timeout        time.Duration
	ConnectTimeout time.Duration
	ReadTimeout    time.Duration
	WriteTimeout   time.Duration
	DefaultTTL     time.Duration
	TokenGenerator TokenGenerator
}

// RedisManager 基于 Redis SET NX PX 和 token 校验脚本提供分布式租约。
type RedisManager struct {
	mu             sync.RWMutex
	name           string
	keyPrefix      string
	defaultTTL     time.Duration
	executor       db.RedisExecutor
	closer         interface{ Close() error }
	tokenGenerator TokenGenerator
	metrics        *Metrics
}

// NewRedis 创建 Redis 分布式锁管理器；未传 Executor 时会按地址配置创建 redigo 执行器。
func NewRedis(options RedisOptions) *RedisManager {
	options = normRedisOpts(options)
	executor := options.Executor
	var closer interface{ Close() error }
	if executor == nil {
		redigo := db.NewRedigoExecutor(db.RedigoOptions{
			Address:        options.Address,
			Password:       options.Password,
			DB:             options.DB,
			PoolSize:       options.PoolSize,
			Timeout:        options.Timeout,
			ConnectTimeout: options.ConnectTimeout,
			ReadTimeout:    options.ReadTimeout,
			WriteTimeout:   options.WriteTimeout,
		})
		executor = redigo
		closer = redigo
	} else if candidate, ok := executor.(interface{ Close() error }); ok {
		closer = candidate
	}
	return &RedisManager{
		name:           options.Name,
		keyPrefix:      options.KeyPrefix,
		defaultTTL:     options.DefaultTTL,
		executor:       executor,
		closer:         closer,
		tokenGenerator: options.TokenGenerator,
	}
}

// SetMetrics 替换 Redis 锁采集器；Acquire 会在读锁内取快照后再记录。
func (m *RedisManager) SetMetrics(metrics *Metrics) {
	if m == nil {
		return
	}
	m.mu.Lock()
	m.metrics = metrics
	m.mu.Unlock()
}

// Acquire 尝试在 Redis 中获取指定 key 的租约。
func (m *RedisManager) Acquire(ctx context.Context, key string, ttl time.Duration) (lease Lease, err error) {
	if m == nil {
		return nil, ErrManagerRequired
	}
	m.mu.RLock()
	metricsRef := m.metrics
	defaultTTL := m.defaultTTL
	executor := m.executor
	tokenGenerator := m.tokenGenerator
	keyPrefix := m.keyPrefix
	m.mu.RUnlock()
	defer func() {
		metricsRef.recordAcquire(classifyAcquire(err))
	}()
	if err := ctxErr(ctx); err != nil {
		return nil, err
	}
	key = normalizeKey(key)
	if key == "" {
		return nil, ErrKeyRequired
	}
	if executor == nil {
		return nil, ErrNilRedisExecutor
	}
	if ttl <= 0 {
		ttl = defaultTTL
		if ttl <= 0 {
			ttl = 10 * time.Second
		}
	}
	if tokenGenerator == nil {
		tokenGenerator = randomToken
	}
	token, err := tokenGenerator()
	if err != nil {
		return nil, err
	}
	ok, err := AcquireRedis(ctx, executor, redisLockKey(keyPrefix, key), token, ttl)
	if errors.Is(err, redis.ErrNil) {
		return nil, ErrLockHeld
	}
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, ErrLockHeld
	}
	return &redisLease{manager: m, key: key, token: token, expiresAt: time.Now().UTC().Add(ttl)}, nil
}

// Snapshot 返回 Redis 锁管理器的配置状态。
func (m *RedisManager) Snapshot() Snapshot {
	if m == nil {
		return Snapshot{Kind: "redis", Closed: true}
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	return Snapshot{
		Name:      m.name,
		Kind:      "redis",
		KeyPrefix: m.keyPrefix,
	}
}

// Close 关闭由管理器持有的 Redis 连接池；外部注入的无 Close 执行器会被忽略。
func (m *RedisManager) Close() error {
	if m == nil || m.closer == nil {
		return nil
	}
	m.mu.RLock()
	closer := m.closer
	m.mu.RUnlock()
	if closer == nil {
		return nil
	}
	return closer.Close()
}

func (m *RedisManager) release(ctx context.Context, key, token string) error {
	if m == nil {
		return ErrManagerRequired
	}
	m.mu.RLock()
	executor := m.executor
	keyPrefix := m.keyPrefix
	m.mu.RUnlock()
	if executor == nil {
		return ErrNilRedisExecutor
	}
	return ReleaseRedis(ctx, executor, redisLockKey(keyPrefix, key), token)
}

// AcquireRedis 使用 Redis 原生命令获取一个带 TTL 的分布式锁。
func AcquireRedis(ctx context.Context, executor db.RedisExecutor, key, token string, ttl time.Duration) (bool, error) {
	if executor == nil {
		return false, ErrNilRedisExecutor
	}
	key = normalizeKey(key)
	token = strings.TrimSpace(token)
	if key == "" {
		return false, ErrKeyRequired
	}
	if token == "" {
		return false, ErrTokenMismatch
	}
	if ttl <= 0 {
		ttl = 10 * time.Second
	}
	result, err := redis.String(executor.Do(ctx, "SET", key, token, "NX", "PX", ttlMillis(ttl)))
	if errors.Is(err, redis.ErrNil) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return strings.EqualFold(strings.TrimSpace(result), "OK"), nil
}

// ReleaseRedis 只在 token 匹配时释放 Redis 锁，避免旧租约删除新持有者。
func ReleaseRedis(ctx context.Context, executor db.RedisExecutor, key, token string) error {
	if executor == nil {
		return ErrNilRedisExecutor
	}
	key = normalizeKey(key)
	token = strings.TrimSpace(token)
	if key == "" {
		return ErrKeyRequired
	}
	if token == "" {
		return ErrTokenMismatch
	}
	result, err := executor.Do(ctx, "EVAL", redisReleaseToken, 1, key, token)
	if err != nil {
		return err
	}
	return redisReleaseResult(result)
}

func redisReleaseResult(result any) error {
	switch value := result.(type) {
	case int64:
		if value == 1 {
			return nil
		}
	case []byte:
		if string(value) == "1" || string(value) == "missing" {
			return nil
		}
	case string:
		if value == "1" || value == "missing" {
			return nil
		}
	}
	if text, err := redis.String(result, nil); err == nil && (text == "1" || text == "missing") {
		return nil
	}
	return ErrTokenMismatch
}

type redisLease struct {
	manager   *RedisManager
	key       string
	token     string
	expiresAt time.Time
}

func (l *redisLease) Key() string {
	return l.key
}

func (l *redisLease) Token() string {
	return l.token
}

func (l *redisLease) ExpiresAt() time.Time {
	return l.expiresAt
}

func (l *redisLease) Release(ctx context.Context) error {
	if l == nil || l.manager == nil {
		return nil
	}
	releaseCtx, cancel := releaseContext(ctx)
	defer cancel()
	return l.manager.release(releaseCtx, l.key, l.token)
}

func redisLockKey(prefix, key string) string {
	key = normalizeKey(key)
	prefix = strings.TrimSpace(prefix)
	if prefix == "" {
		return key
	}
	return prefix + ":" + key
}

func normRedisOpts(options RedisOptions) RedisOptions {
	options.Name = strings.TrimSpace(options.Name)
	if options.Name == "" {
		options.Name = "redis-lock"
	}
	options.KeyPrefix = strings.TrimSpace(options.KeyPrefix)
	if options.KeyPrefix == "" {
		options.KeyPrefix = "lock"
	}
	options.Address = strings.TrimSpace(options.Address)
	if options.Address == "" {
		options.Address = "127.0.0.1:6379"
	}
	if options.PoolSize <= 0 {
		options.PoolSize = 8
	}
	if options.Timeout <= 0 {
		options.Timeout = 2 * time.Second
	}
	if options.ConnectTimeout <= 0 {
		options.ConnectTimeout = options.Timeout
	}
	if options.ReadTimeout <= 0 {
		options.ReadTimeout = options.Timeout
	}
	if options.WriteTimeout <= 0 {
		options.WriteTimeout = options.Timeout
	}
	if options.DefaultTTL <= 0 {
		options.DefaultTTL = 10 * time.Second
	}
	if options.TokenGenerator == nil {
		options.TokenGenerator = randomToken
	}
	return options
}

func ttlMillis(ttl time.Duration) int64 {
	ms := int64(ttl / time.Millisecond)
	if ms <= 0 {
		return 1
	}
	return ms
}