Repository: zeromicro/go-zero
Pull request #3073 (2023-03-28 06:40:18)

refactor(redis): add NonBlock config, disable redis ping by default

codecov[bot]: ## [Codecov](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#3073](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (fd6e27d) into [master](https://codecov.io/gh/zeromicro/go-zero/commit/04434646eb38a713fc5d429290eb9b546f10576f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (0443464) will **decrease** coverage by `0.09%`.
> The diff coverage is `66.66%`.


| [Impacted Files](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [core/stores/redis/conf.go](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvcmVkaXMvY29uZi5nbw==) | `100.00% <ø> (ø)` | |
| [core/stores/redis/redisclustermanager.go](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvcmVkaXMvcmVkaXNjbHVzdGVybWFuYWdlci5nbw==) | `68.57% <0.00%> (-7.19%)` | :arrow_down: |
| [core/stores/redis/redis.go](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvcmVkaXMvcmVkaXMuZ28=) | `98.26% <71.42%> (-0.61%)` | :arrow_down: |
| [core/stores/redis/redisclientmanager.go](https://codecov.io/gh/zeromicro/go-zero/pull/3073?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvcmVkaXMvcmVkaXNjbGllbnRtYW5hZ2VyLmdv) | `88.00% <100.00%> (+1.04%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://codecov.io/gh/zeromicro/go-zero/pull/3073/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)


Patch:
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/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
 	})


Test patch:
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")