fix(serverless-init): eliminate GlobalTags data race in MicroVM trace tags

Codex flagged a data race on PR #53036: MicroVM's /run lifecycle hook calls
serverlessTraceAgent.SetTags from an async goroutine, which mutates
Agent.GlobalTags via SetGlobalTagsUnsafe concurrently with the trace agent's
span-processing loop reading that same field unsynchronized. GlobalTags is
computed once at config-build time and every reader in pkg/trace/agent and
pkg/trace/api assumes it is frozen thereafter; MicroVM's dynamic
lambda_microvm_id update breaks that invariant.

Route the async update through the span modifier instead, which already runs
at the identical point in the span-processing loop and is exclusively
serverless-owned:

- pkg/serverless/trace/span_modifier.go: tags field is now
  atomic.Pointer[map[string]string]; ModifySpan reads it lock-free instead of
  reading a field written by a separate mutator with no synchronization.
- pkg/serverless/trace/trace.go: new UpdateRuntimeTags method that only
  updates the span modifier, never GlobalTags. Existing SetTags (used once
  synchronously at startup, before the trace agent runs) is untouched.
- cmd/serverless-init/main.go: the two async TraceTagSetterFunc closures now
  call UpdateRuntimeTags instead of SetTags.

This is scoped entirely to MicroVM: every other cloud service ignores
LifecycleCtx in Init, so they never exercise the async path and keep calling
the original SetTags at startup unchanged.

Also fixes an issue flagged by Codex review: ModifySpan's new tag-apply loop
unconditionally overwrote _dd.origin whenever the tags map contained it,
undoing the "only fill _dd.origin if absent" guard immediately above it.
Every CloudService.GetTags() sets _dd.origin (not just MicroVM's), and that
value flows into the tags applied here via SetTags/UpdateRuntimeTags at
startup for every cloud service — so this would have silently overwritten a
tracer-supplied span origin (e.g. _dd.origin:rum) for all of them, not just
MicroVM. Fixed by skipping _dd.origin in the loop. Covered by
TestSpanModifierModifySpanPreservesExistingOrigin.

A second Codex finding on this PR — UpdateRuntimeTags not reaching spans
processed via the V1 payload path (ProcessV1 skips SpanModifier) — is left
as a known follow-up; not fixed here.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: litianningdatadog

cmd/serverless-init/lifecycle/server.go

@@ -140,7 +140,7 @@ type LogsTagSetterFunc func([]string)
 func (f LogsTagSetterFunc) SetLogsTags(tags []string) { f(tags) }
 
 // TraceTagSetter can replace the full tag map on the live trace pipeline.
-// Satisfied by trace.ServerlessTraceAgent.SetTags (wrapped via TraceTagSetterFunc).
+// Satisfied by trace.ServerlessTraceAgent.UpdateRuntimeTags (wrapped via TraceTagSetterFunc).
 type TraceTagSetter interface {
 	SetTraceTags(tags map[string]string)
 }

cmd/serverless-init/main.go

@@ -496,7 +496,7 @@ func setup(
 				}),
 				BaseTags: logTagsBase,
 				TraceTagSetter: lifecycle.TraceTagSetterFunc(func(tags map[string]string) {
-					traceAgent.SetTags(tags)
+					traceAgent.UpdateRuntimeTags(tags)
 				}),
 				BaseTraceTags: serverlessInitTag.MakeTraceAgentTags(tagConfig.Tags),
 			},
@@ -532,7 +532,7 @@ func setup(
 			}),
 			BaseTags: logTagsBase,
 			TraceTagSetter: lifecycle.TraceTagSetterFunc(func(tags map[string]string) {
-				traceAgent.SetTags(tags)
+				traceAgent.UpdateRuntimeTags(tags)
 			}),
 			BaseTraceTags: serverlessInitTag.MakeTraceAgentTags(tagConfig.Tags),
 		},

pkg/serverless/trace/span_modifer_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
 	gzip "github.com/DataDog/datadog-agent/comp/trace/compression/impl-gzip"
+	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/trace/agent"
 	"github.com/DataDog/datadog-agent/pkg/trace/api"
 	"github.com/DataDog/datadog-agent/pkg/trace/config"
@@ -109,3 +110,89 @@ func TestSpanModifierDetectsCloudService(t *testing.T) {
 		os.Unsetenv(cloudServiceEnvVar)
 	}
 }
+
+// TestSpanModifierModifySpanBeforeSetTags verifies that ModifySpan only
+// applies the _dd.origin tag when SetTags has never been called, since
+// spanModifier.tags starts as an unset atomic.Pointer.
+func TestSpanModifierModifySpanBeforeSetTags(t *testing.T) {
+	sm := &spanModifier{ddOrigin: "lambda"}
+	span := &pb.Span{Meta: map[string]string{}}
+
+	sm.ModifySpan(&pb.TraceChunk{}, span)
+
+	assert.Equal(t, "lambda", span.Meta[ddOriginTagName])
+	assert.Len(t, span.Meta, 1)
+}
+
+// TestSpanModifierModifySpanAppliesTagsSetDynamically verifies that tags
+// applied via SetTags after construction are picked up by ModifySpan, the
+// mechanism MicroVM uses to deliver the lambda_microvm_id tag once it becomes
+// known at /run time.
+func TestSpanModifierModifySpanAppliesTagsSetDynamically(t *testing.T) {
+	sm := &spanModifier{ddOrigin: "lambda"}
+	sm.SetTags(map[string]string{"lambda_microvm_id": "vm-123"})
+
+	span := &pb.Span{Meta: map[string]string{}}
+	sm.ModifySpan(&pb.TraceChunk{}, span)
+
+	assert.Equal(t, "lambda", span.Meta[ddOriginTagName])
+	assert.Equal(t, "vm-123", span.Meta["lambda_microvm_id"])
+}
+
+// TestSpanModifierModifySpanPreservesExistingOrigin verifies that ModifySpan
+// does not overwrite a span's existing _dd.origin when the dynamically-set
+// tags also contain _dd.origin. Every CloudService.GetTags() sets _dd.origin
+// (e.g. MicroVM sets "lambdamicrovm"), and that value flows into the tags
+// applied here via SetTags/UpdateRuntimeTags — so without this guard, every
+// span would have a tracer-supplied origin (e.g. "rum") silently replaced.
+func TestSpanModifierModifySpanPreservesExistingOrigin(t *testing.T) {
+	sm := &spanModifier{ddOrigin: "lambda"}
+	sm.SetTags(map[string]string{ddOriginTagName: "lambdamicrovm", "lambda_microvm_id": "vm-123"})
+
+	span := &pb.Span{Meta: map[string]string{ddOriginTagName: "rum"}}
+	sm.ModifySpan(&pb.TraceChunk{}, span)
+
+	assert.Equal(t, "rum", span.Meta[ddOriginTagName], "must not overwrite a tracer-supplied origin")
+	assert.Equal(t, "vm-123", span.Meta["lambda_microvm_id"])
+}
+
+// TestSpanModifierModifySpanReflectsLatestSetTags verifies that a later
+// SetTags call replaces the tag set used by subsequent ModifySpan calls.
+func TestSpanModifierModifySpanReflectsLatestSetTags(t *testing.T) {
+	sm := &spanModifier{ddOrigin: "lambda"}
+	sm.SetTags(map[string]string{"lambda_microvm_id": "vm-1"})
+	sm.SetTags(map[string]string{"lambda_microvm_id": "vm-2"})
+
+	span := &pb.Span{Meta: map[string]string{}}
+	sm.ModifySpan(&pb.TraceChunk{}, span)
+
+	assert.Equal(t, "vm-2", span.Meta["lambda_microvm_id"])
+}
+
+// TestSpanModifierSetTagsConcurrentWithModifySpan exercises SetTags and
+// ModifySpan concurrently under the race detector. This is a regression test
+// for the data race Codex flagged on PR #53036: MicroVM's /run hook calls
+// SetTags from a goroutine that runs concurrently with the trace agent's
+// span-processing loop, which calls ModifySpan on every span.
+func TestSpanModifierSetTagsConcurrentWithModifySpan(t *testing.T) {
+	sm := &spanModifier{ddOrigin: "lambda"}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 1000; i++ {
+			sm.SetTags(map[string]string{"lambda_microvm_id": "vm-1"})
+		}
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 1000; i++ {
+			span := &pb.Span{Meta: map[string]string{}}
+			sm.ModifySpan(&pb.TraceChunk{}, span)
+		}
+	}()
+	wg.Wait()
+
+	assert.Equal(t, map[string]string{"lambda_microvm_id": "vm-1"}, *sm.tags.Load())
+}

pkg/serverless/trace/span_modifier.go

@@ -7,6 +7,8 @@
 package trace
 
 import (
+	"go.uber.org/atomic"
+
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/trace/traceutil"
 )
@@ -16,7 +18,7 @@ const (
 )
 
 type spanModifier struct {
-	tags     map[string]string
+	tags     atomic.Pointer[map[string]string]
 	ddOrigin string
 }
 
@@ -26,9 +28,17 @@ func (s *spanModifier) ModifySpan(_ *pb.TraceChunk, span *pb.Span) {
 	if origin := span.Meta[ddOriginTagName]; origin == "" {
 		traceutil.SetMeta(span, ddOriginTagName, s.ddOrigin)
 	}
+	if tags := s.tags.Load(); tags != nil {
+		for k, v := range *tags {
+			if k == ddOriginTagName {
+				continue
+			}
+			traceutil.SetMeta(span, k, v)
+		}
+	}
 }
 
 // SetTags sets the tags to be used by the span modifier.
 func (s *spanModifier) SetTags(tags map[string]string) {
-	s.tags = tags
+	s.tags.Store(&tags)
 }

pkg/serverless/trace/trace.go

@@ -43,6 +43,7 @@ type ServerlessTraceAgent interface {
 	Flush()
 	Process(p *api.Payload)
 	SetTags(map[string]string)
+	UpdateRuntimeTags(map[string]string)
 	SetTargetTPS(float64)
 	SetSpanModifier(agent.SpanModifier)
 	GetSpanModifier() agent.SpanModifier
@@ -204,6 +205,15 @@ func (t *serverlessTraceAgent) SetTags(tags map[string]string) {
 	}
 }
 
+// UpdateRuntimeTags updates the tags applied by the span modifier only. Unlike
+// SetTags, it does not touch the trace agent's GlobalTags, so it is safe to
+// call concurrently with span processing.
+func (t *serverlessTraceAgent) UpdateRuntimeTags(tags map[string]string) {
+	if tagger, ok := t.ta.SpanModifier.(taggable); ok {
+		tagger.SetTags(tags)
+	}
+}
+
 // Stop cancels the trace agent's context and waits for its Run loop to finish.
 // The Run loop handles the full shutdown sequence: draining in-flight traces,
 // flushing stats producers, sending buffered data to the network, and stopping
@@ -300,6 +310,7 @@ func (t noopTraceAgent) Stop()                               {}
 func (t noopTraceAgent) Flush()                              {}
 func (t noopTraceAgent) Process(*api.Payload)                {}
 func (t noopTraceAgent) SetTags(map[string]string)           {}
+func (t noopTraceAgent) UpdateRuntimeTags(map[string]string) {}
 func (t noopTraceAgent) SetTargetTPS(float64)                {}
 func (t noopTraceAgent) SetSpanModifier(agent.SpanModifier)  {}
 func (t noopTraceAgent) GetSpanModifier() agent.SpanModifier { return nil }

pkg/serverless/trace/trace_test.go

@@ -8,6 +8,7 @@
 package trace
 
 import (
+	"context"
 	"errors"
 	"os"
 	"path/filepath"
@@ -19,13 +20,31 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
+	gzip "github.com/DataDog/datadog-agent/comp/trace/compression/impl-gzip"
 	configmock "github.com/DataDog/datadog-agent/pkg/config/mock"
 	pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup"
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
+	"github.com/DataDog/datadog-agent/pkg/trace/agent"
 	"github.com/DataDog/datadog-agent/pkg/trace/config"
+	"github.com/DataDog/datadog-agent/pkg/trace/telemetry"
 	"github.com/DataDog/datadog-agent/pkg/trace/testutil"
+
+	"github.com/DataDog/datadog-go/v5/statsd"
 )
 
+// newTestServerlessTraceAgent builds a minimal serverlessTraceAgent backed by
+// a real *agent.Agent and spanModifier, without starting the agent's Run loop.
+func newTestServerlessTraceAgent(t *testing.T) (*serverlessTraceAgent, *config.AgentConfig) {
+	t.Helper()
+	cfg := config.New()
+	cfg.Endpoints[0].APIKey = "test"
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	ta := agent.NewAgent(ctx, cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, gzip.NewComponent())
+	ta.SpanModifier = &spanModifier{ddOrigin: "lambda"}
+	return &serverlessTraceAgent{ta: ta, cancel: cancel}, cfg
+}
+
 func setupTraceAgentTest(t *testing.T) {
 	// ensure a free port is used for starting the trace agent
 	port, err := testutil.FindTCPPort()
@@ -247,3 +266,53 @@ func TestServerlessTraceAgentDisableTraceStats(t *testing.T) {
 		})
 	}
 }
+
+// TestServerlessTraceAgentSetTagsUpdatesGlobalTagsAndSpanModifier verifies
+// that the synchronous SetTags path (used once at startup, before the trace
+// agent starts processing spans) still updates both GlobalTags and the span
+// modifier, as it did before this fix.
+func TestServerlessTraceAgentSetTagsUpdatesGlobalTagsAndSpanModifier(t *testing.T) {
+	sta, cfg := newTestServerlessTraceAgent(t)
+
+	sta.SetTags(map[string]string{"lambda_microvm_id": "vm-1"})
+
+	assert.Equal(t, map[string]string{"lambda_microvm_id": "vm-1"}, cfg.GlobalTags)
+
+	sm, ok := sta.ta.SpanModifier.(*spanModifier)
+	require.True(t, ok)
+	got := sm.tags.Load()
+	require.NotNil(t, got)
+	assert.Equal(t, map[string]string{"lambda_microvm_id": "vm-1"}, *got)
+}
+
+// TestServerlessTraceAgentUpdateRuntimeTagsDoesNotTouchGlobalTags is the core
+// regression test for the fix: UpdateRuntimeTags (used by MicroVM's async
+// /run hook) must only update the span modifier and must never write to
+// GlobalTags, since GlobalTags is read unsynchronized by the trace agent's
+// span-processing hot path.
+func TestServerlessTraceAgentUpdateRuntimeTagsDoesNotTouchGlobalTags(t *testing.T) {
+	sta, cfg := newTestServerlessTraceAgent(t)
+	cfg.GlobalTags = map[string]string{"env": "prod"}
+
+	sta.UpdateRuntimeTags(map[string]string{"lambda_microvm_id": "vm-2"})
+
+	assert.Equal(t, map[string]string{"env": "prod"}, cfg.GlobalTags)
+
+	sm, ok := sta.ta.SpanModifier.(*spanModifier)
+	require.True(t, ok)
+	got := sm.tags.Load()
+	require.NotNil(t, got)
+	assert.Equal(t, map[string]string{"lambda_microvm_id": "vm-2"}, *got)
+}
+
+// TestServerlessTraceAgentUpdateRuntimeTagsNoSpanModifier verifies
+// UpdateRuntimeTags is safe to call when SpanModifier doesn't implement
+// taggable (e.g. nil, or some other SpanModifier set via SetSpanModifier).
+func TestServerlessTraceAgentUpdateRuntimeTagsNoSpanModifier(t *testing.T) {
+	sta, _ := newTestServerlessTraceAgent(t)
+	sta.ta.SpanModifier = nil
+
+	assert.NotPanics(t, func() {
+		sta.UpdateRuntimeTags(map[string]string{"lambda_microvm_id": "vm-3"})
+	})
+}