github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/redis/redis_config.go (about)

     1  package redis
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"strconv"
     7  	"strings"
     8  	"time"
     9  
    10  	"github.com/gomodule/redigo/redis"
    11  
    12  	"github.com/machinefi/w3bstream/pkg/depends/base/consts"
    13  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    14  )
    15  
    16  type Redis struct {
    17  	Protocol       string
    18  	Host           string
    19  	Port           int
    20  	Password       types.Password
    21  	ConnectTimeout types.Duration
    22  	WriteTimeout   types.Duration
    23  	ReadTimeout    types.Duration
    24  	IdleTimeout    types.Duration
    25  	MaxActive      int
    26  	MaxIdle        int
    27  	Wait           bool
    28  	DB             int
    29  	pool           *redis.Pool
    30  	Prefix         string
    31  }
    32  
    33  func (r *Redis) Acquire() redis.Conn {
    34  	if r.pool != nil {
    35  		return r.pool.Get()
    36  	}
    37  	return nil
    38  }
    39  
    40  func (r *Redis) clone() *Redis {
    41  	cloned := *r
    42  	return &cloned
    43  }
    44  
    45  func (r *Redis) WithPrefix(prefix string) *Redis {
    46  	cloned := r.clone()
    47  	if prefix != "" {
    48  		cloned.Prefix += ":" + prefix
    49  	}
    50  	return cloned
    51  }
    52  
    53  func (r *Redis) WithDBIndex(n int) *Redis {
    54  	cloned := r.clone()
    55  	cloned.DB = n
    56  	return cloned
    57  }
    58  
    59  func (r *Redis) Exec(cmd *Cmd, others ...*Cmd) (interface{}, error) {
    60  	c := r.Acquire()
    61  	defer c.Close()
    62  
    63  	if (len(others)) == 0 {
    64  		return c.Do(cmd.Name, cmd.Args...)
    65  	}
    66  
    67  	err := c.Send("MULTI")
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	err = c.Send(cmd.Name, cmd.Args...)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  
    77  	for i := range others {
    78  		o := others[i]
    79  		if o == nil {
    80  			continue
    81  		}
    82  		err := c.Send(o.Name, o.Args...)
    83  		if err != nil {
    84  			return nil, err
    85  		}
    86  	}
    87  
    88  	return c.Do("EXEC")
    89  }
    90  
    91  func (r *Redis) Key(key string) string {
    92  	return fmt.Sprintf("%s:%s", r.Prefix, key)
    93  }
    94  
    95  func (r *Redis) Name() string { return "redis-cli" }
    96  
    97  func (r *Redis) LivenessCheck() map[string]string {
    98  	m := map[string]string{}
    99  
   100  	conn := r.Acquire()
   101  	defer conn.Close()
   102  
   103  	_, err := conn.Do("PING")
   104  	key := r.Host
   105  	if r.Port != 0 {
   106  		key += ":" + strconv.Itoa(r.Port)
   107  	}
   108  	if err != nil {
   109  		m[key] = err.Error()
   110  	} else {
   111  		m[key] = "ok"
   112  	}
   113  
   114  	return m
   115  }
   116  
   117  func (r *Redis) SetDefault() {
   118  	if r.Protocol == "" {
   119  		r.Protocol = "tcp"
   120  	}
   121  	if r.Host == "" {
   122  		r.Host = "127.0.0.1"
   123  	}
   124  	if r.Port == 0 {
   125  		r.Port = 6379
   126  	}
   127  	if r.ConnectTimeout == 0 {
   128  		r.ConnectTimeout = types.Duration(10 * time.Second)
   129  	}
   130  	if r.ReadTimeout == 0 {
   131  		r.ReadTimeout = types.Duration(10 * time.Second)
   132  	}
   133  	if r.WriteTimeout == 0 {
   134  		r.WriteTimeout = types.Duration(10 * time.Second)
   135  	}
   136  	if r.IdleTimeout == 0 {
   137  		r.IdleTimeout = types.Duration(240 * time.Second)
   138  	}
   139  	if r.MaxActive == 0 {
   140  		r.MaxActive = 5
   141  	}
   142  	if r.MaxIdle == 0 {
   143  		r.MaxIdle = 3
   144  	}
   145  	if !r.Wait {
   146  		r.Wait = true
   147  	}
   148  	if r.DB == 0 {
   149  		r.DB = 1
   150  	}
   151  }
   152  
   153  func (r *Redis) Init() {
   154  	if r.Prefix == "" {
   155  		env := consts.ProduceEnv
   156  		if v := os.Getenv(consts.GoRuntimeEnv); v != "" {
   157  			env = strings.ToLower(v)
   158  		}
   159  		prj := "unknown"
   160  		if v := os.Getenv(consts.EnvProjectName); v != "" {
   161  			prj = strings.ToLower(v)
   162  		}
   163  		r.Prefix = fmt.Sprintf("%s:%s:", env, prj)
   164  	}
   165  	if r.pool == nil {
   166  		r.init()
   167  	}
   168  }
   169  
   170  func (r *Redis) init() {
   171  	dialer := func() (c redis.Conn, err error) {
   172  		c, err = redis.Dial(
   173  			r.Protocol,
   174  			fmt.Sprintf("%s:%d", r.Host, r.Port),
   175  
   176  			redis.DialWriteTimeout(time.Duration(r.WriteTimeout)),
   177  			redis.DialConnectTimeout(time.Duration(r.ConnectTimeout)),
   178  			redis.DialReadTimeout(time.Duration(r.ReadTimeout)),
   179  			redis.DialPassword(r.Password.String()),
   180  			redis.DialDatabase(r.DB),
   181  		)
   182  		return
   183  	}
   184  
   185  	r.pool = &redis.Pool{
   186  		Dial:        dialer,
   187  		MaxIdle:     r.MaxIdle,
   188  		MaxActive:   r.MaxActive,
   189  		IdleTimeout: time.Duration(r.IdleTimeout),
   190  		Wait:        true,
   191  	}
   192  }