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

     1  package redis
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/gomodule/redigo/redis"
    10  
    11  	"github.com/machinefi/w3bstream/pkg/depends/base/consts"
    12  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    13  )
    14  
    15  type Endpoint struct {
    16  	Endpoint types.Endpoint `env:""`
    17  	Wait     bool
    18  	Prefix   string
    19  	pool     *redis.Pool
    20  }
    21  
    22  func (r *Endpoint) Acquire() redis.Conn {
    23  	if r.pool != nil {
    24  		return r.pool.Get()
    25  	}
    26  	return nil
    27  }
    28  
    29  func (r *Endpoint) Key(key string) string {
    30  	return fmt.Sprintf("%s:%s", r.Prefix, key)
    31  }
    32  
    33  func (r *Endpoint) LivenessCheck() map[string]string {
    34  	m := map[string]string{}
    35  
    36  	conn := r.Acquire()
    37  	defer conn.Close()
    38  	_, err := conn.Do("PING")
    39  	if err != nil {
    40  		m[r.Endpoint.Host()] = err.Error()
    41  	} else {
    42  		m[r.Endpoint.Host()] = "ok"
    43  	}
    44  
    45  	return m
    46  }
    47  
    48  func (r *Endpoint) SetDefault() {
    49  	if r.Endpoint.Scheme == "" {
    50  		r.Endpoint.Scheme = "tcp"
    51  	}
    52  	if r.Endpoint.Hostname == "" {
    53  		r.Endpoint.Hostname = "127.0.0.1"
    54  	}
    55  	if r.Endpoint.Port == 0 {
    56  		r.Endpoint.Port = 6379
    57  	}
    58  	if !r.Wait {
    59  		r.Wait = true
    60  	}
    61  	if r.Prefix == "" {
    62  		r.Prefix = fmt.Sprintf("%s:%s:",
    63  			strings.ToLower(os.Getenv(consts.GoRuntimeEnv)),
    64  			strings.ToLower(os.Getenv(consts.EnvProjectName)),
    65  		)
    66  	}
    67  }
    68  
    69  func (r *Endpoint) Init() {
    70  	if r.pool == nil {
    71  		r.init()
    72  	}
    73  }
    74  
    75  func (r *Endpoint) init() {
    76  	opt := struct {
    77  		ConnectTimeout types.Duration `name:"connectTimeout" default:"10s"`
    78  		ReadTimeout    types.Duration `name:"readTimeout"    default:"10s"`
    79  		WriteTimeout   types.Duration `name:"writeTimeout"   default:"10s"`
    80  		IdleTimeout    types.Duration `name:"idleTimeout"    default:"240s"`
    81  		MaxActive      int            `name:"maxActive"      default:"5"`
    82  		MaxIdle        int            `name:"maxIdle"        default:"3"`
    83  		DB             int            `name:"dB"             default:"1"`
    84  	}{}
    85  
    86  	err := types.UnmarshalExtra(r.Endpoint.Param, &opt)
    87  	if err != nil {
    88  		panic(err)
    89  	}
    90  
    91  	dialer := func() (c redis.Conn, err error) {
    92  		options := []redis.DialOption{
    93  			redis.DialDatabase(opt.DB),
    94  			redis.DialConnectTimeout(time.Duration(opt.ConnectTimeout)),
    95  			redis.DialWriteTimeout(time.Duration(opt.WriteTimeout)),
    96  			redis.DialReadTimeout(time.Duration(opt.ReadTimeout)),
    97  		}
    98  
    99  		if r.Endpoint.Password != "" {
   100  			options = append(options, redis.DialPassword(r.Endpoint.Password.String()))
   101  		}
   102  
   103  		return redis.Dial("tcp", r.Endpoint.Host(), options...)
   104  	}
   105  
   106  	r.pool = &redis.Pool{
   107  		Dial:        dialer,
   108  		MaxIdle:     opt.MaxIdle,
   109  		MaxActive:   opt.MaxActive,
   110  		IdleTimeout: time.Duration(opt.IdleTimeout),
   111  		Wait:        true,
   112  	}
   113  }
   114  
   115  func (r *Endpoint) Exec(cmd *Cmd, others ...*Cmd) (interface{}, error) {
   116  	c := r.Acquire()
   117  	defer c.Close()
   118  
   119  	if (len(others)) == 0 {
   120  		return c.Do(cmd.Name, cmd.Args...)
   121  	}
   122  
   123  	err := c.Send("MULTI")
   124  	if err != nil {
   125  		return nil, err
   126  	}
   127  
   128  	err = c.Send(cmd.Name, cmd.Args...)
   129  	if err != nil {
   130  		return nil, err
   131  	}
   132  
   133  	for i := range others {
   134  		o := others[i]
   135  		if o == nil {
   136  			continue
   137  		}
   138  		err := c.Send(o.Name, o.Args...)
   139  		if err != nil {
   140  			return nil, err
   141  		}
   142  	}
   143  
   144  	return c.Do("EXEC")
   145  }