Skip to content

Commit 07de1f5

Browse files
darklight3itdarklight3it
committed
feat: add support for invocation-id
1 parent 8c02edc commit 07de1f5

6 files changed

Lines changed: 217 additions & 34 deletions

File tree

‎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_response.rs‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,10 @@ where
123123
};
124124

125125
let request_id = req.context.request_id.clone();
126+
let invocation_id = req.context.invocation_id.clone();
126127
let lambda_event = match deserializer::deserialize::<EventPayload>(&req.body, req.context) {
127128
Ok(lambda_event) => lambda_event,
128-
Err(err) => match build_event_error_request(&request_id, err) {
129+
Err(err) => match build_event_error_request(request_id, invocation_id, err) {
129130
Ok(request) => return RuntimeApiResponseFuture::Ready(Box::new(Some(Ok(request)))),
130131
Err(err) => {
131132
error!(error = ?err, "failed to build error response for Lambda Runtime API");
@@ -137,23 +138,28 @@ where
137138
// Once the handler input has been generated successfully, pass it through to inner services
138139
// allowing processing both before reaching the handler function and after the handler completes.
139140
let fut = self.inner.call(lambda_event);
140-
RuntimeApiResponseFuture::Future(fut, request_id, PhantomData)
141+
RuntimeApiResponseFuture::Future(fut, request_id, invocation_id, PhantomData)
141142
}
142143
}
143144

144-
fn build_event_error_request<T>(request_id: &str, err: T) -> Result<http::Request<Body>, BoxError>
145+
fn build_event_error_request<T>(
146+
request_id: String,
147+
invocation_id: Option<String>,
148+
err: T,
149+
) -> Result<http::Request<Body>, BoxError>
145150
where
146151
T: Into<Diagnostic> + Debug,
147152
{
148153
error!(error = ?err, "Request payload deserialization into LambdaEvent<T> failed. The handler will not be called. Log at TRACE level to see the payload.");
149-
EventErrorRequest::new(request_id, err).into_req()
154+
EventErrorRequest::new(&request_id, invocation_id.as_deref(), err).into_req()
150155
}
151156

152157
#[pin_project(project = RuntimeApiResponseFutureProj)]
153158
pub enum RuntimeApiResponseFuture<F, Response, BufferedResponse, StreamingResponse, StreamItem, StreamError> {
154159
Future(
155160
#[pin] F,
156161
String,
162+
Option<String>,
157163
PhantomData<(
158164
(),
159165
Response,
@@ -183,9 +189,9 @@ where
183189

184190
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
185191
task::Poll::Ready(match self.as_mut().project() {
186-
RuntimeApiResponseFutureProj::Future(fut, request_id, _) => match ready!(fut.poll(cx)) {
187-
Ok(ok) => EventCompletionRequest::new(request_id, ok).into_req(),
188-
Err(err) => EventErrorRequest::new(request_id, err).into_req(),
192+
RuntimeApiResponseFutureProj::Future(fut, request_id, invocation_id, _) => match ready!(fut.poll(cx)) {
193+
Ok(ok) => EventCompletionRequest::new(request_id, invocation_id.as_deref(), ok).into_req(),
194+
Err(err) => EventErrorRequest::new(request_id, invocation_id.as_deref(), err).into_req(),
189195
},
190196
RuntimeApiResponseFutureProj::Ready(ready) => ready.take().expect("future polled after completion"),
191197
})

‎lambda-runtime/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ pub use tower::{self, service_fn, Service};
2323
#[macro_use]
2424
mod macros;
2525

26+
mod constants;
27+
2628
/// Diagnostic utilities to convert Rust types into Lambda Error types.
2729
pub mod diagnostic;
2830
pub use diagnostic::Diagnostic;

‎lambda-runtime/src/requests.rs‎

Lines changed: 141 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
use crate::{types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse, IntoFunctionResponse};
1+
use crate::{
2+
constants::LAMBDA_RUNTIME_INVOCATION_ID, types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse,
3+
IntoFunctionResponse,
4+
};
25
use bytes::Bytes;
36
use http::{header::CONTENT_TYPE, Method, Request, Uri};
47
use lambda_runtime_api_client::{body::Body, build_request};
@@ -88,6 +91,7 @@ where
8891
E: Into<Error> + Send + Debug,
8992
{
9093
pub(crate) request_id: &'a str,
94+
pub(crate) invocation_id: Option<&'a str>,
9195
pub(crate) body: R,
9296
pub(crate) _unused_b: PhantomData<B>,
9397
pub(crate) _unused_s: PhantomData<S>,
@@ -102,9 +106,14 @@ where
102106
E: Into<Error> + Send + Debug,
103107
{
104108
/// Initialize a new EventCompletionRequest
105-
pub(crate) fn new(request_id: &'a str, body: R) -> EventCompletionRequest<'a, R, B, S, D, E> {
109+
pub(crate) fn new(
110+
request_id: &'a str,
111+
invocation_id: Option<&'a str>,
112+
body: R,
113+
) -> EventCompletionRequest<'a, R, B, S, D, E> {
106114
EventCompletionRequest {
107115
request_id,
116+
invocation_id,
108117
body,
109118
_unused_b: PhantomData::<B>,
110119
_unused_s: PhantomData::<S>,
@@ -129,7 +138,15 @@ where
129138
let body = serde_json::to_vec(&body)?;
130139
let body = Body::from(body);
131140

132-
let req = build_request().method(Method::POST).uri(uri).body(body)?;
141+
let mut req = build_request()
142+
.method(Method::POST)
143+
.uri(uri)
144+
.body(body)?;
145+
146+
if let Some(id) = self.invocation_id {
147+
req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?);
148+
}
149+
133150
Ok(req)
134151
}
135152
FunctionResponse::StreamingResponse(mut response) => {
@@ -145,6 +162,11 @@ where
145162
// See the details in Lambda Developer Doc: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html#runtimes-custom-response-streaming
146163
req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Type".parse()?);
147164
req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Body".parse()?);
165+
166+
if let Some(id) = self.invocation_id {
167+
req_headers.append(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?);
168+
}
169+
148170
req_headers.insert(
149171
"Content-Type",
150172
"application/vnd.awslambda.http-integration-response".parse()?,
@@ -193,29 +215,22 @@ where
193215
}
194216
}
195217

196-
#[test]
197-
fn test_event_completion_request() {
198-
let req = EventCompletionRequest::new("id", "hello, world!");
199-
let req = req.into_req().unwrap();
200-
let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response");
201-
assert_eq!(req.method(), Method::POST);
202-
assert_eq!(req.uri(), &expected);
203-
assert!(match req.headers().get("User-Agent") {
204-
Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
205-
None => false,
206-
});
207-
}
208-
209218
// /runtime/invocation/{AwsRequestId}/error
210219
pub(crate) struct EventErrorRequest<'a> {
211220
pub(crate) request_id: &'a str,
221+
pub(crate) invocation_id: Option<&'a str>,
212222
pub(crate) diagnostic: Diagnostic,
213223
}
214224

215225
impl<'a> EventErrorRequest<'a> {
216-
pub(crate) fn new(request_id: &'a str, diagnostic: impl Into<Diagnostic>) -> EventErrorRequest<'a> {
226+
pub(crate) fn new(
227+
request_id: &'a str,
228+
invocation_id: Option<&'a str>,
229+
diagnostic: impl Into<Diagnostic>,
230+
) -> EventErrorRequest<'a> {
217231
EventErrorRequest {
218232
request_id,
233+
invocation_id,
219234
diagnostic: diagnostic.into(),
220235
}
221236
}
@@ -228,11 +243,16 @@ impl IntoRequest for EventErrorRequest<'_> {
228243
let body = serde_json::to_vec(&self.diagnostic)?;
229244
let body = Body::from(body);
230245

231-
let req = build_request()
246+
let mut req = build_request()
232247
.method(Method::POST)
233248
.uri(uri)
234249
.header("lambda-runtime-function-error-type", "unhandled")
235250
.body(body)?;
251+
252+
if let Some(id) = self.invocation_id {
253+
req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?);
254+
}
255+
236256
Ok(req)
237257
}
238258
}
@@ -253,10 +273,93 @@ mod tests {
253273
});
254274
}
255275

276+
#[test]
277+
fn test_event_completion_request() {
278+
let req = EventCompletionRequest::new("id", Option::Some("invocation_id"), "hello, world!");
279+
let req = req.into_req().unwrap();
280+
let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response");
281+
assert_eq!(req.method(), Method::POST);
282+
assert_eq!(req.uri(), &expected);
283+
284+
assert!(req
285+
.headers()
286+
.get("User-Agent")
287+
.unwrap()
288+
.to_str()
289+
.unwrap()
290+
.starts_with("aws-lambda-rust/"));
291+
292+
assert_eq!(
293+
req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(),
294+
"invocation_id"
295+
);
296+
}
297+
298+
#[test]
299+
fn test_event_completion_request_invocation_id_not_added_when_none() {
300+
let req = EventCompletionRequest::new("id", Option::None, "hello, world!");
301+
let req = req.into_req().unwrap();
302+
303+
assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none());
304+
}
305+
306+
#[test]
307+
fn test_streaming_event_completion_request_with_invocation_id() {
308+
use crate::StreamResponse;
309+
310+
let runtime = tokio::runtime::Builder::new_current_thread()
311+
.enable_all()
312+
.build()
313+
.unwrap();
314+
315+
runtime.block_on(async {
316+
let stream = tokio_stream::iter(vec![Ok::<Bytes, Error>(Bytes::from_static(b"chunk"))]);
317+
let stream_response: StreamResponse<_> = stream.into();
318+
let response = FunctionResponse::StreamingResponse(stream_response);
319+
320+
let req: EventCompletionRequest<'_, _, (), _, _, _> =
321+
EventCompletionRequest::new("id", Some("invocation_id"), response);
322+
323+
let http_req = req.into_req().expect("into_req should succeed");
324+
let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response");
325+
assert_eq!(http_req.method(), Method::POST);
326+
assert_eq!(http_req.uri(), &expected);
327+
328+
assert_eq!(
329+
http_req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(),
330+
"invocation_id"
331+
);
332+
});
333+
}
334+
335+
#[test]
336+
fn test_streaming_event_completion_request_invocation_id_not_added_when_none() {
337+
use crate::StreamResponse;
338+
339+
let runtime = tokio::runtime::Builder::new_current_thread()
340+
.enable_all()
341+
.build()
342+
.unwrap();
343+
344+
runtime.block_on(async {
345+
let stream = tokio_stream::iter(vec![Ok::<Bytes, Error>(Bytes::from_static(b"chunk"))]);
346+
let stream_response: StreamResponse<_> = stream.into();
347+
let response = FunctionResponse::StreamingResponse(stream_response);
348+
349+
let req: EventCompletionRequest<'_, _, (), _, _, _> =
350+
EventCompletionRequest::new("id", None, response);
351+
352+
let http_req = req.into_req().expect("into_req should succeed");
353+
354+
assert!(http_req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none());
355+
});
356+
}
357+
256358
#[test]
257359
fn test_event_error_request() {
258360
let req = EventErrorRequest {
259361
request_id: "id",
362+
invocation_id: Option::Some("invocation_id"),
260363
diagnostic: Diagnostic {
261364
error_type: "InvalidEventDataError".into(),
262365
error_message: "Error parsing event data".into(),
@@ -270,6 +373,26 @@ mod tests {
270373
Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"),
271374
None => false,
272375
});
376+
377+
assert!(match req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID) {
378+
Some(header) => header.to_str().unwrap() == "invocation_id",
379+
None => false,
380+
});
381+
}
382+
383+
#[test]
384+
fn test_event_error_request_invocation_id_not_added_when_none() {
385+
let req = EventErrorRequest {
386+
request_id: "id",
387+
invocation_id: None,
388+
diagnostic: Diagnostic {
389+
error_type: "InvalidEventDataError".into(),
390+
error_message: "Error parsing event data".into(),
391+
},
392+
};
393+
let req = req.into_req().unwrap();
394+
395+
assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none());
273396
}
274397

275398
#[test]

‎lambda-runtime/src/runtime.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -790,7 +790,11 @@ mod endpoint_tests {
790790
let base = server.base_url().parse().expect("Invalid mock server Uri");
791791
let client = Client::builder().with_endpoint(base).build();
792792

793-
let req = EventCompletionRequest::new("156cb537-e2d4-11e8-9b34-d36013741fb9", "{}");
793+
let req = EventCompletionRequest::new(
794+
"156cb537-e2d4-11e8-9b34-d36013741fb9",
795+
Option::Some("invocation_id"),
796+
"{}",
797+
);
794798
let req = req.into_req()?;
795799

796800
let rsp = client.call(req).await?;
@@ -822,6 +826,7 @@ mod endpoint_tests {
822826

823827
let req = EventErrorRequest {
824828
request_id: "156cb537-e2d4-11e8-9b34-d36013741fb9",
829+
invocation_id: Option::Some("invocation_id"),
825830
diagnostic,
826831
};
827832
let req = req.into_req()?;

0 commit comments

Comments
 (0)