Skip to content

Commit b7df6fc

Browse files
calaveracalavera
authored
Extract the request ID without allocating extra memory. (#735)
Changes the way that the Context is initialized to receive the request ID as an argument. This way we also avoid allocating additional memory for it. Signed-off-by: David Calavera <david.calavera@gmail.com>
1 parent 53637e7 commit b7df6fc

File tree

2 files changed

+67
-45
lines changed

2 files changed

+67
-45
lines changed

‎lambda-runtime/src/lib.rs‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ use hyper::{
1717
use lambda_runtime_api_client::Client;
1818
use serde::{Deserialize, Serialize};
1919
use std::{
20-
convert::TryFrom,
2120
env,
2221
fmt::{self, Debug, Display},
2322
future::Future,
@@ -41,6 +40,8 @@ mod types;
4140
use requests::{EventCompletionRequest, EventErrorRequest, IntoRequest, NextEventRequest};
4241
pub use types::{Context, FunctionResponse, IntoFunctionResponse, LambdaEvent, MetadataPrelude, StreamResponse};
4342

43+
use types::invoke_request_id;
44+
4445
/// Error type that lambdas may result in
4546
pub type Error = lambda_runtime_api_client::Error;
4647

@@ -121,6 +122,7 @@ where
121122
trace!("New event arrived (run loop)");
122123
let event = next_event_response?;
123124
let (parts, body) = event.into_parts();
125+
let request_id = invoke_request_id(&parts.headers)?;
124126

125127
#[cfg(debug_assertions)]
126128
if parts.status == http::StatusCode::NO_CONTENT {
@@ -130,19 +132,8 @@ where
130132
continue;
131133
}
132134

133-
let ctx: Context = Context::try_from((self.config.clone(), parts.headers))?;
134-
let request_id = &ctx.request_id.clone();
135-
136-
let request_span = match &ctx.xray_trace_id {
137-
Some(trace_id) => {
138-
env::set_var("_X_AMZN_TRACE_ID", trace_id);
139-
tracing::info_span!("Lambda runtime invoke", requestId = request_id, xrayTraceId = trace_id)
140-
}
141-
None => {
142-
env::remove_var("_X_AMZN_TRACE_ID");
143-
tracing::info_span!("Lambda runtime invoke", requestId = request_id)
144-
}
145-
};
135+
let ctx: Context = Context::new(request_id, self.config.clone(), &parts.headers)?;
136+
let request_span = ctx.request_span();
146137

147138
// Group the handling in one future and instrument it with the span
148139
async {

‎lambda-runtime/src/types.rs‎

Lines changed: 62 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
use crate::{Error, RefConfig};
22
use base64::prelude::*;
33
use bytes::Bytes;
4-
use http::{HeaderMap, HeaderValue, StatusCode};
4+
use http::{header::ToStrError, HeaderMap, HeaderValue, StatusCode};
55
use serde::{Deserialize, Serialize};
66
use std::{
77
collections::HashMap,
8-
convert::TryFrom,
8+
env,
99
fmt::Debug,
1010
time::{Duration, SystemTime},
1111
};
1212
use tokio_stream::Stream;
13+
use tracing::Span;
1314

1415
#[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)]
1516
#[serde(rename_all = "camelCase")]
@@ -120,11 +121,10 @@ pub struct Context {
120121
pub env_config: RefConfig,
121122
}
122123

123-
impl TryFrom<(RefConfig, HeaderMap)> for Context {
124-
type Error = Error;
125-
fn try_from(data: (RefConfig, HeaderMap)) -> Result<Self, Self::Error> {
126-
let env_config = data.0;
127-
let headers = data.1;
124+
impl Context {
125+
/// Create a new [Context] struct based on the fuction configuration
126+
/// and the incoming request data.
127+
pub fn new(request_id: &str, env_config: RefConfig, headers: &HeaderMap) -> Result<Self, Error> {
128128
let client_context: Option<ClientContext> = if let Some(value) = headers.get("lambda-runtime-client-context") {
129129
serde_json::from_str(value.to_str()?)?
130130
} else {
@@ -138,11 +138,7 @@ impl TryFrom<(RefConfig, HeaderMap)> for Context {
138138
};
139139

140140
let ctx = Context {
141-
request_id: headers
142-
.get("lambda-runtime-aws-request-id")
143-
.expect("missing lambda-runtime-aws-request-id header")
144-
.to_str()?
145-
.to_owned(),
141+
request_id: request_id.to_owned(),
146142
deadline: headers
147143
.get("lambda-runtime-deadline-ms")
148144
.expect("missing lambda-runtime-deadline-ms header")
@@ -165,13 +161,37 @@ impl TryFrom<(RefConfig, HeaderMap)> for Context {
165161

166162
Ok(ctx)
167163
}
168-
}
169164

170-
impl Context {
171165
/// The execution deadline for the current invocation.
172166
pub fn deadline(&self) -> SystemTime {
173167
SystemTime::UNIX_EPOCH + Duration::from_millis(self.deadline)
174168
}
169+
170+
/// Create a new [`tracing::Span`] for an incoming invocation.
171+
pub(crate) fn request_span(&self) -> Span {
172+
match &self.xray_trace_id {
173+
Some(trace_id) => {
174+
env::set_var("_X_AMZN_TRACE_ID", trace_id);
175+
tracing::info_span!(
176+
"Lambda runtime invoke",
177+
requestId = &self.request_id,
178+
xrayTraceId = trace_id
179+
)
180+
}
181+
None => {
182+
env::remove_var("_X_AMZN_TRACE_ID");
183+
tracing::info_span!("Lambda runtime invoke", requestId = &self.request_id)
184+
}
185+
}
186+
}
187+
}
188+
189+
/// Extract the invocation request id from the incoming request.
190+
pub(crate) fn invoke_request_id(headers: &HeaderMap) -> Result<&str, ToStrError> {
191+
headers
192+
.get("lambda-runtime-aws-request-id")
193+
.expect("missing lambda-runtime-aws-request-id header")
194+
.to_str()
175195
}
176196

177197
/// Incoming Lambda request containing the event payload and context.
@@ -313,7 +333,7 @@ mod test {
313333
HeaderValue::from_static("arn::myarn"),
314334
);
315335
headers.insert("lambda-runtime-trace-id", HeaderValue::from_static("arn::myarn"));
316-
let tried = Context::try_from((config, headers));
336+
let tried = Context::new("id", config, &headers);
317337
assert!(tried.is_ok());
318338
}
319339

@@ -324,7 +344,7 @@ mod test {
324344
let mut headers = HeaderMap::new();
325345
headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id"));
326346
headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123"));
327-
let tried = Context::try_from((config, headers));
347+
let tried = Context::new("id", config, &headers);
328348
assert!(tried.is_ok());
329349
}
330350

@@ -355,7 +375,7 @@ mod test {
355375
);
356376

357377
let config = Arc::new(Config::default());
358-
let tried = Context::try_from((config, headers));
378+
let tried = Context::new("id", config, &headers);
359379
assert!(tried.is_ok());
360380
let tried = tried.unwrap();
361381
assert!(tried.client_context.is_some());
@@ -369,7 +389,7 @@ mod test {
369389
headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id"));
370390
headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123"));
371391
headers.insert("lambda-runtime-client-context", HeaderValue::from_static("{}"));
372-
let tried = Context::try_from((config, headers));
392+
let tried = Context::new("id", config, &headers);
373393
assert!(tried.is_ok());
374394
assert!(tried.unwrap().client_context.is_some());
375395
}
@@ -390,7 +410,7 @@ mod test {
390410
"lambda-runtime-cognito-identity",
391411
HeaderValue::from_str(&cognito_identity_str).unwrap(),
392412
);
393-
let tried = Context::try_from((config, headers));
413+
let tried = Context::new("id", config, &headers);
394414
assert!(tried.is_ok());
395415
let tried = tried.unwrap();
396416
assert!(tried.identity.is_some());
@@ -412,7 +432,7 @@ mod test {
412432
HeaderValue::from_static("arn::myarn"),
413433
);
414434
headers.insert("lambda-runtime-trace-id", HeaderValue::from_static("arn::myarn"));
415-
let tried = Context::try_from((config, headers));
435+
let tried = Context::new("id", config, &headers);
416436
assert!(tried.is_err());
417437
}
418438

@@ -427,7 +447,7 @@ mod test {
427447
"lambda-runtime-client-context",
428448
HeaderValue::from_static("BAD-Type,not JSON"),
429449
);
430-
let tried = Context::try_from((config, headers));
450+
let tried = Context::new("id", config, &headers);
431451
assert!(tried.is_err());
432452
}
433453

@@ -439,7 +459,7 @@ mod test {
439459
headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id"));
440460
headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123"));
441461
headers.insert("lambda-runtime-cognito-identity", HeaderValue::from_static("{}"));
442-
let tried = Context::try_from((config, headers));
462+
let tried = Context::new("id", config, &headers);
443463
assert!(tried.is_err());
444464
}
445465

@@ -454,14 +474,13 @@ mod test {
454474
"lambda-runtime-cognito-identity",
455475
HeaderValue::from_static("BAD-Type,not JSON"),
456476
);
457-
let tried = Context::try_from((config, headers));
477+
let tried = Context::new("id", config, &headers);
458478
assert!(tried.is_err());
459479
}
460480

461481
#[test]
462482
#[should_panic]
463-
#[allow(unused_must_use)]
464-
fn context_with_missing_request_id_should_panic() {
483+
fn context_with_missing_deadline_should_panic() {
465484
let config = Arc::new(Config::default());
466485

467486
let mut headers = HeaderMap::new();
@@ -471,22 +490,34 @@ mod test {
471490
HeaderValue::from_static("arn::myarn"),
472491
);
473492
headers.insert("lambda-runtime-trace-id", HeaderValue::from_static("arn::myarn"));
474-
Context::try_from((config, headers));
493+
let _ = Context::new("id", config, &headers);
475494
}
476495

477496
#[test]
478-
#[should_panic]
479-
#[allow(unused_must_use)]
480-
fn context_with_missing_deadline_should_panic() {
481-
let config = Arc::new(Config::default());
497+
fn invoke_request_id_should_not_panic() {
498+
let mut headers = HeaderMap::new();
499+
headers.insert("lambda-runtime-aws-request-id", HeaderValue::from_static("my-id"));
500+
headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123"));
501+
headers.insert(
502+
"lambda-runtime-invoked-function-arn",
503+
HeaderValue::from_static("arn::myarn"),
504+
);
505+
headers.insert("lambda-runtime-trace-id", HeaderValue::from_static("arn::myarn"));
506+
507+
let _ = invoke_request_id(&headers);
508+
}
482509

510+
#[test]
511+
#[should_panic]
512+
fn invoke_request_id_should_panic() {
483513
let mut headers = HeaderMap::new();
484514
headers.insert("lambda-runtime-deadline-ms", HeaderValue::from_static("123"));
485515
headers.insert(
486516
"lambda-runtime-invoked-function-arn",
487517
HeaderValue::from_static("arn::myarn"),
488518
);
489519
headers.insert("lambda-runtime-trace-id", HeaderValue::from_static("arn::myarn"));
490-
Context::try_from((config, headers));
520+
521+
let _ = invoke_request_id(&headers);
491522
}
492523
}

0 commit comments

Comments
 (0)