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

cmd/serverless-init/cloudservice/microvm.go

@@ -115,6 +115,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.GetEnhancedMetricTags and CloudRun.GetEnhancedMetricTags use for
+// their own per-instance usage tag.
+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

@@ -128,6 +128,82 @@ 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, nil,
+		(&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, nil,
+		(&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)
+	})
+
+	port := l.Addr().(*net.TCPAddr).Port
+	runPath := "/aws/lambda-microvms/runtime/v1/run"
+	body := strings.NewReader(`{"microvmId":"vm-abc123"}`)
+	resp, err := http.Post("http://127.0.0.1:"+strconv.Itoa(port)+runPath, "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

@@ -310,14 +310,93 @@ func TestCollectorSendsUsageMetricOnCgroupFailure(t *testing.T) {
 
 func TestNewCollectorNilMetricAgent(t *testing.T) {
 	// Untyped nil interface.
-	c, err := NewCollector(nil, metrics.MetricSourceGoogleCloudRunEnhanced, "gcp.run.", "instance", time.Second)
+	c, err := NewCollector(nil, metrics.MetricSourceGoogleCloudRunEnhanced, "gcp.run.", "instance", time.Second, nil)
 	assert.Nil(t, c)
 	assert.Error(t, err)
 
 	// Typed nil implementing EnhancedMetricSender, which is what main.go passes
 	// when metricAgent is a nil *ServerlessMetricAgent (use_dogstatsd disabled).
 	var typedNil *mockEnhancedMetricSender
-	c, err = NewCollector(typedNil, metrics.MetricSourceGoogleCloudRunEnhanced, "gcp.run.", "instance", time.Second)
+	c, err = NewCollector(typedNil, metrics.MetricSourceGoogleCloudRunEnhanced, "gcp.run.", "instance", time.Second, nil)
 	assert.Nil(t, c)
 	assert.Error(t, err)
 }
+
+// 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

@@ -318,6 +318,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

@@ -158,6 +158,26 @@ func TestHandleRun_OversizedBody_Returns500(t *testing.T) {
 	assert.NotContains(t, emitter.getEmitted(), runMetricName, "run metric must not be emitted when the body exceeds the cap")
 }
 
+// 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

@@ -434,6 +434,14 @@ func run(
 	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,
@@ -548,9 +556,14 @@ func setup(
 
 	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

@@ -150,6 +150,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")