mycommon/myredis/redis.go

package myredis

import (
	"context"
	"fmt"
	"git.makemake.in/kzkzzzz/mycommon/myconf"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
	"log"
	"sync"
	"time"
)

const (
	DefaultInstance = "redis"
)

type Client struct {
	*redis.Client
	redisOpt    *redis.Options
	disablePing bool
}

//type Conf struct {
//	Addr            string
//	Password        string
//	Db              int
//	PoolSize        int    // 连接池数量 如果不够用会继续新建连接 可以用 MaxActiveConns 限制
//	MinIdleConn     int    // 最小空闲连接数 预热可能用到
//	MaxIdleConn     int    // 最大空闲连接数 0不限制
//	MaxActiveConn   int    // 0 不限制, 如果不限制, 超出pool size可以继续新建立连接, 但使用完之后不会放回连接池
//	PoolTimeout     string // 等待连接池超时时间 Default is ReadTimeout + 1 second.
//	DialTimeout     string // 拨号超时时间
//	ReadTimeout     string // 读取超时
//	WriteTimeout    string // 写入超时
//	MaxRetries      int    // 重试次数
//	ConnMaxIdleTime string // 连接空闲时间
//}

var (
	defaultClient *redis.Client
	instanceMap   = &sync.Map{}
)

func GetClient(name ...string) *Client {
	var instanceName string
	if len(name) > 0 {
		instanceName = name[0]
	} else {
		instanceName = DefaultInstance
	}

	v, ok := instanceMap.Load(instanceName)
	if !ok {
		panic(fmt.Errorf("redis instance [%s] not init", instanceName))
	}

	return v.(*Client)
}

// InitClient 初始化全局默认client
func InitClient(config *myconf.Config, opts ...Opt) {
	client, err := NewClient(DefaultInstance, config, opts...)
	if err != nil {
		panic(err)
	}
	instanceMap.Store(DefaultInstance, client)
}

// InitClientInstance 初始化全局的client
func InitClientInstance(instanceName string, config *myconf.Config, opts ...Opt) {
	client, err := NewClient(instanceName, config, opts...)
	if err != nil {
		panic(err)
	}
	instanceMap.Store(instanceName, client)
}

func NewClient(instanceName string, config *myconf.Config, opts ...Opt) (*Client, error) {
	cf := &Conf{}
	err := config.UnmarshalKey(instanceName, &cf)
	if err != nil {
		return nil, err
	}
	client, err := NewClientFromConf(cf, opts...)
	if err != nil {
		return nil, err
	}

	instanceMap.Store(instanceMap, client)
	return client, nil
}

type Conf struct {
	Addr            string
	Password        string
	Db              int
	PoolSize        int    // 连接池数量 如果不够用会继续新建连接 可以用 MaxActiveConns 限制
	MinIdleConn     int    // 最小空闲连接数 预热可能用到
	MaxIdleConn     int    // 最大空闲连接数 0不限制
	MaxActiveConn   int    // 0 不限制, 如果不限制, 超出pool size可以继续新建立连接, 但使用完之后不会放回连接池
	DialTimeout     string // 拨号超时时间, 时间单位格式 10ms, 60s, 15m, 2h...
	ReadTimeout     string // 读取超时
	WriteTimeout    string // 写入超时
	ConnMaxIdleTime string // 连接空闲时间
	PoolTimeout     string // 等待连接池超时时间 Default is ReadTimeout + 1 second.
	MaxRetries      int    // 重试次数
}

func NewClientFromConf(cf *Conf, opts ...Opt) (*Client, error) {
	c := &Client{
		redisOpt: DefaultRedisOpt(),
	}

	for _, opt := range opts {
		opt(c)
	}

	err := c.parseConfToRedisOpt(cf)
	if err != nil {
		return nil, err
	}

	client := redis.NewClient(c.redisOpt)

	if c.disablePing == false {
		_, err := client.Ping(context.Background()).Result()
		if err != nil {
			return nil, fmt.Errorf("redis ping err: %s", err)
		}
	}

	c.Client = client

	instanceMap.Store(uuid.New().String(), c)

	log.Printf("connect redis success [addr:%s - db:%d]", cf.Addr, cf.Db)

	return c, nil
}

func DefaultRedisOpt() *redis.Options {
	return &redis.Options{
		Addr:                  "",
		Password:              "",
		DB:                    0,
		MaxRetries:            3,                         // 重试次数
		DialTimeout:           time.Millisecond * 500,    // 拨号超时时间
		ReadTimeout:           time.Second,               // 读取超时
		WriteTimeout:          time.Second * 3,           // 写入超时
		PoolSize:              1024,                      // 连接池数量 如果不够用会继续新建连接 可以用 MaxActiveConns 限制
		PoolTimeout:           time.Second + time.Second, // 等待连接池超时时间 Default is ReadTimeout + 1 second.
		MinIdleConns:          0,                         // 最小空闲连接数 预热可能用到
		MaxIdleConns:          0,                         // 最大空闲连接数 0不限制
		MaxActiveConns:        0,                         // 0 不限制, 如果不限制, 超出pool size可以继续新建立连接, 但使用完之后不会放回连接池
		ConnMaxIdleTime:       time.Minute * 10,          // 连接空闲时间
		ContextTimeoutEnabled: true,                      // context控制超时用到
	}
}

func (c *Client) parseTime(v string) (time.Duration, error) {
	if v == "" {
		return 0, nil
	}

	d, err := time.ParseDuration(v)
	if err != nil {
		return 0, fmt.Errorf("parse time %s err: %s", v, err)
	}

	return d, nil
}

func (c *Client) parseConfToRedisOpt(cf *Conf) error {
	c.redisOpt.Addr = cf.Addr
	c.redisOpt.Password = cf.Password
	c.redisOpt.DB = cf.Db

	if v := cf.PoolSize; v > 0 {
		c.redisOpt.PoolSize = v
	}

	if v := cf.MinIdleConn; v > 0 {
		c.redisOpt.MinIdleConns = v
	}

	if v := cf.MaxIdleConn; v > 0 {
		c.redisOpt.MaxIdleConns = v
	}

	if v := cf.MaxActiveConn; v > 0 {
		c.redisOpt.MaxActiveConns = v
	}

	if v := cf.MaxRetries; v > 0 {
		c.redisOpt.MaxRetries = v
	}

	if v, err := c.parseTime(cf.DialTimeout); err != nil {
		return err
	} else if v > 0 {
		c.redisOpt.DialTimeout = v
	}

	if v, err := c.parseTime(cf.ReadTimeout); err != nil {
		return err
	} else if v > 0 {
		c.redisOpt.ReadTimeout = v
	}

	if v, err := c.parseTime(cf.WriteTimeout); err != nil {
		return err
	} else if v > 0 {
		c.redisOpt.WriteTimeout = v
	}

	if v, err := c.parseTime(cf.ConnMaxIdleTime); err != nil {
		return err
	} else if v > 0 {
		c.redisOpt.ConnMaxIdleTime = v
	}

	if v, err := c.parseTime(cf.PoolTimeout); err != nil {
		return err
	} else if v > 0 {
		c.redisOpt.PoolTimeout = v
	}

	return nil
}

const luaSetOnce = `if redis.call('setnx',KEYS[1],ARGV[1]) == 1 then redis.call('expire',KEYS[1],ARGV[2]) return 1 else return 0 end`

// SetOnce 设置一次并设置过期时间, key不存在则设置成功返回1, key已存在返回0
func (c *Client) SetOnce(ctx context.Context, key, value string, t time.Duration) (int, error) {
	if t < time.Second {
		return 0, fmt.Errorf("time must >= 1s")
	}
	return c.Client.Eval(ctx, luaSetOnce, []string{key}, value, t).Int()
}

func CloseAllClient() {
	instanceMap.Range(func(k, v any) bool {
		v.(*Client).Close()
		return true
	})
}