fix:trace graceful stop,pre loss trace (#2358)

core/service/serviceconf.go

@@ -3,6 +3,8 @@ package service
 import (
 	"log"
 
+	"github.com/zeromicro/go-zero/core/proc"
+
 	"github.com/zeromicro/go-zero/core/load"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/prometheus"
@@ -56,6 +58,9 @@ func (sc ServiceConf) SetUp() error {
 		sc.Telemetry.Name = sc.Name
 	}
 	trace.StartAgent(sc.Telemetry)
+	proc.AddShutdownListener(func() {
+		trace.StopAgent()
+	})
 
 	if len(sc.MetricsUrl) > 0 {
 		stat.SetReportWriter(stat.NewRemoteWriter(sc.MetricsUrl))

core/stores/redis/hook_test.go

@@ -21,6 +21,7 @@ func TestHookProcessCase1(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	writer := log.Writer()
 	var buf strings.Builder
@@ -44,6 +45,7 @@ func TestHookProcessCase2(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	w, restore := injectLog()
 	defer restore()
@@ -108,6 +110,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	w, restore := injectLog()
 	defer restore()

core/trace/agent.go

@@ -1,6 +1,7 @@
 package trace
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
@@ -23,6 +24,7 @@ const (
 var (
 	agents = make(map[string]lang.PlaceholderType)
 	lock   sync.Mutex
+	tp     *sdktrace.TracerProvider
 )
 
 // StartAgent starts a opentelemetry agent.
@@ -43,6 +45,11 @@ func StartAgent(c Config) {
 	agents[c.Endpoint] = lang.Placeholder
 }
 
+// StopAgent shuts down the span processors in the order they were registered.
+func StopAgent() {
+	_ = tp.Shutdown(context.Background())
+}
+
 func createExporter(c Config) (sdktrace.SpanExporter, error) {
 	// Just support jaeger and zipkin now, more for later
 	switch c.Batcher {
@@ -74,7 +81,7 @@ func startAgent(c Config) error {
 		opts = append(opts, sdktrace.WithBatcher(exp))
 	}
 
-	tp := sdktrace.NewTracerProvider(opts...)
+	tp = sdktrace.NewTracerProvider(opts...)
 	otel.SetTracerProvider(tp)
 	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
 		propagation.TraceContext{}, propagation.Baggage{}))

rest/handler/tracinghandler_test.go

@@ -21,6 +21,7 @@ func TestOtelHandler(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	for _, test := range []string{"", "bar"} {
 		t.Run(test, func(t *testing.T) {

rest/httpc/requests_test.go

@@ -22,6 +22,7 @@ func TestDoRequest(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))

zrpc/internal/clientinterceptors/tracinginterceptor_test.go

@@ -23,6 +23,7 @@ func TestOpenTracingInterceptor(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer trace.StopAgent()
 
 	cc := new(grpc.ClientConn)
 	ctx := metadata.NewOutgoingContext(context.Background(), metadata.MD{})

zrpc/internal/rpcserver.go

@@ -6,6 +6,7 @@ import (
 	"github.com/zeromicro/go-zero/core/proc"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/zrpc/internal/serverinterceptors"
+
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
 )

zrpc/internal/serverinterceptors/tracinginterceptor_test.go

@@ -32,6 +32,8 @@ func TestUnaryOpenTracingInterceptor_Enabled(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer trace.StopAgent()
+
 	_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
 		FullMethod: "/package.TestService.GetUser",
 	}, func(ctx context.Context, req interface{}) (interface{}, error) {