chore(lambda-runtime): slightly optimize graceful shutdown helper by using new tokio::try_join biased; api (#1007)

Author: jlizen

examples/extension-internal-flush/Cargo.toml

@@ -9,4 +9,4 @@ aws_lambda_events = { path = "../../lambda-events" }
 lambda-extension = { path = "../../lambda-extension" }
 lambda_runtime = { path = "../../lambda-runtime" }
 serde = "1.0.136"
-tokio = { version = "1", features = ["macros", "sync"] }
+tokio = { version = "1.46", features = ["macros", "sync"] }

examples/extension-internal-flush/src/main.rs

@@ -101,9 +101,11 @@ async fn main() -> Result<(), Error> {
 
     let handler = Arc::new(EventHandler::new(request_done_sender));
 
-    // TODO: add biased! to always poll the handler future first, once supported:
-    // https://github.com/tokio-rs/tokio/issues/7304
     tokio::try_join!(
+        // always poll the handler function first before the flush extension,
+        // this results in a smaller future due to not needing to track which was polled first
+        // each time, and also a tiny latency savings
+        biased;
         lambda_runtime::run(service_fn(|event| {
             let handler = handler.clone();
             async move { handler.invoke(event).await }

lambda-runtime/Cargo.toml

@@ -52,7 +52,7 @@ pin-project = "1"
 serde = { version = "1", features = ["derive", "rc"] }
 serde_json = "^1"
 serde_path_to_error = "0.1.11"
-tokio = { version = "1.0", features = [
+tokio = { version = "1.46", features = [
     "macros",
     "io-util",
     "sync",

lambda-runtime/src/lib.rs

@@ -223,9 +223,11 @@ where
             }
         };
 
-        // TODO: add biased! to always poll the signal handling future first, once supported:
-        // https://github.com/tokio-rs/tokio/issues/7304
-        let _: (_, ()) = tokio::join!(graceful_shutdown_future, async {
+        let _: (_, ()) = tokio::join!(
+            // we always poll the graceful shutdown future first,
+            // which results in a smaller future due to lack of bookkeeping of which was last polled
+            biased;
+            graceful_shutdown_future, async {
             // we suppress extension errors because we don't actually mind if it crashes,
             // all we need to do is kick off the run so that lambda exits the init phase
             let _ = extension.run().await;