Merge pull request #377 from askuy/feature/tracebycontextvalue

优化trace context value
gotomicro · May 3, 2024 · b458510 · b458510
2 parents 59c6238 + c4993ae
commit b458510
Show file tree

Hide file tree

Showing 11 changed files with 58 additions and 55 deletions.
diff --git a/client/egrpc/interceptor.go b/client/egrpc/interceptor.go
@@ -319,7 +319,6 @@ func (c *Container) loggerUnaryClientInterceptor() grpc.UnaryClientInterceptor {
 			if value := tools.ContextValue(ctx, key); value != "" {
 				// 替换context
 				ctx = metadata.AppendToOutgoingContext(ctx, key, value)
-
 				// grpc metadata 存在同一个 key set 多次，客户端可通过日志排查这种错误使用。
 				if md, ok := metadata.FromOutgoingContext(ctx); ok {
 					fields = append(fields, elog.FieldCustomKeyValue(key, strings.Join(md[strings.ToLower(key)], ";")))
@@ -390,7 +389,10 @@ func customHeader(egoLogExtraKeys []string) grpc.UnaryClientInterceptor {
 	return func(ctx context.Context, method string, req, res interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 		for _, key := range egoLogExtraKeys {
 			if value := tools.GrpcHeaderValue(ctx, key); value != "" {
-				ctx = transport.WithValue(ctx, key, value)
+				if ctx.Value(key) != nil {
+					ctx = context.WithValue(ctx, key, value)
+				}
+				//ctx = transport.WithValue(ctx, key, value)
 			}
 		}
 		return invoker(ctx, method, req, res, cc, opts...)

diff --git a/client/ehttp/interceptor.go b/client/ehttp/interceptor.go
@@ -2,8 +2,6 @@ package ehttp
 
 import (
 	"context"
-	"github.com/gotomicro/ego/core/transport"
-	"github.com/spf13/cast"
 	"log"
 	"net/http"
 	"net/url"
@@ -12,6 +10,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/gotomicro/ego/core/transport"
+	"github.com/spf13/cast"
+
 	"github.com/go-resty/resty/v2"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
@@ -58,7 +59,7 @@ func logAccess(name string, config *Config, logger *elog.Component, req *resty.R
 
 	// 支持自定义log
 	for _, key := range loggerKeys {
-		if value := transport.Value(req.Context(), key); value != nil {
+		if value := req.Context().Value(key); value != nil {
 			fields = append(fields, elog.FieldCustomKeyValue(key, cast.ToString(value)))
 		}
 	}

diff --git a/core/transport/transport.go b/core/transport/transport.go
@@ -8,31 +8,22 @@ import (
 
 var customKeyStore = contextKeyStore{
 	keyArr: make([]string, 0),
-	keyMap: make(map[string]*contextKey),
 }
 
 type contextKeyStore struct {
 	keyArr []string
-	keyMap map[string]*contextKey
 	length int
 }
 
 func init() {
 	customKeyStore.keyArr = eapp.EgoLogExtraKeys()
-	for _, value := range eapp.EgoLogExtraKeys() {
-		customKeyStore.keyMap[value] = newContextKey(value)
-	}
 	customKeyStore.length = len(customKeyStore.keyArr)
 }
 
 // Set overrides custom keys with provided array.
 func Set(arr []string) {
 	length := len(arr)
 	customKeyStore.keyArr = arr
-	customKeyStore.keyMap = make(map[string]*contextKey, length)
-	for _, value := range arr {
-		customKeyStore.keyMap[value] = newContextKey(value)
-	}
 	customKeyStore.length = length
 }
 
@@ -48,26 +39,16 @@ func CustomContextKeysLength() int {
 
 // WithValue returns a new context with your key and value
 func WithValue(ctx context.Context, key string, value interface{}) context.Context {
-	return context.WithValue(ctx, getContextKey(key), value)
+	info := ctx.Value(key)
+	if info != nil {
+		return ctx
+	}
+	return context.WithValue(ctx, key, value)
 }
 
 // Value returns value of your key
+// Deprecated
+// Use ctx.Value()
 func Value(ctx context.Context, key string) interface{} {
-	return ctx.Value(getContextKey(key))
-}
-
-func newContextKey(name string) *contextKey {
-	return &contextKey{name: name}
+	return ctx.Value(key)
 }
-
-func getContextKey(key string) *contextKey {
-	return customKeyStore.keyMap[key]
-}
-
-// contextKey is a value for use with context.WithValue. It's used as
-// a pointer so it fits in an interface{} without allocation.
-type contextKey struct {
-	name string
-}
-
-func (k *contextKey) String() string { return "ego context value " + k.name }
diff --git a/core/transport/transport_test.go b/core/transport/transport_test.go
@@ -19,11 +19,20 @@ func TestValue(t *testing.T) {
 	Set([]string{"X-EGO-Test"})
 	ctx := context.Background()
 	ctx = WithValue(ctx, "X-EGO-Test", "hello")
-	val := Value(ctx, "X-EGO-Test")
+	val := ctx.Value("X-EGO-Test")
 	assert.Equal(t, "hello", val)
 }
 
-func Test_newContextKey(t *testing.T) {
-	key := newContextKey("hello")
-	assert.Equal(t, "ego context value hello", key.String())
+//func Test_newContextKey(t *testing.T) {
+//	key := newContextKey("hello")
+//	assert.Equal(t, "ego context value hello", key.String())
+//}
+
+func TestWithValue(t *testing.T) {
+	Set([]string{"X-EGO-Test"})
+	ctx := context.Background()
+	ctx = WithValue(ctx, "X-EGO-Test", "hello1")
+	ctx = WithValue(ctx, "X-EGO-Test", "hello2")
+	val := ctx.Value("X-EGO-Test")
+	assert.Equal(t, "hello1", val)
 }
diff --git a/examples/grpc/header/client/main.go b/examples/grpc/header/client/main.go
@@ -1,10 +1,10 @@
 package main
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/davecgh/go-spew/spew"
-	"github.com/gotomicro/ego/core/transport"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 
@@ -33,7 +33,7 @@ func invokerGrpc() error {
 
 func callGrpc() error {
 	req := http.Request{}
-	parentContext := transport.WithValue(req.Context(), "X-Ego-Uid", 9527)
+	parentContext := context.WithValue(req.Context(), "X-Ego-Uid", 9527)
 	var headers metadata.MD
 	var trailers metadata.MD
 	_, err := grpcComp.SayHello(parentContext, &helloworld.HelloRequest{

diff --git a/examples/http/client/main.go b/examples/http/client/main.go
@@ -3,7 +3,6 @@ package main
 import (
 	"context"
 	"fmt"
-	"github.com/gotomicro/ego/core/transport"
 
 	"github.com/gotomicro/ego"
 	"github.com/gotomicro/ego/client/ehttp"
@@ -53,7 +52,7 @@ func callHTTPWithCustomTrace() error {
 
 	traceID := "123456"
 
-	ctx = transport.WithValue(ctx, "myTraceID", traceID)
+	ctx = context.WithValue(ctx, "myTraceID", traceID)
 
 	req := httpComp.R()
 

diff --git a/examples/http/headeruid/main.go b/examples/http/headeruid/main.go
@@ -1,13 +1,13 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/gin-gonic/gin"
 	"github.com/gotomicro/ego"
 	"github.com/gotomicro/ego/core/elog"
 	"github.com/gotomicro/ego/core/etrace"
-	"github.com/gotomicro/ego/core/transport"
 	"github.com/gotomicro/ego/server/egin"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -29,7 +29,7 @@ func main() {
 		})
 
 		server.GET("/hello", func(ctx *gin.Context) {
-			pCtx := transport.WithValue(ctx.Request.Context(), "X-Ego-Uid", 9527)
+			pCtx := context.WithValue(ctx.Request.Context(), "X-Ego-Uid", 9527)
 			ctx.Request = ctx.Request.WithContext(pCtx)
 			// Get traceId from Request's context
 			// span, _ := etrace.StartSpanFromContext(ctx.Request.Context(), "Handle: /Hello")

diff --git a/internal/tools/tools.go b/internal/tools/tools.go
@@ -6,12 +6,9 @@ import (
 	"go/format"
 	"log"
 	"reflect"
-	"strings"
 
 	"github.com/spf13/cast"
 	"google.golang.org/grpc/metadata"
-
-	"github.com/gotomicro/ego/core/transport"
 )
 
 // GrpcHeaderValue 获取context value
@@ -23,16 +20,18 @@ func GrpcHeaderValue(ctx context.Context, key string) string {
 	if !ok {
 		return ""
 	}
-	// 小写
-	return strings.Join(md.Get(key), ";")
+	if len(md.Get(key)) > 0 {
+		return md.Get(key)[0]
+	}
+	return ""
 }
 
 // ContextValue gRPC日志获取context value
 func ContextValue(ctx context.Context, key string) string {
 	if key == "" {
 		return ""
 	}
-	return cast.ToString(transport.Value(ctx, key))
+	return cast.ToString(ctx.Value(key))
 }
 
 // ToSliceStringMap casts an empty interface to []map[string]interface{} ignoring error

diff --git a/server/egin/interceptor.go b/server/egin/interceptor.go
@@ -416,7 +416,13 @@ func getHeaderValue(c *gin.Context, key string, enableTrustedCustomHeader bool)
 	value := c.GetHeader(key)
 	if value != "" {
 		// 如果信任该Header，将header数据赋值到context上
-		c.Request = c.Request.WithContext(transport.WithValue(c.Request.Context(), key, value))
+		for _, customContextKey := range transport.CustomContextKeys() {
+			// 如果header key和自定义key是一样的，写入到contet value中
+			if customContextKey == key {
+				c.Request = c.Request.WithContext(transport.WithValue(c.Request.Context(), key, value))
+			}
+		}
+
 	}
 	return value
 }

diff --git a/server/egin/interceptor_test.go b/server/egin/interceptor_test.go
@@ -16,11 +16,10 @@ import (
 	"testing"
 
 	"github.com/gin-gonic/gin"
+	"github.com/gotomicro/ego/core/transport"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/stretchr/testify/assert"
 
-	"github.com/gotomicro/ego/core/transport"
-
 	"github.com/gotomicro/ego/core/elog"
 )
 
@@ -174,7 +173,7 @@ func TestPrometheus(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	text, err := ioutil.ReadAll(res.Body)
+	text, err := io.ReadAll(res.Body)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -211,12 +210,12 @@ func Test_getHeaderValue(t *testing.T) {
 
 func Test_getHeaderAssignValue(t *testing.T) {
 	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	transport.Set([]string{"X-Ego-Uid"})
 	c.Request, _ = http.NewRequest("GET", "/chat", nil)
 	c.Request.Header.Set("X-Ego-Uid", "9527")
 	value := getHeaderValue(c, "X-Ego-Uid", true)
 	assert.Equal(t, "9527", value)
-
-	value2 := transport.Value(c.Request.Context(), "X-Ego-Uid")
+	value2 := c.Request.Context().Value("X-Ego-Uid")
 	assert.Equal(t, "9527", value2)
 }
 

diff --git a/server/egrpc/interceptor.go b/server/egrpc/interceptor.go
@@ -238,7 +238,14 @@ func (c *Container) defaultUnaryServerInterceptor() grpc.UnaryServerInterceptor
 		// 必须在defer外层，因为要赋值，替换ctx
 		for _, key := range loggerKeys {
 			if value := tools.GrpcHeaderValue(ctx, key); value != "" {
-				ctx = transport.WithValue(ctx, key, value)
+				// 如果信任该Header，将header数据赋值到context上
+				for _, customContextKey := range transport.CustomContextKeys() {
+					// 如果header key和自定义key是一样的，写入到contet value中
+					if customContextKey == key {
+						ctx = transport.WithValue(ctx, key, value)
+					}
+				}
+
 			}
 		}