chore: code review

Author: darklight3it

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<String>,
+    from: String,
 }
 
 #[derive(Debug)]
@@ -70,12 +69,9 @@ pub(crate) async fn my_handler(event: LambdaEvent<Request>) -> Result<Response,
         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)
+    Ok(Response {
+        from: event.payload._command,
+    })
 }
 
 #[cfg(test)]
@@ -84,46 +80,17 @@ mod tests {
     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,
+    async fn handler_echoes_marker() {
+        let event = LambdaEvent {
+            payload: Request {
+                _command: "invoke-B".into(),
+                sleep: 0,
+            },
+            context: Context::default(),
         };
-        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,
-            }
-        );
+        assert_eq!(result, Response { from: "invoke-B".into() });
     }
 }

lambda-runtime/src/layers/api_response.rs

@@ -1,4 +1,5 @@
 use crate::{
+    constants::LAMBDA_RUNTIME_INVOCATION_ID,
     deserializer,
     requests::{EventCompletionRequest, IntoRequest},
     runtime::LambdaInvocation,
@@ -123,7 +124,17 @@ where
         };
 
         let request_id = req.context.request_id.clone();
-        let invocation_id = req.context.invocation_id.clone();
+
+        // 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());
+
         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, 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<serde_json::Value>| 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")),
+        );
+    }
+}

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<String>,
 }
 
 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()));
-    }
 }

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},
         )],
     ]