fix(serverless-init): attach per-instance tag to MicroVM enhanced usage metric

Codex flagged on PR #53093 (review comment
#53093 (comment))
that MicroVM's enhanced usage metric (aws.lambda.microvm.instance) is
emitted with the same static tag set for every MicroVM booted from a given
image. MicroVM.GetEnhancedMetricTags returns Usage tags at startup, before
the MicroVM's instance ID is known (it only becomes available once the /run
lifecycle hook fires), and nothing ever adds it afterward. Under normal
autoscaling, with multiple concurrent MicroVMs from the same image, their
usage samples become indistinguishable from each other.

Rather than mutating shared, periodically-read state on
ServerlessMetricAgent (which would need a lock/atomic and broke go vet's
copylocks check, since CloudService.Shutdown takes ServerlessMetricAgent by
value across every cloud service), this attaches the tag at the point of
emission instead:

- lifecycle.Server already tracks the instance ID race-free
  (instanceID *atomic.String, captured in handleRun). Added InstanceID(),
  a nil-safe accessor, so callers outside the lifecycle package can read it.
- MicroVM.CurrentUsageMetricTags() turns that into an "instance:<id>" tag,
  or nil before /run has fired.
- enhanced-metrics.Collector gained an optional usageMetricTagsFunc,
  invoked on every collection tick and passed through the existing
  AddEnhancedUsageMetric(..., extraTags ...string) parameter — no new
  shared mutable state required.
- main.go duck-types cloudService against a local usageMetricTagProvider
  interface to wire this hook; every other cloud service doesn't implement
  it, so usageMetricTagsFunc stays nil and their usage metrics are
  unaffected.

This keeps the change scoped to MicroVM and the (nil-safe, additive)
collector plumbing — no other cloud service, ServerlessMetricAgent, or the
Shutdown interface needed to change.

Unit tests:
  dda inv test --targets=./cmd/serverless-init/...

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

Author: litianningdatadog

cmd/serverless-init/cloudservice/microvm.go

@@ -116,6 +116,19 @@ func (m *MicroVM) GetMetricPrefix() string { return MicroVMPrefix }
 // GetUsageMetricSuffix returns the usage metric suffix.
 func (m *MicroVM) GetUsageMetricSuffix() string { return MicroVMUsageMetricSuffix }
 
+// CurrentUsageMetricTags returns the dynamic tags to attach to the enhanced
+// usage metric on each periodic emission: the per-instance tag, once known
+// from /run, or nil before that (see GetEnhancedMetricTags's doc comment on
+// why Usage never carries it directly). "instance" matches the tag key
+// AppService/CloudRun use for their own per-instance usage tag
+// (appservice.go:73, cloudrun.go:156).
+func (m *MicroVM) CurrentUsageMetricTags() []string {
+	if id := m.server.InstanceID(); id != "" {
+		return []string{"instance:" + id}
+	}
+	return nil
+}
+
 // GetOrigin returns the origin tag value.
 func (m *MicroVM) GetOrigin() string { return MicroVMOrigin }
 

cmd/serverless-init/cloudservice/microvm_test.go

@@ -129,6 +129,81 @@ func TestMicroVMGetEnhancedMetricTagsMissingARN(t *testing.T) {
 	assert.Equal(t, result.Base["resource_id"], result.Usage["resource_id"])
 }
 
+// TestMicroVM_CurrentUsageMetricTags_NilServer_ReturnsNil verifies that
+// CurrentUsageMetricTags is safe to call before Init (m.server is nil) — the
+// enhanced-metrics collector may call it before the lifecycle server exists.
+func TestMicroVM_CurrentUsageMetricTags_NilServer_ReturnsNil(t *testing.T) {
+	m := &MicroVM{}
+	assert.Nil(t, m.CurrentUsageMetricTags())
+}
+
+// TestMicroVM_CurrentUsageMetricTags_BeforeRun_ReturnsNil verifies that no
+// instance tag is produced before /run fires, matching GetEnhancedMetricTags'
+// documented behavior that instance_id is unknown until then.
+func TestMicroVM_CurrentUsageMetricTags_BeforeRun_ReturnsNil(t *testing.T) {
+	metricAgent := &serverlessMetrics.ServerlessMetricAgent{}
+	srv := lifecycle.NewServer(
+		0,
+		metricAgent, &noopTraceAgent{}, &noopLogsFlusher{},
+		metricAgent, metricAgent,
+		(&MicroVM{}).GetSource(),
+		time.Second,
+		lifecycle.NewNoopChildHandle(),
+		nil, // no forwarder
+		nil, // no heartbeat
+	)
+	m := &MicroVM{server: srv}
+	assert.Nil(t, m.CurrentUsageMetricTags())
+}
+
+// TestMicroVM_CurrentUsageMetricTags_AfterRun_ReturnsInstanceTag verifies the
+// end-to-end path: once /run has captured the MicroVM instance ID,
+// CurrentUsageMetricTags returns the "instance:<id>" tag the enhanced-metrics
+// collector attaches to the usage metric on every subsequent tick.
+func TestMicroVM_CurrentUsageMetricTags_AfterRun_ReturnsInstanceTag(t *testing.T) {
+	metricAgent := &serverlessMetrics.ServerlessMetricAgent{}
+	srv := lifecycle.NewServer(
+		0,
+		metricAgent, &noopTraceAgent{}, &noopLogsFlusher{},
+		metricAgent, metricAgent,
+		(&MicroVM{}).GetSource(),
+		time.Second,
+		lifecycle.NewNoopChildHandle(),
+		nil, // no forwarder
+		nil, // no heartbeat
+	)
+	l, err := srv.Listen()
+	require.NoError(t, err)
+	go srv.Serve(l)
+	t.Cleanup(func() {
+		shutCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = srv.Stop(shutCtx)
+	})
+
+	launchPath := "/aws/lambda-microvms/runtime/v1/run"
+	body := strings.NewReader(`{"microvmId":"vm-abc123"}`)
+	resp, err := http.Post("http://"+l.Addr().String()+launchPath, "application/json", body)
+	require.NoError(t, err)
+	require.NoError(t, resp.Body.Close())
+
+	m := &MicroVM{server: srv}
+	assert.Equal(t, []string{"instance:vm-abc123"}, m.CurrentUsageMetricTags())
+}
+
+// TestMicroVM_SatisfiesUsageMetricTagProvider is a compile-time guard: main.go
+// duck-types cloudService against an unexported usageMetricTagProvider
+// interface with this exact method set to wire the enhanced-metrics
+// collector's dynamic tag hook. If CurrentUsageMetricTags' signature ever
+// drifts, that type assertion silently stops matching instead of failing to
+// compile — this pins the method set so such drift shows up here instead.
+func TestMicroVM_SatisfiesUsageMetricTagProvider(t *testing.T) {
+	var m any = &MicroVM{}
+	provider, ok := m.(interface{ CurrentUsageMetricTags() []string })
+	require.True(t, ok, "*MicroVM must implement CurrentUsageMetricTags() []string")
+	assert.NotPanics(t, func() { provider.CurrentUsageMetricTags() })
+}
+
 // Compile-time guard: *MicroVM must satisfy the CloudService interface,
 // including the new Run method.
 var _ CloudService = (*MicroVM)(nil)

cmd/serverless-init/enhanced-metrics/collector.go

@@ -78,10 +78,15 @@ type Collector struct {
 	usageMetricSuffix  string
 	// Previous stats for rate calculation
 	previousRateStats ServerlessRateStats
+	// usageMetricTagsFunc, when non-nil, is called on every collection tick to
+	// obtain extra tags for the enhanced usage metric — e.g. MicroVM's
+	// per-instance tag, which is only known once the /run lifecycle hook
+	// fires. nil for every cloud service that has no such dynamic tag.
+	usageMetricTagsFunc func() []string
 }
 
-// NewCollector creates a new Collector
-func NewCollector(metricAgent EnhancedMetricSender, metricSource metrics.MetricSource, metricPrefix string, usageMetricSuffix string, collectionInterval time.Duration) (*Collector, error) {
+// NewCollector creates a new Collector. usageMetricTagsFunc may be nil.
+func NewCollector(metricAgent EnhancedMetricSender, metricSource metrics.MetricSource, metricPrefix string, usageMetricSuffix string, collectionInterval time.Duration, usageMetricTagsFunc func() []string) (*Collector, error) {
 	if metricAgent == nil || reflect.ValueOf(metricAgent).IsNil() {
 		return nil, errors.New("metricAgent cannot be nil")
 	}
@@ -92,13 +97,14 @@ func NewCollector(metricAgent EnhancedMetricSender, metricSource metrics.MetricS
 	}
 
 	return &Collector{
-		metricAgent:        metricAgent,
-		metricSource:       metricSource,
-		cgroupReader:       cgroupReader,
-		collectionInterval: collectionInterval,
-		metricPrefix:       metricPrefix + "enhanced.",
-		usageMetricSuffix:  usageMetricSuffix,
-		previousRateStats:  NullServerlessRateStats,
+		metricAgent:         metricAgent,
+		metricSource:        metricSource,
+		cgroupReader:        cgroupReader,
+		collectionInterval:  collectionInterval,
+		metricPrefix:        metricPrefix + "enhanced.",
+		usageMetricSuffix:   usageMetricSuffix,
+		previousRateStats:   NullServerlessRateStats,
+		usageMetricTagsFunc: usageMetricTagsFunc,
 	}, nil
 }
 
@@ -151,7 +157,11 @@ func (c *Collector) collect() {
 
 	// Always send the usage metric, regardless of cgroup collection success.
 	if c.usageMetricSuffix != "" {
-		c.metricAgent.AddEnhancedUsageMetric(c.metricPrefix+c.usageMetricSuffix, 1, c.metricSource, timestamp)
+		var extraTags []string
+		if c.usageMetricTagsFunc != nil {
+			extraTags = c.usageMetricTagsFunc()
+		}
+		c.metricAgent.AddEnhancedUsageMetric(c.metricPrefix+c.usageMetricSuffix, 1, c.metricSource, timestamp, extraTags...)
 	}
 
 	if err := c.cgroupReader.RefreshCgroups(0); err != nil {

cmd/serverless-init/enhanced-metrics/collector_test.go

@@ -307,3 +307,82 @@ func TestCollectorSendsUsageMetricOnCgroupFailure(t *testing.T) {
 		metrics.MetricSourceGoogleCloudRunEnhanced, mock.Anything, mock.Anything)
 	mockAgent.AssertNotCalled(t, "AddEnhancedMetric", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 }
+
+// TestCollectorUsageMetricIncludesDynamicTags verifies that when
+// usageMetricTagsFunc is set (MicroVM's use case: attaching the per-instance
+// tag once known from /run), its return value is forwarded as extraTags on
+// every AddEnhancedUsageMetric call.
+func TestCollectorUsageMetricIncludesDynamicTags(t *testing.T) {
+	mockAgent := new(mockEnhancedMetricSender)
+	mockAgent.On("AddEnhancedUsageMetric", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+	mockReader := &mockCgroupReader{refreshErr: errors.New("cgroup failure"), version: 1}
+
+	c := &Collector{
+		metricAgent:         mockAgent,
+		metricSource:        metrics.MetricSourceAWSMicroVMEnhanced,
+		cgroupReader:        mockReader,
+		metricPrefix:        "aws.lambda.microvm.enhanced.",
+		usageMetricSuffix:   "instance",
+		previousRateStats:   NullServerlessRateStats,
+		usageMetricTagsFunc: func() []string { return []string{"instance:vm-abc123"} },
+	}
+
+	c.collect()
+
+	mockAgent.AssertCalled(t, "AddEnhancedUsageMetric",
+		"aws.lambda.microvm.enhanced.instance", float64(1),
+		metrics.MetricSourceAWSMicroVMEnhanced, mock.Anything, []string{"instance:vm-abc123"})
+}
+
+// TestCollectorUsageMetricNilTagsFuncSendsNoExtraTags verifies that when
+// usageMetricTagsFunc is nil (every cloud service except MicroVM), no extra
+// tags are added to the usage metric — pinning today's behavior.
+func TestCollectorUsageMetricNilTagsFuncSendsNoExtraTags(t *testing.T) {
+	mockAgent := new(mockEnhancedMetricSender)
+	mockAgent.On("AddEnhancedUsageMetric", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+	mockReader := &mockCgroupReader{refreshErr: errors.New("cgroup failure"), version: 1}
+
+	c := &Collector{
+		metricAgent:       mockAgent,
+		metricSource:      metrics.MetricSourceGoogleCloudRunEnhanced,
+		cgroupReader:      mockReader,
+		metricPrefix:      "gcp.run.container.enhanced.",
+		usageMetricSuffix: "instance",
+		previousRateStats: NullServerlessRateStats,
+	}
+
+	c.collect()
+
+	mockAgent.AssertCalled(t, "AddEnhancedUsageMetric",
+		"gcp.run.container.enhanced.instance", float64(1),
+		metrics.MetricSourceGoogleCloudRunEnhanced, mock.Anything, []string(nil))
+}
+
+// TestNewCollectorWiresUsageMetricTagsFunc verifies that NewCollector stores
+// the provided usageMetricTagsFunc on the returned Collector so collect()
+// picks it up on every tick.
+func TestNewCollectorWiresUsageMetricTagsFunc(t *testing.T) {
+	mockAgent := new(mockEnhancedMetricSender)
+	tagsFunc := func() []string { return []string{"instance:vm-abc123"} }
+
+	c, err := NewCollector(mockAgent, metrics.MetricSourceAWSMicroVMEnhanced, "aws.lambda.microvm.", "instance", time.Second, tagsFunc)
+
+	assert.NoError(t, err)
+	if assert.NotNil(t, c.usageMetricTagsFunc) {
+		assert.Equal(t, []string{"instance:vm-abc123"}, c.usageMetricTagsFunc())
+	}
+}
+
+// TestNewCollectorNilUsageMetricTagsFuncIsAccepted verifies that
+// NewCollector accepts a nil usageMetricTagsFunc — the case for every cloud
+// service except MicroVM.
+func TestNewCollectorNilUsageMetricTagsFuncIsAccepted(t *testing.T) {
+	mockAgent := new(mockEnhancedMetricSender)
+
+	c, err := NewCollector(mockAgent, metrics.MetricSourceGoogleCloudRunEnhanced, "gcp.run.container.", "instance", time.Second, nil)
+
+	assert.NoError(t, err)
+	assert.Nil(t, c.usageMetricTagsFunc)
+}

cmd/serverless-init/enhanced-metrics/collector_unsupported.go

@@ -20,7 +20,7 @@ type EnhancedMetricSender interface{}
 
 type Collector struct{}
 
-func NewCollector(_ EnhancedMetricSender, _ metrics.MetricSource, _ string, _ string, _ time.Duration) (*Collector, error) {
+func NewCollector(_ EnhancedMetricSender, _ metrics.MetricSource, _ string, _ string, _ time.Duration, _ func() []string) (*Collector, error) {
 	return nil, errors.New("Collector is only supported on Linux")
 }
 

cmd/serverless-init/lifecycle/server.go

@@ -305,6 +305,18 @@ func (s *Server) Child() *Child {
 // Exposed for white-box tests in external packages; not part of the stable API.
 func (s *Server) Heartbeat() *Heartbeat { return s.heartbeat }
 
+// InstanceID returns the MicroVM instance ID captured from /run, or "" if
+// /run has not fired yet (or the server is nil). Lets callers outside this
+// package — e.g. the enhanced-metrics collector — attach the current
+// per-instance tag to metrics emitted after Init, without this package
+// needing to know about metric-agent internals.
+func (s *Server) InstanceID() string {
+	if s == nil {
+		return ""
+	}
+	return s.instanceID.Load()
+}
+
 func (s *Server) handler() http.Handler {
 	mux := http.NewServeMux()
 	mux.HandleFunc(postReady, s.handleReady)

cmd/serverless-init/lifecycle/server_test.go

@@ -119,6 +119,26 @@ func TestHandleRunParsesInstanceID(t *testing.T) {
 	assert.Equal(t, "vm-abc123", id, "instance ID must be stored on the server for lifecycle metric tags")
 }
 
+// TestInstanceID_EmptyBeforeRun verifies that InstanceID returns "" before
+// /run fires, and the captured ID afterward — the accessor the enhanced
+// metrics collector uses to attach a per-instance tag to the usage metric.
+func TestInstanceID_EmptyBeforeRun(t *testing.T) {
+	srv, _, _, _, _, _ := newTestServer()
+	assert.Empty(t, srv.InstanceID(), "InstanceID must be empty before /run fires")
+
+	body := strings.NewReader(`{"microvmId":"vm-abc123"}`)
+	srv.handleRun(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, pathRun, body))
+
+	assert.Equal(t, "vm-abc123", srv.InstanceID())
+}
+
+// TestInstanceID_NilServer verifies InstanceID is safe to call on a nil
+// *Server, mirroring the existing nil-safety of Child().
+func TestInstanceID_NilServer(t *testing.T) {
+	var srv *Server
+	assert.Empty(t, srv.InstanceID())
+}
+
 // TestHandleRunWithForwarderParsesInstanceID verifies that when a forwarder is
 // configured, /run still decodes the MicroVM instance ID from the request body
 // before delegating to handleWithForwarder. Without the decode-then-restore fix, the

cmd/serverless-init/main.go

@@ -124,6 +124,14 @@ func run(secretComp secrets.Component, delegatedAuthComp delegatedauth.Component
 	return err
 }
 
+// usageMetricTagProvider is satisfied by cloud services whose enhanced usage
+// metric needs a tag that isn't known until after Init — e.g. MicroVM's
+// per-instance tag, only available once the /run lifecycle hook fires.
+// Ignored (via the type assertion in setup) by every other cloud service.
+type usageMetricTagProvider interface {
+	CurrentUsageMetricTags() []string
+}
+
 func setup(secretComp secrets.Component, delegatedAuthComp delegatedauth.Component, _ mode.Conf, tagger tagger.Component, compression logscompression.Component, hostname hostnameinterface.Component) (cloudservice.CloudService, *serverlessInitLog.Config, *cloudservice.TracingContext, *metrics.ServerlessMetricAgent, logsAgent.ServerlessLogsAgent, *enhancedmetrics.Collector, bool) {
 	tracelog.SetLogger(log.NewWrapper(3))
 
@@ -235,9 +243,14 @@ func setup(secretComp secrets.Component, delegatedAuthComp delegatedauth.Compone
 
 	setupOtlpAgent(metricAgent, tagger)
 
+	var usageMetricTagsFunc func() []string
+	if p, ok := cloudService.(usageMetricTagProvider); ok {
+		usageMetricTagsFunc = p.CurrentUsageMetricTags
+	}
+
 	var enhancedMetricsCollector *enhancedmetrics.Collector
 	if enhancedMetricsEnabled {
-		enhancedMetricsCollector, err = enhancedmetrics.NewCollector(metricAgent, cloudService.GetSource(), cloudService.GetMetricPrefix(), cloudService.GetUsageMetricSuffix(), 3*time.Second)
+		enhancedMetricsCollector, err = enhancedmetrics.NewCollector(metricAgent, cloudService.GetSource(), cloudService.GetMetricPrefix(), cloudService.GetUsageMetricSuffix(), 3*time.Second, usageMetricTagsFunc)
 		if err != nil {
 			log.Warnf("Failed to initialize enhanced metrics collector: %v", err)
 		} else {

cmd/serverless-init/main_test.go

@@ -190,6 +190,37 @@ func TestBaseTraceTagsComputedFromTagConfigTags(t *testing.T) {
 	assert.NotEmpty(t, baseTraceTags)
 }
 
+// TestUsageMetricTagProvider_MicroVMSatisfiesInterface verifies that
+// *cloudservice.MicroVM implements usageMetricTagProvider — the interface
+// setup() type-asserts cloudService against to wire the enhanced-metrics
+// collector's dynamic per-instance usage tag. If CurrentUsageMetricTags'
+// signature ever drifts, the assertion in setup() would silently stop
+// matching (no compile error) instead of failing loudly; this test is what
+// would catch that.
+func TestUsageMetricTagProvider_MicroVMSatisfiesInterface(t *testing.T) {
+	var cloudService cloudservice.CloudService = &cloudservice.MicroVM{}
+	_, ok := cloudService.(usageMetricTagProvider)
+	assert.True(t, ok, "*MicroVM must satisfy usageMetricTagProvider so setup() wires its dynamic usage-metric tags")
+}
+
+// TestUsageMetricTagProvider_OtherServicesDoNotSatisfy documents that cloud
+// services with no dynamic usage-metric tag are intentionally left out of
+// usageMetricTagProvider; setup()'s type assertion falls through to a nil
+// usageMetricTagsFunc for them, and NewCollector treats nil as "no extra tags".
+func TestUsageMetricTagProvider_OtherServicesDoNotSatisfy(t *testing.T) {
+	services := []cloudservice.CloudService{
+		&cloudservice.LocalService{},
+		&cloudservice.AppService{},
+		&cloudservice.CloudRun{},
+		&cloudservice.CloudRunJobs{},
+		&cloudservice.ContainerApp{},
+	}
+	for _, svc := range services {
+		_, ok := svc.(usageMetricTagProvider)
+		assert.False(t, ok, "%T must not satisfy usageMetricTagProvider — it has no dynamic usage-metric tag", svc)
+	}
+}
+
 // TestSetupOtlpAgentNoPanic ensures setupOtlpAgent does not panic when OTLP is enabled.
 func TestSetupOtlpAgentNoPanic(t *testing.T) {
 	t.Setenv("DD_OTLP_CONFIG_LOGS_ENABLED", "true")