From 4860f42637409dbc95ad84e53c512d1dfcdac12d Mon Sep 17 00:00:00 2001
From: longxboy <longxboyhi@gmail.com>
Date: Thu, 3 Jun 2021 16:56:55 +0800
Subject: [PATCH] add ctx for logger (#998)

* add ctx for trace
---
 log/helper.go      |  9 +++++++++
 log/helper_test.go | 18 ++++++++++++++++++
 log/log.go         | 31 ++++++++++++++++++++++++-------
 log/value.go       | 15 ++++++++-------
 4 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/log/helper.go b/log/helper.go
index 2c5e67f9f..2a2aceb43 100644
--- a/log/helper.go
+++ b/log/helper.go
@@ -1,6 +1,7 @@
 package log
 
 import (
+	"context"
 	"fmt"
 )
 
@@ -16,6 +17,14 @@ func NewHelper(logger Logger) *Helper {
 	}
 }
 
+// WithContext returns a shallow copy of h with its context changed
+// to ctx. The provided ctx must be non-nil.
+func (h *Helper) WithContext(ctx context.Context) *Helper {
+	return &Helper{
+		logger: WithContext(ctx, h.logger),
+	}
+}
+
 // Log .
 func (h *Helper) Log(level Level, keyvals ...interface{}) {
 	h.logger.Log(level, keyvals...)
diff --git a/log/helper_test.go b/log/helper_test.go
index beb30c1db..42282a3ca 100644
--- a/log/helper_test.go
+++ b/log/helper_test.go
@@ -1,7 +1,9 @@
 package log
 
 import (
+	"context"
 	"io/ioutil"
+	"os"
 	"testing"
 )
 
@@ -43,3 +45,19 @@ func BenchmarkHelperPrintw(b *testing.B) {
 		log.Debugw("key", "value")
 	}
 }
+
+func TestContext(t *testing.T) {
+	logger := With(NewStdLogger(os.Stdout),
+		"trace", Trace(),
+	)
+	log := NewHelper(logger)
+	ctx := context.WithValue(context.Background(), "trace_id", "2233")
+	log.WithContext(ctx).Info("got trace!")
+}
+
+func Trace() Valuer {
+	return func(ctx context.Context) interface{} {
+		s := ctx.Value("trace_id").(string)
+		return s
+	}
+}
diff --git a/log/log.go b/log/log.go
index 6a27974bf..1a5032ac2 100644
--- a/log/log.go
+++ b/log/log.go
@@ -1,6 +1,7 @@
 package log
 
 import (
+	"context"
 	"log"
 )
 
@@ -14,17 +15,18 @@ type Logger interface {
 	Log(level Level, keyvals ...interface{}) error
 }
 
-type context struct {
+type logger struct {
 	logs      []Logger
 	prefix    []interface{}
 	hasValuer bool
+	ctx       context.Context
 }
 
-func (c *context) Log(level Level, keyvals ...interface{}) error {
+func (c *logger) Log(level Level, keyvals ...interface{}) error {
 	kvs := make([]interface{}, 0, len(c.prefix)+len(keyvals))
 	kvs = append(kvs, c.prefix...)
 	if c.hasValuer {
-		bindValues(kvs)
+		bindValues(c.ctx, kvs)
 	}
 	kvs = append(kvs, keyvals...)
 	for _, l := range c.logs {
@@ -37,20 +39,35 @@ func (c *context) Log(level Level, keyvals ...interface{}) error {
 
 // With with logger fields.
 func With(l Logger, kv ...interface{}) Logger {
-	if c, ok := l.(*context); ok {
+	if c, ok := l.(*logger); ok {
 		kvs := make([]interface{}, 0, len(c.prefix)+len(kv))
 		kvs = append(kvs, kv...)
 		kvs = append(kvs, c.prefix...)
-		return &context{
+		return &logger{
 			logs:      c.logs,
 			prefix:    kvs,
 			hasValuer: containsValuer(kvs),
+			ctx:       c.ctx,
 		}
 	}
-	return &context{logs: []Logger{l}, prefix: kv, hasValuer: containsValuer(kv)}
+	return &logger{logs: []Logger{l}, prefix: kv, hasValuer: containsValuer(kv)}
+}
+
+// WithContext returns a shallow copy of l with its context changed
+// to ctx. The provided ctx must be non-nil.
+func WithContext(ctx context.Context, l Logger) Logger {
+	if c, ok := l.(*logger); ok {
+		return &logger{
+			logs:      c.logs,
+			prefix:    c.prefix,
+			hasValuer: c.hasValuer,
+			ctx:       ctx,
+		}
+	}
+	return &logger{logs: []Logger{l}, ctx: ctx}
 }
 
 // MultiLogger wraps multi logger.
 func MultiLogger(logs ...Logger) Logger {
-	return &context{logs: logs}
+	return &logger{logs: logs}
 }
diff --git a/log/value.go b/log/value.go
index 311d68430..186829e77 100644
--- a/log/value.go
+++ b/log/value.go
@@ -1,6 +1,7 @@
 package log
 
 import (
+	"context"
 	"runtime"
 	"strconv"
 	"strings"
@@ -16,19 +17,19 @@ var (
 )
 
 // Valuer is returns a log value.
-type Valuer func() interface{}
+type Valuer func(ctx context.Context) interface{}
 
 // Value return the function value.
-func Value(v interface{}) interface{} {
+func Value(ctx context.Context, v interface{}) interface{} {
 	if v, ok := v.(Valuer); ok {
-		return v()
+		return v(ctx)
 	}
 	return v
 }
 
 // Caller returns returns a Valuer that returns a pkg/file:line description of the caller.
 func Caller(depth int) Valuer {
-	return func() interface{} {
+	return func(context.Context) interface{} {
 		_, file, line, _ := runtime.Caller(depth)
 		if strings.LastIndex(file, "github.com/go-kratos/kratos/log") > 0 {
 			_, file, line, _ = runtime.Caller(depth + 1)
@@ -40,15 +41,15 @@ func Caller(depth int) Valuer {
 
 // Timestamp returns a timestamp Valuer with a custom time format.
 func Timestamp(layout string) Valuer {
-	return func() interface{} {
+	return func(context.Context) interface{} {
 		return time.Now().Format(layout)
 	}
 }
 
-func bindValues(keyvals []interface{}) {
+func bindValues(ctx context.Context, keyvals []interface{}) {
 	for i := 1; i < len(keyvals); i += 2 {
 		if v, ok := keyvals[i].(Valuer); ok {
-			keyvals[i] = v()
+			keyvals[i] = v(ctx)
 		}
 	}
 }