Repository: zeromicro/go-zero
Pull request #1185 (2021-10-31 13:06:26)

feat: slow threshold customizable in redis


Patch:
diff --git a/core/conf/time.go b/core/conf/time.go
new file mode 100644
index 00000000..be47300c
--- /dev/null
+++ b/core/conf/time.go
@@ -0,0 +1,16 @@
+package conf
+
+import "time"
+
+const minDuration = 100 * time.Microsecond
+
+// CheckedDuration returns the duration that guaranteed to be greater than 100us.
+// Why we need this is because users sometimes intend to use 500 to represent 500ms.
+// In config, duration less than 100us should always be missing ms etc.
+func CheckedDuration(duration time.Duration) time.Duration {
+	if duration > minDuration {
+		return duration
+	}
+
+	return duration * time.Millisecond
+}
diff --git a/core/stores/redis/conf.go b/core/stores/redis/conf.go
index 2d93723b..776bd5ac 100644
--- a/core/stores/redis/conf.go
+++ b/core/stores/redis/conf.go
@@ -1,6 +1,11 @@
 package redis
 
-import "errors"
+import (
+	"errors"
+	"time"
+
+	"github.com/tal-tech/go-zero/core/conf"
+)
 
 var (
 	// ErrEmptyHost is an error that indicates no redis host is set.
@@ -14,10 +19,11 @@ var (
 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:",default=false,options=true|false"`
+		Host          string
+		Type          string        `json:",default=node,options=node|cluster"`
+		Pass          string        `json:",optional"`
+		Tls           bool          `json:",default=false,options=true|false"`
+		SlowThreshold time.Duration `json:",default=100ms"`
 	}
 
 	// A RedisKeyConf is a redis config with key.
@@ -36,6 +42,9 @@ func (rc RedisConf) NewRedis() *Redis {
 	if len(rc.Pass) > 0 {
 		opts = append(opts, WithPass(rc.Pass))
 	}
+	if rc.SlowThreshold > 0 {
+		opts = append(opts, WithSlowThreshold(conf.CheckedDuration(rc.SlowThreshold)))
+	}
 	if rc.Tls {
 		opts = append(opts, WithTLS())
 	}
diff --git a/core/stores/redis/process.go b/core/stores/redis/process.go
index d1081b46..2624acbe 100644
--- a/core/stores/redis/process.go
+++ b/core/stores/redis/process.go
@@ -2,6 +2,7 @@ package redis
 
 import (
 	"strings"
+	"time"
 
 	red "github.com/go-redis/redis"
 	"github.com/tal-tech/go-zero/core/logx"
@@ -9,24 +10,26 @@ import (
 	"github.com/tal-tech/go-zero/core/timex"
 )
 
-func process(proc func(red.Cmder) error) func(red.Cmder) error {
-	return func(cmd red.Cmder) error {
-		start := timex.Now()
+func checkDuration(slowThreshold time.Duration) func(proc func(red.Cmder) error) func(red.Cmder) error {
+	return func(proc func(red.Cmder) error) func(red.Cmder) error {
+		return func(cmd red.Cmder) error {
+			start := timex.Now()
 
-		defer func() {
-			duration := timex.Since(start)
-			if duration > slowThreshold {
-				var buf strings.Builder
-				for i, arg := range cmd.Args() {
-					if i > 0 {
-						buf.WriteByte(' ')
+			defer func() {
+				duration := timex.Since(start)
+				if duration > slowThreshold {
+					var buf strings.Builder
+					for i, arg := range cmd.Args() {
+						if i > 0 {
+							buf.WriteByte(' ')
+						}
+						buf.WriteString(mapping.Repr(arg))
 					}
-					buf.WriteString(mapping.Repr(arg))
+					logx.WithDuration(duration).Slowf("[REDIS] slowcall on executing: %s", buf.String())
 				}
-				logx.WithDuration(duration).Slowf("[REDIS] slowcall on executing: %s", buf.String())
-			}
-		}()
+			}()
 
-		return proc(cmd)
+			return proc(cmd)
+		}
 	}
 }
diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index 87347ce4..63acb68c 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -21,8 +21,7 @@ const (
 
 	blockingQueryTimeout = 5 * time.Second
 	readWriteTimeout     = 2 * time.Second
-
-	slowThreshold = time.Millisecond * 100
+	defaultSlowThreshold = time.Millisecond * 100
 )
 
 // ErrNilNode is an error that indicates a nil redis node.
@@ -40,11 +39,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
+		slowThreshold time.Duration
 	}
 
 	// RedisNode interface represents a redis node.
@@ -78,9 +78,10 @@ type (
 // New returns a Redis with given options.
 func New(addr string, opts ...Option) *Redis {
 	r := &Redis{
-		Addr: addr,
-		Type: NodeType,
-		brk:  breaker.NewBreaker(),
+		Addr:          addr,
+		Type:          NodeType,
+		brk:           breaker.NewBreaker(),
+		slowThreshold: defaultSlowThreshold,
 	}
 
 	for _, opt := range opts {
@@ -1765,6 +1766,13 @@ func WithPass(pass string) Option {
 	}
 }
 
+// WithSlowThreshold sets the slow threshold.
+func WithSlowThreshold(threshold time.Duration) Option {
+	return func(r *Redis) {
+		r.slowThreshold = threshold
+	}
+}
+
 // WithTLS customizes the given Redis with TLS enabled.
 func WithTLS() Option {
 	return func(r *Redis) {
diff --git a/core/stores/redis/redisclientmanager.go b/core/stores/redis/redisclientmanager.go
index 9467f34c..943e85d5 100644
--- a/core/stores/redis/redisclientmanager.go
+++ b/core/stores/redis/redisclientmanager.go
@@ -32,7 +32,7 @@ func getClient(r *Redis) (*red.Client, error) {
 			MinIdleConns: idleConns,
 			TLSConfig:    tlsConfig,
 		})
-		store.WrapProcess(process)
+		store.WrapProcess(checkDuration(r.slowThreshold))
 		return store, nil
 	})
 	if err != nil {
diff --git a/core/stores/redis/redisclustermanager.go b/core/stores/redis/redisclustermanager.go
index 73c6dd05..722fb4fb 100644
--- a/core/stores/redis/redisclustermanager.go
+++ b/core/stores/redis/redisclustermanager.go
@@ -25,7 +25,7 @@ func getCluster(r *Redis) (*red.ClusterClient, error) {
 			MinIdleConns: idleConns,
 			TLSConfig:    tlsConfig,
 		})
-		store.WrapProcess(process)
+		store.WrapProcess(checkDuration(r.slowThreshold))
 
 		return store, nil
 	})


Test patch:
diff --git a/core/conf/time_test.go b/core/conf/time_test.go
new file mode 100644
index 00000000..130e0748
--- /dev/null
+++ b/core/conf/time_test.go
@@ -0,0 +1,14 @@
+package conf
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCheckedDuration(t *testing.T) {
+	assert.Equal(t, time.Second, CheckedDuration(1000))
+	assert.Equal(t, 2*time.Second, CheckedDuration(2000))
+	assert.Equal(t, 2*time.Second, CheckedDuration(time.Second*2))
+}
diff --git a/core/stores/redis/redis_test.go b/core/stores/redis/redis_test.go
index 37e60ec4..8b5ec3f2 100644
--- a/core/stores/redis/redis_test.go
+++ b/core/stores/redis/redis_test.go
@@ -1115,7 +1115,7 @@ func runOnRedisTLS(t *testing.T, fn func(client *Redis)) {
 			client.Close()
 		}
 	}()
-	fn(New(s.Addr(), WithTLS()))
+	fn(New(s.Addr(), WithTLS(), WithSlowThreshold(defaultSlowThreshold/2)))
 }
 
 func badType() Option {