modelRT/util/redis_init.go

// Package util provide some utility functions
package util

import (
	"context"
	"fmt"
	"time"

	"modelRT/config"

	redigo "github.com/gomodule/redigo/redis"
	"github.com/redis/go-redis/v9"
)

// NewRedisClient define func of initialize the Redis client
func NewRedisClient(addr string, opts ...RedisOption) (*redis.Client, error) {
	// default options
	options := RedisOptions{
		redisOptions: &redis.Options{
			Addr: addr,
		},
	}

	// Apply configuration options from config
	for _, opt := range opts {
		opt(&options)
	}

	// create redis client
	client := redis.NewClient(options.redisOptions)

	if options.timeout > 0 {
		// check if the connection is successful
		ctx, cancel := context.WithTimeout(context.Background(), options.timeout)
		defer cancel()
		if err := client.Ping(ctx).Err(); err != nil {
			return nil, fmt.Errorf("can not connect redis:%v", err)
		}
	}
	return client, nil
}

// NewRedigoPool define func of initialize the Redigo pool
func NewRedigoPool(rCfg config.RedisConfig) (*redigo.Pool, error) {
	pool := &redigo.Pool{
		MaxIdle:   rCfg.PoolSize / 2,
		MaxActive: rCfg.PoolSize,
		// TODO optimize IdleTimeout with config parameter
		IdleTimeout: 240 * time.Second,

		// Dial function to create the connection
		Dial: func() (redigo.Conn, error) {
			timeout := time.Duration(rCfg.Timeout) * time.Millisecond // 假设 rCfg.Timeout 是毫秒
			opts := []redigo.DialOption{
				redigo.DialDatabase(rCfg.DB),
				redigo.DialPassword(rCfg.Password),
				redigo.DialConnectTimeout(timeout),
				// redigo.DialReadTimeout(timeout),
				// redigo.DialWriteTimeout(timeout),
				// TODO add redigo.DialUsername when redis open acl
				// redis.DialUsername("username"),
			}

			c, err := redigo.Dial("tcp", rCfg.Addr, opts...)
			if err != nil {
				return nil, fmt.Errorf("redigo dial failed: %w", err)
			}
			return c, nil
		},
	}

	conn := pool.Get()
	defer conn.Close()

	if conn.Err() != nil {
		return nil, fmt.Errorf("failed to get connection from pool: %w", conn.Err())
	}

	_, err := conn.Do("PING")
	if err != nil {
		return nil, fmt.Errorf("redis connection test (PING) failed: %w", err)
	}
	return pool, nil
}