chore: to be reviewed

Author: darklight3it

.gitignore

@@ -49,3 +49,6 @@ __pycache__/
 # Misc
 lambda-runtime/libtest.rmeta
 lambda-integration-tests/target/
+
+# Pre-built RIE binary (built by test-rie.sh)
+aws-lambda-rie

Dockerfile.rie

@@ -4,7 +4,10 @@ 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
+# TODO: Replace COPY with ADD from release URL once a RIE release includes
+#       lambda-runtime-invocation-id support (likely v1.36+):
+#       ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/download/v1.XX/aws-lambda-rie /usr/local/bin/aws-lambda-rie
+COPY aws-lambda-rie /usr/local/bin/aws-lambda-rie
 RUN chmod +x /usr/local/bin/aws-lambda-rie
 
 ARG EXAMPLE=basic-lambda
@@ -22,4 +25,4 @@ RUN cargo build --release
 RUN cp target/release/${EXAMPLE} ${LAMBDA_RUNTIME_DIR}/bootstrap
 
 ENTRYPOINT []
-CMD [ "/usr/local/bin/aws-lambda-rie", "/var/runtime/bootstrap" ]
+CMD [ "/usr/local/bin/aws-lambda-rie", "/var/runtime/bootstrap" ]

scripts/test-rie.sh

@@ -1,19 +1,50 @@
 #!/bin/bash
 set -euo pipefail
 
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Detect container runtime: prefer finch, fall back to docker if it works
+if command -v finch &>/dev/null && finch vm status &>/dev/null; then
+    CONTAINER_CMD=finch
+elif command -v docker &>/dev/null && docker info &>/dev/null 2>&1; then
+    CONTAINER_CMD=docker
+else
+    echo "Error: no working container runtime found (tried finch, docker)" >&2
+    exit 1
+fi
+echo "Using container runtime: $CONTAINER_CMD"
+
 EXAMPLE=${1:-basic-lambda}
 # Optional: set RIE_MAX_CONCURRENCY to enable LMI mode (emulates AWS_LAMBDA_MAX_CONCURRENCY)
 RIE_MAX_CONCURRENCY=${RIE_MAX_CONCURRENCY:-}
 
+echo "Building RIE from source (invocation-id support)..."
+if [ ! -f /tmp/aws-lambda-rie ]; then
+    # Build natively on host to avoid container DNS issues with VPN
+    if ! command -v go &>/dev/null; then
+        echo "Error: Go is required to build RIE. Install with: brew install go" >&2
+        exit 1
+    fi
+    RIE_BUILD_DIR=$(mktemp -d)
+    git clone --depth 50 https://github.com/aws/aws-lambda-runtime-interface-emulator.git "$RIE_BUILD_DIR/rie"
+    cd "$RIE_BUILD_DIR/rie"
+    git checkout ddf2bc4f4c95a1c7fdb15633a80e65b9be8ffe46
+    # Cross-compile for Linux/amd64 (runs inside container)
+    CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o /tmp/aws-lambda-rie ./cmd/aws-lambda-rie
+    cd "$SCRIPT_DIR/.."
+    rm -rf "$RIE_BUILD_DIR"
+fi
+cp /tmp/aws-lambda-rie "$SCRIPT_DIR/../aws-lambda-rie"
+
 echo "Building Docker image with RIE for example: $EXAMPLE..."
-docker build -f Dockerfile.rie --build-arg EXAMPLE=$EXAMPLE -t rust-lambda-rie-test .
+$CONTAINER_CMD build --dns 8.8.8.8 -f Dockerfile.rie --build-arg EXAMPLE=$EXAMPLE -t rust-lambda-rie-test .
 
 echo "Starting RIE container on port 9000..."
 if [ -n "$RIE_MAX_CONCURRENCY" ]; then
     echo "Enabling LMI mode with AWS_LAMBDA_MAX_CONCURRENCY=$RIE_MAX_CONCURRENCY"
-    docker run -p 9000:8080 -e AWS_LAMBDA_MAX_CONCURRENCY="$RIE_MAX_CONCURRENCY" rust-lambda-rie-test &
+    $CONTAINER_CMD run -p 9000:8080 -e AWS_LAMBDA_MAX_CONCURRENCY="$RIE_MAX_CONCURRENCY" rust-lambda-rie-test &
 else
-    docker run -p 9000:8080 rust-lambda-rie-test &
+    $CONTAINER_CMD run -p 9000:8080 rust-lambda-rie-test &
 fi
 CONTAINER_PID=$!
 

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 invoation_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,
+    )]
+