fix(serverless-init): align MicroVM lifecycle hooks with AWS contract

Update MicroVM lifecycle hook handling to match AWS's documented behavior for image-build and runtime hooks.

For /ready and /validate, keep the platform-owned retry loop in control by returning fast 503s when the user app is not reachable or when forwarding times out. Increase the default /validate forward budget to 10s so build-time smoke tests are not capped by the runtime hook default, and include the reachability dial budget in the lifecycle server WriteTimeout.

For /run, fail the hook when the platform request body cannot be read instead of parsing or forwarding a partial runHookPayload.

Validated with dda inv test --targets=./cmd/serverless-init/lifecycle/... and dda inv test --targets=./cmd/serverless-init.

Author: litianningdatadog

cmd/serverless-init/lifecycle/forwarder.go

@@ -27,9 +27,21 @@ const defaultMaxResponseBodyBytes int64 = 1 << 20
 const (
 	defaultForwardTimeout  = 1 * time.Second
 	defaultReadyTimeout    = 60 * time.Second
-	defaultValidateTimeout = 1 * time.Second
+	defaultValidateTimeout = 10 * 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.
@@ -38,7 +50,7 @@ type Forwarder struct {
 	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
+	validateTimeout      time.Duration // default 10s, used for /validate
 	maxResponseBodyBytes int64         // default defaultMaxResponseBodyBytes; cap on user-app body surfaced to platform
 }
 
@@ -76,61 +88,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{}
-	for {
-		conn, err := dialer.DialContext(ctx, "tcp", addr)
-		if err == nil {
-			_ = conn.Close()
-			return nil
-		}
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-time.After(50 * time.Millisecond):
-		}
+	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,

cmd/serverless-init/lifecycle/forwarder_test.go

@@ -135,14 +135,14 @@ func TestNewForwarder_DisableKeepAlives_OpensNewConnectionPerRequest(t *testing.
 }
 
 // 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.
+// a 60s budget, /validate gets a 10s smoke-test 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, 10*time.Second, f.validateTimeout, "validateTimeout default must be 10s")
 	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")
@@ -239,7 +239,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)
@@ -251,8 +252,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.
@@ -301,29 +302,33 @@ 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{},
 	}
+	start := time.Now()
 	resp := f.PassThroughWaiting(150*time.Millisecond, "/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, 100*time.Millisecond,
+		"an unreachable app must fail the single dial attempt fast, not retry until the timeout parameter (150ms) or dialCheckTimeout (200ms) elapse")
 }
 
-// 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)
@@ -353,7 +358,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) {

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:<user-app-port> 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.
@@ -225,16 +227,18 @@ 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 60s, matching platform /validate timeout)
+	//   - /ready: dialCheckTimeout + readyTimeout (default 60s)
+	//   - /validate: dialCheckTimeout + validateTimeout (default 10s)
 	// 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
 		// 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{
@@ -334,7 +338,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;
@@ -362,9 +367,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) {
@@ -454,7 +460,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
@@ -548,8 +555,13 @@ func (s *Server) handleRun(w http.ResponseWriter, r *http.Request) {
 	// original payload to the user app. Without this, the forwarder path would
 	// consume r.Body before the decode, losing the instance_id tag on all
 	// subsequent lifecycle metrics.
-	bodyBytes, _ := io.ReadAll(r.Body)
+	bodyBytes, err := io.ReadAll(r.Body)
 	_ = r.Body.Close()
+	if err != nil {
+		log.Warnf("MicroVM lifecycle: could not read run body: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
 	var body runBody
 	if err := json.Unmarshal(bodyBytes, &body); err != nil {
 		log.Debugf("MicroVM lifecycle: could not parse run body: %v", err)