exchange_go/common/helper/redis_helper.go

package helper

import (
	"context"
	"errors"
	"fmt"
	"log"
	"math"
	"reflect"
	"strconv"
	"time"

	"github.com/bytedance/sonic"
	"github.com/go-redis/redis/v8"
)

// RedisHelper 结构体封装了 Redis 客户端及上下文
type RedisHelper struct {
	client          *redis.Client   // Redis 客户端
	ctx             context.Context // 上下文
	emptyCacheValue string          // 缓存空值的标志
}

var DefaultRedis *RedisHelper

// 初始化默认链接
func InitDefaultRedis(addr, password string, db int) {
	if DefaultRedis == nil {
		DefaultRedis = NewRedisHelper(addr, password, db)
	}

	log.Printf("初始化redis链接")
}

// NewRedisHelper 创建一个新的 RedisHelper 实例
func NewRedisHelper(addr, password string, db int) *RedisHelper {
	rdb := redis.NewClient(&redis.Options{
		Addr:         addr,     // Redis 服务器地址
		Password:     password, // Redis 密码
		DB:           db,       // 使用的数据库编号
		PoolSize:     50,
		MinIdleConns: 10,
		DialTimeout:  10 * time.Second, // 调整连接超时时间
		ReadTimeout:  10 * time.Second, // 调整读超时时间
		WriteTimeout: 10 * time.Second, // 调整写超时时间
	})

	return &RedisHelper{
		client: rdb,
		ctx:    context.Background(), // 创建背景上下文
	}
}

// 测试连接
func (r *RedisHelper) Ping() error {
	return r.client.Ping(r.ctx).Err()
}

// SetString 设置字符串值
func (r *RedisHelper) SetString(key, value string) error {
	return r.client.Set(r.ctx, key, value, 0).Err() // 将值存储到指定的键
}

// 批量设置
func (r *RedisHelper) BatchSet(maps *map[string]string) error {
	pipe := r.client.Pipeline()

	for key, val := range *maps {
		pipe.Set(r.ctx, key, val, 0)
	}

	_, err := pipe.Exec(r.ctx)

	return err
}

// SetString 设置字符串值
func (r *RedisHelper) SetStringExpire(key, value string, expireTime time.Duration) error {
	return r.client.Set(r.ctx, key, value, expireTime).Err() // 将值存储到指定的键
}

// SetString 设置字符串值
func (r *RedisHelper) SetAdd(key, value string, expireTime time.Duration) error {
	// 存储到 SET 中
	result, err := r.client.SAdd(r.ctx, key, value).Result()
	if err != nil {
		return err
	}

	if result == 1 {
		// 设置 SET 的过期时间
		err = r.client.Expire(r.ctx, key, expireTime).Err()
		if err != nil {
			return errors.New("设置过期时间失败:" + err.Error())
		}

		return nil
	} else {
		return errors.New("key已存在")
	}
}

// 设置对象
func SetObjString[T any](r *RedisHelper, key string, value T) error {
	keyValue, err := sonic.Marshal(value)

	if err != nil {
		return err
	}

	return r.SetString(key, string(keyValue))
}

// 获取对象
func GetObjString[T any](r *RedisHelper, key string) (T, error) {
	var result T
	value, err := r.GetString(key)

	if err != nil {
		return result, err
	}

	err = sonic.Unmarshal([]byte(value), &result)

	if err != nil {
		return result, err
	}

	return result, nil
}

func (r *RedisHelper) Get(key string) *redis.StringCmd {
	return r.client.Get(r.ctx, key)
}

/*
获取剩余时间
  - @key redis key
*/
func (r *RedisHelper) TTL(key string) *redis.DurationCmd {
	return r.client.TTL(r.ctx, key)
}

// GetString 获取字符串值
func (r *RedisHelper) GetString(key string) (string, error) {
	return r.client.Get(r.ctx, key).Result() // 从指定的键获取值
}

// DeleteString 删除字符串键
func (r *RedisHelper) DeleteString(key string) error {
	return r.client.Del(r.ctx, key).Err() // 删除指定的键
}

// DeleteString 删除目录下所有key
func (r *RedisHelper) DeleteAll(key string) error {
	keys, err := r.ScanKeys(key)

	if err != nil {
		return err
	}

	_, err = r.BatchDeleteKeys(keys)
	return err
}

/*
递增
  - @key rediskey
*/
func (r *RedisHelper) Incr(key string) *redis.IntCmd {
	return r.client.Incr(r.ctx, key)
}

/*
设置过期时间
  - @key redis key
  - @expiration 过期时间
*/
func (r *RedisHelper) Expire(key string, expiration time.Duration) *redis.BoolCmd {
	return r.client.Expire(r.ctx, key, expiration)
}

/*
批量删除

  - @keys 键数组
*/
func (r *RedisHelper) BatchDeleteKeys(keys []string) (int, error) {
	if r.client == nil {
		return 0, errors.New("Redis client is nil")
	}
	if len(keys) == 0 {
		return 0, nil
	}

	deletedCount := 0
	batchSize := 1000 // 每批次删除的键数量
	for i := 0; i < len(keys); i += batchSize {
		end := i + batchSize
		if end > len(keys) {
			end = len(keys)
		}
		batch := keys[i:end]

		_, err := r.client.Pipelined(r.ctx, func(pipe redis.Pipeliner) error {
			for _, key := range batch {
				pipe.Del(r.ctx, key)
				deletedCount++
			}
			return nil
		})
		if err != nil {
			return deletedCount, fmt.Errorf("failed to delete keys in batch: %v", err)
		}
	}

	return deletedCount, nil
}

// DeleteKeysByPrefix 删除指定前缀的键
func (r *RedisHelper) DeleteKeysByPrefix(prefixes ...string) error {
	ctx := context.Background()
	// 遍历每个前缀
	for _, prefix := range prefixes {
		var cursor uint64
		var keys []string

		// 使用 SCAN 命令查找匹配的键
		for {
			var err error
			keys, cursor, err = r.client.Scan(ctx, cursor, prefix+"*", 1000).Result()
			if err != nil {
				return err
			}

			// 删除匹配的键
			if len(keys) > 0 {
				_, err := r.client.Del(ctx, keys...).Result()
				if err != nil {
					return err
				}
				fmt.Printf("Deleted keys with prefix '%s': %v\n", prefix, keys)
			}

			// 如果游标为 0，表示迭代结束
			if cursor == 0 {
				break
			}
		}
	}
	return nil
}

// 查找所有value
func (r *RedisHelper) GetAllKeysAndValues(pattern string) ([]string, error) {
	var cursor uint64
	var result = []string{}

	for {
		// 使用 SCAN 命令获取匹配的键
		keys, nextCursor, err := r.client.Scan(r.ctx, cursor, pattern+"*", 1000).Result()
		if err != nil {
			log.Printf("Error scanning keys: %v", err)
			return nil, err
		}

		// 处理匹配到的键
		for _, key := range keys {
			value, err := r.client.Get(r.ctx, key).Result()
			if err != nil {
				if err == redis.Nil {
					fmt.Printf("Key %s does not exist\n", key)
				} else {
					fmt.Printf("Error getting value for key %s: %v", key, err)
				}
			} else {
				result = append(result, value)
			}
		}

		// 如果 cursor 为 0，表示扫描完成
		if nextCursor == 0 {
			break
		}
		cursor = nextCursor
	}

	return result, nil
}

// LPushList 将一个或多个值插入到列表的头部
func (r *RedisHelper) LPushList(key string, values ...string) error {
	return r.client.LPush(r.ctx, key, values).Err() // 将值插入到列表的头部
}

// RPushList 将一个或多个值插入到列表的尾部
func (r *RedisHelper) RPushList(key string, values ...string) error {
	return r.client.RPush(r.ctx, key, values).Err() // 将值插入到列表的尾部
}

// LPopList 从列表的头部弹出一个元素
func (r *RedisHelper) LPopList(key string) (string, error) {
	return r.client.LPop(r.ctx, key).Result() // 从列表的头部移除并返回第一个元素
}

// RPopList 从列表的尾部弹出一个元素
func (r *RedisHelper) RPopList(key string) (string, error) {
	return r.client.RPop(r.ctx, key).Result() // 从列表的尾部移除并返回最后一个元素
}

// LRangeList 获取列表中指定范围的元素
func (r *RedisHelper) LRangeList(key string, start, stop int64) ([]string, error) {
	return r.client.LRange(r.ctx, key, start, stop).Result() // 获取列表中指定范围的元素
}

// GetAllList 获取列表中的所有元素
func (r *RedisHelper) GetAllList(key string) ([]string, error) {
	values, err := r.client.LRange(r.ctx, key, 0, -1).Result()
	if err == redis.Nil {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	// 检查是否包含空值标志
	if len(values) == 1 && values[0] == r.emptyCacheValue {
		return nil, nil
	}

	return values, nil
}

func (r *RedisHelper) LRem(key, val string) (int64, error) {
	count := 0 // 删除所有与 valueToRemove 相等的元素
	result, err := r.client.LRem(r.ctx, key, int64(count), val).Result()
	if err != nil {
		fmt.Printf("删除元素失败: %v\n", err)
	}

	return result, nil
}

func (r *RedisHelper) IsElementInList(key string, element string) (bool, error) {
	var cursor int64 = 0
	const batchSize int64 = 1000 // 每批次获取的元素数量

	for {
		// 分批次获取列表元素
		elements, err := r.client.LRange(r.ctx, key, cursor, cursor+batchSize-1).Result()
		if err != nil {
			return false, err
		}
		if len(elements) == 0 {
			break // 没有更多数据
		}

		// 遍历当前批次的元素
		for _, e := range elements {
			if e == element {
				return true, nil
			}
		}

		cursor += batchSize // 移动到下一批次
	}

	return false, nil
}

/*
SetListCache 重新设置列表缓存
  - @expiration 0-过期 1-过期时间
*/
func (r *RedisHelper) SetListCache(key string, expiration time.Duration, values ...string) error {
	tempKey := key + ":temp"

	// 使用事务来确保操作的原子性
	pipe := r.client.TxPipeline()

	// 将新数据插入到临时列表中
	pipe.RPush(r.ctx, tempKey, values)

	// 重命名临时列表为目标列表
	pipe.Rename(r.ctx, tempKey, key)

	if expiration > 0 {
		// 设置目标列表的过期时间
		pipe.Expire(r.ctx, key, expiration)
	}

	// 执行事务
	_, err := pipe.Exec(r.ctx)
	return err
}

// SetEmptyListCache 设置空值缓存
func (r *RedisHelper) SetEmptyListCache(key string, expiration time.Duration) error {
	// 使用一个特殊标志值表示列表为空
	_, err := r.client.RPush(r.ctx, key, r.emptyCacheValue).Result()
	if err != nil {
		return err
	}

	// 设置列表的过期时间
	return r.client.Expire(r.ctx, key, expiration).Err()
}

// scanKeys 使用 SCAN 命令获取所有匹配的键
func (r *RedisHelper) ScanKeys(pattern string) ([]string, error) {

	var cursor uint64
	var keys []string
	for {
		var newKeys []string
		var err error

		// SCAN 命令每次返回部分匹配的键
		newKeys, cursor, err = r.client.Scan(r.ctx, cursor, pattern, 1000).Result()
		if err != nil {
			return nil, err
		}
		keys = append(keys, newKeys...)
		if cursor == 0 {
			break
		}
	}
	return keys, nil
}

// 泛型函数，用于获取所有键的列表数据并合并为一个数组
func GetAndMergeLists[T any](r *RedisHelper, keys []string) ([]T, error) {
	var combinedList []T
	for _, key := range keys {
		// 获取每个键的列表数据
		listData, err := r.client.LRange(r.ctx, key, 0, -1).Result()
		if err != nil {
			return nil, err
		}

		// 解码每个数据项为类型 T，并添加到结果列表中
		for _, data := range listData {
			var item T
			if err := sonic.Unmarshal([]byte(data), &item); err != nil {
				return nil, err
			}
			combinedList = append(combinedList, item)
		}
	}
	return combinedList, nil
}

// SetNX 实现类似于 Redis 的 SETNX 命令
func (r *RedisHelper) SetNX(key string, value interface{}, expiration time.Duration) (bool, error) {
	result, err := r.client.Set(r.ctx, key, value, expiration).Result()
	if err != nil {
		return false, err
	}

	// 如果键不存在则 result 会等于 "OK"
	return result == "OK", nil
}

func getFieldsFromStruct(obj interface{}) map[string]interface{} {
	fields := make(map[string]interface{})
	val := reflect.ValueOf(obj)
	typ := reflect.TypeOf(obj)

	for i := 0; i < val.NumField(); i++ {
		field := typ.Field(i)
		tag := field.Tag.Get("redis")
		if tag != "" {
			fieldVal := val.Field(i)
			if fieldVal.Kind() == reflect.Slice || fieldVal.Kind() == reflect.Map {
				// 处理切片或映射类型
				// 对于切片，使用索引作为字段名
				if fieldVal.Kind() == reflect.Slice {
					for j := 0; j < fieldVal.Len(); j++ {
						elem := fieldVal.Index(j).Interface()
						fields[fmt.Sprintf("%s_%d", tag, j)] = elem
					}
				} else if fieldVal.Kind() == reflect.Map {
					// 对于映射，使用键作为字段名
					for _, key := range fieldVal.MapKeys() {
						elem := fieldVal.MapIndex(key).Interface()
						fields[fmt.Sprintf("%s_%v", tag, key.Interface())] = elem
					}
				}
			} else {
				fields[tag] = fieldVal.Interface()
			}
		}
	}

	return fields
}

func (r *RedisHelper) SetHashWithTags(key string, obj interface{}) error {
	fields := getFieldsFromStruct(obj)
	_, err := r.client.HSet(r.ctx, key, fields).Result()
	return err
}

// HSetField 设置哈希中的一个字段
func (r *RedisHelper) HSetField(key, field string, value interface{}) error {
	_, err := r.client.HSet(r.ctx, key, field, value).Result()
	if err != nil {
		log.Printf("Error setting field %s in hash %s: %v", field, key, err)
		return err
	}
	return nil
}

// HSetMultipleFields 设置哈希中的多个字段
func (r *RedisHelper) HSetMultipleFields(key string, fields map[string]interface{}) error {
	_, err := r.client.HSet(r.ctx, key, fields).Result()
	if err != nil {
		log.Printf("Error setting multiple fields in hash %s: %v", key, err)
		return err
	}
	return nil
}

// HGetField 获取哈希中某个字段的值
func (r *RedisHelper) HGetField(key, field string) (string, error) {
	val, err := r.client.HGet(r.ctx, key, field).Result()
	if err != nil {
		if err == redis.Nil {
			return "", nil // Field does not exist
		}
		log.Printf("Error getting field %s from hash %s: %v", field, key, err)
		return "", err
	}
	return val, nil
}

// HGetAllFields 获取哈希中所有字段的值
func (r *RedisHelper) HGetAllFields(key string) (map[string]string, error) {
	fields, err := r.client.HGetAll(r.ctx, key).Result()
	if err != nil {
		log.Printf("Error getting all fields from hash %s: %v", key, err)
		return nil, err
	}
	return fields, nil
}

// HDelField 删除哈希中的某个字段
func (r *RedisHelper) HDelField(key, field string) error {
	_, err := r.client.HDel(r.ctx, key, field).Result()
	if err != nil {
		log.Printf("Error deleting field %s from hash %s: %v", field, key, err)
		return err
	}
	return nil
}

// 删除哈希
func (r *RedisHelper) HDelAll(key string) error {
	_, err := r.client.Del(r.ctx, key).Result()
	if err != nil {
		log.Printf("Error deleting from hash %s: %v", key, err)
		return err
	}
	return nil
}

// HKeys 获取哈希中所有字段的名字
func (r *RedisHelper) HKeys(key string) ([]string, error) {
	fields, err := r.client.HKeys(r.ctx, key).Result()
	if err != nil {
		log.Printf("Error getting keys from hash %s: %v", key, err)
		return nil, err
	}
	return fields, nil
}

// DelSet 从集合中删除元素
func (r *RedisHelper) DelSet(key string, value string) error {
	_, err := r.client.SRem(r.ctx, key, value).Result()
	if err != nil {
		log.Printf("Error del value from set %s: %v", key, err)
		return err
	}
	return nil
}

func (r *RedisHelper) Sismember(key string, value string) (bool, error) {
	result, err := r.client.SIsMember(r.ctx, key, value).Result()
	if err != nil {
		log.Printf("Error Sismember value from set %s: %v", key, err)
	}
	return result, err
}

// sort set start
// 批量添加
func (r *RedisHelper) BatchSortSet(key string, array []*redis.Z) error {
	pipe := r.client.Pipeline()
	for _, val := range array {
		pipe.ZAdd(r.ctx, key, val)
	}

	_, err := pipe.Exec(r.ctx)
	return err
}

// 单一写入 sort set
func (e *RedisHelper) SignelAdd(key string, score float64, member string) error {
	// 先删除具有相同 score 的所有成员
	scoreStr := strconv.FormatFloat(score, 'g', -1, 64)
	_, err := e.client.ZRemRangeByScore(e.ctx, key, scoreStr, scoreStr).Result()

	if err != nil {
		fmt.Printf("删除score失败,err:%s", err.Error())
	}
	_, err = e.client.ZAdd(e.ctx, key, &redis.Z{
		Score:  score,
		Member: member,
	}).Result()

	if err != nil {
		return err
	}

	return nil
}

// 写入数据
func (e *RedisHelper) AddSortSet(key string, score float64, member string) error {
	_, err := e.client.ZAdd(e.ctx, key, &redis.Z{
		Score:  score,
		Member: member,
	}).Result()

	if err != nil {
		return err
	}

	return nil
}

// 删除指定元素
func (e *RedisHelper) DelSortSet(key, member string) error {
	return e.client.ZRem(e.ctx, key, member).Err()
}

// RemoveBeforeScore 移除 Sorted Set 中分数小于等于指定值的数据
// key: Sorted Set 的键
// score: 分数上限，所有小于等于此分数的元素将被移除
// 返回值: 移除的元素数量和可能的错误
func (e *RedisHelper) RemoveBeforeScore(key string, score float64) (int64, error) {
	if key == "" {
		return 0, errors.New("key 不能为空")
	}
	if math.IsNaN(score) || math.IsInf(score, 0) {
		return 0, errors.New("score 必须是有效数字")
	}

	// 使用 ZRemRangeByScore 移除数据
	count, err := e.client.ZRemRangeByScore(e.ctx, key, "-inf", strconv.FormatFloat(score, 'f', -1, 64)).Result()
	if err != nil {
		return 0, fmt.Errorf("移除 Sorted Set 数据失败, key: %s, score: %f, err: %v", key, score, err)
	}

	return count, nil
}

// GetNextAfterScore 获取指定分数及之后的第一条数据（包含指定分数）
func (e *RedisHelper) GetNextAfterScore(key string, score float64) (string, error) {
	// 使用 ZRangeByScore 获取大于等于 score 的第一条数据
	zs, err := e.client.ZRangeByScoreWithScores(e.ctx, key, &redis.ZRangeBy{
		Min:    fmt.Sprintf("%f", score), // 包含指定分数
		Max:    "+inf",                   // 上限为正无穷
		Offset: 0,                        // 从第 0 条开始
		Count:  1,                        // 只取 1 条
	}).Result()

	if err != nil {
		return "", fmt.Errorf("获取数据失败: %v", err)
	}
	if len(zs) == 0 {
		return "", nil // 没有符合条件的元素
	}
	return zs[0].Member.(string), nil
}

/*
获取sort set 所有数据
*/
func (e *RedisHelper) GetAllSortSet(key string) ([]string, error) {
	return e.client.ZRange(e.ctx, key, 0, -1).Result()
}

/*
获取sort set 所有数据和score
*/
func (e *RedisHelper) GetRevRangeScoresSortSet(key string) ([]redis.Z, error) {
	return e.client.ZRevRangeWithScores(e.ctx, key, 0, -1).Result()
}

// 获取指定区间数据
func (e *RedisHelper) GetSortSetMembers(key string, start, stop int64) ([]string, error) {
	return e.client.ZRange(e.ctx, key, start, stop).Result()
}

// 获取最后N条数据
func (e *RedisHelper) GetLastSortSetMembers(key string, num int64) ([]string, error) {
	return e.client.ZRevRange(e.ctx, key, 0, num).Result()
}

// func (e *RedisHelper) DelSortSet(key,)

// 根据索引范围删除
func (e *RedisHelper) DelByRank(key string, start, stop int64) error {
	return e.client.ZRemRangeByRank(e.ctx, key, start, stop).Err()
}

// sort set end

// GetUserLoginPwdErrFre 获取用户登录密码错误频次
func (e *RedisHelper) GetUserLoginPwdErrFre(key string) (total int, wait time.Duration, err error) {
	total, _ = e.client.Get(e.ctx, key).Int()
	wait = e.client.TTL(e.ctx, key).Val()
	return
}

// SetUserLoginPwdErrFre 设置用户登录密码错误频次
func (e *RedisHelper) SetUserLoginPwdErrFre(key string, expire time.Duration) (val int64, err error) {
	val, err = e.client.Incr(e.ctx, key).Result()
	if err != nil {
		return
	}
	if err = e.client.Expire(e.ctx, key, expire).Err(); err != nil {
		return
	}
	return
}