feat: fix grpc trace (#978)

mosn · Aug 2, 2023 · 77d045b · 77d045b
1 parent c484b6c
commit 77d045b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 17 deletions.
diff --git a/components/trace/generator.go b/components/trace/generator.go
@@ -29,7 +29,6 @@ var (
 
 // Generator  is used to get or generate traceId/spanId/context
 type Generator interface {
-	Init(ctx context.Context)
 	GetTraceId(ctx context.Context) string
 	GetSpanId(ctx context.Context) string
 	GenerateNewContext(ctx context.Context, span api.Span) context.Context

diff --git a/components/trace/generator_test.go b/components/trace/generator_test.go
@@ -24,7 +24,6 @@ import (
 type MockGenerator struct {
 }
 
-func (m *MockGenerator) Init(ctx context.Context) {}
 func (m *MockGenerator) GetTraceId(ctx context.Context) string {
 	return "mock"
 }

diff --git a/diagnostics/genetator.go b/diagnostics/genetator.go
@@ -21,18 +21,12 @@ func init() {
 
 // OpenGenerator is the default implementation of Generator
 type OpenGenerator struct {
-	md metadata.MD
 }
 
-func (o *OpenGenerator) Init(ctx context.Context) {
-	o.md = map[string][]string{}
-	if md, ok := metadata.FromIncomingContext(ctx); ok {
-		o.md = md
-	}
-}
 func (o *OpenGenerator) GetTraceId(ctx context.Context) string {
 	var traceId string
-	if v, ok := o.md[strings.ToLower(sofa.TRACER_ID_KEY)]; ok {
+	md, _ := metadata.FromIncomingContext(ctx)
+	if v, ok := md[strings.ToLower(sofa.TRACER_ID_KEY)]; ok {
 		traceId = v[0]
 	} else {
 		traceId = mtrace.IdGen().GenerateTraceId()
@@ -42,7 +36,8 @@ func (o *OpenGenerator) GetTraceId(ctx context.Context) string {
 
 func (o *OpenGenerator) GetSpanId(ctx context.Context) string {
 	var spanId string
-	if v, ok := o.md[strings.ToLower(sofa.RPC_ID_KEY)]; ok {
+	md, _ := metadata.FromIncomingContext(ctx)
+	if v, ok := md[strings.ToLower(sofa.RPC_ID_KEY)]; ok {
 		spanId = v[0]
 	} else {
 		spanId = "0"
@@ -55,7 +50,8 @@ func (o *OpenGenerator) GetSpanId(ctx context.Context) string {
 func (o *OpenGenerator) GetParentSpanId(ctx context.Context) string {
 	// TODO: need some design to get the parent id
 	var spanId string
-	if v, ok := o.md[strings.ToLower(sofa.RPC_ID_KEY)]; ok {
+	md, _ := metadata.FromIncomingContext(ctx)
+	if v, ok := md[strings.ToLower(sofa.RPC_ID_KEY)]; ok {
 		spanId = v[0]
 	} else {
 		spanId = "0"
@@ -64,13 +60,14 @@ func (o *OpenGenerator) GetParentSpanId(ctx context.Context) string {
 }
 
 func (o *OpenGenerator) GenerateNewContext(ctx context.Context, span api.Span) context.Context {
-	newMd := o.md.Copy()
+	md, _ := metadata.FromIncomingContext(ctx)
+	newMd := md.Copy()
 	newMd[strings.ToLower(sofa.TRACER_ID_KEY)] = []string{span.TraceId()}
 	newMd[strings.ToLower(sofa.RPC_ID_KEY)] = []string{span.SpanId()}
-	if v, ok := o.md[strings.ToLower(sofa.APP_NAME_KEY)]; ok && len(v) > 0 {
+	if v, ok := md[strings.ToLower(sofa.APP_NAME_KEY)]; ok && len(v) > 0 {
 		span.SetTag(trace.LAYOTTO_APP_NAME, v[0])
 	}
-	if v, ok := o.md[strings.ToLower(sofa.SOFA_TRACE_BAGGAGE_DATA)]; ok && len(v) > 0 {
+	if v, ok := md[strings.ToLower(sofa.SOFA_TRACE_BAGGAGE_DATA)]; ok && len(v) > 0 {
 		span.SetTag(trace.LAYOTTO_ATTRS_CONTENT, v[0])
 	}
 	ctx = metadata.NewIncomingContext(ctx, newMd)

diff --git a/diagnostics/tracing.go b/diagnostics/tracing.go
@@ -65,7 +65,6 @@ func NewSpan(ctx context.Context, startTime time.Time, config map[string]interfa
 		log.DefaultLogger.Errorf("not support trace type: %+v", generatorName)
 		return nil
 	}
-	ge.Init(ctx)
 	// use generator to extract the span/trace/parentSpan IDs
 	spanId := ge.GetSpanId(ctx)
 	traceId := ge.GetTraceId(ctx)
@@ -85,7 +84,6 @@ func GetNewContext(ctx context.Context, span api.Span) context.Context {
 	if ge == nil {
 		return ctx
 	}
-	ge.Init(ctx)
 	newCtx := ge.GenerateNewContext(ctx, span)
 	return newCtx
 }