From f9619328f25ba4733d3bf88bf54ecf21f3c07faf Mon Sep 17 00:00:00 2001
From: cong <zhangcong1992@gmail.com>
Date: Thu, 12 Jan 2023 09:40:18 +0800
Subject: [PATCH] refactor(rest): use static config for trace ignore paths.
 (#2773)

---
 rest/config.go                      |  2 ++
 rest/engine.go                      |  8 +++++--
 rest/handler/tracinghandler.go      | 35 +++++++++++++++++++++--------
 rest/handler/tracinghandler_test.go |  5 +----
 4 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/rest/config.go b/rest/config.go
index 48738001..ece29bfd 100644
--- a/rest/config.go
+++ b/rest/config.go
@@ -57,5 +57,7 @@ type (
 		Signature    SignatureConf `json:",optional"`
 		// There are default values for all the items in Middlewares.
 		Middlewares MiddlewaresConf
+		// TraceIgnorePaths is paths blacklist for trace middleware.
+		TraceIgnorePaths []string `json:",optional"`
 	}
 )
diff --git a/rest/engine.go b/rest/engine.go
index 8beb5ad0..6f12ee23 100644
--- a/rest/engine.go
+++ b/rest/engine.go
@@ -118,7 +118,9 @@ func (ng *engine) buildChainWithNativeMiddlewares(fr featuredRoutes, route Route
 	chn := chain.New()
 
 	if ng.conf.Middlewares.Trace {
-		chn = chn.Append(handler.TracingHandler(ng.conf.Name, route.Path))
+		chn = chn.Append(handler.TracingHandler(ng.conf.Name,
+			route.Path,
+			handler.WithTraceIgnorePaths(ng.conf.TraceIgnorePaths)))
 	}
 	if ng.conf.Middlewares.Log {
 		chn = chn.Append(ng.getLogHandler())
@@ -202,7 +204,9 @@ func (ng *engine) getShedder(priority bool) load.Shedder {
 func (ng *engine) notFoundHandler(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		chn := chain.New(
-			handler.TracingHandler(ng.conf.Name, ""),
+			handler.TracingHandler(ng.conf.Name,
+				"",
+				handler.WithTraceIgnorePaths(ng.conf.TraceIgnorePaths)),
 			ng.getLogHandler(),
 		)
 
diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go
index d3151569..a664cd78 100644
--- a/rest/handler/tracinghandler.go
+++ b/rest/handler/tracinghandler.go
@@ -2,9 +2,8 @@ package handler
 
 import (
 	"net/http"
-	"sync"
 
-	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/trace"
 	"github.com/zeromicro/go-zero/rest/internal/response"
 	"go.opentelemetry.io/otel"
@@ -13,15 +12,26 @@ import (
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
-var notTracingSpans sync.Map
+type (
+	// TracingOption defines the method to customize an tracingOptions.
+	TracingOption func(options *tracingOptions)
 
-// DontTraceSpan disable tracing for the specified span name.
-func DontTraceSpan(spanName string) {
-	notTracingSpans.Store(spanName, lang.Placeholder)
-}
+	// tracingOptions is TracingHandler options.
+	tracingOptions struct {
+		traceIgnorePaths []string
+	}
+)
 
 // TracingHandler return a middleware that process the opentelemetry.
-func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
+func TracingHandler(serviceName, path string, opts ...TracingOption) func(http.Handler) http.Handler {
+	var tracingOpts tracingOptions
+	for _, opt := range opts {
+		opt(&tracingOpts)
+	}
+
+	ignorePaths := collection.NewSet()
+	ignorePaths.AddStr(tracingOpts.traceIgnorePaths...)
+
 	return func(next http.Handler) http.Handler {
 		propagator := otel.GetTextMapPropagator()
 		tracer := otel.GetTracerProvider().Tracer(trace.TraceName)
@@ -32,7 +42,7 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 				spanName = r.URL.Path
 			}
 
-			if _, ok := notTracingSpans.Load(spanName); ok {
+			if ignorePaths.Contains(spanName) {
 				next.ServeHTTP(w, r)
 				return
 			}
@@ -58,3 +68,10 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 		})
 	}
 }
+
+// WithTraceIgnorePaths specifies the traceIgnorePaths option for TracingHandler.
+func WithTraceIgnorePaths(traceIgnorePaths []string) TracingOption {
+	return func(options *tracingOptions) {
+		options.traceIgnorePaths = traceIgnorePaths
+	}
+}
diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go
index 79a94c4d..0e110010 100644
--- a/rest/handler/tracinghandler_test.go
+++ b/rest/handler/tracinghandler_test.go
@@ -63,12 +63,9 @@ func TestDontTracingSpan(t *testing.T) {
 	})
 	defer ztrace.StopAgent()
 
-	DontTraceSpan("bar")
-	defer notTracingSpans.Delete("bar")
-
 	for _, test := range []string{"", "bar", "foo"} {
 		t.Run(test, func(t *testing.T) {
-			h := chain.New(TracingHandler("foo", test)).Then(
+			h := chain.New(TracingHandler("foo", test, WithTraceIgnorePaths([]string{"bar"}))).Then(
 				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 					span := trace.SpanFromContext(r.Context())
 					spanCtx := span.SpanContext()