feat: add lambda-runtime-invocation-id header

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"] }

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<String>,
+}
+
+#[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<HandlerError> 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<Request>) -> Result<Response, HandlerError> {
+    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,
+            }
+        );
+    }
+}

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";

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::<EventPayload>(&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,23 +138,28 @@ 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<T>(request_id: &str, err: T) -> Result<http::Request<Body>, BoxError>
+fn build_event_error_request<T>(
+    request_id: String,
+    invocation_id: Option<String>,
+    err: T,
+) -> Result<http::Request<Body>, BoxError>
 where
     T: Into<Diagnostic> + Debug,
 {
     error!(error = ?err, "Request payload deserialization into LambdaEvent<T> 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)]
 pub enum RuntimeApiResponseFuture<F, Response, BufferedResponse, StreamingResponse, StreamItem, StreamError> {
     Future(
         #[pin] F,
         String,
+        Option<String>,
         PhantomData<(
             (),
             Response,
@@ -183,9 +189,9 @@ where
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
         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"),
         })

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;