refactor(snapstart): drop the redundant pre-snapshot readiness gate

With AWS_LWA_ASYNC_INIT ignored under SnapStart, check_init_health takes its
blocking path and `main` awaits it before `run()` — so the application is already
ready by the time before_snapshot runs. The gate only duplicated that invariant in a
second place.

The one case it still covered was a library consumer calling Adapter::new() then
run() while skipping check_init_health(). That is not a coherent guarantee for them
anyway: it would protect the snapshot while leaving their invocations unguarded,
since they also skip ready_at_init and so never get the first-request re-check. If
someone re-enables async_init under SnapStart, the test on effective_async_init
fails, which is an earlier and louder signal than a runtime wait.

Removes the gate, the two tests written for it, and a duplicate happy-path test it
left behind. `ensure_ready`'s `phase` parameter goes too — after_restore is now its
only caller, so it was a parameter with one possible value.

Docs reconciled with the removal: the guide's lifecycle step 1 no longer claims a
pre-snapshot wait (initialization already did it),
AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS bounds two waits again rather than three, and
the CHANGELOG's "hooks firing before the application is ready" fix folds into the
async_init entry, which is where the actual fix lives.

Author: bnusunny

CHANGELOG.md

@@ -24,12 +24,11 @@
   app. Default: 4 seconds. A value that is set but unusable falls back to the
   default and logs a warning.
 - Add `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` to bound the readiness check
-  (fractional seconds allowed, e.g. `0.5`), applied to the initial cold-start
-  readiness wait, the pre-snapshot wait, and the post-SnapStart-restore readiness
-  check. When set and the app does not become ready within it, the adapter
-  **refuses to serve**: cold-start init fails (the runtime never starts), and taking
-  a snapshot or completing a restore fails, rather than admitting traffic to an
-  app that never reported ready. When unset (the default) the wait is
+  (fractional seconds allowed, e.g. `0.5`), applied to both the initial cold-start
+  readiness wait and the post-SnapStart-restore readiness check. When set and the app
+  does not become ready within it, the adapter **refuses to serve**: cold-start init
+  fails (the runtime never starts) and a restore fails, rather than admitting traffic
+  to an app that never reported ready. When unset (the default) the wait is
   **unbounded**, matching the previous behavior, so existing slow-cold-start apps
   are unaffected unless they opt in. On-demand cold starts using `async_init` keep
   that path's own fixed ~9.8s bound (non-fatal) and are not affected by this variable.
@@ -45,14 +44,6 @@
   is normalized so it behaves like `/api`. **Upgrade note:** this changes the path
   forwarded to your app for those inputs — deployments that relied on the old
   repeated/partial stripping should verify their routes.
-- Fix the hooks firing before the application is ready. With `AWS_LWA_ASYNC_INIT=true`
-  the adapter finishes initialization after 9.8 seconds even if the app has not bound
-  its port yet, so the snapshot could capture a still-booting app and each hook `POST`
-  then failed immediately with a connection error. The adapter now waits for the
-  readiness check before the snapshot is taken — bounded by
-  `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` when it is set — which covers both hooks,
-  since the after-restore hook deliberately runs before its own readiness check so the
-  application can reconnect before its health is judged.
 - `AWS_LWA_ASYNC_INIT` is now ignored under SnapStart and Provisioned Concurrency, with
   a warning. It exists to work around the short initialization limit for on-demand cold
   starts by reporting init complete before the application is ready; neither of those
@@ -63,7 +54,8 @@
   `AWS_LWA_ASYNC_INIT=true` together with SnapStart or provisioned concurrency now waits
   for its readiness check during initialization instead of finishing early; bound that
   wait with `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` if you want it to fail rather than
-  block.
+  block. Under SnapStart this is also what guarantees the snapshot is taken of a fully
+  initialized application.
 
 ### Dependencies
 

docs/guide/src/features/snapstart.md

@@ -15,13 +15,13 @@ Both hooks are opt-in and independent — each fires only when its variable is s
 
 The adapter always registers for the SnapStart lifecycle; the Lambda runtime invokes the hooks only when your function runs under SnapStart. When it does, the adapter participates as follows:
 
-1. **Before checkpoint** — the adapter first waits for your application to pass the readiness check, so a snapshot is never taken of a still-booting app. Then, if `AWS_LWA_SNAPSTART_BEFORE_CHECKPOINT_PATH` is set, it sends an empty `POST` to that path, and signals Lambda that it is ready for the snapshot.
+1. **Before checkpoint** — if `AWS_LWA_SNAPSTART_BEFORE_CHECKPOINT_PATH` is set, the adapter sends an empty `POST` to that path on your application, then signals Lambda that it is ready for the snapshot. Initialization has already waited for your application to report ready by this point, so the snapshot is never taken of a still-booting app.
 
 2. **After restore** — Lambda restores the environment. The adapter first refreshes its own HTTP connection to your application (so it never reuses a connection captured in the snapshot); then, if `AWS_LWA_SNAPSTART_AFTER_RESTORE_PATH` is set, sends an empty `POST` to that path; and finally re-runs the readiness check before admitting traffic.
 
 > **Note:** `AWS_LWA_ASYNC_INIT` is ignored under SnapStart (and under Provisioned Concurrency). It works around the short initialization limit for on-demand cold starts by reporting init complete before the application is ready; neither of those environments has that limit, and finishing early would snapshot — or serve — a half-initialized application. The adapter logs a warning when it ignores the setting.
 
-Each hook is an empty `POST`, and your application must respond with a `2xx` status. A non-`2xx` response or a connection failure fails the SnapStart phase — initialization for the before-checkpoint hook, restore for the after-restore hook — rather than serving traffic against an improperly prepared application. The adapter does not impose its own deadline on a hook: Lambda already bounds both phases, and the after-restore hook in particular must complete within your function timeout, so keep that in mind when a hook does slow work such as draining a large connection pool. The final readiness check runs on every restore (whether or not an after-restore path is configured). By default this readiness wait is unbounded; set `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` (fractional seconds allowed, e.g. `0.5`) to bound it, in which case a restore whose application does not report ready within that timeout fails. The same variable also bounds the initial cold-start readiness check and the pre-snapshot wait in step 1: when set and the application does not report ready within the timeout, initialization fails (the Lambda runtime never starts) rather than serving traffic — or snapshotting — an app that never came up.
+Each hook is an empty `POST`, and your application must respond with a `2xx` status. A non-`2xx` response or a connection failure fails the SnapStart phase — initialization for the before-checkpoint hook, restore for the after-restore hook — rather than serving traffic against an improperly prepared application. The adapter does not impose its own deadline on a hook: Lambda already bounds both phases, and the after-restore hook in particular must complete within your function timeout, so keep that in mind when a hook does slow work such as draining a large connection pool. The final readiness check runs on every restore (whether or not an after-restore path is configured). By default this readiness wait is unbounded; set `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` (fractional seconds allowed, e.g. `0.5`) to bound it, in which case a restore whose application does not report ready within that timeout fails. The same variable also bounds the initial cold-start readiness check, which under SnapStart is what guarantees the application is ready before the snapshot: when set and the application does not report ready within the timeout, initialization fails (the Lambda runtime never starts) rather than snapshotting or serving an app that never came up.
 
 ## Why you need the hooks
 

src/snapstart.rs

@@ -119,16 +119,11 @@ impl SnapStartHooks {
 impl SnapStartResource for SnapStartHooks {
     fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(async move {
-            // Wait for the app unconditionally, before any hook, so a snapshot is
-            // never taken of a still-booting app. `check_init_health` already blocks
-            // until ready under SnapStart (AWS_LWA_ASYNC_INIT is ignored there, see
-            // `effective_async_init`), so this is defense in depth — it also covers a
-            // consumer that drives the `Service` impl without calling
-            // `check_init_health`. It has to be unconditional: `after_restore` POSTs
-            // before its own readiness check (step 2 before step 3, deliberately), so
-            // a half-booted snapshot would fail every restore even with no
-            // before-checkpoint hook configured.
-            self.ensure_ready(&self.client, "before-checkpoint").await?;
+            // No readiness wait here: the app is already ready. `AWS_LWA_ASYNC_INIT` is
+            // ignored under SnapStart (see `effective_async_init`), so
+            // `check_init_health` takes its blocking path and `main` awaits it before
+            // `run()`. Re-checking would only duplicate that invariant in a second
+            // place.
             if let Some(path) = self.before_checkpoint_path.as_deref() {
                 Self::post_hook(&self.client, &self.domain, path).await?;
             }
@@ -160,24 +155,22 @@ impl SnapStartResource for SnapStartHooks {
             // 3. Confirm the app is serving again before traffic is admitted.
             //    A configured timeout bounds the wait and fails the restore on
             //    expiry; when unset the wait is unbounded (historical behavior).
-            self.ensure_ready(&fresh, "after-restore").await?;
+            self.ensure_ready(&fresh).await?;
 
             Ok(())
         })
     }
 }
 
 impl SnapStartHooks {
-    /// Waits for the app to report ready, bounded by
-    /// `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` when set. `phase` names the SnapStart
-    /// phase in the timeout error so an init failure is distinguishable from a restore
-    /// failure.
+    /// Waits for the app to report ready after a restore, bounded by
+    /// `AWS_LWA_READINESS_CHECK_TIMEOUT_SECONDS` when set.
     ///
-    /// Unset means unbounded, which cannot fail — only block until Lambda's own phase
+    /// Unset means unbounded, which cannot fail — only block until Lambda's own restore
     /// timeout, with the escalating `app is not ready after {}ms` log as the signal.
-    async fn ensure_ready(&self, client: &Client<HttpConnector, Body>, phase: &str) -> Result<(), Error> {
+    async fn ensure_ready(&self, client: &Client<HttpConnector, Body>) -> Result<(), Error> {
         match self.readiness_timeout {
-            Some(t) => self.ensure_ready_with_timeout(client, t, phase).await,
+            Some(t) => self.ensure_ready_with_timeout(client, t).await,
             None => {
                 self.wait_ready(client).await;
                 Ok(())
@@ -191,11 +184,10 @@ impl SnapStartHooks {
         &self,
         client: &Client<HttpConnector, Body>,
         readiness_timeout: Duration,
-        phase: &str,
     ) -> Result<(), Error> {
         timeout(readiness_timeout, self.wait_ready(client)).await.map_err(|_| {
             Error::from(format!(
-                "SnapStart {phase} readiness check timed out after {readiness_timeout:?}"
+                "SnapStart after-restore readiness check timed out after {readiness_timeout:?}"
             ))
         })
     }
@@ -265,85 +257,8 @@ mod tests {
         m.assert();
     }
 
-    /// `before_snapshot` must gate the hook POST on readiness: an app that has not
-    /// bound its port yet would fail the init phase with `ECONNREFUSED`.
-    ///
-    /// The hook route here would answer 200, but readiness never passes — so the hook
-    /// must not be called at all.
-    #[tokio::test]
-    async fn before_snapshot_waits_for_readiness_before_posting() {
-        let server = MockServer::start();
-        let hook = server.mock(|when, then| {
-            when.method(httpmock::Method::POST).path("/before");
-            then.status(200);
-        });
-        // Readiness target answers 503 forever, so the app is never ready.
-        server.mock(|when, then| {
-            when.path("/never-ready");
-            then.status(503);
-        });
-        let mut h = hooks_with_health(&server, Some("/before"), None, "/never-ready");
-        h.readiness_timeout = Some(Duration::from_millis(150));
-
-        let err = h
-            .before_snapshot()
-            .await
-            .expect_err("an app that is not ready must fail the before-checkpoint phase");
-        assert!(
-            err.to_string().contains("before-checkpoint") && err.to_string().contains("readiness"),
-            "error must name the before-checkpoint readiness check, got: {err}"
-        );
-        hook.assert_calls(0);
-    }
-
-    /// The readiness gate must be UNCONDITIONAL, not tied to the before-checkpoint
-    /// hook being configured.
-    ///
-    /// The hooks are independent, so a function may set only
-    /// `AWS_LWA_SNAPSTART_AFTER_RESTORE_PATH`. With the gate inside
-    /// `if let Some(before_checkpoint_path)`, that configuration snapshots whatever
-    /// state the app is in — and `after_restore` POSTs its hook before its own
-    /// readiness check (step 2 before step 3, deliberately, so the reconnect happens
-    /// before health is judged), so the POST would hit a process that is not listening
-    /// and every restore would fail.
-    #[tokio::test]
-    async fn before_snapshot_waits_for_readiness_with_no_hook_configured() {
-        let server = MockServer::start();
-        server.mock(|when, then| {
-            when.path("/never-ready");
-            then.status(503);
-        });
-        // Only the AFTER-restore hook is configured.
-        let mut h = hooks_with_health(&server, None, Some("/after"), "/never-ready");
-        h.readiness_timeout = Some(Duration::from_millis(150));
-
-        let err = h
-            .before_snapshot()
-            .await
-            .expect_err("the snapshot must not be taken against an app that is not ready");
-        assert!(
-            err.to_string().contains("before-checkpoint") && err.to_string().contains("readiness"),
-            "error must name the before-checkpoint readiness check, got: {err}"
-        );
-    }
-
-    /// The gate must not change the happy path: a ready app still gets the POST.
-    #[tokio::test]
-    async fn before_snapshot_posts_once_app_is_ready() {
-        let server = MockServer::start();
-        let m = server.mock(|when, then| {
-            when.method(httpmock::Method::POST).path("/before");
-            then.status(200);
-        });
-        let h = hooks(&server, Some("/before"), None);
-        assert!(h.before_snapshot().await.is_ok());
-        m.assert();
-    }
-
     #[tokio::test]
-    /// With no hook configured there is nothing to POST, but the readiness wait still
-    /// runs (`hooks` mocks a healthy `/health`), so this is not a full no-op.
-    async fn before_snapshot_posts_nothing_when_unset() {
+    async fn before_snapshot_noop_when_unset() {
         let server = MockServer::start();
         let h = hooks(&server, None, None);
         assert!(h.before_snapshot().await.is_ok());
@@ -451,9 +366,7 @@ mod tests {
         let h = hooks_with_health(&server, None, None, "/never");
         let client = build_client(Duration::from_secs(4), Pooling::Enabled);
 
-        let result = h
-            .ensure_ready_with_timeout(&client, Duration::from_millis(100), "after-restore")
-            .await;
+        let result = h.ensure_ready_with_timeout(&client, Duration::from_millis(100)).await;
 
         let err = result.expect_err("unready app should fail the readiness check");
         assert!(err.to_string().contains("timed out"), "unexpected error: {err}");