feat(tracing): support trace info LevelDisabled

hertz-contrib · Nov 16, 2023 · 4777abf · 4777abf
1 parent 78fb3de
commit 4777abf
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 3 deletions.
diff --git 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
@@ -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:6666"),
+		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:6666/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
@@ -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