support async init for long initialization lambda functions (#53)

* support async init for long initialization lambda functions

Cargo.toml

@@ -11,7 +11,7 @@ lambda-extension = "0.6.0"
 lambda_http = "0.6.0"
 log = "0.4.14"
 reqwest = { version = "0.11", default-features = false, features = ["blocking", "json"] }
-tokio = { version = "1.20.0", features = ["macros", "io-util", "sync", "rt-multi-thread"] }
+tokio = { version = "1.20.0", features = ["macros", "io-util", "sync", "rt-multi-thread", "time"] }
 tokio-retry = "0.3"
 
 [[bin]]

README.md

@@ -67,12 +67,20 @@ After passing readiness check, Lambda Web Adapter will start Lambda Runtime and
 
 The readiness check port/path and traffic port can be configured using environment variables. These environment variables can be defined either within docker file or as Lambda function configuration.
 
-|Environment Variable| Description                                                | Default |
-|--------------------|------------------------------------------------------------|---------|
-|PORT                | traffic port                                               | "8080"  |
-|READINESS_CHECK_PORT| readiness check port, default to the fraffic port          | PORT    |
-|READINESS_CHECK_PATH| readiness check path                                       | "/"     |
-|REMOVE_BASE_PATH    | (optional) the base path to be removed from request path   | None    |
+| Environment Variable | Description                                                          | Default |
+|----------------------|----------------------------------------------------------------------|---------|
+| PORT                 | traffic port                                                         | "8080"  |
+| READINESS_CHECK_PORT | readiness check port, default to the traffic port                    | PORT    |
+| READINESS_CHECK_PATH | readiness check path                                                 | "/"     |
+| ASYNC_INIT           | enable asynchronous initialization for long initialization functions | "false" |
+| REMOVE_BASE_PATH     | (optional) the base path to be removed from request path             | None    |
+
+**ASYNC_INIT** Lambda managed runtimes offer up to 10 seconds for function initialization. During this period of time, Lambda functions have burst of CPU to accelerate initialization, and it is free. 
+If a lambda function couldn't complete the initialization within 10 seconds, Lambda will restart the function, and bill for the initialization. 
+To help functions to use this 10 seconds free initialization time and avoid the restart, Lambda Web Adapter supports asynchronous initialization. 
+When this feature is enabled, Lambda Web Adapter performs readiness check up to 9.8 seconds. If the web app is not ready by then, 
+Lambda Web Adapter signals to Lambda service that the init is completed, and continues readiness check in the handler. 
+This feature is disabled by default. Enable it by setting environment variable `ASYNC_INIT` to `true`. 
 
 **REMOVE_BASE_PATH** - The value of this environment variable tells the adapter whether the application is running under a base path.
 For example, you could have configured your API Gateway to have a /orders/{proxy+} and a /catalog/{proxy+} resource.

src/main.rs

@@ -7,6 +7,7 @@ use log::*;
 use reqwest::{redirect, Client};
 use std::time::Duration;
 use std::{env, future, mem};
+use tokio::time::timeout;
 use tokio_retry::{strategy::FixedInterval, Retry};
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
@@ -17,6 +18,7 @@ struct AdapterOptions {
     readiness_check_port: String,
     readiness_check_path: String,
     base_path: Option<String>,
+    async_init: bool,
 }
 
 #[tokio::main]
@@ -31,6 +33,10 @@ async fn main() -> Result<(), Error> {
             .unwrap_or_else(|_| env::var("PORT").unwrap_or_else(|_| "8080".to_string())),
         readiness_check_path: env::var("READINESS_CHECK_PATH").unwrap_or_else(|_| "/".to_string()),
         base_path: env::var("REMOVE_BASE_PATH").ok(),
+        async_init: env::var("ASYNC_INIT")
+            .unwrap_or_else(|_| "false".to_string())
+            .parse()
+            .unwrap_or(false),
     };
 
     // register as an external extension
@@ -43,7 +49,36 @@ async fn main() -> Result<(), Error> {
             .expect("extension thread error");
     });
 
-    // check if the application is ready every 10 milliseconds
+    // check if the application is ready
+    let is_ready = if options.async_init {
+        timeout(Duration::from_secs_f32(9.8), check_readiness(options))
+            .await
+            .is_ok()
+    } else {
+        check_readiness(options).await
+    };
+
+    // start lambda runtime
+    let http_client = &Client::builder()
+        .redirect(redirect::Policy::none())
+        .pool_idle_timeout(Duration::from_secs(4))
+        .build()
+        .unwrap();
+    let mut first_invoke = true;
+    lambda_http::run(http_handler(|event: Request| {
+        let is_app_ready = if first_invoke {
+            first_invoke = false;
+            is_ready
+        } else {
+            true // the app must be ready after the first invoke
+        };
+        async move { http_proxy_handler(event, http_client, options, is_app_ready).await }
+    }))
+    .await?;
+    Ok(())
+}
+
+async fn check_readiness(options: &AdapterOptions) -> bool {
     Retry::spawn(FixedInterval::from_millis(10), || {
         let readiness_check_url = format!(
             "http://{}:{}{}",
@@ -55,26 +90,20 @@ async fn main() -> Result<(), Error> {
         }
     })
     .await
-    .expect("application server is not ready");
-
-    // start lambda runtime
-    let http_client = &Client::builder()
-        .redirect(redirect::Policy::none())
-        .pool_idle_timeout(Duration::from_secs(4))
-        .build()
-        .unwrap();
-    lambda_http::run(http_handler(|event: Request| async move {
-        http_proxy_handler(event, http_client, options).await
-    }))
-    .await?;
-    Ok(())
+    .is_ok()
 }
 
 async fn http_proxy_handler(
     event: Request,
     http_client: &Client,
     options: &AdapterOptions,
+    is_app_ready: bool,
 ) -> Result<Response<Body>, Error> {
+    // continue checking readiness if async_init is configured and the app is not ready
+    if options.async_init && !is_app_ready {
+        check_readiness(options).await;
+    }
+
     let raw_path = event.raw_http_path();
     let (parts, body) = event.into_parts();