From f2075c18b1204396845d67ca08b09114b246b902 Mon Sep 17 00:00:00 2001 From: jlizen Date: Thu, 3 Jul 2025 17:45:05 +0000 Subject: [PATCH] chore(lambda-runtime): slightly optimize graceful shutdown helper by using new tokio::try_join biased; api --- examples/extension-internal-flush/Cargo.toml | 2 +- examples/extension-internal-flush/src/main.rs | 6 ++++-- lambda-runtime/Cargo.toml | 2 +- lambda-runtime/src/lib.rs | 8 +++++--- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/examples/extension-internal-flush/Cargo.toml b/examples/extension-internal-flush/Cargo.toml index daadd0eb..70e27ebf 100644 --- a/examples/extension-internal-flush/Cargo.toml +++ b/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"] } diff --git a/examples/extension-internal-flush/src/main.rs b/examples/extension-internal-flush/src/main.rs index c728030b..9a515ff8 100644 --- a/examples/extension-internal-flush/src/main.rs +++ b/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 } diff --git a/lambda-runtime/Cargo.toml b/lambda-runtime/Cargo.toml index 6950e1ba..c4787d56 100644 --- a/lambda-runtime/Cargo.toml +++ b/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", diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index e1dd408f..cbcd0a9e 100644 --- a/lambda-runtime/src/lib.rs +++ b/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;