Skip to content

Commit 31be935

Browse files
authored
feat: add lambda-runtime-invocation-id header (#1159)
* feat: add support for invocation-id * test: add multiconcurrency testing * chore: additional fixes * cohre: other changes * chore: code review * chore: add serde default to example * chore: fix the possible problem with malformed bytes * chore: fix append * chore: change message for 410 * chore: add rate limiting capability to the loggin. * chore: fmt * chore: taking rate_limiter from Dial 9 * chore: last fixes
1 parent 8c02edc commit 31be935

19 files changed

Lines changed: 634 additions & 94 deletions

File tree

.env

Lines changed: 0 additions & 16 deletions
This file was deleted.

.github/workflows/dockerized-test.yml

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,10 @@ jobs:
1919

2020
- uses: dtolnay/rust-toolchain@stable
2121
- uses: Swatinem/rust-cache@v2
22-
- name: Load environment variables
23-
run: |
24-
if [ -f .env ]; then
25-
set -a
26-
source .env
27-
set +a
28-
echo "HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD}" >> $GITHUB_ENV
29-
echo "OUTPUT_DIR=${OUTPUT_DIR}" >> $GITHUB_ENV
30-
fi
31-
3222
- name: Build Lambda artifacts for testing
3323
run: |
3424
mkdir -p test/dockerized/tasks
35-
OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples
25+
HANDLERS_TO_BUILD="basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh
3626
ls -la test/dockerized/tasks/
3727
3828
- name: Build base test image with RIE and custom entrypoint
@@ -57,7 +47,7 @@ jobs:
5747
- name: Build Lambda artifacts for testing
5848
run: |
5949
mkdir -p test/dockerized/tasks
60-
HANDLERS_TO_BUILD="basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples
50+
HANDLERS_TO_BUILD="basic-lambda-concurrent invocation-id-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh
6151
ls -la test/dockerized/tasks/
6252
6353
- name: Build base test image with RIE and custom entrypoint
@@ -68,6 +58,8 @@ jobs:
6858

6959
- name: Run concurrent scenarios
7060
uses: aws/containerized-test-runner-for-aws-lambda@main
61+
env:
62+
CONTAINER_READY_DELAY_SECS: 5
7163
with:
7264
suiteFileArray: '[]'
7365
dockerImageName: 'local/test-base'

Dockerfile.rie

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,20 @@ RUN dnf install -y gcc
44
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
55
ENV PATH="/root/.cargo/bin:${PATH}"
66

7-
ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie
8-
RUN chmod +x /usr/local/bin/aws-lambda-rie
7+
ARG TARGETARCH
8+
ENV RIE_VERSION=1.36 \
9+
RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \
10+
RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \
11+
RIE_PATH=/usr/local/bin/aws-lambda-rie
12+
13+
COPY scripts/download-rie.sh /tmp/download-rie.sh
14+
RUN sh /tmp/download-rie.sh \
15+
"${TARGETARCH}" \
16+
"${RIE_VERSION}" \
17+
"${RIE_SHA256_AMD64}" \
18+
"${RIE_SHA256_ARM64}" \
19+
"${RIE_PATH}" \
20+
&& rm /tmp/download-rie.sh
921

1022
ARG EXAMPLE=basic-lambda
1123

Dockerfile.test

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,19 @@
11
FROM public.ecr.aws/lambda/provided:al2023
22

3-
ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie
4-
RUN chmod +x /usr/local/bin/aws-lambda-rie
3+
ARG TARGETARCH
4+
ENV RIE_VERSION=1.36 \
5+
RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \
6+
RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \
7+
RIE_PATH=/usr/local/bin/aws-lambda-rie
8+
9+
COPY scripts/download-rie.sh /tmp/download-rie.sh
10+
RUN sh /tmp/download-rie.sh \
11+
"${TARGETARCH}" \
12+
"${RIE_VERSION}" \
13+
"${RIE_SHA256_AMD64}" \
14+
"${RIE_SHA256_ARM64}" \
15+
"${RIE_PATH}" \
16+
&& rm /tmp/download-rie.sh
517

618
COPY scripts/custom-lambda-entrypoint.sh /usr/local/bin/lambda-entrypoint
719
RUN chmod +x /usr/local/bin/lambda-entrypoint

Makefile

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,10 @@ INTEG_EXTENSIONS := extension-fn extension-trait logs-trait
77
INTEG_ARCH := x86_64-unknown-linux-musl
88
RIE_MAX_CONCURRENCY ?= 4
99
TEST_RUNNER_BRANCH ?= main
10+
CONTAINER_READY_DELAY_SECS ?= 5
1011
OUTPUT_DIR ?= test/dockerized/tasks
11-
HANDLERS_TO_BUILD ?=
12-
HANDLER ?=
13-
14-
# Load environment variables from .env file if it exists
15-
-include .env
16-
export
12+
HANDLERS_TO_BUILD ?= basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent
13+
HANDLER ?= basic-lambda
1714

1815
.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
1916

@@ -125,7 +122,7 @@ fmt:
125122
cargo +nightly fmt --all
126123

127124
build-examples:
128-
HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD} OUTPUT_DIR=${OUTPUT_DIR} ./scripts/build-examples.sh
125+
HANDLERS_TO_BUILD="$(subst ",,$(HANDLERS_TO_BUILD))" OUTPUT_DIR="$(OUTPUT_DIR)" ./scripts/build-examples.sh
129126

130127
nuke:
131128
docker kill $$(docker ps -q)
@@ -145,6 +142,7 @@ build-test-runner: build-examples
145142
@echo "Building test runner Docker image..."
146143
@docker build -t test-runner:local -f .test-runner/Dockerfile .test-runner
147144

145+
test-dockerized-concurrent: HANDLERS_TO_BUILD := basic-lambda-concurrent invocation-id-concurrent
148146
test-dockerized-concurrent: build-test-runner
149147
@echo "Running concurrent scenarios in Docker..."
150148
@docker network rm concurrent-test-net 2>/dev/null || true
@@ -156,6 +154,7 @@ test-dockerized-concurrent: build-test-runner
156154
-e TASK_FOLDER=./test/dockerized/tasks \
157155
-e GITHUB_WORKSPACE=/workspace \
158156
-e DOCKER_SHARED_NETWORK=concurrent-test-net \
157+
-e CONTAINER_READY_DELAY_SECS=$(CONTAINER_READY_DELAY_SECS) \
159158
-v /var/run/docker.sock:/var/run/docker.sock \
160159
-v "$(CURDIR):/workspace" \
161160
-w /workspace \
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
[package]
2+
name = "invocation-id-concurrent"
3+
version = "0.1.0"
4+
edition = "2021"
5+
6+
[dependencies]
7+
lambda_runtime = { path = "../../lambda-runtime", features = ["concurrency-tokio"] }
8+
serde = "1"
9+
tokio = { version = "1", features = ["macros", "rt", "time"] }
10+
11+
[dev-dependencies]
12+
serde_json = "1.0"
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// This example requires the following input to succeed:
2+
// { "command": "do something" }
3+
4+
use lambda_runtime::{service_fn, tracing, Diagnostic, Error, LambdaEvent};
5+
use serde::{Deserialize, Serialize};
6+
7+
#[derive(Deserialize)]
8+
struct Request {
9+
#[serde(default, rename = "command")]
10+
_command: String,
11+
sleep: u32,
12+
}
13+
14+
#[derive(Serialize, Debug, PartialEq)]
15+
struct Response {
16+
from: String,
17+
}
18+
19+
#[derive(Debug)]
20+
struct HandlerError(String);
21+
22+
impl std::fmt::Display for HandlerError {
23+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24+
write!(f, "{}", self.0)
25+
}
26+
}
27+
28+
impl From<HandlerError> for Diagnostic {
29+
fn from(e: HandlerError) -> Diagnostic {
30+
Diagnostic {
31+
error_type: "HandlerError".into(),
32+
error_message: e.0,
33+
}
34+
}
35+
}
36+
37+
38+
/**
39+
* Cross-wiring protection: duplicate request-id after timeout.
40+
41+
Timeline:
42+
t=0: Invoke A starts, handler sleeps 7s
43+
t=5: A times out (timeout=5s). Batch 1 completes with timeout error.
44+
t=5: Invoke B starts (same request-id), handler sleeps 4s
45+
t=7: A's handler wakes up, posts stale /response/{same-id}
46+
t=9: B's handler wakes up, posts correct /response/{same-id}
47+
48+
With invocation-id: A's stale post at t=7 gets 410 Gone. B responds at t=9 correctly.
49+
Without: A's stale response at t=7 is accepted for B (cross-wired).
50+
*/
51+
52+
#[tokio::main]
53+
async fn main() -> Result<(), Error> {
54+
// required to enable CloudWatch error logging by the runtime
55+
tracing::init_default_subscriber();
56+
let max_concurrency = std::env::var("AWS_LAMBDA_MAX_CONCURRENCY").unwrap_or_else(|_| "not set".to_string());
57+
tracing::info!(AWS_LAMBDA_MAX_CONCURRENCY = %max_concurrency, "starting concurrent handler");
58+
59+
let func = service_fn(my_handler);
60+
if let Err(err) = lambda_runtime::run_concurrent(func).await {
61+
tracing::error!(error = %err, "run error");
62+
return Err(err);
63+
}
64+
Ok(())
65+
}
66+
67+
pub(crate) async fn my_handler(event: LambdaEvent<Request>) -> Result<Response, HandlerError> {
68+
if event.payload.sleep > 0 {
69+
tokio::time::sleep(tokio::time::Duration::from_secs(event.payload.sleep.into())).await;
70+
}
71+
72+
Ok(Response {
73+
from: event.payload._command,
74+
})
75+
}
76+
77+
#[cfg(test)]
78+
mod tests {
79+
use super::*;
80+
use lambda_runtime::{Context, LambdaEvent};
81+
82+
#[tokio::test]
83+
async fn handler_echoes_marker() {
84+
let event = LambdaEvent {
85+
payload: Request {
86+
_command: "invoke-B".into(),
87+
sleep: 0,
88+
},
89+
context: Context::default(),
90+
};
91+
92+
let result = my_handler(event).await.unwrap();
93+
94+
assert_eq!(result, Response { from: "invoke-B".into() });
95+
}
96+
97+
#[test]
98+
fn request_defaults_missing_command() {
99+
let request: Request = serde_json::from_str(r#"{"sleep": 0}"#).unwrap();
100+
101+
assert_eq!(request._command, "");
102+
}
103+
}

lambda-runtime/src/constants.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
/// Header names used in the Lambda Runtime API.
2+
pub(crate) const LAMBDA_RUNTIME_REQUEST_ID: &str = "lambda-runtime-aws-request-id";
3+
pub(crate) const LAMBDA_RUNTIME_DEADLINE_MS: &str = "lambda-runtime-deadline-ms";
4+
pub(crate) const LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN: &str = "lambda-runtime-invoked-function-arn";
5+
pub(crate) const LAMBDA_RUNTIME_TRACE_ID: &str = "lambda-runtime-trace-id";
6+
pub(crate) const LAMBDA_RUNTIME_CLIENT_CONTEXT: &str = "lambda-runtime-client-context";
7+
pub(crate) const LAMBDA_RUNTIME_COGNITO_IDENTITY: &str = "lambda-runtime-cognito-identity";
8+
pub(crate) const LAMBDA_RUNTIME_TENANT_ID: &str = "lambda-runtime-aws-tenant-id";
9+
pub(crate) const LAMBDA_RUNTIME_INVOCATION_ID: &str = "lambda-runtime-invocation-id";

lambda-runtime/src/layers/api_client.rs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,12 @@ where
114114
// Adding more information on top of 410 Gone, to make it more clear since we cannot access the body of the message
115115
if status == 410 {
116116
log_or_print!(
117-
tracing: tracing::error!("Lambda function timeout!"),
118-
fallback: eprintln!("Lambda function timeout!")
117+
tracing: tracing::error!(
118+
"Lambda Runtime API rejected the response: invocation timed out or response was stale"
119+
),
120+
fallback: eprintln!(
121+
"Lambda Runtime API rejected the response: invocation timed out or response was stale"
122+
)
119123
);
120124
}
121125

@@ -219,7 +223,9 @@ mod tests {
219223

220224
// Verify the error was logged
221225
assert!(logs_contain("Lambda Runtime API returned non-200 response"));
222-
assert!(logs_contain("Lambda function timeout!"));
226+
assert!(logs_contain(
227+
"Lambda Runtime API rejected the response: invocation timed out or response was stale"
228+
));
223229
}
224230

225231
#[tokio::test]

0 commit comments

Comments
 (0)