Set up tracing a bit better

cmd/todoapp/main.go

@@ -19,7 +19,6 @@ import (
 	"github.com/craigpastro/todoapp/internal/postgres"
 	"github.com/craigpastro/todoapp/internal/server"
 	"github.com/sethvargo/go-envconfig"
-	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"golang.org/x/exp/slog"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
@@ -61,18 +60,16 @@ func run(ctx context.Context, cfg *config) {
 	logger := slog.New(handler)
 	slog.SetDefault(logger)
 
-	tp := sdktrace.NewTracerProvider()
-	if cfg.TraceEnabled {
-		tp = instrumentation.MustNewTracerProvider(instrumentation.TracerConfig{
+	tpShutdown := instrumentation.MustNewTracerProvider(
+		cfg.TraceEnabled,
+		instrumentation.TracerConfig{
 			ServiceName:    cfg.ServiceName,
 			ServiceVersion: cfg.ServiceVersion,
 			Environment:    cfg.ServiceEnvironment,
 			Endpoint:       cfg.TraceProviderURL,
 			SampleRatio:    cfg.TraceSampleRatio,
-		})
-
-		slog.Info(fmt.Sprintf("sending traces to '%s'", cfg.TraceProviderURL))
-	}
+		},
+	)
 
 	pool := postgres.MustNew(&postgres.Config{
 		ConnString:        cfg.PostgresConnString,
@@ -118,16 +115,15 @@ func run(ctx context.Context, cfg *config) {
 
 	slog.Info("todoapp attempting to shutdown gracefully")
 
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
 	if err := srv.Shutdown(ctx); err != nil {
 		slog.Error("todoapp shutdown failed", slog.String("error", err.Error()))
 		os.Exit(1)
 	}
 
-	_ = tp.ForceFlush(ctx)
-	_ = tp.Shutdown(ctx)
+	_ = tpShutdown(ctx)
 	pool.Close()
 
 	slog.Info("todoapp shutdown gracefully. bye 👋")

internal/instrumentation/trace.go

@@ -2,19 +2,23 @@ package instrumentation
 
 import (
 	"context"
+	"fmt"
+	"log/slog"
 	"time"
 
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/codes"
-	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
+	"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.20.0"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
 )
 
+type shutdownFunc func(context.Context) error
+
 type TracerConfig struct {
 	ServiceName    string
 	ServiceVersion string
@@ -23,37 +27,53 @@ type TracerConfig struct {
 	SampleRatio    float64
 }
 
-func MustNewTracerProvider(cfg TracerConfig) *sdktrace.TracerProvider {
-	client := otlptracegrpc.NewClient(
-		otlptracegrpc.WithInsecure(),
-		otlptracegrpc.WithEndpoint(cfg.Endpoint),
-		otlptracegrpc.WithDialOption(grpc.WithBlock()),
-	)
+func MustNewTracerProvider(enabled bool, cfg TracerConfig) shutdownFunc {
+	if !enabled {
+		return func(_ context.Context) error {
+			return nil
+		}
+	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-	defer cancel()
+	prop := propagation.NewCompositeTextMapPropagator(
+		propagation.TraceContext{},
+		propagation.Baggage{},
+	)
+	otel.SetTextMapPropagator(prop)
 
-	exp, err := otlptrace.New(ctx, client)
+	resource, err := resource.Merge(
+		resource.Default(),
+		resource.NewWithAttributes(
+			semconv.SchemaURL,
+			semconv.ServiceNameKey.String(cfg.ServiceName),
+			semconv.ServiceVersionKey.String(cfg.ServiceVersion),
+			semconv.DeploymentEnvironmentKey.String(cfg.Environment),
+		),
+	)
 	if err != nil {
 		panic(err)
 	}
 
-	resource := resource.NewWithAttributes(
-		semconv.SchemaURL,
-		semconv.ServiceNameKey.String(cfg.ServiceName),
-		semconv.ServiceVersionKey.String(cfg.ServiceVersion),
-		semconv.DeploymentEnvironmentKey.String(cfg.Environment),
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	traceExporter, err := otlptracegrpc.New(
+		timeoutCtx,
+		otlptracegrpc.WithEndpoint(cfg.Endpoint),
+		otlptracegrpc.WithInsecure(),
+		otlptracegrpc.WithDialOption(grpc.WithBlock()),
 	)
 
-	tp := sdktrace.NewTracerProvider(
+	tracerProvider := sdktrace.NewTracerProvider(
 		sdktrace.WithSampler(sdktrace.TraceIDRatioBased(cfg.SampleRatio)),
+		sdktrace.WithBatcher(traceExporter),
 		sdktrace.WithResource(resource),
-		sdktrace.WithSpanProcessor(sdktrace.NewBatchSpanProcessor(exp)),
 	)
 
-	otel.SetTracerProvider(tp)
+	otel.SetTracerProvider(tracerProvider)
+
+	slog.Info(fmt.Sprintf("sending traces to '%s'", cfg.Endpoint))
 
-	return tp
+	return tracerProvider.Shutdown
 }
 
 func TraceError(span trace.Span, err error) {