From 07de1f501faf6544a18c13d3f131b77737f5dfbd Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 9 Aug 2026 19:13:34 +0100 Subject: [PATCH 01/11] feat: add support for invocation-id --- lambda-runtime/src/constants.rs | 9 ++ lambda-runtime/src/layers/api_response.rs | 20 ++- lambda-runtime/src/lib.rs | 2 + lambda-runtime/src/requests.rs | 159 +++++++++++++++++++--- lambda-runtime/src/runtime.rs | 7 +- lambda-runtime/src/types.rs | 54 ++++++-- 6 files changed, 217 insertions(+), 34 deletions(-) create mode 100644 lambda-runtime/src/constants.rs diff --git a/lambda-runtime/src/constants.rs b/lambda-runtime/src/constants.rs new file mode 100644 index 000000000..98c789d0f --- /dev/null +++ b/lambda-runtime/src/constants.rs @@ -0,0 +1,9 @@ +/// Header names used in the Lambda Runtime API. +pub(crate) const LAMBDA_RUNTIME_REQUEST_ID: &str = "lambda-runtime-aws-request-id"; +pub(crate) const LAMBDA_RUNTIME_DEADLINE_MS: &str = "lambda-runtime-deadline-ms"; +pub(crate) const LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN: &str = "lambda-runtime-invoked-function-arn"; +pub(crate) const LAMBDA_RUNTIME_TRACE_ID: &str = "lambda-runtime-trace-id"; +pub(crate) const LAMBDA_RUNTIME_CLIENT_CONTEXT: &str = "lambda-runtime-client-context"; +pub(crate) const LAMBDA_RUNTIME_COGNITO_IDENTITY: &str = "lambda-runtime-cognito-identity"; +pub(crate) const LAMBDA_RUNTIME_TENANT_ID: &str = "lambda-runtime-aws-tenant-id"; +pub(crate) const LAMBDA_RUNTIME_INVOCATION_ID: &str = "lambda-runtime-invocation-id"; diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index 5bb3c96f8..378199e42 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -123,9 +123,10 @@ where }; let request_id = req.context.request_id.clone(); + let invocation_id = req.context.invocation_id.clone(); let lambda_event = match deserializer::deserialize::(&req.body, req.context) { Ok(lambda_event) => lambda_event, - Err(err) => match build_event_error_request(&request_id, err) { + Err(err) => match build_event_error_request(request_id, invocation_id, err) { Ok(request) => return RuntimeApiResponseFuture::Ready(Box::new(Some(Ok(request)))), Err(err) => { error!(error = ?err, "failed to build error response for Lambda Runtime API"); @@ -137,16 +138,20 @@ where // Once the handler input has been generated successfully, pass it through to inner services // allowing processing both before reaching the handler function and after the handler completes. let fut = self.inner.call(lambda_event); - RuntimeApiResponseFuture::Future(fut, request_id, PhantomData) + RuntimeApiResponseFuture::Future(fut, request_id, invocation_id, PhantomData) } } -fn build_event_error_request(request_id: &str, err: T) -> Result, BoxError> +fn build_event_error_request( + request_id: String, + invocation_id: Option, + err: T, +) -> Result, BoxError> where T: Into + Debug, { error!(error = ?err, "Request payload deserialization into LambdaEvent failed. The handler will not be called. Log at TRACE level to see the payload."); - EventErrorRequest::new(request_id, err).into_req() + EventErrorRequest::new(&request_id, invocation_id.as_deref(), err).into_req() } #[pin_project(project = RuntimeApiResponseFutureProj)] @@ -154,6 +159,7 @@ pub enum RuntimeApiResponseFuture, PhantomData<( (), Response, @@ -183,9 +189,9 @@ where fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll { task::Poll::Ready(match self.as_mut().project() { - RuntimeApiResponseFutureProj::Future(fut, request_id, _) => match ready!(fut.poll(cx)) { - Ok(ok) => EventCompletionRequest::new(request_id, ok).into_req(), - Err(err) => EventErrorRequest::new(request_id, err).into_req(), + RuntimeApiResponseFutureProj::Future(fut, request_id, invocation_id, _) => match ready!(fut.poll(cx)) { + Ok(ok) => EventCompletionRequest::new(request_id, invocation_id.as_deref(), ok).into_req(), + Err(err) => EventErrorRequest::new(request_id, invocation_id.as_deref(), err).into_req(), }, RuntimeApiResponseFutureProj::Ready(ready) => ready.take().expect("future polled after completion"), }) diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 69c04ebc6..59344194d 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -23,6 +23,8 @@ pub use tower::{self, service_fn, Service}; #[macro_use] mod macros; +mod constants; + /// Diagnostic utilities to convert Rust types into Lambda Error types. pub mod diagnostic; pub use diagnostic::Diagnostic; diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index b03f14c79..91ce16345 100644 --- a/lambda-runtime/src/requests.rs +++ b/lambda-runtime/src/requests.rs @@ -1,4 +1,7 @@ -use crate::{types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse, IntoFunctionResponse}; +use crate::{ + constants::LAMBDA_RUNTIME_INVOCATION_ID, types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse, + IntoFunctionResponse, +}; use bytes::Bytes; use http::{header::CONTENT_TYPE, Method, Request, Uri}; use lambda_runtime_api_client::{body::Body, build_request}; @@ -88,6 +91,7 @@ where E: Into + Send + Debug, { pub(crate) request_id: &'a str, + pub(crate) invocation_id: Option<&'a str>, pub(crate) body: R, pub(crate) _unused_b: PhantomData, pub(crate) _unused_s: PhantomData, @@ -102,9 +106,14 @@ where E: Into + Send + Debug, { /// Initialize a new EventCompletionRequest - pub(crate) fn new(request_id: &'a str, body: R) -> EventCompletionRequest<'a, R, B, S, D, E> { + pub(crate) fn new( + request_id: &'a str, + invocation_id: Option<&'a str>, + body: R, + ) -> EventCompletionRequest<'a, R, B, S, D, E> { EventCompletionRequest { request_id, + invocation_id, body, _unused_b: PhantomData::, _unused_s: PhantomData::, @@ -129,7 +138,15 @@ where let body = serde_json::to_vec(&body)?; let body = Body::from(body); - let req = build_request().method(Method::POST).uri(uri).body(body)?; + let mut req = build_request() + .method(Method::POST) + .uri(uri) + .body(body)?; + + if let Some(id) = self.invocation_id { + req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + Ok(req) } FunctionResponse::StreamingResponse(mut response) => { @@ -145,6 +162,11 @@ where // See the details in Lambda Developer Doc: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html#runtimes-custom-response-streaming req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Type".parse()?); req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Body".parse()?); + + if let Some(id) = self.invocation_id { + req_headers.append(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + req_headers.insert( "Content-Type", "application/vnd.awslambda.http-integration-response".parse()?, @@ -193,29 +215,22 @@ where } } -#[test] -fn test_event_completion_request() { - let req = EventCompletionRequest::new("id", "hello, world!"); - let req = req.into_req().unwrap(); - let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); - assert_eq!(req.method(), Method::POST); - assert_eq!(req.uri(), &expected); - assert!(match req.headers().get("User-Agent") { - Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"), - None => false, - }); -} - // /runtime/invocation/{AwsRequestId}/error pub(crate) struct EventErrorRequest<'a> { pub(crate) request_id: &'a str, + pub(crate) invocation_id: Option<&'a str>, pub(crate) diagnostic: Diagnostic, } impl<'a> EventErrorRequest<'a> { - pub(crate) fn new(request_id: &'a str, diagnostic: impl Into) -> EventErrorRequest<'a> { + pub(crate) fn new( + request_id: &'a str, + invocation_id: Option<&'a str>, + diagnostic: impl Into, + ) -> EventErrorRequest<'a> { EventErrorRequest { request_id, + invocation_id, diagnostic: diagnostic.into(), } } @@ -228,11 +243,16 @@ impl IntoRequest for EventErrorRequest<'_> { let body = serde_json::to_vec(&self.diagnostic)?; let body = Body::from(body); - let req = build_request() + let mut req = build_request() .method(Method::POST) .uri(uri) .header("lambda-runtime-function-error-type", "unhandled") .body(body)?; + + if let Some(id) = self.invocation_id { + req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + Ok(req) } } @@ -253,10 +273,93 @@ mod tests { }); } + #[test] + fn test_event_completion_request() { + let req = EventCompletionRequest::new("id", Option::Some("invocation_id"), "hello, world!"); + let req = req.into_req().unwrap(); + let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); + assert_eq!(req.method(), Method::POST); + assert_eq!(req.uri(), &expected); + + assert!(req + .headers() + .get("User-Agent") + .unwrap() + .to_str() + .unwrap() + .starts_with("aws-lambda-rust/")); + + assert_eq!( + req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(), + "invocation_id" + ); + } + + #[test] + fn test_event_completion_request_invocation_id_not_added_when_none() { + let req = EventCompletionRequest::new("id", Option::None, "hello, world!"); + let req = req.into_req().unwrap(); + + assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); + } + + #[test] + fn test_streaming_event_completion_request_with_invocation_id() { + use crate::StreamResponse; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + runtime.block_on(async { + let stream = tokio_stream::iter(vec![Ok::(Bytes::from_static(b"chunk"))]); + let stream_response: StreamResponse<_> = stream.into(); + let response = FunctionResponse::StreamingResponse(stream_response); + + let req: EventCompletionRequest<'_, _, (), _, _, _> = + EventCompletionRequest::new("id", Some("invocation_id"), response); + + let http_req = req.into_req().expect("into_req should succeed"); + let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); + assert_eq!(http_req.method(), Method::POST); + assert_eq!(http_req.uri(), &expected); + + assert_eq!( + http_req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(), + "invocation_id" + ); + }); + } + + #[test] + fn test_streaming_event_completion_request_invocation_id_not_added_when_none() { + use crate::StreamResponse; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + runtime.block_on(async { + let stream = tokio_stream::iter(vec![Ok::(Bytes::from_static(b"chunk"))]); + let stream_response: StreamResponse<_> = stream.into(); + let response = FunctionResponse::StreamingResponse(stream_response); + + let req: EventCompletionRequest<'_, _, (), _, _, _> = + EventCompletionRequest::new("id", None, response); + + let http_req = req.into_req().expect("into_req should succeed"); + + assert!(http_req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); + }); + } + #[test] fn test_event_error_request() { let req = EventErrorRequest { request_id: "id", + invocation_id: Option::Some("invocation_id"), diagnostic: Diagnostic { error_type: "InvalidEventDataError".into(), error_message: "Error parsing event data".into(), @@ -270,6 +373,26 @@ mod tests { Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"), None => false, }); + + assert!(match req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID) { + Some(header) => header.to_str().unwrap() == "invocation_id", + None => false, + }); + } + + #[test] + fn test_event_error_request_invocation_id_not_added_when_none() { + let req = EventErrorRequest { + request_id: "id", + invocation_id: None, + diagnostic: Diagnostic { + error_type: "InvalidEventDataError".into(), + error_message: "Error parsing event data".into(), + }, + }; + let req = req.into_req().unwrap(); + + assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); } #[test] diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs index ae00bc20b..b73355768 100644 --- a/lambda-runtime/src/runtime.rs +++ b/lambda-runtime/src/runtime.rs @@ -790,7 +790,11 @@ mod endpoint_tests { let base = server.base_url().parse().expect("Invalid mock server Uri"); let client = Client::builder().with_endpoint(base).build(); - let req = EventCompletionRequest::new("156cb537-e2d4-11e8-9b34-d36013741fb9", "{}"); + let req = EventCompletionRequest::new( + "156cb537-e2d4-11e8-9b34-d36013741fb9", + Option::Some("invocation_id"), + "{}", + ); let req = req.into_req()?; let rsp = client.call(req).await?; @@ -822,6 +826,7 @@ mod endpoint_tests { let req = EventErrorRequest { request_id: "156cb537-e2d4-11e8-9b34-d36013741fb9", + invocation_id: Option::Some("invocation_id"), diagnostic, }; let req = req.into_req()?; diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs index 2f8b36986..20481fd92 100644 --- a/lambda-runtime/src/types.rs +++ b/lambda-runtime/src/types.rs @@ -1,4 +1,11 @@ -use crate::{Error, RefConfig}; +use crate::{ + constants::{ + LAMBDA_RUNTIME_CLIENT_CONTEXT, LAMBDA_RUNTIME_COGNITO_IDENTITY, LAMBDA_RUNTIME_DEADLINE_MS, + LAMBDA_RUNTIME_INVOCATION_ID, LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN, LAMBDA_RUNTIME_REQUEST_ID, + LAMBDA_RUNTIME_TENANT_ID, LAMBDA_RUNTIME_TRACE_ID, + }, + Error, RefConfig, +}; use base64::prelude::*; use bytes::Bytes; use http::{header::ToStrError, HeaderMap, HeaderValue, StatusCode}; @@ -85,6 +92,11 @@ pub struct Context { /// Includes information such as the function name, memory allocation, /// version, and log streams. pub env_config: RefConfig, + /// The invocation ID assigned by the Lambda runtime for cross-wiring protection. + /// Echoed back on `/response` and `/error` to allow RAPID to reject stale responses + /// from timed-out invocations. `None` when running against older RAPID versions + /// that don't send this header. + pub invocation_id: Option, } impl Default for Context { @@ -98,6 +110,7 @@ impl Default for Context { identity: None, tenant_id: None, env_config: std::sync::Arc::new(crate::Config::default()), + invocation_id: None, } } } @@ -106,7 +119,7 @@ impl Context { /// Create a new [Context] struct based on the function configuration /// and the incoming request data. pub fn new(request_id: &str, env_config: RefConfig, headers: &HeaderMap) -> Result { - let client_context: Option = if let Some(value) = headers.get("lambda-runtime-client-context") { + let client_context: Option = if let Some(value) = headers.get(LAMBDA_RUNTIME_CLIENT_CONTEXT) { let raw = value.to_str()?; if raw.is_empty() { None @@ -117,7 +130,7 @@ impl Context { None }; - let identity: Option = if let Some(value) = headers.get("lambda-runtime-cognito-identity") { + let identity: Option = if let Some(value) = headers.get(LAMBDA_RUNTIME_COGNITO_IDENTITY) { let raw = value.to_str()?; if raw.is_empty() { None @@ -131,26 +144,29 @@ impl Context { let ctx = Context { request_id: request_id.to_owned(), deadline: headers - .get("lambda-runtime-deadline-ms") + .get(LAMBDA_RUNTIME_DEADLINE_MS) .expect("missing lambda-runtime-deadline-ms header") .to_str()? .parse::()?, invoked_function_arn: headers - .get("lambda-runtime-invoked-function-arn") + .get(LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN) .unwrap_or(&HeaderValue::from_static( "No header lambda-runtime-invoked-function-arn found.", )) .to_str()? .to_owned(), xray_trace_id: headers - .get("lambda-runtime-trace-id") + .get(LAMBDA_RUNTIME_TRACE_ID) .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), client_context, identity, tenant_id: headers - .get("lambda-runtime-aws-tenant-id") + .get(LAMBDA_RUNTIME_TENANT_ID) .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), env_config, + invocation_id: headers + .get(LAMBDA_RUNTIME_INVOCATION_ID) + .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), }; Ok(ctx) @@ -165,7 +181,7 @@ impl Context { /// Extract the invocation request id from the incoming request. pub(crate) fn invoke_request_id(headers: &HeaderMap) -> Result<&str, ToStrError> { headers - .get("lambda-runtime-aws-request-id") + .get(LAMBDA_RUNTIME_REQUEST_ID) .expect("missing lambda-runtime-aws-request-id header") .to_str() } @@ -291,6 +307,8 @@ where #[cfg(test)] mod test { + use http::HeaderName; + use super::*; use crate::Config; use std::sync::Arc; @@ -535,4 +553,24 @@ mod test { let context = Context::new("id", config, &headers).unwrap(); assert_eq!(context.tenant_id, None); } + + #[test] + fn context_with_invocation_id_resolves() { + let config = Arc::new(Config::default()); + let mut headers = HeaderMap::new(); + + let context = Context::new("id", config, &headers).unwrap(); + + assert_eq!(context.invocation_id, None); + + let config = Arc::new(Config::default()); + headers.insert( + "lambda-runtime-invocation-id", + HeaderValue::from_static("invocation-123"), + ); + + let context = Context::new("id", config, &headers).unwrap(); + + assert_eq!(context.invocation_id, Some("invocation-123".to_string())); + } } From de1f2e408bf8efa5dba4e165906c589ac1dc6c9a Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 9 Aug 2026 19:15:30 +0100 Subject: [PATCH 02/11] test: add multiconcurrency testing --- Dockerfile.rie | 16 ++- Dockerfile.test | 16 ++- examples/invocation-id-concurrent/Cargo.toml | 9 ++ examples/invocation-id-concurrent/src/main.rs | 128 ++++++++++++++++++ scripts/download-rie.sh | 46 +++++++ scripts/test-rie.sh | 2 +- .../scenarios/concurrent_scenarios.py | 32 +++++ 7 files changed, 244 insertions(+), 5 deletions(-) create mode 100644 examples/invocation-id-concurrent/Cargo.toml create mode 100644 examples/invocation-id-concurrent/src/main.rs create mode 100644 scripts/download-rie.sh diff --git a/Dockerfile.rie b/Dockerfile.rie index 1a46b5771..b55545b22 100644 --- a/Dockerfile.rie +++ b/Dockerfile.rie @@ -4,8 +4,20 @@ RUN dnf install -y gcc RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y ENV PATH="/root/.cargo/bin:${PATH}" -ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie -RUN chmod +x /usr/local/bin/aws-lambda-rie +ARG TARGETARCH +ENV RIE_VERSION=1.36 \ + RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \ + RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \ + RIE_PATH=/usr/local/bin/aws-lambda-rie + +COPY scripts/download-rie.sh /tmp/download-rie.sh +RUN sh /tmp/download-rie.sh \ + "${TARGETARCH}" \ + "${RIE_VERSION}" \ + "${RIE_SHA256_AMD64}" \ + "${RIE_SHA256_ARM64}" \ + "${RIE_PATH}" \ + && rm /tmp/download-rie.sh ARG EXAMPLE=basic-lambda diff --git a/Dockerfile.test b/Dockerfile.test index b36b1f28c..ece55bf63 100644 --- a/Dockerfile.test +++ b/Dockerfile.test @@ -1,7 +1,19 @@ FROM public.ecr.aws/lambda/provided:al2023 -ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie -RUN chmod +x /usr/local/bin/aws-lambda-rie +ARG TARGETARCH +ENV RIE_VERSION=1.36 \ + RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \ + RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \ + RIE_PATH=/usr/local/bin/aws-lambda-rie + +COPY scripts/download-rie.sh /tmp/download-rie.sh +RUN sh /tmp/download-rie.sh \ + "${TARGETARCH}" \ + "${RIE_VERSION}" \ + "${RIE_SHA256_AMD64}" \ + "${RIE_SHA256_ARM64}" \ + "${RIE_PATH}" \ + && rm /tmp/download-rie.sh COPY scripts/custom-lambda-entrypoint.sh /usr/local/bin/lambda-entrypoint RUN chmod +x /usr/local/bin/lambda-entrypoint diff --git a/examples/invocation-id-concurrent/Cargo.toml b/examples/invocation-id-concurrent/Cargo.toml new file mode 100644 index 000000000..8d0841ced --- /dev/null +++ b/examples/invocation-id-concurrent/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "invocation-id-concurrent" +version = "0.1.0" +edition = "2021" + +[dependencies] +lambda_runtime = { path = "../../lambda-runtime", features = ["concurrency-tokio"] } +serde = "1.0.219" +tokio = { version = "1", features = ["macros", "rt", "time"] } diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs new file mode 100644 index 000000000..da695a164 --- /dev/null +++ b/examples/invocation-id-concurrent/src/main.rs @@ -0,0 +1,128 @@ +// This example requires the following input to succeed: +// { "command": "do something" } + +use lambda_runtime::{service_fn, tracing, Diagnostic, Error, LambdaEvent}; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize)] +struct Request { + command: String, + sleep: u32 +} + +#[derive(Serialize, Debug, PartialEq)] +struct Response { + req_id: String, + inv_id: Option, +} + +#[derive(Debug)] +struct HandlerError(String); + +impl std::fmt::Display for HandlerError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl From for Diagnostic { + fn from(e: HandlerError) -> Diagnostic { + Diagnostic { + error_type: "HandlerError".into(), + error_message: e.0, + } + } +} + + +/** + * Cross-wiring protection: duplicate request-id after timeout. + + Timeline: + t=0: Invoke A starts, handler sleeps 7s + t=5: A times out (timeout=5s). Batch 1 completes with timeout error. + t=5: Invoke B starts (same request-id), handler sleeps 4s + t=7: A's handler wakes up, posts stale /response/{same-id} + t=9: B's handler wakes up, posts correct /response/{same-id} + + With invocation-id: A's stale post at t=7 gets 410 Gone. B responds at t=9 correctly. + Without: A's stale response at t=7 is accepted for B (cross-wired). + */ + +#[tokio::main] +async fn main() -> Result<(), Error> { + // required to enable CloudWatch error logging by the runtime + tracing::init_default_subscriber(); + let max_concurrency = std::env::var("AWS_LAMBDA_MAX_CONCURRENCY").unwrap_or_else(|_| "not set".to_string()); + tracing::info!(AWS_LAMBDA_MAX_CONCURRENCY = %max_concurrency, "starting concurrent handler"); + + let func = service_fn(my_handler); + if let Err(err) = lambda_runtime::run_concurrent(func).await { + tracing::error!(error = %err, "run error"); + return Err(err); + } + Ok(()) +} + +pub(crate) async fn my_handler(event: LambdaEvent) -> Result { + if event.payload.sleep > 0 { + tokio::time::sleep(tokio::time::Duration::from_secs(event.payload.sleep.into())).await; + } + + let resp = Response { + req_id: event.context.request_id, + inv_id: event.context.invocation_id, + }; + + Ok(resp) +} + +#[cfg(test)] +mod tests { + use super::*; + use lambda_runtime::{Context, LambdaEvent}; + + #[tokio::test] + async fn handler_returns_request_and_invocation_ids() { + let mut context = Context::default(); + context.request_id = "req-123".to_string(); + context.invocation_id = Some("inv-456".to_string()); + + let payload = Request { + command: "test".to_string(), + sleep: 0, + }; + let event = LambdaEvent { payload, context }; + let result = my_handler(event).await.unwrap(); + + assert_eq!( + result, + Response { + req_id: "req-123".to_string(), + inv_id: Some("inv-456".to_string()), + } + ); + } + + #[tokio::test] + async fn handler_works_without_invocation_id() { + let mut context = Context::default(); + context.request_id = "req-789".to_string(); + // invocation_id defaults to None + + let payload = Request { + command: "test".to_string(), + sleep: 0, + }; + let event = LambdaEvent { payload, context }; + let result = my_handler(event).await.unwrap(); + + assert_eq!( + result, + Response { + req_id: "req-789".to_string(), + inv_id: None, + } + ); + } +} diff --git a/scripts/download-rie.sh b/scripts/download-rie.sh new file mode 100644 index 000000000..0d22509aa --- /dev/null +++ b/scripts/download-rie.sh @@ -0,0 +1,46 @@ +#!/bin/sh + +set -eu + +if [ "$#" -ne 5 ]; then + echo "Usage: $0 TARGETARCH RIE_VERSION RIE_SHA256_AMD64 RIE_SHA256_ARM64 RIE_PATH" >&2 + exit 1 +fi + +TARGETARCH=$1 +RIE_VERSION=$2 +RIE_SHA256_AMD64=$3 +RIE_SHA256_ARM64=$4 +RIE_PATH=$5 + +case "${TARGETARCH}" in + amd64) + RIE_ASSET=aws-lambda-rie + RIE_SHA256=${RIE_SHA256_AMD64} + ;; + arm64) + RIE_ASSET=aws-lambda-rie-arm64 + RIE_SHA256=${RIE_SHA256_ARM64} + ;; + *) + echo "Unsupported target architecture: ${TARGETARCH}" >&2 + exit 1 + ;; +esac + +: "${RIE_PATH:?RIE_PATH must be set}" +RIE_TMP=$(mktemp) +trap 'rm -f "${RIE_TMP}"' EXIT + +curl \ + --fail \ + --location \ + --silent \ + --show-error \ + --retry 3 \ + --retry-all-errors \ + --output "${RIE_TMP}" \ + "https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/download/v${RIE_VERSION}/${RIE_ASSET}" + +echo "${RIE_SHA256} ${RIE_TMP}" | sha256sum --check --status +install -m 0755 "${RIE_TMP}" "${RIE_PATH}" diff --git a/scripts/test-rie.sh b/scripts/test-rie.sh index c5949fe8f..9de6feddd 100755 --- a/scripts/test-rie.sh +++ b/scripts/test-rie.sh @@ -18,7 +18,7 @@ fi CONTAINER_PID=$! echo "Container started. Test with:" -if [ "$EXAMPLE" = "basic-lambda" ] || [ "$EXAMPLE" = "basic-lambda-concurrent" ]; then +if [ "$EXAMPLE" = "basic-lambda" ] || [ "$EXAMPLE" = "basic-lambda-concurrent" ] || [ "$EXAMPLE" = "invocation-id-concurrent" ]; then echo "curl -XPOST 'http://localhost:9000/2015-03-31/functions/function/invocations' -d '{\"command\": \"test from RIE\"}' -H 'Content-Type: application/json'" else echo "For example '$EXAMPLE', check examples/$EXAMPLE/src/main.rs for the expected payload format." diff --git a/test/dockerized/scenarios/concurrent_scenarios.py b/test/dockerized/scenarios/concurrent_scenarios.py index 8c2ce842b..86b147ae1 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -11,6 +11,7 @@ HANDLER = "basic-lambda-concurrent" IMAGE = os.environ.get("TEST_IMAGE", "local/test-base") DEFAULT_CONCURRENCY = 10 +TIMEOUT = 5 def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: @@ -21,6 +22,12 @@ def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: } +def _invocation_id_env(concurrency: int = DEFAULT_CONCURRENCY, timeout: int = TIMEOUT) -> dict: + return _make_env | { + "AWS_LAMBDA_FUNCTION_TIMEOUT": str(timeout), + } + + def get_concurrent_scenarios(): scenarios = [] @@ -62,3 +69,28 @@ def get_concurrent_scenarios(): )) return scenarios + + +def invocation_id_scenarios(): + batches = [ + [Request.create( + payload={"name": "invoke-A", "sleep": TIMEOUT + 2}, + assertions=[{"transform": ".errorType", "error": "Sandbox.Timedout"}], + headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, + )], + [Request.create( + payload={"name": "invoke-B", "sleep": TIMEOUT - 1}, + assertions={"response": {"from": "invoke-B"}}, + headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, + )], + ] + + + return [ConcurrentTest( + name="invocation_id", + handler="invocation-id-concurrent", + environment_variables=_invocation_id_env(timeout=1), + request_batches=batches, + image=IMAGE, + )] + From a75e5f18d335b92c6c6d6f37152c2571e6173ca2 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Thu, 27 Aug 2026 11:03:06 +0000 Subject: [PATCH 03/11] chore: additional fixes --- .github/workflows/dockerized-test.yml | 6 ++++- Makefile | 5 ++++- examples/invocation-id-concurrent/src/main.rs | 9 ++++---- lambda-runtime/src/requests.rs | 10 +++------ lambda-runtime/src/types.rs | 3 ++- scripts/build-examples.sh | 22 ++++++++++++++----- .../scenarios/concurrent_scenarios.py | 16 ++++++++------ 7 files changed, 45 insertions(+), 26 deletions(-) diff --git a/.github/workflows/dockerized-test.yml b/.github/workflows/dockerized-test.yml index cb6cf495d..d6bf11567 100644 --- a/.github/workflows/dockerized-test.yml +++ b/.github/workflows/dockerized-test.yml @@ -57,7 +57,9 @@ jobs: - name: Build Lambda artifacts for testing run: | mkdir -p test/dockerized/tasks - HANDLERS_TO_BUILD="basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples + HANDLERS_TO_BUILD="basic-lambda-concurrent invocation-id-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples + test -x test/dockerized/tasks/basic-lambda-concurrent + test -x test/dockerized/tasks/invocation-id-concurrent ls -la test/dockerized/tasks/ - name: Build base test image with RIE and custom entrypoint @@ -68,6 +70,8 @@ jobs: - name: Run concurrent scenarios uses: aws/containerized-test-runner-for-aws-lambda@main + env: + CONTAINER_READY_DELAY_SECS: 5 with: suiteFileArray: '[]' dockerImageName: 'local/test-base' diff --git a/Makefile b/Makefile index 2b652e12e..c3e43286d 100644 --- a/Makefile +++ b/Makefile @@ -7,6 +7,7 @@ INTEG_EXTENSIONS := extension-fn extension-trait logs-trait INTEG_ARCH := x86_64-unknown-linux-musl RIE_MAX_CONCURRENCY ?= 4 TEST_RUNNER_BRANCH ?= main +CONTAINER_READY_DELAY_SECS ?= 5 OUTPUT_DIR ?= test/dockerized/tasks HANDLERS_TO_BUILD ?= HANDLER ?= @@ -125,7 +126,7 @@ fmt: cargo +nightly fmt --all build-examples: - HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD} OUTPUT_DIR=${OUTPUT_DIR} ./scripts/build-examples.sh + HANDLERS_TO_BUILD="$(subst ",,$(HANDLERS_TO_BUILD))" OUTPUT_DIR="$(OUTPUT_DIR)" ./scripts/build-examples.sh nuke: docker kill $$(docker ps -q) @@ -145,6 +146,7 @@ build-test-runner: build-examples @echo "Building test runner Docker image..." @docker build -t test-runner:local -f .test-runner/Dockerfile .test-runner +test-dockerized-concurrent: HANDLERS_TO_BUILD := basic-lambda-concurrent invocation-id-concurrent test-dockerized-concurrent: build-test-runner @echo "Running concurrent scenarios in Docker..." @docker network rm concurrent-test-net 2>/dev/null || true @@ -156,6 +158,7 @@ test-dockerized-concurrent: build-test-runner -e TASK_FOLDER=./test/dockerized/tasks \ -e GITHUB_WORKSPACE=/workspace \ -e DOCKER_SHARED_NETWORK=concurrent-test-net \ + -e CONTAINER_READY_DELAY_SECS=$(CONTAINER_READY_DELAY_SECS) \ -v /var/run/docker.sock:/var/run/docker.sock \ -v "$(CURDIR):/workspace" \ -w /workspace \ diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs index da695a164..7497dd1c6 100644 --- a/examples/invocation-id-concurrent/src/main.rs +++ b/examples/invocation-id-concurrent/src/main.rs @@ -6,8 +6,9 @@ use serde::{Deserialize, Serialize}; #[derive(Deserialize)] struct Request { - command: String, - sleep: u32 + #[serde(rename = "command")] + _command: String, + sleep: u32, } #[derive(Serialize, Debug, PartialEq)] @@ -89,7 +90,7 @@ mod tests { context.invocation_id = Some("inv-456".to_string()); let payload = Request { - command: "test".to_string(), + _command: "test".to_string(), sleep: 0, }; let event = LambdaEvent { payload, context }; @@ -111,7 +112,7 @@ mod tests { // invocation_id defaults to None let payload = Request { - command: "test".to_string(), + _command: "test".to_string(), sleep: 0, }; let event = LambdaEvent { payload, context }; diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index 91ce16345..a4b637029 100644 --- a/lambda-runtime/src/requests.rs +++ b/lambda-runtime/src/requests.rs @@ -138,10 +138,7 @@ where let body = serde_json::to_vec(&body)?; let body = Body::from(body); - let mut req = build_request() - .method(Method::POST) - .uri(uri) - .body(body)?; + let mut req = build_request().method(Method::POST).uri(uri).body(body)?; if let Some(id) = self.invocation_id { req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); @@ -346,8 +343,7 @@ mod tests { let stream_response: StreamResponse<_> = stream.into(); let response = FunctionResponse::StreamingResponse(stream_response); - let req: EventCompletionRequest<'_, _, (), _, _, _> = - EventCompletionRequest::new("id", None, response); + let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", None, response); let http_req = req.into_req().expect("into_req should succeed"); @@ -466,7 +462,7 @@ mod tests { let stream_response: StreamResponse<_> = stream.into(); let response = FunctionResponse::StreamingResponse(stream_response); - let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", response); + let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", None, response); let http_req = req.into_req().expect("into_req should succeed"); diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs index 20481fd92..100aa5bd5 100644 --- a/lambda-runtime/src/types.rs +++ b/lambda-runtime/src/types.rs @@ -307,7 +307,6 @@ where #[cfg(test)] mod test { - use http::HeaderName; use super::*; use crate::Config; @@ -558,6 +557,8 @@ mod test { fn context_with_invocation_id_resolves() { let config = Arc::new(Config::default()); let mut headers = HeaderMap::new(); + headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id")); + headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123")); let context = Context::new("id", config, &headers).unwrap(); diff --git a/scripts/build-examples.sh b/scripts/build-examples.sh index b479059dc..2e577c3a0 100755 --- a/scripts/build-examples.sh +++ b/scripts/build-examples.sh @@ -11,13 +11,25 @@ echo "Building handlers: ${HANDLERS_TO_BUILD}" for handler in ${HANDLERS_TO_BUILD}; do dir="examples/$handler" - [ ! -f "$dir/Cargo.toml" ] && echo "✗ $handler not found" && continue - + if [ ! -f "$dir/Cargo.toml" ]; then + echo "✗ $handler not found" + continue + fi + echo "Building $handler..." - (cd "$dir" && cargo build --release) || continue - - [ -f "$dir/target/release/$handler" ] && cp "$dir/target/release/$handler" "$OUTPUT_DIR/" && echo "✓ $handler" + if ! (cd "$dir" && cargo build --release); then + continue + fi + + if [ ! -f "$dir/target/release/$handler" ]; then + echo "✗ $handler artifact not found" + continue + fi + + cp "$dir/target/release/$handler" "$OUTPUT_DIR/" + echo "✓ $handler" done echo "" ls -lh "$OUTPUT_DIR/" 2>/dev/null || echo "No binaries built" +exit 0 diff --git a/test/dockerized/scenarios/concurrent_scenarios.py b/test/dockerized/scenarios/concurrent_scenarios.py index 86b147ae1..2ed0f7630 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -9,7 +9,9 @@ from containerized_test_runner.models import Request, ConcurrentTest HANDLER = "basic-lambda-concurrent" +INVOCATION_ID_HANDLER = "invocation-id-concurrent" IMAGE = os.environ.get("TEST_IMAGE", "local/test-base") +SAME_REQUEST_ID = "shared-request-id" DEFAULT_CONCURRENCY = 10 TIMEOUT = 5 @@ -23,7 +25,7 @@ def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: def _invocation_id_env(concurrency: int = DEFAULT_CONCURRENCY, timeout: int = TIMEOUT) -> dict: - return _make_env | { + return _make_env(concurrency) | { "AWS_LAMBDA_FUNCTION_TIMEOUT": str(timeout), } @@ -71,16 +73,16 @@ def get_concurrent_scenarios(): return scenarios -def invocation_id_scenarios(): +def get_invocation_id_scenarios(): batches = [ [Request.create( - payload={"name": "invoke-A", "sleep": TIMEOUT + 2}, + payload={"command": "invoke-A", "sleep": TIMEOUT + 2}, assertions=[{"transform": ".errorType", "error": "Sandbox.Timedout"}], headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, )], [Request.create( - payload={"name": "invoke-B", "sleep": TIMEOUT - 1}, - assertions={"response": {"from": "invoke-B"}}, + payload={"command": "invoke-B", "sleep": TIMEOUT - 1}, + assertions=[{"transform": ".req_id", "response": SAME_REQUEST_ID}], headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, )], ] @@ -88,8 +90,8 @@ def invocation_id_scenarios(): return [ConcurrentTest( name="invocation_id", - handler="invocation-id-concurrent", - environment_variables=_invocation_id_env(timeout=1), + handler=INVOCATION_ID_HANDLER, + environment_variables=_invocation_id_env(timeout=TIMEOUT), request_batches=batches, image=IMAGE, )] From 2cabdac443495956f7907f75e58d129013ebee78 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Thu, 27 Aug 2026 13:11:38 +0000 Subject: [PATCH 04/11] cohre: other changes --- .env | 16 ---------------- .github/workflows/dockerized-test.yml | 16 ++-------------- Makefile | 8 ++------ 3 files changed, 4 insertions(+), 36 deletions(-) delete mode 100644 .env diff --git a/.env b/.env deleted file mode 100644 index 1873f15dd..000000000 --- a/.env +++ /dev/null @@ -1,16 +0,0 @@ -# Test configuration for RIE and dockerized tests -# Customize these values as needed for testing both local and on github - -# Handlers to build -HANDLERS_TO_BUILD="basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent" - -HANDLER=basic-lambda - -# Output directory for built binaries -OUTPUT_DIR=test/dockerized/tasks - -# Max concurrent Lambda invocations for LMI mode -RIE_MAX_CONCURRENCY=4 - -# Branch of containerized-test-runner-for-aws-lambda to clone -TEST_RUNNER_BRANCH=main diff --git a/.github/workflows/dockerized-test.yml b/.github/workflows/dockerized-test.yml index d6bf11567..be86e1470 100644 --- a/.github/workflows/dockerized-test.yml +++ b/.github/workflows/dockerized-test.yml @@ -19,20 +19,10 @@ jobs: - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 - - name: Load environment variables - run: | - if [ -f .env ]; then - set -a - source .env - set +a - echo "HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD}" >> $GITHUB_ENV - echo "OUTPUT_DIR=${OUTPUT_DIR}" >> $GITHUB_ENV - fi - - name: Build Lambda artifacts for testing run: | mkdir -p test/dockerized/tasks - OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples + HANDLERS_TO_BUILD="basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh ls -la test/dockerized/tasks/ - name: Build base test image with RIE and custom entrypoint @@ -57,9 +47,7 @@ jobs: - name: Build Lambda artifacts for testing run: | mkdir -p test/dockerized/tasks - HANDLERS_TO_BUILD="basic-lambda-concurrent invocation-id-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples - test -x test/dockerized/tasks/basic-lambda-concurrent - test -x test/dockerized/tasks/invocation-id-concurrent + HANDLERS_TO_BUILD="basic-lambda-concurrent invocation-id-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh ls -la test/dockerized/tasks/ - name: Build base test image with RIE and custom entrypoint diff --git a/Makefile b/Makefile index c3e43286d..ccc415a72 100644 --- a/Makefile +++ b/Makefile @@ -9,12 +9,8 @@ RIE_MAX_CONCURRENCY ?= 4 TEST_RUNNER_BRANCH ?= main CONTAINER_READY_DELAY_SECS ?= 5 OUTPUT_DIR ?= test/dockerized/tasks -HANDLERS_TO_BUILD ?= -HANDLER ?= - -# Load environment variables from .env file if it exists --include .env -export +HANDLERS_TO_BUILD ?= basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent +HANDLER ?= basic-lambda .PHONY: help pr-check integration-tests check-event-features fmt build-examples build-test-runner test-rie test-rie-lmi nuke test-dockerized test-dockerized-concurrent From 1152b7c6f4dd9eba1cd1dbc6ea1e15caa6a1ec51 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Thu, 27 Aug 2026 15:58:20 +0000 Subject: [PATCH 05/11] chore: code review --- examples/invocation-id-concurrent/src/main.rs | 57 ++++--------------- lambda-runtime/src/layers/api_response.rs | 49 +++++++++++++++- lambda-runtime/src/types.rs | 35 +----------- .../scenarios/concurrent_scenarios.py | 2 +- 4 files changed, 63 insertions(+), 80 deletions(-) diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs index 7497dd1c6..fd4af07f1 100644 --- a/examples/invocation-id-concurrent/src/main.rs +++ b/examples/invocation-id-concurrent/src/main.rs @@ -13,8 +13,7 @@ struct Request { #[derive(Serialize, Debug, PartialEq)] struct Response { - req_id: String, - inv_id: Option, + from: String, } #[derive(Debug)] @@ -70,12 +69,9 @@ pub(crate) async fn my_handler(event: LambdaEvent) -> Result(&req.body, req.context) { Ok(lambda_event) => lambda_event, Err(err) => match build_event_error_request(request_id, invocation_id, err) { @@ -197,3 +208,39 @@ where }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{constants::LAMBDA_RUNTIME_INVOCATION_ID, runtime::LambdaInvocation, Context}; + use http::{HeaderValue, Response}; + use serde_json::json; + use tower::{service_fn, Service}; + + #[tokio::test] + async fn forwards_invocation_id_from_next_response_headers() { + let mut response = Response::new(()); + response + .headers_mut() + .insert(LAMBDA_RUNTIME_INVOCATION_ID, HeaderValue::from_static("invocation-123")); + let (parts, _) = response.into_parts(); + + let mut service = RuntimeApiResponseService::new(service_fn(|_event: LambdaEvent| async { + Ok::<_, Diagnostic>(json!({"ok": true})) + })); + + let request = service + .call(LambdaInvocation { + parts, + body: bytes::Bytes::from_static(b"{}"), + context: Context::default(), + }) + .await + .expect("response request should be created"); + + assert_eq!( + request.headers().get(LAMBDA_RUNTIME_INVOCATION_ID), + Some(&HeaderValue::from_static("invocation-123")), + ); + } +} diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs index 100aa5bd5..edd932402 100644 --- a/lambda-runtime/src/types.rs +++ b/lambda-runtime/src/types.rs @@ -1,8 +1,8 @@ use crate::{ constants::{ LAMBDA_RUNTIME_CLIENT_CONTEXT, LAMBDA_RUNTIME_COGNITO_IDENTITY, LAMBDA_RUNTIME_DEADLINE_MS, - LAMBDA_RUNTIME_INVOCATION_ID, LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN, LAMBDA_RUNTIME_REQUEST_ID, - LAMBDA_RUNTIME_TENANT_ID, LAMBDA_RUNTIME_TRACE_ID, + LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN, LAMBDA_RUNTIME_REQUEST_ID, LAMBDA_RUNTIME_TENANT_ID, + LAMBDA_RUNTIME_TRACE_ID, }, Error, RefConfig, }; @@ -92,11 +92,6 @@ pub struct Context { /// Includes information such as the function name, memory allocation, /// version, and log streams. pub env_config: RefConfig, - /// The invocation ID assigned by the Lambda runtime for cross-wiring protection. - /// Echoed back on `/response` and `/error` to allow RAPID to reject stale responses - /// from timed-out invocations. `None` when running against older RAPID versions - /// that don't send this header. - pub invocation_id: Option, } impl Default for Context { @@ -110,7 +105,6 @@ impl Default for Context { identity: None, tenant_id: None, env_config: std::sync::Arc::new(crate::Config::default()), - invocation_id: None, } } } @@ -164,9 +158,6 @@ impl Context { .get(LAMBDA_RUNTIME_TENANT_ID) .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), env_config, - invocation_id: headers - .get(LAMBDA_RUNTIME_INVOCATION_ID) - .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), }; Ok(ctx) @@ -552,26 +543,4 @@ mod test { let context = Context::new("id", config, &headers).unwrap(); assert_eq!(context.tenant_id, None); } - - #[test] - fn context_with_invocation_id_resolves() { - let config = Arc::new(Config::default()); - let mut headers = HeaderMap::new(); - headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id")); - headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123")); - - let context = Context::new("id", config, &headers).unwrap(); - - assert_eq!(context.invocation_id, None); - - let config = Arc::new(Config::default()); - headers.insert( - "lambda-runtime-invocation-id", - HeaderValue::from_static("invocation-123"), - ); - - let context = Context::new("id", config, &headers).unwrap(); - - assert_eq!(context.invocation_id, Some("invocation-123".to_string())); - } } diff --git a/test/dockerized/scenarios/concurrent_scenarios.py b/test/dockerized/scenarios/concurrent_scenarios.py index 2ed0f7630..c7a1127f0 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -82,7 +82,7 @@ def get_invocation_id_scenarios(): )], [Request.create( payload={"command": "invoke-B", "sleep": TIMEOUT - 1}, - assertions=[{"transform": ".req_id", "response": SAME_REQUEST_ID}], + assertions=[{"response": {"from": "invoke-B"}}], headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, )], ] From fe3e866ac6cfd36255b4074e37c43c1da5f90541 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 30 Aug 2026 17:17:33 +0000 Subject: [PATCH 06/11] chore: add serde default to example --- examples/invocation-id-concurrent/Cargo.toml | 3 +++ examples/invocation-id-concurrent/src/main.rs | 9 ++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/examples/invocation-id-concurrent/Cargo.toml b/examples/invocation-id-concurrent/Cargo.toml index 8d0841ced..4d61e6820 100644 --- a/examples/invocation-id-concurrent/Cargo.toml +++ b/examples/invocation-id-concurrent/Cargo.toml @@ -7,3 +7,6 @@ edition = "2021" lambda_runtime = { path = "../../lambda-runtime", features = ["concurrency-tokio"] } serde = "1.0.219" tokio = { version = "1", features = ["macros", "rt", "time"] } + +[dev-dependencies] +serde_json = "=1.0.151" diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs index fd4af07f1..acca09a05 100644 --- a/examples/invocation-id-concurrent/src/main.rs +++ b/examples/invocation-id-concurrent/src/main.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; #[derive(Deserialize)] struct Request { - #[serde(rename = "command")] + #[serde(default, rename = "command")] _command: String, sleep: u32, } @@ -93,4 +93,11 @@ mod tests { assert_eq!(result, Response { from: "invoke-B".into() }); } + + #[test] + fn request_defaults_missing_command() { + let request: Request = serde_json::from_str(r#"{"sleep": 0}"#).unwrap(); + + assert_eq!(request._command, ""); + } } From 69fc5d0ba20ada05909612d7da12c218cc9d245f Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 30 Aug 2026 18:40:06 +0000 Subject: [PATCH 07/11] chore: fix the possible problem with malformed bytes --- lambda-runtime/src/layers/api_response.rs | 48 +++++++++++++++++++---- 1 file changed, 41 insertions(+), 7 deletions(-) diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index e7976c15b..9bf782822 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -11,7 +11,7 @@ use pin_project::pin_project; use serde::{Deserialize, Serialize}; use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task}; use tower::Service; -use tracing::{error, trace}; +use tracing::{error, trace, warn}; /// Tower service that turns the result or an error of a handler function into a Lambda Runtime API /// response. @@ -128,12 +128,17 @@ where // The invocation ID assigned by the Lambda runtime for cross-wiring protection. // Echoed back on `/response` and `/error` to allow RAPID to reject stale responses // from timed-out invocations. `None` when running against older RAPID versions - // that don't send this header - let invocation_id = req - .parts - .headers - .get(LAMBDA_RUNTIME_INVOCATION_ID) - .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()); + // that don't send this header. + let invocation_id = match req.parts.headers.get(LAMBDA_RUNTIME_INVOCATION_ID) { + Some(value) => match value.to_str() { + Ok(value) => Some(value.to_owned()), + Err(error) => { + warn!(error = ?error, "Ignoring malformed Lambda runtime invocation ID header"); + None + } + }, + None => None, + }; let lambda_event = match deserializer::deserialize::(&req.body, req.context) { Ok(lambda_event) => lambda_event, @@ -243,4 +248,33 @@ mod tests { Some(&HeaderValue::from_static("invocation-123")), ); } + + #[tokio::test] + async fn malformed_invocation_id_does_not_block_deserialization_error_response() { + let mut response = Response::new(()); + response.headers_mut().insert( + LAMBDA_RUNTIME_INVOCATION_ID, + HeaderValue::from_bytes(&[0xff]).expect("header value should accept opaque bytes"), + ); + let (parts, _) = response.into_parts(); + + let mut service = RuntimeApiResponseService::new(service_fn(|_event: LambdaEvent| async { + Ok::<_, Diagnostic>(json!({"ok": true})) + })); + + let request = service + .call(LambdaInvocation { + parts, + body: bytes::Bytes::from_static(b"{"), + context: Context { + request_id: "request-123".to_owned(), + ..Context::default() + }, + }) + .await + .expect("deserialization errors should produce an error request"); + + assert_eq!(request.uri().path(), "/2018-06-01/runtime/invocation/request-123/error",); + assert!(request.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); + } } From 68434addb643628cb3ef7e8f1e6e387e221c7b36 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Mon, 31 Aug 2026 08:05:14 +0000 Subject: [PATCH 08/11] chore: fix append --- lambda-runtime/src/requests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index a4b637029..933e3100a 100644 --- a/lambda-runtime/src/requests.rs +++ b/lambda-runtime/src/requests.rs @@ -161,7 +161,7 @@ where req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Body".parse()?); if let Some(id) = self.invocation_id { - req_headers.append(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + req_headers.insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); } req_headers.insert( From 5beb7606f32a1cdfafecf5e5d686db27acf9378c Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Mon, 31 Aug 2026 10:51:28 +0000 Subject: [PATCH 09/11] chore: change message for 410 --- lambda-runtime/src/layers/api_client.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lambda-runtime/src/layers/api_client.rs b/lambda-runtime/src/layers/api_client.rs index bc2bcd0f3..d91ea78f7 100644 --- a/lambda-runtime/src/layers/api_client.rs +++ b/lambda-runtime/src/layers/api_client.rs @@ -114,8 +114,12 @@ where // Adding more information on top of 410 Gone, to make it more clear since we cannot access the body of the message if status == 410 { log_or_print!( - tracing: tracing::error!("Lambda function timeout!"), - fallback: eprintln!("Lambda function timeout!") + tracing: tracing::error!( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + ), + fallback: eprintln!( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + ) ); } @@ -219,7 +223,9 @@ mod tests { // Verify the error was logged assert!(logs_contain("Lambda Runtime API returned non-200 response")); - assert!(logs_contain("Lambda function timeout!")); + assert!(logs_contain( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + )); } #[tokio::test] From e03fa4298938ff1cd4a2f567f36f50a3c762cc36 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 30 Aug 2026 18:57:42 +0000 Subject: [PATCH 10/11] chore: add rate limiting capability to the loggin. --- lambda-runtime/src/layers/api_response.rs | 20 ++-- lambda-runtime/src/lib.rs | 2 + lambda-runtime/src/rate_limiter.rs | 109 ++++++++++++++++++++++ 3 files changed, 124 insertions(+), 7 deletions(-) create mode 100644 lambda-runtime/src/rate_limiter.rs diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index 9bf782822..bfff12572 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -1,18 +1,17 @@ use crate::{ - constants::LAMBDA_RUNTIME_INVOCATION_ID, - deserializer, - requests::{EventCompletionRequest, IntoRequest}, - runtime::LambdaInvocation, - Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, + Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, constants::LAMBDA_RUNTIME_INVOCATION_ID, deserializer, rate_limiter::RateLimiter, requests::{EventCompletionRequest, IntoRequest}, runtime::LambdaInvocation, }; use futures::{ready, Stream}; use lambda_runtime_api_client::{body::Body, BoxError}; use pin_project::pin_project; use serde::{Deserialize, Serialize}; -use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task}; +use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task, time::Duration}; use tower::Service; use tracing::{error, trace, warn}; + +static MALFORMED_INVOCATION_ID_LIMITER: RateLimiter = RateLimiter::new(Duration::from_secs(60)); + /// Tower service that turns the result or an error of a handler function into a Lambda Runtime API /// response. /// @@ -133,7 +132,14 @@ where Some(value) => match value.to_str() { Ok(value) => Some(value.to_owned()), Err(error) => { - warn!(error = ?error, "Ignoring malformed Lambda runtime invocation ID header"); + + if MALFORMED_INVOCATION_ID_LIMITER.allow() { + warn!( + error = ?error, + rate_limit_interval_ms = MALFORMED_INVOCATION_ID_LIMITER.interval().as_millis(), + "Ignoring malformed Lambda runtime invocation ID header; this warning is rate limited" + ); + } None } }, diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 59344194d..ce6688baa 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -29,6 +29,8 @@ mod constants; pub mod diagnostic; pub use diagnostic::Diagnostic; +mod rate_limiter; + mod deserializer; /// Tower middleware to be applied to runtime invocations. pub mod layers; diff --git a/lambda-runtime/src/rate_limiter.rs b/lambda-runtime/src/rate_limiter.rs new file mode 100644 index 000000000..7b86b175b --- /dev/null +++ b/lambda-runtime/src/rate_limiter.rs @@ -0,0 +1,109 @@ +//! Thread-safe, process-local rate limiting for infrequent runtime events. + +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +/// Allows an operation at most once during each configured interval. +/// +/// A `RateLimiter` is intended to be shared by concurrent runtime tasks. When +/// stored in a `static`, it is initialized once per Lambda execution environment +/// and retains its state across warm invocations. A new cold-started environment +/// receives a new limiter. +pub(crate) struct RateLimiter { + /// Minimum duration between allowed operations. + interval: Duration, + /// Timestamp of the most recent allowed operation. + last_allowed: Mutex>, +} + +impl RateLimiter { + /// Creates a rate limiter with the specified minimum interval. + pub(crate) const fn new(interval: Duration) -> RateLimiter { + RateLimiter { + interval, + last_allowed: Mutex::new(None), + } + } + + /// Returns the minimum duration between allowed operations. + pub(crate) const fn interval(&self) -> Duration { + self.interval + } + + /// + /// The first call is allowed. Subsequent calls are rejected until the + /// configured interval has elapsed since the previous allowed call. + /// Concurrent callers are serialized while checking and updating the last + /// allowed timestamp so only one caller crosses the interval boundary. + pub(crate) fn allow(&self) -> bool { + let mut last_allowed = match self.last_allowed.lock() { + Ok(guard) => guard, + Err(poisoned) => { + // The limiter state is disposable, so reset it and recover instead of + // allowing a poisoned mutex to crash the runtime or suppress future logs. + let mut guard = poisoned.into_inner(); + *guard = None; + self.last_allowed.clear_poison(); + guard + } + }; + + if last_allowed + .as_ref() + .is_some_and(|value| value.elapsed() < self.interval) + { + return false; + } + + *last_allowed = Some(Instant::now()); + + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::panic::{catch_unwind, AssertUnwindSafe}; + use std::thread; + + #[test] + fn allows_first_call() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + assert!(limiter.allow()); + } + + #[test] + fn rejects_calls_inside_interval() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + assert!(limiter.allow()); + assert!(!limiter.allow()); + } + + #[test] + fn allows_call_after_interval() { + let limiter = RateLimiter::new(Duration::from_millis(10)); + + assert!(limiter.allow()); + thread::sleep(Duration::from_millis(15)); + + assert!(limiter.allow()); + } + + #[test] + fn recovers_from_poisoned_mutex() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + let _ = catch_unwind(AssertUnwindSafe(|| { + let _guard = limiter.last_allowed.lock().unwrap(); + panic!("poison the limiter mutex"); + })); + + assert!(limiter.allow()); + assert!(!limiter.allow()); + } +} From eddf4c577363aecfdcd2bda33ffae192dc00dd59 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Mon, 31 Aug 2026 07:48:07 +0000 Subject: [PATCH 11/11] chore: fmt --- lambda-runtime/src/layers/api_response.rs | 9 +++++--- lambda-runtime/src/rate_limiter.rs | 6 +++-- .../scenarios/concurrent_scenarios.py | 23 +++++++++---------- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index bfff12572..11d9c5b32 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -1,5 +1,10 @@ use crate::{ - Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, constants::LAMBDA_RUNTIME_INVOCATION_ID, deserializer, rate_limiter::RateLimiter, requests::{EventCompletionRequest, IntoRequest}, runtime::LambdaInvocation, + constants::LAMBDA_RUNTIME_INVOCATION_ID, + deserializer, + rate_limiter::RateLimiter, + requests::{EventCompletionRequest, IntoRequest}, + runtime::LambdaInvocation, + Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, }; use futures::{ready, Stream}; use lambda_runtime_api_client::{body::Body, BoxError}; @@ -9,7 +14,6 @@ use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task, time: use tower::Service; use tracing::{error, trace, warn}; - static MALFORMED_INVOCATION_ID_LIMITER: RateLimiter = RateLimiter::new(Duration::from_secs(60)); /// Tower service that turns the result or an error of a handler function into a Lambda Runtime API @@ -132,7 +136,6 @@ where Some(value) => match value.to_str() { Ok(value) => Some(value.to_owned()), Err(error) => { - if MALFORMED_INVOCATION_ID_LIMITER.allow() { warn!( error = ?error, diff --git a/lambda-runtime/src/rate_limiter.rs b/lambda-runtime/src/rate_limiter.rs index 7b86b175b..f2fa6514f 100644 --- a/lambda-runtime/src/rate_limiter.rs +++ b/lambda-runtime/src/rate_limiter.rs @@ -66,8 +66,10 @@ impl RateLimiter { #[cfg(test)] mod tests { use super::*; - use std::panic::{catch_unwind, AssertUnwindSafe}; - use std::thread; + use std::{ + panic::{catch_unwind, AssertUnwindSafe}, + thread, + }; #[test] fn allows_first_call() { diff --git a/test/dockerized/scenarios/concurrent_scenarios.py b/test/dockerized/scenarios/concurrent_scenarios.py index c7a1127f0..8370289b3 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -1,12 +1,7 @@ -""" -Multi-concurrency test scenarios for basic-lambda-concurrent. - -The handler expects: { "command": "" } -and responds with: { "req_id": "", "msg": "Command executed." } -""" +"""Multi-concurrency test scenarios.""" import os -from containerized_test_runner.models import Request, ConcurrentTest +from containerized_test_runner.models import ConcurrentTest, Request HANDLER = "basic-lambda-concurrent" INVOCATION_ID_HANDLER = "invocation-id-concurrent" @@ -16,16 +11,20 @@ TIMEOUT = 5 -def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: +def _make_env(handler: str = HANDLER, concurrency: int = DEFAULT_CONCURRENCY) -> dict: return { - "_HANDLER": HANDLER, + "_HANDLER": handler, "AWS_LAMBDA_MAX_CONCURRENCY": str(concurrency), "AWS_LAMBDA_LOG_FORMAT": "JSON", } -def _invocation_id_env(concurrency: int = DEFAULT_CONCURRENCY, timeout: int = TIMEOUT) -> dict: - return _make_env(concurrency) | { +def _invocation_id_env( + handler: str = INVOCATION_ID_HANDLER, + concurrency: int = DEFAULT_CONCURRENCY, + timeout: int = TIMEOUT, +) -> dict: + return _make_env(handler, concurrency) | { "AWS_LAMBDA_FUNCTION_TIMEOUT": str(timeout), } @@ -91,7 +90,7 @@ def get_invocation_id_scenarios(): return [ConcurrentTest( name="invocation_id", handler=INVOCATION_ID_HANDLER, - environment_variables=_invocation_id_env(timeout=TIMEOUT), + environment_variables=_invocation_id_env(handler=INVOCATION_ID_HANDLER, timeout=TIMEOUT), request_batches=batches, image=IMAGE, )]