test: add multiconcurrency testing

Author: darklight3it

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

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