diff --git a/core/stores/redis/conf.go b/core/stores/redis/conf.go index bdf9afef..c5421087 100644 --- a/core/stores/redis/conf.go +++ b/core/stores/redis/conf.go @@ -1,6 +1,9 @@ package redis -import "errors" +import ( + "errors" + "time" +) var ( // ErrEmptyHost is an error that indicates no redis host is set. @@ -9,17 +12,18 @@ var ( ErrEmptyType = errors.New("empty redis type") // ErrEmptyKey is an error that indicates no redis key is set. ErrEmptyKey = errors.New("empty redis key") - // ErrPing is an error that indicates ping failed. - ErrPing = errors.New("ping redis failed") ) type ( // A RedisConf is a redis config. RedisConf struct { - Host string - Type string `json:",default=node,options=node|cluster"` - Pass string `json:",optional"` - Tls bool `json:",optional"` + Host string + Type string `json:",default=node,options=node|cluster"` + Pass string `json:",optional"` + Tls bool `json:",optional"` + NonBlock bool `json:",default=true"` + // PingTimeout is the timeout for ping redis. + PingTimeout time.Duration `json:",default=1s"` } // A RedisKeyConf is a redis config with key. diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go index 31c7e514..e8eb8964 100644 --- a/core/stores/redis/redis.go +++ b/core/stores/redis/redis.go @@ -10,6 +10,7 @@ import ( red "github.com/go-redis/redis/v8" "github.com/zeromicro/go-zero/core/breaker" + "github.com/zeromicro/go-zero/core/errorx" "github.com/zeromicro/go-zero/core/mapping" "github.com/zeromicro/go-zero/core/syncx" ) @@ -25,6 +26,7 @@ const ( blockingQueryTimeout = 5 * time.Second readWriteTimeout = 2 * time.Second defaultSlowThreshold = time.Millisecond * 100 + defaultPingTimeout = time.Second ) var ( @@ -51,11 +53,12 @@ type ( // Redis defines a redis node/cluster. It is thread-safe. Redis struct { - Addr string - Type string - Pass string - tls bool - brk breaker.Breaker + Addr string + Type string + Pass string + tls bool + brk breaker.Breaker + hooks []red.Hook } // RedisNode interface represents a redis node. @@ -119,8 +122,10 @@ func NewRedis(conf RedisConf, opts ...Option) (*Redis, error) { } rds := newRedis(conf.Host, opts...) - if !rds.Ping() { - return nil, ErrPing + if !conf.NonBlock { + if err := rds.checkConnection(conf.PingTimeout); err != nil { + return nil, errorx.Wrap(err, fmt.Sprintf("redis connect error, addr: %s", conf.Host)) + } } return rds, nil @@ -2769,6 +2774,23 @@ func (s *Redis) ZunionstoreCtx(ctx context.Context, dest string, store *ZStore) return } +func (s *Redis) checkConnection(pingTimeout time.Duration) error { + conn, err := getRedis(s) + if err != nil { + return err + } + + timeout := defaultPingTimeout + if pingTimeout > 0 { + timeout = pingTimeout + } + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + return conn.Ping(ctx).Err() +} + // Cluster customizes the given Redis as a cluster. func Cluster() Option { return func(r *Redis) { @@ -2795,6 +2817,14 @@ func WithTLS() Option { } } +// withHook customizes the given Redis with given hook, only for private use now, +// maybe expose later. +func withHook(hook red.Hook) Option { + return func(r *Redis) { + r.hooks = append(r.hooks, hook) + } +} + func acceptable(err error) bool { return err == nil || err == red.Nil || err == context.Canceled } diff --git a/core/stores/redis/redis_test.go b/core/stores/redis/redis_test.go index ffd4d00d..d23b9127 100644 --- a/core/stores/redis/redis_test.go +++ b/core/stores/redis/redis_test.go @@ -16,6 +16,25 @@ import ( "github.com/zeromicro/go-zero/core/stringx" ) +type myHook struct { + red.Hook + includePing bool +} + +var _ red.Hook = myHook{} + +func (m myHook) BeforeProcess(ctx context.Context, cmd red.Cmder) (context.Context, error) { + return ctx, nil +} + +func (m myHook) AfterProcess(ctx context.Context, cmd red.Cmder) error { + // skip ping cmd + if cmd.Name() == "ping" && !m.includePing { + return nil + } + return errors.New("hook error") +} + func TestNewRedis(t *testing.T) { r1, err := miniredis.Run() assert.NoError(t, err) @@ -126,6 +145,31 @@ func TestNewRedis(t *testing.T) { } } +func TestRedis_NonBlock(t *testing.T) { + logx.Disable() + + t.Run("nonBlock true", func(t *testing.T) { + s := miniredis.RunT(t) + // use hook to simulate redis ping error + _, err := NewRedis(RedisConf{ + Host: s.Addr(), + NonBlock: true, + Type: NodeType, + }, withHook(myHook{includePing: true})) + assert.NoError(t, err) + }) + + t.Run("nonBlock false", func(t *testing.T) { + s := miniredis.RunT(t) + _, err := NewRedis(RedisConf{ + Host: s.Addr(), + NonBlock: false, + Type: NodeType, + }, withHook(myHook{includePing: true})) + assert.ErrorContains(t, err, "redis connect error") + }) +} + func TestRedis_Decr(t *testing.T) { runOnRedis(t, func(client *Redis) { _, err := New(client.Addr, badType()).Decr("a") diff --git a/core/stores/redis/redisclientmanager.go b/core/stores/redis/redisclientmanager.go index fa4b72c9..8311e8c0 100644 --- a/core/stores/redis/redisclientmanager.go +++ b/core/stores/redis/redisclientmanager.go @@ -33,6 +33,9 @@ func getClient(r *Redis) (*red.Client, error) { TLSConfig: tlsConfig, }) store.AddHook(durationHook) + for _, hook := range r.hooks { + store.AddHook(hook) + } return store, nil }) diff --git a/core/stores/redis/redisclustermanager.go b/core/stores/redis/redisclustermanager.go index dda1bcdb..d4a489e6 100644 --- a/core/stores/redis/redisclustermanager.go +++ b/core/stores/redis/redisclustermanager.go @@ -29,6 +29,9 @@ func getCluster(r *Redis) (*red.ClusterClient, error) { TLSConfig: tlsConfig, }) store.AddHook(durationHook) + for _, hook := range r.hooks { + store.AddHook(hook) + } return store, nil })