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