zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: refactor trace in redis & sql & mongo (#1865)

Browse Source 
* refactor: refactor tracing in redis & sql & mongo

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* fix: fix some tests

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: add missing content

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: adjust `log` and `return`

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: reformat code

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: reformat code

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: reformat code

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: simpler span name

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: fix a bug

Signed-off-by: chenquan <chenquan.dev@gmail.com>

* refactor: fix a bug

Signed-off-by: chenquan <chenquan.dev@gmail.com>
This commit is contained in:
chen quan
2022-05-12 23:32:34 -05:00
committed by GitHub
parent 52f060caae
commit e567a0c718
11 changed files with 434 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
core/stores/redis/hook.go
View File

@@ -6,35 +6,40 @@ import (
"time"
red "github.com/go-redis/redis/v8"
"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"
"github.com/zeromicro/go-zero/core/trace"
"go.opentelemetry.io/otel"
tracestd "go.opentelemetry.io/otel/trace"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
oteltrace "go.opentelemetry.io/otel/trace"
)
// spanName is the span name of the redis calls.
const spanName = "redis"
var (
startTimeKey = contextKey("startTime")
durationHook = hook{tracer: otel.GetTracerProvider().Tracer(trace.TraceName)}
startTimeKey = contextKey("startTime")
durationHook = hook{tracer: otel.GetTracerProvider().Tracer(trace.TraceName)}
redisCmdsAttributeKey = attribute.Key("redis.cmds")
)
type (
contextKey string
hook struct {
tracer tracestd.Tracer
tracer oteltrace.Tracer
}
)
func (h hook) BeforeProcess(ctx context.Context, _ red.Cmder) (context.Context, error) {
return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now())), nil
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) AfterProcess(ctx context.Context, cmd red.Cmder) error {
h.endSpan(ctx)
err := cmd.Err()
h.endSpan(ctx, err)
val := ctx.Value(startTimeKey)
if val == nil {
@@ -54,17 +59,30 @@ func (h hook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
return nil
}
func (h hook) BeforeProcessPipeline(ctx context.Context, _ []red.Cmder) (context.Context, error) {
return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now())), 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 {
h.endSpan(ctx)
if len(cmds) == 0 {
return nil
}
batchError := errorx.BatchError{}
for _, cmd := range cmds {
err := cmd.Err()
if err == nil {
continue
}
batchError.Add(err)
}
h.endSpan(ctx, batchError.Err())
val := ctx.Value(startTimeKey)
if val == nil {
return nil
@@ -94,11 +112,30 @@ func logDuration(ctx context.Context, cmd 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) context.Context {
ctx, _ = h.tracer.Start(ctx, spanName)
func (h hook) startSpan(ctx context.Context, cmds ...red.Cmder) context.Context {
ctx, span := h.tracer.Start(ctx,
spanName,
oteltrace.WithSpanKind(oteltrace.SpanKindClient),
)
cmdStrs := make([]string, 0, len(cmds))
for _, cmd := range cmds {
cmdStrs = append(cmdStrs, cmd.Name())
}
span.SetAttributes(redisCmdsAttributeKey.StringSlice(cmdStrs))
return ctx
}
func (h hook) endSpan(ctx context.Context) {
tracestd.SpanFromContext(ctx).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
}
span.SetStatus(codes.Error, err.Error())
span.RecordError(err)
}

10
core/stores/redis/hook_test.go
View File

@@ -27,7 +27,7 @@ func TestHookProcessCase1(t *testing.T) {
log.SetOutput(&buf)
defer log.SetOutput(writer)
ctx, err := durationHook.BeforeProcess(context.Background(), nil)
ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
if err != nil {
t.Fatal(err)
}
@@ -48,7 +48,7 @@ func TestHookProcessCase2(t *testing.T) {
w, restore := injectLog()
defer restore()
ctx, err := durationHook.BeforeProcess(context.Background(), nil)
ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
if err != nil {
t.Fatal(err)
}
@@ -89,7 +89,7 @@ func TestHookProcessPipelineCase1(t *testing.T) {
log.SetOutput(&buf)
defer log.SetOutput(writer)
ctx, err := durationHook.BeforeProcessPipeline(context.Background(), nil)
ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{red.NewCmd(context.Background())})
if err != nil {
t.Fatal(err)
}
@@ -112,7 +112,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
w, restore := injectLog()
defer restore()
ctx, err := durationHook.BeforeProcessPipeline(context.Background(), nil)
ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{red.NewCmd(context.Background())})
if err != nil {
t.Fatal(err)
}
@@ -156,7 +156,7 @@ func TestHookProcessPipelineCase5(t *testing.T) {
defer log.SetOutput(writer)
ctx := context.WithValue(context.Background(), startTimeKey, "foo")
assert.Nil(t, durationHook.AfterProcessPipeline(ctx, nil))
assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{red.NewCmd(context.Background())}))
assert.True(t, buf.Len() == 0)
}