feat(serverless-init): wire MicroVM CloudService into main.go and register it (#53096)

### What does this PR do?

Completes the MicroVM integration started in #53092/#53093/#53094:

- **`cloudservice/service.go`** — `GetCloudServiceType` now returns
`&MicroVM{}` when `DD_AWS_MICROVM_IMAGE_ARN` is set. This check runs
before the amd64-only arch gate, since MicroVM also supports arm64.
- **`main.go`** — `setup()` builds a `LifecycleContext` from the
telemetry dependencies it already constructs (metric/trace/log flushers,
base log tags, base trace tags) and passes it through
`TracingContext.LifecycleCtx` so `MicroVM.Init` can start the lifecycle
server. This happens on both the normal path and the no-API-key path —
MicroVM's lifecycle handshake with the platform must complete regardless
of whether trace/metric collection is enabled, unlike other cloud
services, which are intentionally not initialized on the no-API-key path
(doing so would create trace spans with tracing disabled and unset span
tags, leading to a nil-map panic on shutdown for e.g. Cloud Run Jobs).
- **`lifecycle/server.go`** — `flushAll` no longer calls
`s.logsFlusher.Flush(...)` unconditionally. The `logsAgent` passed
through as `LifecycleContext.LogsFlusher` can be nil when the logs agent
fails to start (`SetupLogAgent`'s error is discarded in
`cmd/serverless-init/log/log.go`), which would panic during MicroVM's
`/suspend` and `/terminate` handshake. Guarded with a nil check,
matching the existing precedent in `main.go`'s `lastFlush` and
`flushAll`'s own `sampleDrainer` handling.
- **`main_test.go`** — `TestRun_LocalService_SidecarMode` started the
real `RunSidecar` signal-wait path but never sent a signal or cleaned
up, leaking a goroutine that would intercept SIGTERM for the rest of the
test binary's life. It now registers its own SIGTERM listener first (to
override the default terminate disposition before signaling), sends
itself a real SIGTERM, and asserts `RunSidecar` returns.

### Motivation

This is the atomic step that turns on the feature: before this PR,
`MicroVM` (PR 2, fully tested by PR 3/PR 4) exists in the binary but is
unreachable — `GetCloudServiceType` never selects it and `main.go` never
gives it a `LifecycleContext`, so `DD_AWS_MICROVM_IMAGE_ARN` has no
effect. Registration and `main.go` wiring are deliberately kept together
in one PR (rather than splitting further) because registering `MicroVM`
without also wiring `LifecycleContext` would leave `Init` a no-op and
`MicroVM.Run` dereferencing a nil child handle — this PR is the only
point in the stack where that combination is safe to land.

The `flushAll` nil-guard and the sidecar-test signal leak were both
flagged by Codex's automated review of this PR and are fixed here rather
than as follow-ups, since both are one-line-scoped and directly touch
code this PR introduces.

### Update (amended)

This branch was originally built on top of an older, unrestacked copy of
`forwarder.go`/`wire.go`/`server.go` before being reset onto the current
`07-02` base. The reset didn't refresh the working tree, so the original
commit silently reintroduced pre-fix versions of five things (each had a
test pinning it that was deleted alongside it), none of which were
mentioned in the original PR description above:

- **`forwarder.go`** — restores the `CheckRedirect` handler on the
forwarder's `http.Client` so a 3xx from the user app is mirrored to the
platform as-is instead of silently followed (which would replay a POST
hook as a body-dropped GET). See
`TestForwarder_PassThrough_DoesNotFollowRedirects`.
- **`wire.go`** — restores the `sidecarMode` early-return before parsing
`userAppPort` and the forward/ready/validate timeouts, so a stale or
colliding value inherited from an init-mode config produces a warning
instead of failing setup in sidecar mode. See the three
`TestSetupComponents_SidecarMode_*` tests.
- **`lifecycle/server.go`** — restores `writeTimeoutHeadroom`
(`heartbeatStopTimeout + mirrorResponseTimeout`) in the `WriteTimeout`
calculation; restores unconditional response-body buffering in
`handleWithForwarder` for both `flushParallel` (`/suspend`) and
`flushSequential` (`/terminate`) — the `/suspend` path had silently lost
its buffering, risking a partial mirrored body if the parallel flush
outlives `forwardTimeout`; and restores the nil-map guard before writing
`lambda_microvm_id` into a cloned `baseTraceTags` map. See
`TestHandleSuspend_WithForwarder_BodyBufferedBeforeFlush` and
`TestHandleRun_NilBaseTraceTags_DoesNotPanic`.
- **`server_test.go`** — restores `go.uber.org/atomic` in place of
`sync/atomic`, per this repo's `codereview_guideline.md` (this had
already been fixed once, per Codex review on an earlier PR in this
stack, before being reintroduced by the same reset).

Verified via `git merge-tree` that these restorations merge cleanly with
the PRs already stacked on top of this branch (#53104, #53230, #53231).

### Describe how you validated your changes

```
dda inv test --targets=./cmd/serverless-init/...
```

279 tests, 275 passed, 4 skipped (pre-existing platform skips). New
coverage:
- `TestFlushAllNilLogsFlusherDoesNotPanic` — verifies `flushAll`
tolerates a nil `logsFlusher` and still completes promptly (asserts it
returns via the normal "all workers done" path, not the `flushTimeout`
fallback, since a panic in `flushAll`'s spawned goroutine wouldn't be
caught by `assert.NotPanics` alone).
- `TestRun_LocalService_SidecarMode` (revised) — now drives `RunSidecar`
through a real SIGTERM instead of a fixed timeout, and asserts it
returns cleanly with no leaked signal handler.
- `main_test.go` also verifies the `logTagsBase`/`baseTraceTags` values
threaded into `LifecycleContext`, that `metrics.ServerlessMetricAgent`
methods are nil-safe without a started `Demux`, and that
`CloudService.Run` correctly dispatches to both init-container and
sidecar paths.
- `service_test.go` verifies `GetCloudServiceType` selects `MicroVM` and
that it takes priority over `CloudRunJobs` when env vars for both happen
to be set.
- `TestForwarder_PassThrough_DoesNotFollowRedirects`,
`TestSetupComponents_SidecarMode_*` (×3), and
`TestHandleSuspend_WithForwarder_BodyBufferedBeforeFlush`/`TestHandleRun_NilBaseTraceTags_DoesNotPanic`
— restored, per above.

Also ran the linter:

```
dda inv linter.go --targets=./cmd/serverless-init/...
```

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

Author: litianningdatadog

cmd/serverless-init/cloudservice/microvm_test.go

@@ -10,7 +10,6 @@ import (
 	"net"
 	"net/http"
 	"os"
-	"runtime"
 	"strconv"
 	"strings"
 	"sync"
@@ -138,9 +137,6 @@ var _ CloudService = (*MicroVM)(nil)
 // init-container mode threads m.child into RunInit so the lifecycle server's
 // /ready alive-check reflects the user app's actual state.
 func TestMicroVM_Run_InitMode_ThreadsChildLiveness(t *testing.T) {
-	if runtime.GOOS == "windows" {
-		t.Skip("serverless-init is not supported on windows")
-	}
 	if testing.Short() {
 		t.Skip("spawns a subprocess")
 	}
@@ -151,23 +147,12 @@ func TestMicroVM_Run_InitMode_ThreadsChildLiveness(t *testing.T) {
 	child := lifecycle.NewChild()
 	m := &MicroVM{child: child}
 
-	// Poll for the alive transition instead of sleeping a fixed delay: a
-	// single sample after a guessed delay can land before cmd.Start has
-	// invoked MarkAlive, or after the short-lived subprocess has already
-	// exited, on a loaded host. Polling records the alive state the instant
-	// it appears, so the race window is eliminated regardless of host load.
 	var midRunAlive atomic.Bool
 	probeDone := make(chan struct{})
 	go func() {
 		defer close(probeDone)
-		deadline := time.Now().Add(2 * time.Second)
-		for time.Now().Before(deadline) {
-			if child.IsAlive() {
-				midRunAlive.Store(true)
-				return
-			}
-			time.Sleep(time.Millisecond)
-		}
+		time.Sleep(100 * time.Millisecond)
+		midRunAlive.Store(child.IsAlive())
 	}()
 
 	err := m.Run(mode.Conf{SidecarMode: false}, &serverlessInitLog.Config{})

cmd/serverless-init/cloudservice/service.go

@@ -214,6 +214,10 @@ func (l *LocalService) ShouldForceFlushAllOnForceFlushToSerializer() bool {
 func GetCloudServiceType() CloudService {
 	arch := runtime.GOARCH
 
+	if isMicroVM() {
+		return &MicroVM{}
+	}
+
 	if arch != archAMD64 {
 		log.Errorf(unsupportedArchMsg, arch)
 	}

cmd/serverless-init/cloudservice/service_test.go

@@ -10,6 +10,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	serverlessenv "github.com/DataDog/datadog-agent/pkg/serverless/env"
 )
 
 func TestGetCloudServiceType(t *testing.T) {
@@ -36,3 +38,20 @@ func TestGetCloudServiceTypeForCloudRunJob(t *testing.T) {
 	_, ok := cloudService.(*CloudRunJobs)
 	assert.True(t, ok)
 }
+
+func TestGetCloudServiceTypeMicroVM(t *testing.T) {
+	t.Setenv(serverlessenv.MicroVMImageARNEnvVar, "arn:aws:lambda:us-east-1:123456789012:microvm-image:my-image")
+	svc := GetCloudServiceType()
+	_, ok := svc.(*MicroVM)
+	assert.True(t, ok, "expected MicroVM CloudService")
+}
+
+func TestGetCloudServiceTypeMicroVMTakesPriorityOverCloudRun(t *testing.T) {
+	// MicroVM is checked first — both would never be set in practice,
+	// but the ordering must be explicit.
+	t.Setenv(ServiceNameEnvVar, "my-service")
+	t.Setenv(serverlessenv.MicroVMImageARNEnvVar, "arn:aws:lambda:us-east-1:123456789012:microvm-image:my-image")
+	svc := GetCloudServiceType()
+	_, ok := svc.(*MicroVM)
+	assert.True(t, ok, "MicroVM should take priority over CloudRun")
+}

cmd/serverless-init/lifecycle/server.go

@@ -223,14 +223,13 @@ func NewServer(
 		s.child = c
 	}
 	// WriteTimeout must cover the full handler wall-clock for every path:
-	//   - No forwarder: flushTimeout (flush budget)
+	//   - 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 60s, matching platform /validate timeout)
-	// plus writeTimeoutHeadroom (heartbeat.Stop() before /suspend and
-	// /terminate, and the final mirrored-response write). Use the largest of
-	// all applicable budgets so the HTTP server does not close the
-	// platform-facing connection before the handler writes the response.
+	// 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.
 	maxTimeout := s.flushTimeout
 	if s.fwd != nil {
 		// /terminate uses flushSequential: flush runs after the forward, so its
@@ -427,7 +426,12 @@ func (s *Server) flushAll(flushCtx context.Context) {
 	flushDone := make(chan struct{}, flushWorkerCount)
 	go func() { s.metricFlusher.Flush(); flushDone <- struct{}{} }()
 	go func() { s.traceFlusher.Flush(); flushDone <- struct{}{} }()
-	go func() { s.logsFlusher.Flush(flushCtx); flushDone <- struct{}{} }()
+	go func() {
+		if s.logsFlusher != nil {
+			s.logsFlusher.Flush(flushCtx)
+		}
+		flushDone <- struct{}{}
+	}()
 	s.waitForFlushes(flushCtx, flushDone)
 }
 

cmd/serverless-init/lifecycle/server_test.go

@@ -1034,6 +1034,32 @@ func TestFlushAllDrainTimeoutDoesNotBlock(t *testing.T) {
 	assert.Less(t, time.Since(start), 500*time.Millisecond, "flushAll must return within flushTimeout even when drainer blocks")
 }
 
+// TestFlushAllNilLogsFlusherDoesNotPanic verifies that flushAll tolerates a nil
+// logsFlusher without panicking, and still flushes the metric and trace
+// flushers promptly. Production setup() can pass a nil logsFlusher into the
+// LifecycleContext when the logs agent failed to start (SetupLogAgent
+// returns nil on error), so MicroVM's /suspend and /terminate handshake must
+// not crash on that value. A nil-interface method call in the flushAll
+// goroutine wouldn't be caught by assert.NotPanics (it panics in a different
+// goroutine), so this also asserts flushAll returns promptly via the normal
+// "all workers done" path rather than falling back to the flushTimeout path,
+// which is what happens when the logs-flush goroutine never signals
+// completion.
+func TestFlushAllNilLogsFlusherDoesNotPanic(t *testing.T) {
+	metric := &mockFlusher{}
+	trace := &mockFlusher{}
+	drainer := &mockSampleDrainer{}
+	srv := NewServer(0, metric, trace, nil, &mockMetricEmitter{}, drainer, metrics.MetricSourceAWSMicroVMEnhanced, 2*time.Second, nil, nil, nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), srv.flushTimeout)
+	defer cancel()
+	start := time.Now()
+	assert.NotPanics(t, func() { srv.flushAll(ctx) })
+	assert.Less(t, time.Since(start), 500*time.Millisecond, "flushAll must not fall back to the flushTimeout path when logsFlusher is nil")
+	assert.Equal(t, int32(1), metric.count.Load())
+	assert.Equal(t, int32(1), trace.count.Load())
+}
+
 // withFakeHeartbeat installs a real *Heartbeat with a long interval that
 // will never tick during the test, lets us observe Start/Stop side effects
 // via running goroutine count, and returns a teardown that ensures cleanup.

cmd/serverless-init/main.go

@@ -9,6 +9,7 @@ package main
 
 import (
 	"context"
+	"github.com/DataDog/datadog-agent/comp/logs-library/processor"
 	"os"
 	"strings"
 	"sync"
@@ -44,11 +45,13 @@ import (
 
 	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
 	enhancedmetrics "github.com/DataDog/datadog-agent/cmd/serverless-init/enhanced-metrics"
+	"github.com/DataDog/datadog-agent/cmd/serverless-init/lifecycle"
 	serverlessInitTag "github.com/DataDog/datadog-agent/cmd/serverless-init/tag"
 	logsAgent "github.com/DataDog/datadog-agent/comp/logs/agent/def"
 	"github.com/DataDog/datadog-agent/pkg/config/model"
 	pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup"
 	configUtils "github.com/DataDog/datadog-agent/pkg/config/utils"
+	serverlessLogs "github.com/DataDog/datadog-agent/pkg/serverless/logs"
 	"github.com/DataDog/datadog-agent/pkg/serverless/metrics"
 	"github.com/DataDog/datadog-agent/pkg/serverless/otlp"
 	serverlessTag "github.com/DataDog/datadog-agent/pkg/serverless/tags"
@@ -150,6 +153,10 @@ func setup(secretComp secrets.Component, delegatedAuthComp delegatedauth.Compone
 	origin := cloudService.GetOrigin()
 	// Note: we do not modify tags for the LogsAgent.
 	logsAgent := serverlessInitLog.SetupLogAgent(agentLogConfig, tagConfig.Tags, tagger, compression, hostname, origin)
+	// Snapshot the startup log tags so the lifecycle server can append lambda_microvm_id
+	// at /run without losing the base set. Must be computed after SetupLogAgent
+	// since MapToArray normalises the same map that SetupLogAgent passes to SetLogsTags.
+	logTagsBase := serverlessTag.MapToArray(tagConfig.Tags)
 
 	// When no API key is configured, skip trace and metric agent initialization
 	// to avoid noisy error logs. The process wrapper and logs agent still function normally.
@@ -159,26 +166,68 @@ func setup(secretComp secrets.Component, delegatedAuthComp delegatedauth.Compone
 	if apiKey == "" && apmAPIKey == "" {
 		log.Warnf("DD_API_KEY is not set; trace and metric collection are disabled. Set DD_API_KEY to enable monitoring.")
 		traceAgent := trace.NewNoopTraceAgent()
-		tracingCtx := &cloudservice.TracingContext{TraceAgent: traceAgent}
-		metricAgent := &metrics.ServerlessMetricAgent{
-			Tagger: tagger,
+		metricAgent := &metrics.ServerlessMetricAgent{Tagger: tagger}
+		tracingCtx := &cloudservice.TracingContext{
+			TraceAgent: traceAgent,
+			LifecycleCtx: &cloudservice.LifecycleContext{
+				MetricFlusher: metricAgent,
+				LogsFlusher:   logsAgent,
+				MetricEmitter: metricAgent,
+				SampleDrainer: metricAgent,
+				FlushTimeout:  agentLogConfig.FlushTimeout,
+				SidecarMode:   modeConf.SidecarMode,
+				LogsTagSetter: lifecycle.LogsTagSetterFunc(func(tags []string) {
+					serverlessLogs.SetLogsTags(tags)
+					processor.SetServerlessInitTagCache(tags)
+				}),
+				BaseTags: logTagsBase,
+				TraceTagSetter: lifecycle.TraceTagSetterFunc(func(tags map[string]string) {
+					traceAgent.SetTags(tags)
+				}),
+				BaseTraceTags: serverlessInitTag.MakeTraceAgentTags(tagConfig.Tags),
+			},
+		}
+		// Only MicroVM needs initialization without an API key: its Init starts the
+		// lifecycle hook server so the platform can complete lifecycle handshakes.
+		// Initializing other services here would create trace spans even though
+		// tracing is disabled and span tags are unset, leading to a nil-map panic
+		// on shutdown (e.g. Cloud Run Jobs writing error tags into a nil span Meta).
+		if origin == cloudservice.MicroVMOrigin {
+			_ = cloudService.Init(tracingCtx)
 		}
 		return cloudService, agentLogConfig, tracingCtx, metricAgent, logsAgent, nil, false
 	}
 
 	traceTags := serverlessInitTag.MakeTraceAgentTags(tagConfig.Tags)
 	traceAgent := setupTraceAgent(traceTags, tagConfig.ConfiguredTags, tagger, origin)
 
+	metricAgent := setupMetricAgent(tagConfig.Tags, tagConfig.EnhancedMetricTags, tagConfig.EnhancedUsageMetricTags, tagger, cloudService.ShouldForceFlushAllOnForceFlushToSerializer())
+
 	tracingCtx := &cloudservice.TracingContext{
 		TraceAgent: traceAgent,
 		SpanTags:   traceTags,
+		LifecycleCtx: &cloudservice.LifecycleContext{
+			MetricFlusher: metricAgent,
+			LogsFlusher:   logsAgent,
+			MetricEmitter: metricAgent,
+			SampleDrainer: metricAgent,
+			FlushTimeout:  agentLogConfig.FlushTimeout,
+			SidecarMode:   modeConf.SidecarMode,
+			LogsTagSetter: lifecycle.LogsTagSetterFunc(func(tags []string) {
+				serverlessLogs.SetLogsTags(tags)
+				processor.SetServerlessInitTagCache(tags)
+			}),
+			BaseTags: logTagsBase,
+			TraceTagSetter: lifecycle.TraceTagSetterFunc(func(tags map[string]string) {
+				traceAgent.SetTags(tags)
+			}),
+			BaseTraceTags: serverlessInitTag.MakeTraceAgentTags(tagConfig.Tags),
+		},
 	}
 
 	// TODO check for errors and exit
 	_ = cloudService.Init(tracingCtx)
 
-	metricAgent := setupMetricAgent(tagConfig.Tags, tagConfig.EnhancedMetricTags, tagConfig.EnhancedUsageMetricTags, tagger, cloudService.ShouldForceFlushAllOnForceFlushToSerializer())
-
 	enhancedMetricsEnabled := pkgconfigsetup.Datadog().GetBool("enhanced_metrics")
 	if enhancedMetricsEnabled {
 		cloudService.AddStartMetric(metricAgent)
@@ -197,6 +246,7 @@ func setup(secretComp secrets.Component, delegatedAuthComp delegatedauth.Compone
 	}
 
 	go flushMetricsAgent(metricAgent)
+
 	return cloudService, agentLogConfig, tracingCtx, metricAgent, logsAgent, enhancedMetricsCollector, enhancedMetricsEnabled
 }
 
@@ -289,6 +339,7 @@ func setupTraceAgent(tags map[string]string, configuredTags []string, tagger tag
 		FunctionTags:          functionTags,
 	})
 	traceAgent.SetTags(tags)
+
 	go func() {
 		for range time.Tick(3 * time.Second) {
 			traceAgent.Flush()