90CNServer/server/internal/platform/db/redis_executor.go

package db

import (
	"context"
	"errors"
	"time"

	"github.com/gomodule/redigo/redis"
)

// ErrRedisExecutorClosed 表示 Redis 执行器已关闭或未初始化。
var ErrRedisExecutorClosed = errors.New("redis executor is closed")

// closeRedisConnErr 在 Redis 操作本身成功时保留连接关闭错误，避免池连接异常被静默吞掉。
func closeRedisConnErr(conn redis.Conn, err *error) {
	if conn == nil || err == nil {
		return
	}
	if closeErr := conn.Close(); closeErr != nil && *err == nil {
		*err = closeErr
	}
}

// RedigoOptions 配置 redigo 连接池。
type RedigoOptions struct {
	Address           string
	Password          string
	DB                int
	PoolSize          int
	MaxActive         int
	IdleTimeout       time.Duration
	Timeout           time.Duration
	ConnectTimeout    time.Duration
	ReadTimeout       time.Duration
	WriteTimeout      time.Duration
	TestOnBorrowAfter time.Duration
}

// RedisPoolStats 描述 Redis 连接池状态。
type RedisPoolStats struct {
	Address      string        `json:"address"`
	DB           int           `json:"db"`
	MaxIdle      int           `json:"max_idle"`
	MaxActive    int           `json:"max_active"`
	IdleTimeout  time.Duration `json:"idle_timeout"`
	ActiveCount  int           `json:"active_count"`
	IdleCount    int           `json:"idle_count"`
	WaitCount    int64         `json:"wait_count"`
	WaitDuration time.Duration `json:"wait_duration"`
}

// RedigoExecutor 用 redigo 连接池实现 RedisExecutor。
type RedigoExecutor struct {
	options RedigoOptions
	pool    *redis.Pool
}

// NormalizeRedigoOptions 补齐 redigo 默认连接参数。
func NormalizeRedigoOptions(options RedigoOptions) RedigoOptions {
	if options.Address == "" {
		options.Address = "127.0.0.1:6379"
	}
	if options.PoolSize <= 0 {
		options.PoolSize = 8
	}
	if options.MaxActive <= 0 {
		options.MaxActive = options.PoolSize * 2
	}
	if options.IdleTimeout <= 0 {
		options.IdleTimeout = time.Minute
	}
	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.TestOnBorrowAfter < 0 {
		options.TestOnBorrowAfter = 0
	}
	if options.TestOnBorrowAfter == 0 {
		options.TestOnBorrowAfter = time.Minute
	}
	return options
}

// NewRedigoExecutor 创建 redigo Redis 执行器。
func NewRedigoExecutor(options RedigoOptions) *RedigoExecutor {
	options = NormalizeRedigoOptions(options)
	return &RedigoExecutor{options: options, pool: &redis.Pool{
		MaxIdle:     options.PoolSize,
		MaxActive:   options.MaxActive,
		IdleTimeout: options.IdleTimeout,
		Wait:        true,
		Dial: func() (redis.Conn, error) {
			dialOptions := []redis.DialOption{
				redis.DialConnectTimeout(options.ConnectTimeout),
				redis.DialReadTimeout(options.ReadTimeout),
				redis.DialWriteTimeout(options.WriteTimeout),
				redis.DialDatabase(options.DB),
			}
			if options.Password != "" {
				dialOptions = append(dialOptions, redis.DialPassword(options.Password))
			}
			return redis.Dial("tcp", options.Address, dialOptions...)
		},
		TestOnBorrow: func(conn redis.Conn, lastUsed time.Time) error {
			if time.Since(lastUsed) < options.TestOnBorrowAfter {
				return nil
			}
			_, err := conn.Do("PING")
			return err
		},
	}}
}

// Do 执行单条 Redis 命令。
func (e *RedigoExecutor) Do(ctx context.Context, command string, args ...any) (reply any, err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if e == nil || e.pool == nil {
		return nil, ErrRedisExecutorClosed
	}
	conn, err := e.pool.GetContext(ctx)
	if err != nil {
		return nil, err
	}
	defer closeRedisConnErr(conn, &err)
	return conn.Do(command, args...)
}

// DoBatch 使用同一连接批量发送 Redis 命令。
func (e *RedigoExecutor) DoBatch(ctx context.Context, commands []RedisCommand) (err error) {
	if len(commands) == 0 {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if e == nil || e.pool == nil {
		return ErrRedisExecutorClosed
	}
	conn, err := e.pool.GetContext(ctx)
	if err != nil {
		return err
	}
	defer closeRedisConnErr(conn, &err)
	return RunRedigoBatch(ctx, conn, commands)
}

// RunRedigoBatch 在已获取的 redigo 连接上发送批量命令，并在 Flush 后完整接收响应。
func RunRedigoBatch(ctx context.Context, conn redis.Conn, commands []RedisCommand) error {
	if len(commands) == 0 {
		return nil
	}
	if conn == nil {
		return ErrRedisExecutorClosed
	}
	sent := 0
	for _, command := range commands {
		if err := ctx.Err(); err != nil {
			return err
		}
		if command.Name == "" {
			continue
		}
		if err := conn.Send(command.Name, command.Args...); err != nil {
			return err
		}
		sent++
	}
	if sent == 0 {
		return nil
	}
	if err := conn.Flush(); err != nil {
		return err
	}
	// Flush 成功后命令可能已经被 Redis 执行；此时必须排空响应，不能再让调用方 ctx
	// 中途打断 Receive，否则异步刷盘会把已提交写误判成失败并重复入队。
	for idx := 0; idx < sent; idx++ {
		if _, err := conn.Receive(); err != nil {
			return err
		}
	}
	return nil
}

// PoolStats 返回连接池统计信息。
func (e *RedigoExecutor) PoolStats() RedisPoolStats {
	if e == nil || e.pool == nil {
		return RedisPoolStats{}
	}
	stats := e.pool.Stats()
	return RedisPoolStats{
		Address:      e.options.Address,
		DB:           e.options.DB,
		MaxIdle:      e.options.PoolSize,
		MaxActive:    e.options.MaxActive,
		IdleTimeout:  e.options.IdleTimeout,
		ActiveCount:  stats.ActiveCount,
		IdleCount:    stats.IdleCount,
		WaitCount:    stats.WaitCount,
		WaitDuration: stats.WaitDuration,
	}
}

// Close 关闭 Redis 连接池。
func (e *RedigoExecutor) Close() error {
	if e == nil || e.pool == nil {
		return nil
	}
	return e.pool.Close()
}