From efeef9076c989a49e640b399106aa88bc385a7ae Mon Sep 17 00:00:00 2001
From: James Blair <mail@jamesblair.net>
Date: Sun, 26 Nov 2023 09:46:05 +1300
Subject: [PATCH] Backport server: Have tracingExporter own resources it
 initialises.

Signed-off-by: James Blair <mail@jamesblair.net>
---
 server/embed/config_tracing.go | 50 +++++++++++++++++++++++++---------
 server/embed/etcd.go           |  9 +++---
 2 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/server/embed/config_tracing.go b/server/embed/config_tracing.go
index 75a5febaad63..880479e515b3 100644
--- a/server/embed/config_tracing.go
+++ b/server/embed/config_tracing.go
@@ -40,13 +40,19 @@ func validateTracingConfig(samplingRate int) error {
 	return nil
 }
 
-func setupTracingExporter(ctx context.Context, cfg *Config) (exporter tracesdk.SpanExporter, options []otelgrpc.Option, err error) {
-	exporter, err = otlptracegrpc.New(ctx,
+type tracingExporter struct {
+	exporter tracesdk.SpanExporter
+	opts     []otelgrpc.Option
+	provider *tracesdk.TracerProvider
+}
+
+func newTracingExporter(ctx context.Context, cfg *Config) (*tracingExporter, error) {
+	exporter, err := otlptracegrpc.New(ctx,
 		otlptracegrpc.WithInsecure(),
 		otlptracegrpc.WithEndpoint(cfg.ExperimentalDistributedTracingAddress),
 	)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	res, err := resource.New(ctx,
@@ -55,7 +61,7 @@ func setupTracingExporter(ctx context.Context, cfg *Config) (exporter tracesdk.S
 		),
 	)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	if resWithIDKey := determineResourceWithIDKey(cfg.ExperimentalDistributedTracingServiceInstanceID); resWithIDKey != nil {
@@ -63,11 +69,19 @@ func setupTracingExporter(ctx context.Context, cfg *Config) (exporter tracesdk.S
 		// resource in case of duplicates.
 		res, err = resource.Merge(res, resWithIDKey)
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 	}
 
-	options = append(options,
+		traceProvider := tracesdk.NewTracerProvider(
+		tracesdk.WithBatcher(exporter),
+		tracesdk.WithResource(res),
+		tracesdk.WithSampler(
+			tracesdk.ParentBased(determineSampler(cfg.ExperimentalDistributedTracingSamplingRatePerMillion)),
+		),
+	)
+
+	options := []otelgrpc.Option{
 		otelgrpc.WithPropagators(
 			propagation.NewCompositeTextMapPropagator(
 				propagation.TraceContext{},
@@ -75,13 +89,9 @@ func setupTracingExporter(ctx context.Context, cfg *Config) (exporter tracesdk.S
 			),
 		),
 		otelgrpc.WithTracerProvider(
-			tracesdk.NewTracerProvider(
-				tracesdk.WithBatcher(exporter),
-				tracesdk.WithResource(res),
-				tracesdk.WithSampler(tracesdk.ParentBased(tracesdk.NeverSample())),
-			),
+			traceProvider,
 		),
-	)
+	}
 
 	cfg.logger.Debug(
 		"distributed tracing enabled",
@@ -90,7 +100,21 @@ func setupTracingExporter(ctx context.Context, cfg *Config) (exporter tracesdk.S
 		zap.String("service-instance-id", cfg.ExperimentalDistributedTracingServiceInstanceID),
 	)
 
-	return exporter, options, err
+	return &tracingExporter{
+		exporter: exporter,
+		opts:     options,
+		provider: traceProvider,
+	}, nil
+}
+
+func (te *tracingExporter) Close(ctx context.Context) {
+	if te.provider != nil {
+		te.provider.Shutdown(ctx)
+	}
+
+	if te.exporter != nil {
+		te.exporter.Shutdown(ctx)
+	}
 }
 
 func determineSampler(samplingRate int) tracesdk.Sampler {
diff --git a/server/embed/etcd.go b/server/embed/etcd.go
index b4f4defe54e7..67764968fbf9 100644
--- a/server/embed/etcd.go
+++ b/server/embed/etcd.go
@@ -228,15 +228,14 @@ func StartEtcd(inCfg *Config) (e *Etcd, err error) {
 
 	if srvcfg.ExperimentalEnableDistributedTracing {
 		tctx := context.Background()
-		tracingExporter, opts, err := setupTracingExporter(tctx, cfg)
+		tracingExporter, err := newTracingExporter(tctx, cfg)
 		if err != nil {
 			return e, err
 		}
-		if tracingExporter == nil || len(opts) == 0 {
-			return e, fmt.Errorf("error setting up distributed tracing")
+		e.tracingExporterShutdown = func() {
+			tracingExporter.Close(tctx)
 		}
-		e.tracingExporterShutdown = func() { tracingExporter.Shutdown(tctx) }
-		srvcfg.ExperimentalTracerOptions = opts
+		srvcfg.ExperimentalTracerOptions = tracingExporter.opts
 
 		e.cfg.logger.Info("distributed tracing setup enabled")
 	}