github.com/erda-project/erda-infra@v1.0.9/providers/redis/redis.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package redis
    16  
    17  import (
    18  	"fmt"
    19  	"reflect"
    20  	"strings"
    21  	"sync"
    22  	"time"
    23  
    24  	"github.com/go-redis/redis"
    25  
    26  	"github.com/erda-project/erda-infra/base/logs"
    27  	"github.com/erda-project/erda-infra/base/servicehub"
    28  )
    29  
    30  // Interface .
    31  type Interface interface {
    32  	DB() *redis.Client
    33  	Open(db int) (*redis.Client, error)
    34  }
    35  
    36  var (
    37  	interfaceType = reflect.TypeOf((*Interface)(nil)).Elem()
    38  	clientType    = reflect.TypeOf((*redis.Client)(nil))
    39  )
    40  
    41  type config struct {
    42  	Addr          string `file:"addr" env:"REDIS_ADDR"`
    43  	MasterName    string `file:"master_name" env:"REDIS_MASTER_NAME"`
    44  	SentinelsAddr string `file:"sentinels_addr" env:"REDIS_SENTINELS_ADDR"`
    45  	Password      string `file:"password" env:"REDIS_PASSWORD"`
    46  	DB            int    `file:"db" env:"REDIS_DB" default:"0"`
    47  
    48  	MaxRetries int `file:"max_retries" env:"REDIS_MAX_RETRIES"`
    49  
    50  	DialTimeout  time.Duration `file:"dial_timeout" env:"REDIS_DIAL_TIMEOUT"`
    51  	ReadTimeout  time.Duration `file:"read_timeout" env:"REDIS_READ_TIMEOUT"`
    52  	WriteTimeout time.Duration `file:"write_timeout" env:"REDIS_WRITE_TIMEOUT"`
    53  
    54  	PoolSize           int           `file:"pool_size" env:"REDIS_POOL_SIZE"`
    55  	PoolTimeout        time.Duration `file:"pool_timeout" env:"REDIS_POOL_TIMEOUT"`
    56  	IdleTimeout        time.Duration `file:"idle_timeout" env:"REDIS_IDLE_TIMEOUT"`
    57  	IdleCheckFrequency time.Duration `file:"idle_check_frequency" env:"REDIS_IDLE_CHECK_FREQUENCY"`
    58  }
    59  
    60  // provider .
    61  type provider struct {
    62  	Cfg     *config
    63  	Log     logs.Logger
    64  	client  *redis.Client
    65  	clients map[int]*redis.Client
    66  	lock    sync.Mutex
    67  }
    68  
    69  // Init .
    70  func (p *provider) Init(ctx servicehub.Context) error {
    71  	c, err := p.Open(p.Cfg.DB)
    72  	if err != nil {
    73  		return err
    74  	}
    75  	p.client = c
    76  	return nil
    77  }
    78  
    79  func (p *provider) DB() *redis.Client {
    80  	// already initialized at provider.Init
    81  	return p.client
    82  }
    83  
    84  func (p *provider) Open(db int) (*redis.Client, error) {
    85  	p.lock.Lock()
    86  	defer p.lock.Unlock()
    87  	if c, ok := p.clients[db]; ok {
    88  		return c, nil
    89  	}
    90  	var c *redis.Client
    91  	if p.Cfg.MasterName != "" && p.Cfg.SentinelsAddr != "" {
    92  		addrs := strings.Split(p.Cfg.SentinelsAddr, ",")
    93  		c = redis.NewFailoverClient(&redis.FailoverOptions{
    94  			MasterName:         p.Cfg.MasterName,
    95  			SentinelAddrs:      addrs,
    96  			Password:           p.Cfg.Password,
    97  			DB:                 db,
    98  			MaxRetries:         p.Cfg.MaxRetries,
    99  			DialTimeout:        p.Cfg.DialTimeout,
   100  			ReadTimeout:        p.Cfg.ReadTimeout,
   101  			WriteTimeout:       p.Cfg.WriteTimeout,
   102  			PoolSize:           p.Cfg.PoolSize,
   103  			PoolTimeout:        p.Cfg.PoolTimeout,
   104  			IdleTimeout:        p.Cfg.IdleTimeout,
   105  			IdleCheckFrequency: p.Cfg.IdleCheckFrequency,
   106  		})
   107  	} else if p.Cfg.Addr != "" {
   108  		c = redis.NewClient(&redis.Options{
   109  			Addr:               p.Cfg.Addr,
   110  			Password:           p.Cfg.Password,
   111  			DB:                 db,
   112  			MaxRetries:         p.Cfg.MaxRetries,
   113  			DialTimeout:        p.Cfg.DialTimeout,
   114  			ReadTimeout:        p.Cfg.ReadTimeout,
   115  			WriteTimeout:       p.Cfg.WriteTimeout,
   116  			PoolSize:           p.Cfg.PoolSize,
   117  			PoolTimeout:        p.Cfg.PoolTimeout,
   118  			IdleTimeout:        p.Cfg.IdleTimeout,
   119  			IdleCheckFrequency: p.Cfg.IdleCheckFrequency,
   120  		})
   121  	} else {
   122  		err := fmt.Errorf("redis config error: no addr or sentinel")
   123  		p.Log.Error(err)
   124  		return nil, err
   125  	}
   126  
   127  	if pong, err := c.Ping().Result(); err != nil {
   128  		p.Log.Errorf("redis ping error: %s", err)
   129  		return nil, err
   130  	} else if pong != "PONG" {
   131  		err := fmt.Errorf("redis ping result '%s' is not PONG", pong)
   132  		p.Log.Error(err)
   133  		return nil, err
   134  	} else {
   135  		p.Log.Infof("open redis db %d and ping ok", db)
   136  	}
   137  
   138  	p.clients[db] = c
   139  	return c, nil
   140  }
   141  
   142  func (p *provider) Provide(ctx servicehub.DependencyContext, args ...interface{}) interface{} {
   143  	if ctx.Service() == "redis-client" || ctx.Type() == clientType {
   144  		return p.DB()
   145  	}
   146  	return p
   147  }
   148  
   149  func init() {
   150  	servicehub.Register("redis", &servicehub.Spec{
   151  		Services: []string{"redis", "redis-client"},
   152  		Types: []reflect.Type{
   153  			interfaceType, clientType,
   154  		},
   155  		Description: "redis",
   156  		ConfigFunc:  func() interface{} { return &config{} },
   157  		Creator: func() servicehub.Provider {
   158  			return &provider{clients: make(map[int]*redis.Client)}
   159  		},
   160  	})
   161  }