diff --git a/core/trace/constants.go b/core/trace/constants.go index e840de6be..03a9ca677 100644 --- a/core/trace/constants.go +++ b/core/trace/constants.go @@ -1,6 +1,6 @@ package trace const ( - traceIdKey = "X-Trace-ID" + TraceIdKey = "X-Trace-ID" spanIdKey = "X-Span-ID" ) diff --git a/core/trace/propagator_test.go b/core/trace/propagator_test.go index 92c140c52..4dd4e60f7 100644 --- a/core/trace/propagator_test.go +++ b/core/trace/propagator_test.go @@ -11,11 +11,11 @@ import ( func TestHttpPropagator_Extract(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) - req.Header.Set(traceIdKey, "trace") + req.Header.Set(TraceIdKey, "trace") req.Header.Set(spanIdKey, "span") carrier, err := Extract(HttpFormat, req.Header) assert.Nil(t, err) - assert.Equal(t, "trace", carrier.Get(traceIdKey)) + assert.Equal(t, "trace", carrier.Get(TraceIdKey)) assert.Equal(t, "span", carrier.Get(spanIdKey)) _, err = Extract(HttpFormat, req) @@ -24,11 +24,11 @@ func TestHttpPropagator_Extract(t *testing.T) { func TestHttpPropagator_Inject(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) - req.Header.Set(traceIdKey, "trace") + req.Header.Set(TraceIdKey, "trace") req.Header.Set(spanIdKey, "span") carrier, err := Inject(HttpFormat, req.Header) assert.Nil(t, err) - assert.Equal(t, "trace", carrier.Get(traceIdKey)) + assert.Equal(t, "trace", carrier.Get(TraceIdKey)) assert.Equal(t, "span", carrier.Get(spanIdKey)) _, err = Inject(HttpFormat, req) @@ -37,12 +37,12 @@ func TestHttpPropagator_Inject(t *testing.T) { func TestGrpcPropagator_Extract(t *testing.T) { md := metadata.New(map[string]string{ - traceIdKey: "trace", + TraceIdKey: "trace", spanIdKey: "span", }) carrier, err := Extract(GrpcFormat, md) assert.Nil(t, err) - assert.Equal(t, "trace", carrier.Get(traceIdKey)) + assert.Equal(t, "trace", carrier.Get(TraceIdKey)) assert.Equal(t, "span", carrier.Get(spanIdKey)) _, err = Extract(GrpcFormat, 1) @@ -53,12 +53,12 @@ func TestGrpcPropagator_Extract(t *testing.T) { func TestGrpcPropagator_Inject(t *testing.T) { md := metadata.New(map[string]string{ - traceIdKey: "trace", + TraceIdKey: "trace", spanIdKey: "span", }) carrier, err := Inject(GrpcFormat, md) assert.Nil(t, err) - assert.Equal(t, "trace", carrier.Get(traceIdKey)) + assert.Equal(t, "trace", carrier.Get(TraceIdKey)) assert.Equal(t, "span", carrier.Get(spanIdKey)) _, err = Inject(GrpcFormat, 1) diff --git a/core/trace/span.go b/core/trace/span.go index bdbf0af03..1174283e2 100644 --- a/core/trace/span.go +++ b/core/trace/span.go @@ -34,7 +34,7 @@ type Span struct { func newServerSpan(carrier Carrier, serviceName, operationName string) tracespec.Trace { traceId := stringx.TakeWithPriority(func() string { if carrier != nil { - return carrier.Get(traceIdKey) + return carrier.Get(TraceIdKey) } return "" }, stringx.RandId) diff --git a/core/trace/span_test.go b/core/trace/span_test.go index ab0fa14d7..8323b226d 100644 --- a/core/trace/span_test.go +++ b/core/trace/span_test.go @@ -57,7 +57,7 @@ func TestServerSpan(t *testing.T) { func TestServerSpan_WithCarrier(t *testing.T) { md := metadata.New(map[string]string{ - traceIdKey: "a", + TraceIdKey: "a", spanIdKey: "0.1", }) ctx, span := StartServerSpan(context.Background(), grpcCarrier(md), "service", "operation") @@ -99,7 +99,7 @@ func TestSpan_Follow(t *testing.T) { for _, test := range tests { t.Run(stringx.RandId(), func(t *testing.T) { md := metadata.New(map[string]string{ - traceIdKey: "a", + TraceIdKey: "a", spanIdKey: test.span, }) ctx, span := StartServerSpan(context.Background(), grpcCarrier(md), diff --git a/core/trace/spancontext.go b/core/trace/spancontext.go index dd78d9712..7fc822570 100644 --- a/core/trace/spancontext.go +++ b/core/trace/spancontext.go @@ -14,6 +14,6 @@ func (sc spanContext) SpanId() string { } func (sc spanContext) Visit(fn func(key, val string) bool) { - fn(traceIdKey, sc.traceId) + fn(TraceIdKey, sc.traceId) fn(spanIdKey, sc.spanId) } diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go index b06cb25f7..7e250d2b9 100644 --- a/rest/handler/tracinghandler.go +++ b/rest/handler/tracinghandler.go @@ -21,6 +21,8 @@ func TracingHandler(next http.Handler) http.Handler { defer span.Finish() r = r.WithContext(ctx) + // Conveniently track error messages + w.Header().Set(trace.TraceIdKey, span.TraceId()) next.ServeHTTP(w, r) }) } diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go index 2422ba953..69fc0da48 100644 --- a/rest/handler/tracinghandler_test.go +++ b/rest/handler/tracinghandler_test.go @@ -1,6 +1,8 @@ package handler import ( + "github.com/tal-tech/go-zero/core/stringx" + "github.com/tal-tech/go-zero/core/trace" "net/http" "net/http/httptest" "testing" @@ -11,14 +13,19 @@ import ( func TestTracingHandler(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) - req.Header.Set("X-Trace-ID", "theid") + + traceId := stringx.RandId() + req.Header.Set(trace.TraceIdKey, traceId) + handler := TracingHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { span, ok := r.Context().Value(tracespec.TracingKey).(tracespec.Trace) assert.True(t, ok) - assert.Equal(t, "theid", span.TraceId()) + assert.Equal(t, traceId, span.TraceId()) })) resp := httptest.NewRecorder() handler.ServeHTTP(resp, req) + assert.Equal(t, http.StatusOK, resp.Code) + assert.Equal(t, traceId, resp.Header().Get(trace.TraceIdKey)) }