feat(provisioner): propagate trace info (#17166)

coryb · spikecurtis · web-flow · commit 12e5718b99bb · 2025-04-08T11:58:28.000+04:00
If tracing is enabled, propagate the trace information to the terraform
provisioner via environment variables. This sets the `TRACEPARENT`
environment variable using the default W3C trace propagators. Users can
choose to continue the trace by adding new spans in the provisioner by
reading from the environment like:

ctx := env.ContextWithRemoteSpanContext(context.Background(),
os.Environ())

---------

Co-authored-by: Spike Curtis &lt;spike@spikecurtis.com&gt;
diff --git a/provisioner/terraform/otelenv.go b/provisioner/terraform/otelenv.go
@@ -0,0 +1,88 @@
+package terraform
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"strings"
+	"unicode"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+)
+
+// TODO: replace this with the upstream OTEL env propagation when it is
+// released.
+
+// envCarrier is a propagation.TextMapCarrier that is used to extract or
+// inject tracing environment variables. This is used with a
+// propagation.TextMapPropagator
+type envCarrier struct {
+	Env []string
+}
+
+var _ propagation.TextMapCarrier = (*envCarrier)(nil)
+
+func toKey(key string) string {
+	key = strings.ToUpper(key)
+	key = strings.ReplaceAll(key, "-", "_")
+	return strings.Map(func(r rune) rune {
+		if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' {
+			return r
+		}
+		return -1
+	}, key)
+}
+
+func (c *envCarrier) Set(key, value string) {
+	if c == nil {
+		return
+	}
+	key = toKey(key)
+	for i, e := range c.Env {
+		if strings.HasPrefix(e, key+"=") {
+			// don't directly update the slice so we don't modify the slice
+			// passed in
+			c.Env = slices.Clone(c.Env)
+			c.Env[i] = fmt.Sprintf("%s=%s", key, value)
+			return
+		}
+	}
+	c.Env = append(c.Env, fmt.Sprintf("%s=%s", key, value))
+}
+
+func (c *envCarrier) Get(key string) string {
+	if c == nil {
+		return ""
+	}
+	key = toKey(key)
+	for _, e := range c.Env {
+		if strings.HasPrefix(e, key+"=") {
+			return strings.TrimPrefix(e, key+"=")
+		}
+	}
+	return ""
+}
+
+func (c *envCarrier) Keys() []string {
+	if c == nil {
+		return nil
+	}
+	keys := make([]string, len(c.Env))
+	for i, e := range c.Env {
+		k, _, _ := strings.Cut(e, "=")
+		keys[i] = k
+	}
+	return keys
+}
+
+// otelEnvInject will add add any necessary environment variables for the span
+// found in the Context.  If environment variables are already present
+// in `environ` then they will be updated.  If no variables are found the
+// new ones will be appended.  The new environment will be returned, `environ`
+// will never be modified.
+func otelEnvInject(ctx context.Context, environ []string) []string {
+	c := &envCarrier{Env: environ}
+	otel.GetTextMapPropagator().Inject(ctx, c)
+	return c.Env
+}
diff --git a/provisioner/terraform/otelenv_internal_test.go b/provisioner/terraform/otelenv_internal_test.go
@@ -0,0 +1,85 @@
+package terraform
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type testIDGenerator struct{}
+
+var _ sdktrace.IDGenerator = (*testIDGenerator)(nil)
+
+func (testIDGenerator) NewIDs(_ context.Context) (trace.TraceID, trace.SpanID) {
+	traceID, _ := trace.TraceIDFromHex("60d19e9e9abf2197c1d6d8f93e28ee2a")
+	spanID, _ := trace.SpanIDFromHex("a028bd951229a46f")
+	return traceID, spanID
+}
+
+func (testIDGenerator) NewSpanID(_ context.Context, _ trace.TraceID) trace.SpanID {
+	spanID, _ := trace.SpanIDFromHex("a028bd951229a46f")
+	return spanID
+}
+
+func TestOtelEnvInject(t *testing.T) {
+	t.Parallel()
+	testTraceProvider := sdktrace.NewTracerProvider(
+		sdktrace.WithSampler(sdktrace.AlwaysSample()),
+		sdktrace.WithIDGenerator(testIDGenerator{}),
+	)
+
+	tracer := testTraceProvider.Tracer("example")
+	ctx, span := tracer.Start(context.Background(), "testing")
+	defer span.End()
+
+	input := []string{"PATH=/usr/bin:/bin"}
+
+	otel.SetTextMapPropagator(propagation.TraceContext{})
+	got := otelEnvInject(ctx, input)
+	require.Equal(t, []string{
+		"PATH=/usr/bin:/bin",
+		"TRACEPARENT=00-60d19e9e9abf2197c1d6d8f93e28ee2a-a028bd951229a46f-01",
+	}, got)
+
+	// verify we update rather than append
+	input = []string{
+		"PATH=/usr/bin:/bin",
+		"TRACEPARENT=origTraceParent",
+		"TERM=xterm",
+	}
+
+	otel.SetTextMapPropagator(propagation.TraceContext{})
+	got = otelEnvInject(ctx, input)
+	require.Equal(t, []string{
+		"PATH=/usr/bin:/bin",
+		"TRACEPARENT=00-60d19e9e9abf2197c1d6d8f93e28ee2a-a028bd951229a46f-01",
+		"TERM=xterm",
+	}, got)
+}
+
+func TestEnvCarrierSet(t *testing.T) {
+	t.Parallel()
+	c := &envCarrier{
+		Env: []string{"PATH=/usr/bin:/bin", "TERM=xterm"},
+	}
+	c.Set("PATH", "/usr/local/bin")
+	c.Set("NEWVAR", "newval")
+	require.Equal(t, []string{
+		"PATH=/usr/local/bin",
+		"TERM=xterm",
+		"NEWVAR=newval",
+	}, c.Env)
+}
+
+func TestEnvCarrierKeys(t *testing.T) {
+	t.Parallel()
+	c := &envCarrier{
+		Env: []string{"PATH=/usr/bin:/bin", "TERM=xterm"},
+	}
+	require.Equal(t, []string{"PATH", "TERM"}, c.Keys())
+}
diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
@@ -156,6 +156,7 @@ func (s *server) Plan(
 	if err != nil {
 		return provisionersdk.PlanErrorf("setup env: %s", err)
 	}
+	env = otelEnvInject(ctx, env)
 
 	vars, err := planVars(request)
 	if err != nil {
@@ -208,6 +209,7 @@ func (s *server) Apply(
 	if err != nil {
 		return provisionersdk.ApplyErrorf("provision env: %s", err)
 	}
+	env = otelEnvInject(ctx, env)
 	resp, err := e.apply(
 		ctx, killCtx, env, sess,
 	)

Original file line number	Diff line number	Diff line change
`@@ -156,6 +156,7 @@ func (s *server) Plan(`
`156`	`156`	`if err != nil {`
`157`	`157`	`return provisionersdk.PlanErrorf("setup env: %s", err)`
`158`	`158`	`}`
	`159`	`+ env = otelEnvInject(ctx, env)`
`159`	`160`
`160`	`161`	`vars, err := planVars(request)`
`161`	`162`	`if err != nil {`
`@@ -208,6 +209,7 @@ func (s *server) Apply(`
`208`	`209`	`if err != nil {`
`209`	`210`	`return provisionersdk.ApplyErrorf("provision env: %s", err)`
`210`	`211`	`}`
	`212`	`+ env = otelEnvInject(ctx, env)`
`211`	`213`	`resp, err := e.apply(`
`212`	`214`	`ctx, killCtx, env, sess,`
`213`	`215`	`)`