From 2f13607f020d5f4d00f4886638feaaa892022ccb Mon Sep 17 00:00:00 2001 From: askuy Date: Fri, 3 May 2024 16:07:09 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E4=BC=98=E5=8C=96trace=20context=20value?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- client/egrpc/interceptor.go | 6 +++-- client/ehttp/interceptor.go | 7 +++--- core/transport/transport.go | 35 +++++++---------------------- core/transport/transport_test.go | 17 ++++++++++---- examples/grpc/header/client/main.go | 4 ++-- examples/http/client/main.go | 3 +-- examples/http/headeruid/main.go | 4 ++-- internal/tools/tools.go | 11 ++++----- server/egin/interceptor.go | 8 ++++++- server/egin/interceptor_test.go | 9 ++++---- server/egrpc/interceptor.go | 9 +++++++- 11 files changed, 59 insertions(+), 54 deletions(-) diff --git a/client/egrpc/interceptor.go b/client/egrpc/interceptor.go index cadfbc62..3ef422a6 100644 --- 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 index 641208e4..0a3388cc 100644 --- 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 index e97604ab..c30b0be7 100644 --- a/core/transport/transport.go +++ b/core/transport/transport.go @@ -8,20 +8,15 @@ 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) } @@ -29,10 +24,6 @@ func init() { 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 index f4597ce0..42fd0018 100644 --- 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 index 00b75a92..101060d4 100644 --- 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 index 3ffc3861..f2f122c9 100644 --- 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 index 95f5759b..b1fdc38d 100644 --- 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 index 90eb7cc7..6faf740b 100644 --- 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,8 +20,12 @@ func GrpcHeaderValue(ctx context.Context, key string) string { if !ok { return "" } + + if len(md) > 0 { + return md.Get(key)[0] + } // 小写 - return strings.Join(md.Get(key), ";") + return "" } // ContextValue gRPC日志获取context value @@ -32,7 +33,7 @@ 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 index 4ceafd33..6dd6bdf7 100644 --- 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 index d74fd61b..94106e68 100644 --- 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 index c5dde6c2..8035e3da 100644 --- 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) + } + } + } } From c4993ae7054c7b4bc5dbb7680e1b3b3d60e42fa6 Mon Sep 17 00:00:00 2001 From: askuy Date: Fri, 3 May 2024 16:23:53 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E4=BC=98=E5=8C=96trace=20context=20value?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/tools/tools.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/internal/tools/tools.go b/internal/tools/tools.go index 6faf740b..ae5cd224 100644 --- a/internal/tools/tools.go +++ b/internal/tools/tools.go @@ -20,11 +20,9 @@ func GrpcHeaderValue(ctx context.Context, key string) string { if !ok { return "" } - - if len(md) > 0 { + if len(md.Get(key)) > 0 { return md.Get(key)[0] } - // 小写 return "" }