laic-backend/common/redis.go

package common

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"laic-backend/logger"

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

var _redis *redis.Pool

func InitRedis(source string) {
	_redis = &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial:        func() (redis.Conn, error) { return redis.DialURL(source) },
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
	}
	logger.INFO("redis ready")
}

func GetRedisPool() *redis.Pool {
	return _redis
}

func GetString(key string) (string, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.String(c.Do("GET", key))
}

func GetBytes(key string) ([]byte, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.Bytes(c.Do("GET", key))
}

func GetInt64(key string) (int64, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.Int64(c.Do("GET", key))
}

func GetJSON[T any](key string) (*T, error) {
	c := _redis.Get()
	defer c.Close()
	r, err := redis.Bytes(c.Do("GET", key))
	if err != nil {
		return nil, err
	}
	var v T
	err = json.Unmarshal(r, &v)
	return &v, err
}

func SetValue(key string, n any) error {
	c := _redis.Get()
	defer c.Close()
	_, err := c.Do("SET", key, n)
	return err
}

func SetJSON(key string, n any) error {
	c := _redis.Get()
	defer c.Close()
	data, err := json.Marshal(n)
	if err != nil {
		return err
	}
	_, err = c.Do("SET", key, data)
	return err
}

func SetJSONWithExpired(key string, n any, second uint) error {
	c := _redis.Get()
	defer c.Close()
	data, err := json.Marshal(n)
	if err != nil {
		return err
	}
	_, err = c.Do("SET", key, data, "EX", second)
	return err
}

func SetValueWithExpired(key string, n any, second uint) error {
	c := _redis.Get()
	defer c.Close()
	_, err := c.Do("SET", key, n, "EX", second)
	return err
}

func SetRemoveIfKeyMissing(setKey, presenceKey, member string) (bool, error) {
	const script = `if redis.call('EXISTS', KEYS[2]) == 0 then return redis.call('SREM', KEYS[1], ARGV[1]) end return 0`
	n, err := redis.Int(exec_lua(script, []string{setKey, presenceKey}, member))
	return n == 1, err
}

// SetNX 仅当 key 不存在时写入，返回是否写入成功
func SetNX(key string, value any) (bool, error) {
	c := _redis.Get()
	defer c.Close()
	reply, err := c.Do("SET", key, value, "NX")
	return setNXResult(reply, err)
}

// SET ... NX returns the simple string "OK" when it writes and nil when the
// key already exists. It is not an integer reply.
func setNXResult(reply any, err error) (bool, error) {
	if err != nil {
		return false, err
	}
	if reply == nil {
		return false, nil
	}
	status, err := redis.String(reply, nil)
	if err != nil {
		return false, err
	}
	if strings.EqualFold(status, "OK") {
		return true, nil
	}
	return false, fmt.Errorf("redis: unexpected SET NX reply %q", status)
}

// IncrBy 原子自增，返回自增后的值（计费充值用）
func IncrBy(key string, delta int64) (int64, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.Int64(c.Do("INCRBY", key, delta))
}

// DecrBy 原子自减，返回自减后的值（计费扣减用）
func DecrBy(key string, delta int64) (int64, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.Int64(c.Do("DECRBY", key, delta))
}

func HasKey(key string) bool {
	c := _redis.Get()
	defer c.Close()
	exists, err := redis.Int(c.Do("EXISTS", key))
	if err != nil {
		return false
	}
	return exists == 1
}

func HashGet(key string, field string) (string, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.String(c.Do("HGET", key, field))
}

func HashGetAll(key string) (map[string]string, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.StringMap(c.Do("HGETALL", key))
}

func HashSetValue(key string, field string, n any) error {
	c := _redis.Get()
	defer c.Close()
	_, err := c.Do("HSET", key, field, n)
	return err
}

func HashSetValues(key string, values map[string]any) error {
	c := _redis.Get()
	defer c.Close()
	args := make([]any, 0, len(values)*2+1)
	args = append(args, key)
	for k, v := range values {
		args = append(args, k, v)
	}
	_, err := c.Do("HMSET", args...)
	return err
}

func SetMembers(key string) ([]string, error) {
	c := _redis.Get()
	defer c.Close()
	return redis.Strings(c.Do("SMEMBERS", key))
}

func SetAdd(key string, n any) error {
	c := _redis.Get()
	defer c.Close()
	_, err := c.Do("SADD", key, n)
	return err
}

func SetRemove(key string, n any) error {
	c := _redis.Get()
	defer c.Close()
	_, err := c.Do("SREM", key, n)
	return err
}

func SetMemberExists(key string, value string) (bool, error) {
	c := _redis.Get()
	defer c.Close()
	exists, err := redis.Int(c.Do("SISMEMBER", key, value))
	if err != nil {
		return false, err
	}
	return exists == 1, nil
}

func exec_lua(lua string, keys []string, args ...any) (any, error) {
	c := _redis.Get()
	defer c.Close()
	keylen := 0
	if keys != nil {
		keylen = len(keys)
	}
	arr := make([]any, 0, keylen+len(args)+2)
	arr = append(arr, lua)
	arr = append(arr, keylen)
	for _, k := range keys {
		arr = append(arr, k)
	}
	arr = append(arr, args...)
	return c.Do("EVAL", arr...)
}

func ExecLua(lua string, keys []string, args ...any) error {
	_, err := exec_lua(lua, keys, args...)
	return err
}

func GetLuaInt64(lua string, keys []string, args ...any) (int64, error) {
	return redis.Int64(exec_lua(lua, keys, args...))
}

func GetLuaStrings(lua string, keys []string, args ...any) ([]string, error) {
	return redis.Strings(exec_lua(lua, keys, args...))
}

func GetLuaInt64s(lua string, keys []string, args ...any) ([]int64, error) {
	return redis.Int64s(exec_lua(lua, keys, args...))
}

func Delete(key ...string) error {
	c := _redis.Get()
	defer c.Close()
	param := make([]any, len(key))
	for i, k := range key {
		param[i] = k
	}
	_, err := c.Do("DEL", param...)
	return err
}

func GetKeysWithPrefix(prefix string) ([]string, error) {
	c := _redis.Get()
	defer c.Close()
	cursor := 0
	var keys []string
	for {
		values, err := redis.Values(c.Do("SCAN", cursor, "MATCH", prefix+"*"))
		if err != nil {
			return nil, err
		}
		var items []string
		if _, err := redis.Scan(values, &cursor, &items); err != nil {
			return nil, err
		}
		keys = append(keys, items...)
		if cursor == 0 {
			break
		}
	}
	return keys, nil
}