From 07de1f501faf6544a18c13d3f131b77737f5dfbd Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Sun, 9 Aug 2026 19:13:34 +0100 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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}, )], ]