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

Codex flagged this on PR #53093 ([review
comment](#53093 (comment))):
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` computes the `Usage` tag set once at
startup — before the MicroVM's instance ID is known, since the platform
only reveals it via the `/run` lifecycle hook — and nothing ever updates
it afterward. Under normal autoscaling, with multiple concurrent
MicroVMs running from the same image, their periodic usage samples are
indistinguishable from one another: instance-level usage collapses into
a single series.

The obvious fix — mutate the frozen tag set on `ServerlessMetricAgent`
once `/run` reveals the instance ID, mirroring the existing
`LogsTagSetter`/`TraceTagSetter` pattern — turned out to require a
lock/atomic on `ServerlessMetricAgent`. That struct is passed **by
value** into `CloudService.Shutdown` across every cloud service
(`AppService`, `CloudRun`, `CloudRunJobs`, `ContainerApp`,
`LocalService`, `MicroVM`), so adding any lock-bearing field trips `go
vet`'s copylocks check everywhere, not just for MicroVM.

Instead, this attaches the tag at the point of *emission* rather than
mutating shared state:

- `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.
- `enhancedmetrics.Collector` gained an optional `usageMetricTagsFunc`,
invoked on every collection tick and passed straight through the
existing `AddEnhancedUsageMetric(..., extraTags ...string)` parameter —
no new shared mutable state needed.
- `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` for
them and their usage metrics are unaffected.

This keeps the change scoped to MicroVM plus small, nil-safe, additive
plumbing in the shared collector — no other cloud service,
`ServerlessMetricAgent`, or the `Shutdown` interface needed to change.

- [x] `dda inv test --targets=./cmd/serverless-init/...` — all 281 tests
pass (12 new: `InstanceID()` nil/before/after `/run`,
`CurrentUsageMetricTags()` nil-server/before-run/after-run end-to-end,
collector dynamic-tag forwarding + nil-func regression guard,
`NewCollector` wiring, and the `usageMetricTagProvider` type-assertion
pinning `*MicroVM` in vs. every other cloud service out)
- [x] `gofmt -l` clean on all changed files
- [x] `dda inv linter.go --targets=./cmd/serverless-init/...` — 0 issues
- [ ]
`cmd/serverless-init/enhanced-metrics/collector.go`/`collector_test.go`
carry a pre-existing `//go:build linux` tag; this session ran on macOS
without a Linux cross-toolchain, so those specific new tests are
syntax-checked (`gofmt -e`) and reviewed but not yet executed — first
real run will be in Linux CI

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: litianningdatadog

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,80 @@ 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
+	)
+	require.NoError(t, srv.ListenAndServe(nil))
+	t.Cleanup(func() {
+		shutCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = srv.Stop(shutCtx)
+	})
+
+	port := srv.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/forwarder.go

@@ -27,20 +27,34 @@ import (
 const defaultMaxResponseBodyBytes int64 = 1 << 20
 
 const (
+	// Per AWS, Run/Resume/Suspend/Terminate: 1 second
+	// Ready/Validate: 30 seconds
 	defaultForwardTimeout  = 1 * time.Second
-	defaultReadyTimeout    = 60 * time.Second
-	defaultValidateTimeout = 1 * time.Second
+	defaultReadyTimeout    = 30 * time.Second
+	defaultValidateTimeout = 30 * time.Second
 )
 
+// dialCheckTimeout bounds the single TCP dial attempt PassThroughWaiting uses
+// to check user-app reachability for /ready and /validate. Per the hook
+// contract, the platform — not the hook — owns the retry loop for these two
+// hooks, so the check must answer fast rather than block until the app is up.
+// A refused loopback connection fails in microseconds (the kernel sends RST
+// immediately), so this timeout only matters for a slow/hung connect (e.g.
+// host scheduling jitter on an oversubscribed MicroVM host). 200ms is
+// generous enough to absorb that jitter while staying well inside the "answer
+// fast" contract, and costs at most one extra platform retry (a 503) on the
+// rare call where it's hit.
+const dialCheckTimeout = 200 * time.Millisecond
+
 // Forwarder POSTs lifecycle hooks to the user app. It is constructed only
 // when DD_AWS_MICROVM_USER_APP_PORT is set and we're in init-container
 // mode + MicroVM origin.
 type Forwarder struct {
 	target               string        // e.g. "http://127.0.0.1:8080"
 	client               *http.Client  // shared; no client-level Timeout (per-call deadlines via ctx)
 	forwardTimeout       time.Duration // default 1s, used for suspend/terminate/run/resume
-	readyTimeout         time.Duration // default 60s, used for /ready
-	validateTimeout      time.Duration // default 1s, used for /validate
+	readyTimeout         time.Duration // default 30s, used for /ready
+	validateTimeout      time.Duration // default 30s, used for /validate
 	maxResponseBodyBytes int64         // default defaultMaxResponseBodyBytes; cap on user-app body surfaced to platform
 }
 
@@ -83,63 +97,58 @@ func NewForwarder(port int, forwardTimeout, readyTimeout, validateTimeout time.D
 	}
 }
 
-// PassThroughWaiting waits for the user app to accept TCP connections bounded
-// by timeout, then forwards the request and mirrors the response. Used for:
+// PassThroughWaiting checks user-app reachability with a single fast TCP dial
+// (bounded by dialCheckTimeout), then forwards the request bounded by timeout
+// and mirrors the response. Used for:
 //   - /ready    ("I am booted and ready to be snapshotted"): the platform
-//     retries on non-200, so the TCP wait absorbs the startup race.
+//     retries on non-200 until its own configured timeout, so the hook
+//     answers fast rather than blocking on the startup race.
 //   - /validate ("I was resumed from a snapshot and everything is good"): the
-//     TCP wait handles a crash-then-restart between resume and this
-//     call.
+//     same fast-answer contract applies to the crash-then-restart window
+//     between resume and this call.
 //
-// Body is buffered before the TCP wait so the bytes survive waitForUserApp.
-// Deadline exceeded maps to 504. Body Close contract documented on PassThrough.
+// Per the /ready and /validate hook contract, only 200 and 503 are meaningful
+// responses — the platform retries on 503 until its own configured timeout,
+// while any other non-200 (including 504) fails the build. So unlike
+// PassThrough, an unreachable app or a deadline exceeded here always maps to
+// 503, never 504. Body Close contract documented on PassThrough.
 func (f *Forwarder) PassThroughWaiting(timeout time.Duration, path string, headers http.Header, body io.Reader) *http.Response {
 	var bodyBytes []byte
 	if body != nil {
 		var err error
-		// Read the full inbound body before the TCP wait. A read error is a
-		// server-side failure (network, OS, memory) — not a client mistake —
-		// so return 500 rather than 400. We still forward nothing to the user
-		// app to avoid passing a partial body (which could make it answer
-		// /validate "healthy" off incomplete data).
+		// Read the full inbound body before the reachability check. A read
+		// error is a server-side failure (network, OS, memory) — not a client
+		// mistake — so return 500 rather than 400. We still forward nothing to
+		// the user app to avoid passing a partial body (which could make it
+		// answer /validate "healthy" off incomplete data).
 		if bodyBytes, err = io.ReadAll(body); err != nil {
 			return statusOnlyResponse(http.StatusInternalServerError)
 		}
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	if err := f.waitForUserApp(ctx); err != nil {
-		cancel()
-		return statusOnlyResponse(mapErrToStatus(err))
+	if !f.reachable(dialCheckTimeout) {
+		return statusOnlyResponse(http.StatusServiceUnavailable)
 	}
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	resp, err := f.do(ctx, path, headers, bytes.NewReader(bodyBytes))
 	if err != nil {
 		cancel()
-		return statusOnlyResponse(mapErrToStatus(err))
+		return statusOnlyResponse(http.StatusServiceUnavailable)
 	}
 	resp.Body = wrapResponseBody(resp.Body, f.maxResponseBodyBytes, cancel)
 	return resp
 }
 
-// waitForUserApp polls the user app's TCP port until a connection succeeds or
-// ctx is cancelled. Returns nil when the port is reachable, ctx.Err() when
-// the deadline is exceeded. Polls every 50ms.
-func (f *Forwarder) waitForUserApp(ctx context.Context) error {
+// reachable performs a single TCP dial attempt to the user app, bounded by
+// timeout. No polling/retrying: per the hook contract, the platform (not the
+// hook) owns the retry loop for /ready and /validate.
+func (f *Forwarder) reachable(timeout time.Duration) bool {
 	addr := strings.TrimPrefix(f.target, "http://")
-	dialer := &net.Dialer{}
-	ticker := time.NewTicker(50 * time.Millisecond)
-	defer ticker.Stop()
-	for {
-		conn, err := dialer.DialContext(ctx, "tcp", addr)
-		if err == nil {
-			_ = conn.Close()
-			return nil
-		}
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-ticker.C:
-		}
+	conn, err := net.DialTimeout("tcp", addr, timeout)
+	if err != nil {
+		return false
 	}
+	_ = conn.Close()
+	return true
 }
 
 // PassThrough forwards a per-MicroVM lifecycle hook (/run, /resume,