pkg/log(dm):support reading fields from Ctx (#5518)

close #5517

dm/pkg/log/ctx.go

@@ -0,0 +1,47 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package log
+
+import (
+	"context"
+
+	"go.uber.org/zap"
+)
+
+type logCtxKeyType string
+
+var logCtxKey = logCtxKeyType("zap-fields")
+
+func getZapFieldsFromCtx(ctx context.Context) []zap.Field {
+	// Use a pointer to key to save one allocation.
+	fields := ctx.Value(&logCtxKey)
+	if fields == nil {
+		return nil
+	}
+	// Note that the value is a pointer to slice, to save allocation.
+	return *fields.(*[]zap.Field)
+}
+
+// AppendZapFieldToCtx attaches new fields to the context, which can be then
+// used with log.WithCtx().
+func AppendZapFieldToCtx(ctx context.Context, newFields ...zap.Field) context.Context {
+	fields := getZapFieldsFromCtx(ctx)
+	if fields == nil {
+		// 8 fields should be enough for most situations.
+		fields = make([]zap.Field, 0, 8)
+	}
+	fields = append(fields, newFields...)
+
+	return context.WithValue(ctx, &logCtxKey, &fields)
+}

dm/pkg/log/ctx_test.go

@@ -0,0 +1,57 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package log
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+func TestLogCtxBasics(t *testing.T) {
+	ctx := context.Background()
+	ctx1 := AppendZapFieldToCtx(ctx, zap.String("key1", "value1"))
+	ctx2 := AppendZapFieldToCtx(ctx1, zap.String("key2", "value2"))
+
+	require.Nil(t, getZapFieldsFromCtx(ctx))
+	require.Equal(t, []zap.Field{zap.String("key1", "value1")}, getZapFieldsFromCtx(ctx1))
+	require.Equal(t, []zap.Field{
+		zap.String("key1", "value1"),
+		zap.String("key2", "value2"),
+	}, getZapFieldsFromCtx(ctx2))
+}
+
+func TestLogCtxTooManyFields(t *testing.T) {
+	ctx := context.Background()
+	var (
+		ctxs   []context.Context
+		fields [][]zap.Field
+	)
+
+	ctxs = append(ctxs, ctx)
+	fields = append(fields, nil)
+	for i := 1; i < 100; i++ {
+		k := fmt.Sprintf("key%d", i)
+		v := fmt.Sprintf("value%d", i)
+		ctxs = append(ctxs, AppendZapFieldToCtx(ctxs[i-1], zap.String(k, v)))
+		fields = append(fields, append(fields[i-1], zap.String(k, v)))
+	}
+
+	for i := 0; i < 100; i++ {
+		require.Equal(t, fields[i], getZapFieldsFromCtx(ctxs[i]), "failed, index = %d", i)
+	}
+}

dm/pkg/log/log.go

@@ -183,3 +183,8 @@ func WrapStringerField(message string, object fmt.Stringer) zap.Field {
 
 	return zap.Stringer(message, object)
 }
+
+// WithCtx adds fields from ctx to the logger.
+func WithCtx(ctx context.Context) Logger {
+	return Logger{appLogger.With(getZapFieldsFromCtx(ctx)...)}
+}

dm/pkg/log/log_test.go

@@ -150,3 +150,37 @@ func (s *testLogSuite) TestInitSlowQueryLoggerNotInDebugLevel(c *C) {
 	c.Assert(output[0], Matches, ".*this is from applogger.*")
 	c.Assert(output[1], Equals, "") // no output
 }
+
+func (s *testLogSuite) TestWithCtx(c *C) {
+	// test slow query logger can write debug log
+	logLevel := "debug"
+	cfg := &Config{Level: logLevel, Format: "json"}
+	cfg.Adjust()
+
+	ctx := context.Background()
+	ctx = AppendZapFieldToCtx(ctx, zap.String("key1", "value1"))
+	ctx = AppendZapFieldToCtx(ctx, zap.String("key2", "value2"))
+
+	output, err := captureStdout(func() {
+		c.Assert(InitLogger(cfg), IsNil)
+		WithCtx(ctx).Info("test1")
+	})
+	c.Assert(err, IsNil)
+	c.Assert(output[0], Matches, ".*test1.*key1.*value1.*key2.*value2.*")
+}
+
+func BenchmarkBaseline(b *testing.B) {
+	logger := L().With(zap.String("key1", "value1"))
+	for i := 0; i < b.N; i++ {
+		subLogger := logger.With(zap.String("key2", "value2"))
+		subLogger.Info("test-test-test")
+	}
+}
+
+func BenchmarkWithCtx(b *testing.B) {
+	ctx := AppendZapFieldToCtx(context.Background(), zap.String("key1", "value1"))
+	for i := 0; i < b.N; i++ {
+		subCtx := AppendZapFieldToCtx(ctx, zap.String("key2", "value2"))
+		WithCtx(subCtx).Info("test-test-test")
+	}
+}