add gracefully shutdown in otel_setup.go (#315)

add gracefully shutdown in otel_setup.go

Author: 123liuziming

pkg/data/default.json

@@ -29,6 +29,12 @@
     "OnEnter": "defer func(){ retVal0.otel_trace_context = contextPropagate(callergp.otel_trace_context); retVal0.otel_baggage_container = contextPropagate(callergp.otel_baggage_container); }()",
     "UseRaw": true
   },
+  {
+    "ImportPath": "runtime",
+    "Function": "runExitHooks",
+    "OnEnter": "if ExitHook != nil { ExitHook(); }",
+    "UseRaw": true
+  },
   {
     "ImportPath": "runtime",
     "FileName": "runtime_linker.go",

pkg/data/test_runtime.json

@@ -0,0 +1,43 @@
+[
+  {
+    "ImportPath": "runtime",
+    "FileName": "reorder_init.go",
+    "Path": "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/runtime"
+  },
+  {
+    "ImportPath": "runtime",
+    "Function":"doInit",
+    "OnEnter": "reorderInitTasks(ts)",
+    "UseRaw": true,
+    "GoVersion": "[1.21.0,)"
+  },
+  {
+    "ImportPath": "runtime",
+    "StructType": "g",
+    "FieldName": "otel_trace_context",
+    "FieldType": "interface{}"
+  },
+  {
+    "ImportPath": "runtime",
+    "StructType": "g",
+    "FieldName": "otel_baggage_container",
+    "FieldType": "interface{}"
+  },
+  {
+    "ImportPath": "runtime",
+    "Function": "newproc1",
+    "OnEnter": "defer func(){ retVal0.otel_trace_context = contextPropagate(callergp.otel_trace_context); retVal0.otel_baggage_container = contextPropagate(callergp.otel_baggage_container); }()",
+    "UseRaw": true
+  },
+  {
+    "ImportPath": "runtime",
+    "Function": "runExitHooks",
+    "OnEnter": "if ExitHook != nil { ExitHook(); }",
+    "UseRaw": true
+  },
+  {
+    "ImportPath": "runtime",
+    "FileName": "runtime_linker.go",
+    "Path": "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules/runtime"
+  }
+]

pkg/otel_setup.go

@@ -1,3 +1,5 @@
+//go:build ignore
+
 // Copyright (c) 2024 Alibaba Group Holding Ltd.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,14 +28,14 @@ import (
 	"log"
 	http2 "net/http"
 	"os"
+	"runtime"
 	"strings"
 
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/core/meter"
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/inst-api-semconv/instrumenter/http"
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/test/verifier"
-	"go.opentelemetry.io/contrib/instrumentation/runtime"
+	otelruntime "go.opentelemetry.io/contrib/instrumentation/runtime"
 	"go.opentelemetry.io/otel"
-	_ "go.opentelemetry.io/otel"
 	_ "go.opentelemetry.io/otel/baggage"
 	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc"
 	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp"
@@ -43,7 +45,6 @@ import (
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/sdk/metric"
 	"go.opentelemetry.io/otel/sdk/trace"
-	_ "go.opentelemetry.io/otel/sdk/trace"
 )
 
 // set the following environment variables based on https://opentelemetry.io/docs/specs/otel/configuration/sdk-environment-variables
@@ -57,7 +58,19 @@ const metrics_exporter = "OTEL_METRICS_EXPORTER"
 const prometheus_exporter_port = "OTEL_EXPORTER_PROMETHEUS_PORT"
 const default_prometheus_exporter_port = "9464"
 
+var (
+	spanExporter       trace.SpanExporter
+	traceProvider      *trace.TracerProvider
+	metricsProvider    *metric.MeterProvider
+	batchSpanProcessor trace.SpanProcessor
+)
+
 func init() {
+	ctx := context.Background()
+	// graceful shutdown
+	runtime.ExitHook = func() {
+		gracefullyShutdown(ctx)
+	}
 	path, err := os.Executable()
 	if err != nil {
 		panic(err)
@@ -66,7 +79,7 @@ func init() {
 	if strings.HasSuffix(path, exec_name) {
 		return
 	}
-	if err = initOpenTelemetry(); err != nil {
+	if err = initOpenTelemetry(ctx); err != nil {
 		log.Fatalf("%s: %v", "Failed to initialize opentelemetry resource", err)
 	}
 }
@@ -78,32 +91,27 @@ func newSpanProcessor(ctx context.Context) trace.SpanProcessor {
 		simpleProcessor := trace.NewSimpleSpanProcessor(traceExporter)
 		return simpleProcessor
 	} else {
-		var traceExporter trace.SpanExporter
 		var err error
 		if os.Getenv(report_protocol) == "grpc" || os.Getenv(trace_report_protocol) == "grpc" {
-			traceExporter, err = otlptrace.New(ctx, otlptracegrpc.NewClient())
+			spanExporter, err = otlptrace.New(ctx, otlptracegrpc.NewClient())
 			if err != nil {
 				log.Fatalf("%s: %v", "Failed to create the OpenTelemetry trace exporter", err)
 			}
 		} else {
-			traceExporter, err = otlptrace.New(ctx, otlptracehttp.NewClient())
+			spanExporter, err = otlptrace.New(ctx, otlptracehttp.NewClient())
 			if err != nil {
 				log.Fatalf("%s: %v", "Failed to create the OpenTelemetry trace exporter", err)
 			}
 		}
-		batchSpanProcessor := trace.NewBatchSpanProcessor(traceExporter)
+		batchSpanProcessor = trace.NewBatchSpanProcessor(spanExporter)
 		return batchSpanProcessor
 	}
 }
 
-func initOpenTelemetry() error {
-	ctx := context.Background()
-
-	var batchSpanProcessor trace.SpanProcessor
+func initOpenTelemetry(ctx context.Context) error {
 
 	batchSpanProcessor = newSpanProcessor(ctx)
 
-	var traceProvider *trace.TracerProvider
 	if batchSpanProcessor != nil {
 		traceProvider = trace.NewTracerProvider(
 			trace.WithSpanProcessor(batchSpanProcessor))
@@ -118,10 +126,9 @@ func initOpenTelemetry() error {
 
 func initMetrics() error {
 	ctx := context.Background()
-	var mp *metric.MeterProvider
 	// TODO: abstract the if-else
 	if verifier.IsInTest() {
-		mp = metric.NewMeterProvider(
+		metricsProvider = metric.NewMeterProvider(
 			metric.WithReader(verifier.ManualReader),
 		)
 	} else {
@@ -130,7 +137,7 @@ func initMetrics() error {
 			if err != nil {
 				log.Fatalf("new otlp metric prometheus exporter failed: %v", err)
 			}
-			mp = metric.NewMeterProvider(
+			metricsProvider = metric.NewMeterProvider(
 				metric.WithReader(exporter),
 			)
 			go serveMetrics()
@@ -139,24 +146,24 @@ func initMetrics() error {
 			if err != nil {
 				log.Fatalf("new otlp metric grpc exporter failed: %v", err)
 			}
-			mp = metric.NewMeterProvider(
+			metricsProvider = metric.NewMeterProvider(
 				metric.WithReader(metric.NewPeriodicReader(exporter)),
 			)
 		} else {
 			exporter, err := otlpmetrichttp.New(ctx)
 			if err != nil {
 				log.Fatalf("new otlp metric http exporter failed: %v", err)
 			}
-			mp = metric.NewMeterProvider(
+			metricsProvider = metric.NewMeterProvider(
 				metric.WithReader(metric.NewPeriodicReader(exporter)),
 			)
 		}
 	}
-	if mp == nil {
+	if metricsProvider == nil {
 		return errors.New("No MeterProvider is provided")
 	}
-	otel.SetMeterProvider(mp)
-	m := mp.Meter("opentelemetry-global-meter")
+	otel.SetMeterProvider(metricsProvider)
+	m := metricsProvider.Meter("opentelemetry-global-meter")
 	meter.SetMeter(m)
 	// init http metrics
 	http.InitHttpMetrics(m)
@@ -167,7 +174,7 @@ func initMetrics() error {
 	// nacos experimental metrics
 	experimental.InitNacosExperimentalMetrics(m)
 	// DefaultMinimumReadMemStatsInterval is 15 second
-	return runtime.Start(runtime.WithMeterProvider(mp))
+	return otelruntime.Start(otelruntime.WithMeterProvider(metricsProvider))
 }
 
 func serveMetrics() {
@@ -183,3 +190,26 @@ func serveMetrics() {
 		return
 	}
 }
+
+func gracefullyShutdown(ctx context.Context) {
+	if metricsProvider != nil {
+		if err := metricsProvider.Shutdown(ctx); err != nil {
+			log.Printf("%s: %v", "Failed to shutdown the OpenTelemetry metric provider", err)
+		}
+	}
+	if traceProvider != nil {
+		if err := traceProvider.Shutdown(ctx); err != nil {
+			log.Printf("%s: %v", "Failed to shutdown the OpenTelemetry trace provider", err)
+		}
+	}
+	if spanExporter != nil {
+		if err := spanExporter.Shutdown(ctx); err != nil {
+			log.Printf("%s: %v", "Failed to shutdown the OpenTelemetry span exporter", err)
+		}
+	}
+	if batchSpanProcessor != nil {
+		if err := batchSpanProcessor.Shutdown(ctx); err != nil {
+			log.Printf("%s: %v", "Failed to shutdown the OpenTelemetry batch span processor", err)
+		}
+	}
+}

pkg/rules/runtime/runtime_linker.go

@@ -18,6 +18,10 @@ import (
 	_ "unsafe"
 )
 
+// When the program exit, we should call the exit hook to shutdown gracefully.
+// See https://github.com/alibaba/opentelemetry-go-auto-instrumentation/blob/main/pkg/otel_setup.go
+var ExitHook func()
+
 //go:linkname otel_get_trace_context_from_gls otel_get_trace_context_from_gls
 var otel_get_trace_context_from_gls = _otel_gls_get_trace_context_impl
 

test/build_test.go

@@ -58,7 +58,7 @@ func TestBuildProject4(t *testing.T) {
 	RunGoBuildFallible(t, "go", "build", "m1")
 	RunSet(t, "-rule=../../pkg/data/test_error.json,../../pkg/data/test_fmt.json")
 	RunGoBuild(t, "go", "build", "m1")
-	RunSet(t, "-disabledefault=true", "-rule=../../pkg/data/test_error.json,../../pkg/data/test_fmt.json")
+	RunSet(t, "-disabledefault=true", "-rule=../../pkg/data/test_error.json,../../pkg/data/test_fmt.json,../../pkg/data/test_runtime.json")
 	RunGoBuild(t, "go", "build", "m1")
 	RunSet(t, "-disabledefault=true", "-rule=../../pkg/data/default.json,../../pkg/data/test_fmt.json")
 	RunGoBuild(t, "go", "build", "m1")
@@ -80,7 +80,7 @@ func TestBuildProject6(t *testing.T) {
 	const AppName = "build"
 	UseApp(AppName)
 
-	RunSet(t, "-disabledefault=true", "-rule=../../pkg/data/test_fmt.json", "-verbose")
+	RunSet(t, "-disabledefault=true", "-rule=../../pkg/data/test_fmt.json,../../pkg/data/test_runtime.json", "-verbose")
 	RunGoBuild(t, "go", "build", "m1")
 	// only test_fmt.json should be available because -disabledefault is set
 	ExpectPreprocessContains(t, util.DebugLogFile, "fmt")