fix: properly shutdown tracers (#4127)

cli/server.go

@@ -161,7 +161,7 @@ func Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, error))
 			)
 
 			if traceEnable || telemetryEnable {
-				sdkTracerProvider, err := tracing.TracerProvider(ctx, "coderd", tracing.TracerOpts{
+				sdkTracerProvider, closeTracing, err := tracing.TracerProvider(ctx, "coderd", tracing.TracerOpts{
 					Default: traceEnable,
 					Coder:   telemetryEnable && !isTest(),
 				})
@@ -170,7 +170,7 @@ func Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, error))
 				} else {
 					// allow time for traces to flush even if command context is canceled
 					defer func() {
-						_ = shutdownWithTimeout(sdkTracerProvider, 5*time.Second)
+						_ = shutdownWithTimeout(closeTracing, 5*time.Second)
 					}()
 
 					d, err := tracing.PostgresDriver(sdkTracerProvider, "coderd.database")
@@ -545,7 +545,7 @@ func Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, error))
 				},
 			}
 			defer func() {
-				_ = shutdownWithTimeout(server, 5*time.Second)
+				_ = shutdownWithTimeout(server.Shutdown, 5*time.Second)
 			}()
 
 			eg := errgroup.Group{}
@@ -633,7 +633,7 @@ func Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, error))
 			// in-flight requests, give in-flight requests 5 seconds to
 			// complete.
 			cmd.Println("Shutting down API server...")
-			err = shutdownWithTimeout(server, 5*time.Second)
+			err = shutdownWithTimeout(server.Shutdown, 5*time.Second)
 			if err != nil {
 				cmd.Printf("API server shutdown took longer than 5s: %s", err)
 			} else {
@@ -655,7 +655,7 @@ func Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, error))
 					if verbose {
 						cmd.Printf("Shutting down provisioner daemon %d...\n", id)
 					}
-					err := shutdownWithTimeout(provisionerDaemon, 5*time.Second)
+					err := shutdownWithTimeout(provisionerDaemon.Shutdown, 5*time.Second)
 					if err != nil {
 						cmd.PrintErrf("Failed to shutdown provisioner daemon %d: %s\n", id, err)
 						return
@@ -903,10 +903,10 @@ func isLocalURL(ctx context.Context, u *url.URL) (bool, error) {
 	return false, nil
 }
 
-func shutdownWithTimeout(s interface{ Shutdown(context.Context) error }, timeout time.Duration) error {
+func shutdownWithTimeout(shutdown func(context.Context) error, timeout time.Duration) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
-	return s.Shutdown(ctx)
+	return shutdown(ctx)
 }
 
 // nolint:revive

coderd/tracing/exporter.go

@@ -3,9 +3,11 @@ package tracing
 import (
 	"context"
 
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
+	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/sdk/resource"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
@@ -24,34 +26,53 @@ type TracerOpts struct {
 
 // TracerProvider creates a grpc otlp exporter and configures a trace provider.
 // Caller is responsible for calling TracerProvider.Shutdown to ensure all data is flushed.
-func TracerProvider(ctx context.Context, service string, opts TracerOpts) (*sdktrace.TracerProvider, error) {
+func TracerProvider(ctx context.Context, service string, opts TracerOpts) (*sdktrace.TracerProvider, func(context.Context) error, error) {
 	res := resource.NewWithAttributes(
 		semconv.SchemaURL,
 		// the service name used to display traces in backends
 		semconv.ServiceNameKey.String(service),
 	)
 
-	tracerOpts := []sdktrace.TracerProviderOption{
+	var (
-		sdktrace.WithResource(res),
+		tracerOpts = []sdktrace.TracerProviderOption{
-	}
+			sdktrace.WithResource(res),
+		}
+		closers = []func(context.Context) error{}
+	)
+
 	if opts.Default {
 		exporter, err := DefaultExporter(ctx)
 		if err != nil {
-			return nil, xerrors.Errorf("default exporter: %w", err)
+			return nil, nil, xerrors.Errorf("default exporter: %w", err)
 		}
+		closers = append(closers, exporter.Shutdown)
 		tracerOpts = append(tracerOpts, sdktrace.WithBatcher(exporter))
 	}
 	if opts.Coder {
 		exporter, err := CoderExporter(ctx)
 		if err != nil {
-			return nil, xerrors.Errorf("coder exporter: %w", err)
+			return nil, nil, xerrors.Errorf("coder exporter: %w", err)
 		}
+		closers = append(closers, exporter.Shutdown)
 		tracerOpts = append(tracerOpts, sdktrace.WithBatcher(exporter))
 	}
 
 	tracerProvider := sdktrace.NewTracerProvider(tracerOpts...)
+	otel.SetTracerProvider(tracerProvider)
+	otel.SetTextMapPropagator(
+		propagation.NewCompositeTextMapPropagator(
+			propagation.TraceContext{},
+			propagation.Baggage{},
+		),
+	)
 
-	return tracerProvider, nil
+	return tracerProvider, func(ctx context.Context) error {
+		for _, close := range closers {
+			_ = close(ctx)
+		}
+		_ = tracerProvider.Shutdown(ctx)
+		return nil
+	}, nil
 }
 
 func DefaultExporter(ctx context.Context) (*otlptrace.Exporter, error) {