diff --git a/cmd/serverless-init/cloudservice/microvm.go b/cmd/serverless-init/cloudservice/microvm.go index fe03579806be..abd970a3c2cd 100644 --- a/cmd/serverless-init/cloudservice/microvm.go +++ b/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 } diff --git a/cmd/serverless-init/cloudservice/microvm_test.go b/cmd/serverless-init/cloudservice/microvm_test.go index 0d4e65ded269..8432eefb0903 100644 --- a/cmd/serverless-init/cloudservice/microvm_test.go +++ b/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:" 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) diff --git a/cmd/serverless-init/enhanced-metrics/collector.go b/cmd/serverless-init/enhanced-metrics/collector.go index 8250f3dc7609..604bb64d688a 100644 --- a/cmd/serverless-init/enhanced-metrics/collector.go +++ b/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 { diff --git a/cmd/serverless-init/enhanced-metrics/collector_test.go b/cmd/serverless-init/enhanced-metrics/collector_test.go index 12e30fdd3087..d4f38ff5c8b2 100644 --- a/cmd/serverless-init/enhanced-metrics/collector_test.go +++ b/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) +} diff --git a/cmd/serverless-init/enhanced-metrics/collector_unsupported.go b/cmd/serverless-init/enhanced-metrics/collector_unsupported.go index 0305193cc087..d10eaee7f06d 100644 --- a/cmd/serverless-init/enhanced-metrics/collector_unsupported.go +++ b/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") } diff --git a/cmd/serverless-init/lifecycle/forwarder.go b/cmd/serverless-init/lifecycle/forwarder.go index 55730faf375a..452c46b60d56 100644 --- a/cmd/serverless-init/lifecycle/forwarder.go +++ b/cmd/serverless-init/lifecycle/forwarder.go @@ -27,11 +27,25 @@ 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. @@ -39,8 +53,8 @@ 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, diff --git a/cmd/serverless-init/lifecycle/forwarder_test.go b/cmd/serverless-init/lifecycle/forwarder_test.go index 5665bb845ff2..ee61a17dc8ee 100644 --- a/cmd/serverless-init/lifecycle/forwarder_test.go +++ b/cmd/serverless-init/lifecycle/forwarder_test.go @@ -134,15 +134,15 @@ func TestNewForwarder_DisableKeepAlives_OpensNewConnectionPerRequest(t *testing. "DisableKeepAlives must open a fresh TCP connection per request; with keep-alives enabled only 1 connection would be accepted") } -// NewForwarder defaults reflect the configured wire.go constants. /ready keeps -// a 60s budget (matches the platform hook timeout); the remaining hooks default -// to 1s. Changing these defaults is a deliberate behavior change — the test -// failure is intentional. +// NewForwarder defaults reflect the configured wire.go constants. /ready and +// /validate each keep a 30s budget, and runtime hooks default to 1s. Changing +// these defaults is a deliberate behavior change — the test failure is +// intentional. func TestNewForwarder_Defaults(t *testing.T) { f := NewForwarder(8080, defaultForwardTimeout, defaultReadyTimeout, defaultValidateTimeout) assert.Equal(t, 1*time.Second, f.forwardTimeout, "forwardTimeout default must be 1s") - assert.Equal(t, 60*time.Second, f.readyTimeout, "readyTimeout default must be 60s (matches platform /ready hook timeout)") - assert.Equal(t, 1*time.Second, f.validateTimeout, "validateTimeout default must be 1s") + assert.Equal(t, 30*time.Second, f.readyTimeout, "readyTimeout default must be 30s (matches platform /ready hook timeout)") + assert.Equal(t, 30*time.Second, f.validateTimeout, "validateTimeout default must be 30s") assert.Equal(t, defaultMaxResponseBodyBytes, f.maxResponseBodyBytes, "maxResponseBodyBytes default must be 1 MiB") tr, ok := f.client.Transport.(*http.Transport) require.True(t, ok, "transport must be *http.Transport") @@ -267,7 +267,8 @@ func TestForwarder_PassThroughWaiting_MirrorsStatusAndBody(t *testing.T) { // PassThroughWaiting must honor the timeout it receives. The caller (server.go) // passes readyTimeout or validateTimeout — PassThroughWaiting must not apply // any other field. This test uses a 50ms timeout against a 200ms upstream: the -// context expires and the call returns 504. +// context expires and the call returns 503 (per the /ready and /validate hook +// contract, only 200 and 503 are meaningful — never 504). func TestForwarder_PassThroughWaiting_HonorsProvidedTimeout(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { time.Sleep(200 * time.Millisecond) @@ -279,8 +280,8 @@ func TestForwarder_PassThroughWaiting_HonorsProvidedTimeout(t *testing.T) { resp := f.PassThroughWaiting(50*time.Millisecond, "/ready", nil, nil) require.NotNil(t, resp) defer resp.Body.Close() - assert.Equal(t, http.StatusGatewayTimeout, resp.StatusCode, - "PassThroughWaiting must time out on the provided timeout") + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode, + "PassThroughWaiting must time out on the provided timeout and return 503, never 504") } // On the happy path the response body is wrapped by cancelOnCloseReader. @@ -329,29 +330,36 @@ func TestWrapResponseBody_CapZero_DisablesCap(t *testing.T) { "cap=0 must read the full body — LimitReader must NOT be applied") } -// PassThroughWaiting retries TCP dials until the context expires when the port -// is unbound, so the response is always 504 (deadline exceeded) — never 503. -// This is distinct from PassThrough which returns 503 on the very first -// connect-refused error. The behavioral difference is intentional: /ready and -// /validate must wait for the user app to start, not fail-fast on a transient -// dial error. -func TestForwarder_PassThroughWaiting_UnboundPort_RetriesUntilTimeout504(t *testing.T) { +// PassThroughWaiting makes a single reachability dial attempt, bounded by +// dialCheckTimeout — it must not retry/poll, and must not wait out the +// timeout parameter, before answering. Per the hook contract the platform +// (not the hook) owns the retry loop for /ready and /validate, so an +// unreachable app must return 503 fast, the same fail-fast contract +// PassThrough already applies to connect-refused errors. +func TestForwarder_PassThroughWaiting_UnboundPort_FailsFastWith503(t *testing.T) { f := &Forwarder{ - target: "http://127.0.0.1:1", // unbound — every dial attempt is refused + target: "http://127.0.0.1:1", // unbound — dial attempt is refused client: &http.Client{}, } - resp := f.PassThroughWaiting(150*time.Millisecond, "/ready", nil, nil) + start := time.Now() + // timeout is deliberately much larger than dialCheckTimeout so a passing + // assertion below can only be explained by the fast dial-check path, not + // by this parameter. + resp := f.PassThroughWaiting(2*time.Second, "/ready", nil, nil) + elapsed := time.Since(start) require.NotNil(t, resp) defer resp.Body.Close() - assert.Equal(t, http.StatusGatewayTimeout, resp.StatusCode, - "unbound port must retry until timeout (504), not immediately 503 like PassThrough") + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode, + "unreachable app must return 503, not 504") + assert.Less(t, elapsed, dialCheckTimeout+100*time.Millisecond, + "an unreachable app must fail within dialCheckTimeout, not retry until the much larger timeout parameter (2s) elapses") } -// passThroughWaiting buffers the request body before the TCP wait so it is -// still available after waitForUserApp returns. Without buffering, the body -// reader would be exhausted during the wait loop and f.do would send an empty -// body to the user app. This pins the buffer-then-forward contract. -func TestForwarder_PassThroughWaiting_ForwardsBodyAfterTCPWait(t *testing.T) { +// passThroughWaiting buffers the request body before the reachability check +// so it is still available once that check returns. Without buffering, the +// body reader would be exhausted during the dial attempt and f.do would send +// an empty body to the user app. This pins the buffer-then-forward contract. +func TestForwarder_PassThroughWaiting_ForwardsBodyAfterReachabilityCheck(t *testing.T) { received := make(chan []byte, 1) srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { b, _ := io.ReadAll(r.Body) @@ -381,7 +389,8 @@ func (errReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } // A failed read of the inbound body must NOT forward a partial body to the // user app. PassThroughWaiting returns 500 (server-side read failure) and -// short-circuits before the TCP wait so no partial body reaches the user app. +// short-circuits before the reachability check so no partial body reaches the +// user app. func TestForwarder_PassThroughWaiting_BodyReadError_Returns500(t *testing.T) { var reached atomic.Int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { diff --git a/cmd/serverless-init/lifecycle/server.go b/cmd/serverless-init/lifecycle/server.go index d8b9fff3d092..ba0e8551a961 100644 --- a/cmd/serverless-init/lifecycle/server.go +++ b/cmd/serverless-init/lifecycle/server.go @@ -32,10 +32,12 @@ // - When the env var is set: the agent forwards each hook to // 127.0.0.1: on the same path and mirrors the user app's // response (status, body, Content-Type) back to the platform. /ready and -// /validate wait for TCP reachability before forwarding. For /run, -// /resume, /suspend, and /terminate the agent's own work — metric -// emission, /suspend and /terminate telemetry flush — runs in a goroutine -// in parallel with the pass-through. +// /validate check TCP reachability with a single fast dial before +// forwarding, answering 503 immediately if the app isn't up yet rather +// than blocking — the platform owns the retry loop for those two hooks. +// For /run, /resume, /suspend, and /terminate the agent's own work — +// metric emission, /suspend and /terminate telemetry flush — runs in a +// goroutine in parallel with the pass-through. // // /terminate does NOT synthesize SIGTERM. The platform owns process // termination via OS signals delivered independently of this HTTP event. @@ -238,8 +240,8 @@ func NewServer( // WriteTimeout must cover the full handler wall-clock for every path: // - No forwarder: flushTimeout (flush budget + write headroom) // - /run, /resume, /suspend, /terminate: forwardTimeout (default 1s) - // - /ready: readyTimeout (default 60s, matching platform /ready timeout) - // - /validate: validateTimeout (default 1s) + // - /ready: dialCheckTimeout + readyTimeout (default 30s) + // - /validate: dialCheckTimeout + validateTimeout (default 30s) // Use the largest of all applicable budgets so the HTTP server does not // close the platform-facing connection before the handler writes the // mirrored response. @@ -247,7 +249,9 @@ func NewServer( if s.fwd != nil { // /terminate uses flushSequential: flush runs after the forward, so its // wall-clock is forwardTimeout+flushTimeout, not max of the two. - maxTimeout = max(maxTimeout, s.fwd.forwardTimeout+s.flushTimeout, s.fwd.readyTimeout, s.fwd.validateTimeout) + readyBudget := dialCheckTimeout + s.fwd.readyTimeout + validateBudget := dialCheckTimeout + s.fwd.validateTimeout + maxTimeout = max(maxTimeout, s.fwd.forwardTimeout+s.flushTimeout, readyBudget, validateBudget) } writeTimeout := maxTimeout + writeTimeoutHeadroom s.httpServer = &http.Server{ @@ -318,6 +322,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) @@ -335,7 +351,8 @@ func (s *Server) handler() http.Handler { // // Dispatcher: // - If a Forwarder is configured (env-var opt-in), pass-through to the user -// app with TCP-wait: dial errors map to 503, deadline to 504. +// app with a fast reachability check: dial errors and deadline exceeded +// both map to 503. // - Otherwise, alive-check via ChildHandle: child alive → 200, anything // else (not yet started, already exited, or nil handle) → 503. The // pre-spawn race is absorbed by the platform's /ready retry behavior; @@ -363,9 +380,10 @@ func (s *Server) passThroughReady(w http.ResponseWriter, r *http.Request) { // lifecycle of a production MicroVM. // // When a Forwarder is configured (DD_AWS_MICROVM_USER_APP_PORT set): -// pass-through to the user app with TCP-wait, mirroring the response, so the -// user app's own smoke test drives the build's validity decision. The TCP-wait -// absorbs the window before the app is reachable on the test run. Without a +// pass-through to the user app with a fast reachability check, mirroring the +// response, so the user app's own smoke test drives the build's validity +// decision. A 503 while the app isn't yet reachable on the test run relies on +// the platform's own /validate retry to absorb the window. Without a // forwarder the agent returns 200 directly; the user app is not required to // implement /validate in that mode. func (s *Server) handleValidate(w http.ResponseWriter, r *http.Request) { @@ -455,7 +473,8 @@ func (s *Server) flushAll(flushCtx context.Context) { // // Used for /run and /resume (noFlush), /suspend (flushParallel), and // /terminate (flushSequential). /ready and /validate use passThroughReady -// and passThroughValidate directly (TCP-wait path, not this function). +// and passThroughValidate directly (fast-reachability-check path, not this +// function). // // flushParallel: flush runs concurrently with the forward; wall-clock is // max(forwardTimeout, flushTimeout)+ε. Known limitation: telemetry produced @@ -558,7 +577,7 @@ func (s *Server) handleRun(w http.ResponseWriter, r *http.Request) { bodyBytes, err := io.ReadAll(r.Body) _ = r.Body.Close() if err != nil { - log.Debugf("MicroVM lifecycle: could not read run body: %v", err) + log.Warnf("MicroVM lifecycle: could not read run body: %v", err) w.WriteHeader(http.StatusInternalServerError) return } diff --git a/cmd/serverless-init/lifecycle/server_test.go b/cmd/serverless-init/lifecycle/server_test.go index 7c26cb80bf56..697a39bb5bb3 100644 --- a/cmd/serverless-init/lifecycle/server_test.go +++ b/cmd/serverless-init/lifecycle/server_test.go @@ -74,6 +74,10 @@ type neverDrainer struct{} func (n *neverDrainer) WaitForPendingSamples() { select {} } +type serverErrReader struct{} + +func (serverErrReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } + func newTestServer() (*Server, *mockFlusher, *mockFlusher, *mockLogsAgent, *mockMetricEmitter, *mockSampleDrainer) { metric := &mockFlusher{} trace := &mockFlusher{} @@ -158,6 +162,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 @@ -196,6 +220,44 @@ func TestHandleRunEmptyBodyDoesNotSetInstanceID(t *testing.T) { assert.Empty(t, id, "empty body must not set instance ID") } +func TestHandleRunBodyReadErrorReturns500(t *testing.T) { + srv, _, _, _, emitter, _ := newTestServer() + + req := httptest.NewRequest(http.MethodPost, pathRun, serverErrReader{}) + rec := httptest.NewRecorder() + srv.handleRun(rec, req) + + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.Empty(t, srv.InstanceID(), "failed body read must not set instance ID") + assert.NotContains(t, emitter.getEmitted(), runMetricName, "failed body read must not emit run metric") +} + +func TestHandleRunWithForwarderBodyReadErrorDoesNotForward(t *testing.T) { + var reached atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + reached.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + srv, _, _, _, emitter, _ := newTestServer() + srv.fwd = &Forwarder{ + target: upstream.URL, + client: &http.Client{}, + forwardTimeout: 2 * time.Second, + maxResponseBodyBytes: defaultMaxResponseBodyBytes, + } + + req := httptest.NewRequest(http.MethodPost, pathRun, serverErrReader{}) + rec := httptest.NewRecorder() + srv.handleRun(rec, req) + + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.Equal(t, int32(0), reached.Load(), "user app must not receive a partial runHookPayload") + assert.Empty(t, srv.InstanceID(), "failed body read must not set instance ID") + assert.NotContains(t, emitter.getEmitted(), runMetricName, "failed body read must not emit run metric") +} + func TestHandleSuspendFlushesBeforeResponding(t *testing.T) { srv, metric, trace, logs, emitter, drainer := newTestServer() req := httptest.NewRequest(http.MethodPost, pathSuspend, nil) @@ -391,6 +453,32 @@ func TestHandleReady_WithForwarder_PassesThrough(t *testing.T) { assert.Equal(t, `{"ready":false,"reason":"warming"}`, rec.Body.String()) } +// /validate shares passThroughReady's PassThroughWaiting code path (same +// reachability-check-then-forward shape, just bounded by validateTimeout +// instead of readyTimeout) but had no server-level pass-through coverage of +// its own — this pins the mirror contract for /validate specifically. +func TestHandleValidate_WithForwarder_PassesThrough(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/x-validate") + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(`{"valid":false,"reason":"warming"}`)) + })) + defer upstream.Close() + + srv, _, _, _, _, _ := newTestServer() + srv.fwd = &Forwarder{ + target: upstream.URL, + client: &http.Client{}, + validateTimeout: 200 * time.Millisecond, + maxResponseBodyBytes: defaultMaxResponseBodyBytes, + } + rec := httptest.NewRecorder() + srv.handleValidate(rec, httptest.NewRequest(http.MethodPost, pathValidate, nil)) + assert.Equal(t, http.StatusServiceUnavailable, rec.Code) + assert.Equal(t, "application/x-validate", rec.Header().Get("Content-Type")) + assert.Equal(t, `{"valid":false,"reason":"warming"}`, rec.Body.String()) +} + // /run with a forwarder configured mirrors the user-app's status code, // body, and Content-Type, and emits the run metric. Replaces the prior // fire-and-forget contract. @@ -925,6 +1013,24 @@ func TestNewServerWithForwarderWriteTimeoutCoversForwardBudget(t *testing.T) { "WriteTimeout must cover forwardTimeout+flushTimeout (terminate sequential-flush path)") } +// TestNewServerWithForwarderWriteTimeoutCoversReadinessBudgets verifies that +// WriteTimeout covers the fast reachability dial plus the forwarded /ready or +// /validate request. Without the dial budget, the HTTP server can close the +// platform-facing connection just before the handler mirrors a valid response. +func TestNewServerWithForwarderWriteTimeoutCoversReadinessBudgets(t *testing.T) { + flushTimeout := 5 * time.Second + fwd := &Forwarder{ + forwardTimeout: time.Second, + readyTimeout: 2 * time.Second, + validateTimeout: 30 * time.Second, + client: &http.Client{}, + maxResponseBodyBytes: defaultMaxResponseBodyBytes, + } + srv := NewServer(0, &mockFlusher{}, &mockFlusher{}, &mockLogsAgent{}, &mockMetricEmitter{}, &mockSampleDrainer{}, metrics.MetricSourceAWSMicroVMEnhanced, flushTimeout, nil, fwd, nil) + assert.Equal(t, dialCheckTimeout+fwd.validateTimeout+writeTimeoutHeadroom, srv.httpServer.WriteTimeout, + "WriteTimeout must cover dialCheckTimeout+validateTimeout for /validate") +} + // TestInstanceIDTagAppearsInMetricsAfterRun verifies that once /run stores a // MicroVM instance ID, subsequent lifecycle metrics include lambda_microvm_id: as // an extra tag. This is the primary tagging path for identifying individual MicroVM diff --git a/cmd/serverless-init/main.go b/cmd/serverless-init/main.go index d8236f78bd16..3bb1fe21ec5a 100644 --- a/cmd/serverless-init/main.go +++ b/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 { diff --git a/cmd/serverless-init/main_test.go b/cmd/serverless-init/main_test.go index 46d23961bff7..8fab60188b2f 100644 --- a/cmd/serverless-init/main_test.go +++ b/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")