From 2d78558234018b1cd6d42e35a3cba05024f927ef Mon Sep 17 00:00:00 2001 From: CoderPoet <coderpoet@outlook.com> Date: Thu, 16 Nov 2023 23:43:16 +0800 Subject: [PATCH] feat(tracing): support trace info LevelDisabled --- tracing/middleware.go | 6 +++++- tracing/middleware_test.go | 25 +++++++++++++++++++++++-- tracing/tracer_server.go | 4 ++++ 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/tracing/middleware.go b/tracing/middleware.go index 712a48d..e5ecaed 100644 --- a/tracing/middleware.go +++ b/tracing/middleware.go @@ -18,6 +18,7 @@ import ( "context" "time" + "github.com/cloudwego/hertz/pkg/common/tracer/stats" "go.opentelemetry.io/otel/metric" "github.com/cloudwego/hertz/pkg/app" @@ -109,8 +110,11 @@ func ServerMiddleware(cfg *Config) app.HandlerFunc { } sTracer := tc.Tracer() - ti := c.GetTraceInfo() + if ti.Stats().Level() == stats.LevelDisabled { + c.Next(ctx) + return + } opts := []oteltrace.SpanStartOption{ oteltrace.WithTimestamp(getStartTimeOrNow(ti)), diff --git a/tracing/middleware_test.go b/tracing/middleware_test.go index dbbc3d9..dbcbd70 100644 --- a/tracing/middleware_test.go +++ b/tracing/middleware_test.go @@ -20,10 +20,10 @@ import ( "testing" "time" - "github.com/cloudwego/hertz/pkg/common/test/assert" - "github.com/cloudwego/hertz/pkg/app" "github.com/cloudwego/hertz/pkg/app/server" + "github.com/cloudwego/hertz/pkg/common/test/assert" + "github.com/cloudwego/hertz/pkg/common/tracer/stats" "go.opentelemetry.io/otel" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" @@ -47,3 +47,24 @@ func TestServerMiddleware(t *testing.T) { assert.Nil(t, err) assert.True(t, len(resp.Header.Get("trace-id")) != 0) } + +func TestServerMiddlewareDisableTrace(t *testing.T) { + sr := tracetest.NewSpanRecorder() + otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))) + tracer, cfg := NewServerTracer(WithCustomResponseHandler(func(c context.Context, ctx *app.RequestContext) { + ctx.Header("trace-id", oteltrace.SpanFromContext(c).SpanContext().TraceID().String()) + })) + h := server.Default(tracer, + server.WithHostPorts("127.0.0.1:16666"), + server.WithTraceLevel(stats.LevelDisabled), + ) + h.Use(ServerMiddleware(cfg)) + h.GET("/ping", func(c context.Context, ctx *app.RequestContext) { + }) + + go h.Spin() + time.Sleep(100 * time.Millisecond) + resp, err := http.Get("http://127.0.0.1:16666/ping") + assert.Nil(t, err) + assert.True(t, len(resp.Header.Get("trace-id")) == 0) +} diff --git a/tracing/tracer_server.go b/tracing/tracer_server.go index b46c6f8..d2e2468 100644 --- a/tracing/tracer_server.go +++ b/tracing/tracer_server.go @@ -77,6 +77,10 @@ func (s *serverTracer) Finish(ctx context.Context, c *app.RequestContext) { ti := c.GetTraceInfo() st := ti.Stats() + if st.Level() == stats.LevelDisabled { + return + } + httpStart := st.GetEvent(stats.HTTPStart) if httpStart == nil { return