feat: upgrade go-redis to v9 (#3088)

Co-authored-by: cong <zhangcong1992@gmail.com>

core/stores/redis/hook.go

@@ -7,9 +7,8 @@ import (
 	"strings"
 	"time"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/breaker"
-	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/mapping"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -23,93 +22,65 @@ import (
 const spanName = "redis"
 
 var (
-	startTimeKey          = contextKey("startTime")
 	durationHook          = hook{}
 	redisCmdsAttributeKey = attribute.Key("redis.cmds")
 )
 
-type (
-	contextKey string
-	hook       struct{}
-)
+type hook struct{}
 
-func (h hook) BeforeProcess(ctx context.Context, cmd red.Cmder) (context.Context, error) {
-	return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now()), cmd), nil
+func (h hook) DialHook(next red.DialHook) red.DialHook {
+	// no need to implement
+	return next
 }
 
-func (h hook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
-	err := cmd.Err()
-	h.endSpan(ctx, err)
+func (h hook) ProcessHook(next red.ProcessHook) red.ProcessHook {
+	return func(ctx context.Context, cmd red.Cmder) error {
+		start := timex.Now()
+		ctx, endSpan := h.startSpan(ctx, cmd)
 
-	val := ctx.Value(startTimeKey)
-	if val == nil {
-		return nil
-	}
+		err := next(ctx, cmd)
 
-	start, ok := val.(time.Duration)
-	if !ok {
-		return nil
-	}
+		endSpan(err)
+		duration := timex.Since(start)
 
-	duration := timex.Since(start)
-	if duration > slowThreshold.Load() {
-		logDuration(ctx, []red.Cmder{cmd}, duration)
-		metricSlowCount.Inc(cmd.Name())
-	}
-
-	metricReqDur.ObserveFloat(float64(duration)/float64(time.Millisecond), cmd.Name())
-	if msg := formatError(err); len(msg) > 0 {
-		metricReqErr.Inc(cmd.Name(), msg)
-	}
-
-	return nil
-}
-
-func (h hook) BeforeProcessPipeline(ctx context.Context, cmds []red.Cmder) (context.Context, error) {
-	if len(cmds) == 0 {
-		return ctx, nil
-	}
-
-	return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now()), cmds...), nil
-}
-
-func (h hook) AfterProcessPipeline(ctx context.Context, cmds []red.Cmder) error {
-	if len(cmds) == 0 {
-		return nil
-	}
-
-	batchError := errorx.BatchError{}
-	for _, cmd := range cmds {
-		err := cmd.Err()
-		if err == nil {
-			continue
+		if duration > slowThreshold.Load() {
+			logDuration(ctx, []red.Cmder{cmd}, duration)
+			metricSlowCount.Inc(cmd.Name())
 		}
 
-		batchError.Add(err)
-	}
-	h.endSpan(ctx, batchError.Err())
+		metricReqDur.Observe(duration.Milliseconds(), cmd.Name())
+		if msg := formatError(err); len(msg) > 0 {
+			metricReqErr.Inc(cmd.Name(), msg)
+		}
 
-	val := ctx.Value(startTimeKey)
-	if val == nil {
-		return nil
+		return err
 	}
+}
 
-	start, ok := val.(time.Duration)
-	if !ok {
-		return nil
+func (h hook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
+	return func(ctx context.Context, cmds []red.Cmder) error {
+		if len(cmds) == 0 {
+			return next(ctx, cmds)
+		}
+
+		start := timex.Now()
+		ctx, endSpan := h.startSpan(ctx, cmds...)
+
+		err := next(ctx, cmds)
+
+		endSpan(err)
+		duration := timex.Since(start)
+		if duration > slowThreshold.Load()*time.Duration(len(cmds)) {
+			logDuration(ctx, cmds, duration)
+		}
+
+		metricReqDur.Observe(duration.Milliseconds(), "Pipeline")
+		if msg := formatError(err); len(msg) > 0 {
+			metricReqErr.Inc("Pipeline", msg)
+		}
+
+		return err
 	}
-
-	duration := timex.Since(start)
-	if duration > slowThreshold.Load()*time.Duration(len(cmds)) {
-		logDuration(ctx, cmds, duration)
-	}
-
-	metricReqDur.Observe(duration.Milliseconds(), "Pipeline")
-	if msg := formatError(batchError.Err()); len(msg) > 0 {
-		metricReqErr.Inc("Pipeline", msg)
-	}
-
-	return nil
 }
 
 func formatError(err error) string {
@@ -152,7 +123,7 @@ func logDuration(ctx context.Context, cmds []red.Cmder, duration time.Duration)
 	logx.WithContext(ctx).WithDuration(duration).Slowf("[REDIS] slowcall on executing: %s", buf.String())
 }
 
-func (h hook) startSpan(ctx context.Context, cmds ...red.Cmder) context.Context {
+func (h hook) startSpan(ctx context.Context, cmds ...red.Cmder) (context.Context, func(err error)) {
 	tracer := trace.TracerFromContext(ctx)
 
 	ctx, span := tracer.Start(ctx,
@@ -166,18 +137,15 @@ func (h hook) startSpan(ctx context.Context, cmds ...red.Cmder) context.Context
 	}
 	span.SetAttributes(redisCmdsAttributeKey.StringSlice(cmdStrs))
 
-	return ctx
-}
+	return ctx, func(err error) {
+		defer span.End()
 
-func (h hook) endSpan(ctx context.Context, err error) {
-	span := oteltrace.SpanFromContext(ctx)
-	defer span.End()
+		if err == nil || err == red.Nil {
+			span.SetStatus(codes.Ok, "")
+			return
+		}
 
-	if err == nil || err == red.Nil {
-		span.SetStatus(codes.Ok, "")
-		return
+		span.SetStatus(codes.Error, err.Error())
+		span.RecordError(err)
 	}
-
-	span.SetStatus(codes.Error, err.Error())
-	span.RecordError(err)
 }

core/stores/redis/hook_test.go

@@ -4,167 +4,103 @@ import (
 	"context"
 	"errors"
 	"io"
-	"log"
 	"net"
 	"strings"
 	"testing"
 	"time"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/logx/logtest"
-	ztrace "github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
 	tracesdk "go.opentelemetry.io/otel/trace"
 )
 
 func TestHookProcessCase1(t *testing.T) {
-	ztrace.StartAgent(ztrace.Config{
-		Name:     "go-zero-test",
-		Endpoint: "http://localhost:14268/api/traces",
-		Batcher:  "jaeger",
-		Sampler:  1.0,
-	})
-	defer ztrace.StopAgent()
+	tracetest.NewInMemoryExporter(t)
+	w := logtest.NewCollector(t)
 
-	writer := log.Writer()
-	var buf strings.Builder
-	log.SetOutput(&buf)
-	defer log.SetOutput(writer)
-
-	ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
+	err := durationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
+		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+		return nil
+	})(context.Background(), red.NewCmd(context.Background()))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background())))
-	assert.False(t, strings.Contains(buf.String(), "slow"))
-	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+	assert.False(t, strings.Contains(w.String(), "slow"))
 }
 
 func TestHookProcessCase2(t *testing.T) {
-	ztrace.StartAgent(ztrace.Config{
-		Name:     "go-zero-test",
-		Endpoint: "http://localhost:14268/api/traces",
-		Batcher:  "jaeger",
-		Sampler:  1.0,
-	})
-	defer ztrace.StopAgent()
-
+	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
-	ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
+
+	err := durationHook.ProcessHook(func(ctx context.Context, cmd red.Cmder) error {
+		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+		time.Sleep(slowThreshold.Load() + time.Millisecond)
+		return nil
+	})(context.Background(), red.NewCmd(context.Background()))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	time.Sleep(slowThreshold.Load() + time.Millisecond)
-
-	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background(), "foo", "bar")))
 	assert.True(t, strings.Contains(w.String(), "slow"))
-}
-
-func TestHookProcessCase3(t *testing.T) {
-	writer := log.Writer()
-	var buf strings.Builder
-	log.SetOutput(&buf)
-	defer log.SetOutput(writer)
-
-	assert.Nil(t, durationHook.AfterProcess(context.Background(), red.NewCmd(context.Background())))
-	assert.True(t, buf.Len() == 0)
-}
-
-func TestHookProcessCase4(t *testing.T) {
-	writer := log.Writer()
-	var buf strings.Builder
-	log.SetOutput(&buf)
-	defer log.SetOutput(writer)
-
-	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
-	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background())))
-	assert.True(t, buf.Len() == 0)
+	assert.True(t, strings.Contains(w.String(), "trace"))
+	assert.True(t, strings.Contains(w.String(), "span"))
 }
 
 func TestHookProcessPipelineCase1(t *testing.T) {
-	ztrace.StartAgent(ztrace.Config{
-		Name:     "go-zero-test",
-		Endpoint: "http://localhost:14268/api/traces",
-		Batcher:  "jaeger",
-		Sampler:  1.0,
-	})
-	defer ztrace.StopAgent()
+	tracetest.NewInMemoryExporter(t)
+	w := logtest.NewCollector(t)
 
-	writer := log.Writer()
-	var buf strings.Builder
-	log.SetOutput(&buf)
-	defer log.SetOutput(writer)
-
-	_, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{})
+	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+		return nil
+	})(context.Background(), nil)
 	assert.NoError(t, err)
-	ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{
+
+	err = durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+		return nil
+	})(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background()),
 	})
 	assert.NoError(t, err)
 
-	assert.NoError(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{}))
-	assert.NoError(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
-		red.NewCmd(context.Background()),
-	}))
-	assert.False(t, strings.Contains(buf.String(), "slow"))
+	assert.False(t, strings.Contains(w.String(), "slow"))
 }
 
 func TestHookProcessPipelineCase2(t *testing.T) {
-	ztrace.StartAgent(ztrace.Config{
-		Name:     "go-zero-test",
-		Endpoint: "http://localhost:14268/api/traces",
-		Batcher:  "jaeger",
-		Sampler:  1.0,
-	})
-	defer ztrace.StopAgent()
-
+	tracetest.NewInMemoryExporter(t)
 	w := logtest.NewCollector(t)
-	ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{
+
+	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+		time.Sleep(slowThreshold.Load() + time.Millisecond)
+		return nil
+	})(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background()),
 	})
 	assert.NoError(t, err)
 
-	time.Sleep(slowThreshold.Load() + time.Millisecond)
-
-	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
-		red.NewCmd(context.Background(), "foo", "bar"),
-	}))
 	assert.True(t, strings.Contains(w.String(), "slow"))
+	assert.True(t, strings.Contains(w.String(), "trace"))
+	assert.True(t, strings.Contains(w.String(), "span"))
 }
 
 func TestHookProcessPipelineCase3(t *testing.T) {
-	w := logtest.NewCollector(t)
+	te := tracetest.NewInMemoryExporter(t)
 
-	assert.Nil(t, durationHook.AfterProcessPipeline(context.Background(), []red.Cmder{
+	err := durationHook.ProcessPipelineHook(func(ctx context.Context, cmds []red.Cmder) error {
+		assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
+		return assert.AnError
+	})(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background()),
-	}))
-	assert.True(t, len(w.String()) == 0)
-}
-
-func TestHookProcessPipelineCase4(t *testing.T) {
-	w := logtest.NewCollector(t)
-
-	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
-	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
-		red.NewCmd(context.Background()),
-	}))
-	assert.True(t, len(w.String()) == 0)
-}
-
-func TestHookProcessPipelineCase5(t *testing.T) {
-	writer := log.Writer()
-	var buf strings.Builder
-	log.SetOutput(&buf)
-	defer log.SetOutput(writer)
-
-	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
-	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
-		red.NewCmd(context.Background()),
-	}))
-	assert.True(t, buf.Len() == 0)
+	})
+	assert.ErrorIs(t, err, assert.AnError)
+	traceLogs := te.GetSpans().Snapshots()[0]
+	assert.Equal(t, "redis", traceLogs.Name())
+	assert.Equal(t, assert.AnError.Error(), traceLogs.Events()[0].Attributes[1].Value.AsString(), "trace should record error")
 }
 
 func TestLogDuration(t *testing.T) {

core/stores/redis/metrics.go

@@ -3,8 +3,8 @@ package redis
 import (
 	"sync"
 
-	red "github.com/go-redis/redis/v8"
 	"github.com/prometheus/client_golang/prometheus"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/metric"
 )
 

core/stores/redis/metrics_test.go

@@ -7,9 +7,9 @@ import (
 	"testing"
 	"time"
 
-	red "github.com/go-redis/redis/v8"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/testutil"
+	red "github.com/redis/go-redis/v9"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/conf"
 	"github.com/zeromicro/go-zero/internal/devserver"

core/stores/redis/redis.go

@@ -7,7 +7,7 @@ import (
 	"strconv"
 	"time"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/logx"
@@ -64,6 +64,7 @@ type (
 	// RedisNode interface represents a redis node.
 	RedisNode interface {
 		red.Cmdable
+		red.BitMapCmdable
 	}
 
 	// GeoLocation is used with GeoAdd to add geospatial location.
@@ -2058,7 +2059,7 @@ func (s *Redis) ZaddFloatCtx(ctx context.Context, key string, score float64, val
 			return err
 		}
 
-		v, err := conn.ZAdd(ctx, key, &red.Z{
+		v, err := conn.ZAdd(ctx, key, red.Z{
 			Score:  score,
 			Member: value,
 		}).Result()
@@ -2086,9 +2087,9 @@ func (s *Redis) ZaddsCtx(ctx context.Context, key string, ps ...Pair) (val int64
 			return err
 		}
 
-		var zs []*red.Z
+		var zs []red.Z
 		for _, p := range ps {
-			z := &red.Z{Score: float64(p.Score), Member: p.Key}
+			z := red.Z{Score: float64(p.Score), Member: p.Key}
 			zs = append(zs, z)
 		}
 

core/stores/redis/redis_test.go

@@ -10,29 +10,34 @@ import (
 	"time"
 
 	"github.com/alicebob/miniredis/v2"
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/logx"
 	"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) DialHook(next red.DialHook) red.DialHook {
+	return next
 }
 
-func (m myHook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
-	// skip ping cmd
-	if cmd.Name() == "ping" && !m.includePing {
-		return nil
+func (m myHook) ProcessPipelineHook(next red.ProcessPipelineHook) red.ProcessPipelineHook {
+	return next
+}
+
+func (m myHook) ProcessHook(next red.ProcessHook) red.ProcessHook {
+	return func(ctx context.Context, cmd red.Cmder) error {
+		// skip ping cmd
+		if cmd.Name() == "ping" && !m.includePing {
+			return next(ctx, cmd)
+		}
+		return errors.New("hook error")
 	}
-	return errors.New("hook error")
 }
 
 func TestNewRedis(t *testing.T) {
@@ -1609,7 +1614,7 @@ func TestRedis_Pipelined(t *testing.T) {
 			func(pipe Pipeliner) error {
 				pipe.Incr(context.Background(), "pipelined_counter")
 				pipe.Expire(context.Background(), "pipelined_counter", time.Hour)
-				pipe.ZAdd(context.Background(), "zadd", &Z{Score: 12, Member: "zadd"})
+				pipe.ZAdd(context.Background(), "zadd", Z{Score: 12, Member: "zadd"})
 				return nil
 			},
 		)

core/stores/redis/redisblockingnode.go

@@ -3,7 +3,7 @@ package redis
 import (
 	"fmt"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/logx"
 )
 

core/stores/redis/redisclientmanager.go

@@ -5,7 +5,7 @@ import (
 	"io"
 	"runtime"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/syncx"
 )
 

core/stores/redis/redisclustermanager.go

@@ -6,7 +6,7 @@ import (
 	"runtime"
 	"strings"
 
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/syncx"
 )
 

core/stores/redis/redisclustermanager_test.go

@@ -4,7 +4,7 @@ import (
 	"testing"
 
 	"github.com/alicebob/miniredis/v2"
-	red "github.com/go-redis/redis/v8"
+	red "github.com/redis/go-redis/v9"
 	"github.com/stretchr/testify/assert"
 )
 

core/stores/redis/redislock.go

@@ -7,8 +7,7 @@ import (
 	"sync/atomic"
 	"time"
 
-	red "github.com/go-redis/redis/v8"
-
+	red "github.com/redis/go-redis/v9"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stringx"
 )