diff --git a/lambda-extension/examples/basic.rs b/lambda-extension/examples/basic.rs index 6ed92652..aedc046b 100644 --- a/lambda-extension/examples/basic.rs +++ b/lambda-extension/examples/basic.rs @@ -1,7 +1,7 @@ -use lambda_extension::{extension_fn, Error, NextEvent}; +use lambda_extension::{extension_fn, Error, LambdaEvent, NextEvent}; -async fn my_extension(event: NextEvent) -> Result<(), Error> { - match event { +async fn my_extension(event: LambdaEvent) -> Result<(), Error> { + match event.next { NextEvent::Shutdown(_e) => { // do something with the shutdown event } diff --git a/lambda-extension/examples/custom_events.rs b/lambda-extension/examples/custom_events.rs index cd3dae46..560fa295 100644 --- a/lambda-extension/examples/custom_events.rs +++ b/lambda-extension/examples/custom_events.rs @@ -1,7 +1,7 @@ -use lambda_extension::{extension_fn, Error, NextEvent, Runtime}; +use lambda_extension::{extension_fn, Error, LambdaEvent, NextEvent, Runtime}; -async fn my_extension(event: NextEvent) -> Result<(), Error> { - match event { +async fn my_extension(event: LambdaEvent) -> Result<(), Error> { + match event.next { NextEvent::Shutdown(_e) => { // do something with the shutdown event } diff --git a/lambda-extension/examples/custom_trait_implementation.rs b/lambda-extension/examples/custom_trait_implementation.rs index 910a7e07..7f056957 100644 --- a/lambda-extension/examples/custom_trait_implementation.rs +++ b/lambda-extension/examples/custom_trait_implementation.rs @@ -1,4 +1,4 @@ -use lambda_extension::{run, Error, Extension, InvokeEvent, NextEvent}; +use lambda_extension::{run, Error, Extension, InvokeEvent, LambdaEvent, NextEvent}; use std::{ future::{ready, Future}, pin::Pin, @@ -11,8 +11,8 @@ struct MyExtension { impl Extension for MyExtension { type Fut = Pin>>>; - fn call(&mut self, event: NextEvent) -> Self::Fut { - match event { + fn call(&mut self, event: LambdaEvent) -> Self::Fut { + match event.next { NextEvent::Shutdown(_e) => { self.data.clear(); } diff --git a/lambda-extension/src/lib.rs b/lambda-extension/src/lib.rs index f1d34a18..c38947e5 100644 --- a/lambda-extension/src/lib.rs +++ b/lambda-extension/src/lib.rs @@ -94,12 +94,21 @@ impl NextEvent { } } +/// Wrapper with information about the next +/// event that the Lambda Runtime is going to process +pub struct LambdaEvent { + /// ID assigned to this extension by the Lambda Runtime + pub extension_id: String, + /// Next incoming event + pub next: NextEvent, +} + /// A trait describing an asynchronous extension. pub trait Extension { /// Response of this Extension. type Fut: Future>; /// Handle the incoming event. - fn call(&mut self, event: NextEvent) -> Self::Fut; + fn call(&mut self, event: LambdaEvent) -> Self::Fut; } /// Returns a new [`ExtensionFn`] with the given closure. @@ -119,11 +128,11 @@ pub struct ExtensionFn { impl Extension for ExtensionFn where - F: Fn(NextEvent) -> Fut, + F: Fn(LambdaEvent) -> Fut, Fut: Future>, { type Fut = Fut; - fn call(&mut self, event: NextEvent) -> Self::Fut { + fn call(&mut self, event: LambdaEvent) -> Self::Fut { (self.f)(event) } } @@ -173,6 +182,11 @@ where let event: NextEvent = serde_json::from_slice(&body)?; let is_invoke = event.is_invoke(); + let event = LambdaEvent { + extension_id: self.extension_id.clone(), + next: event, + }; + let res = extension.call(event).await; if let Err(error) = res { let req = if is_invoke { diff --git a/lambda-integration-tests/src/bin/extension-fn.rs b/lambda-integration-tests/src/bin/extension-fn.rs index fabdb4cd..5835b0f7 100644 --- a/lambda-integration-tests/src/bin/extension-fn.rs +++ b/lambda-integration-tests/src/bin/extension-fn.rs @@ -1,8 +1,8 @@ -use lambda_extension::{extension_fn, Error, NextEvent}; +use lambda_extension::{extension_fn, Error, LambdaEvent, NextEvent}; use tracing::info; -async fn my_extension(event: NextEvent) -> Result<(), Error> { - match event { +async fn my_extension(event: LambdaEvent) -> Result<(), Error> { + match event.next { NextEvent::Shutdown(e) => { info!("[extension-fn] Shutdown event received: {:?}", e); } diff --git a/lambda-integration-tests/src/bin/extension-trait.rs b/lambda-integration-tests/src/bin/extension-trait.rs index 62ab2dc7..bc4b3b32 100644 --- a/lambda-integration-tests/src/bin/extension-trait.rs +++ b/lambda-integration-tests/src/bin/extension-trait.rs @@ -1,4 +1,4 @@ -use lambda_extension::{Error, Extension, NextEvent}; +use lambda_extension::{Error, Extension, LambdaEvent, NextEvent}; use std::{ future::{ready, Future}, pin::Pin, @@ -13,8 +13,8 @@ struct MyExtension { impl Extension for MyExtension { type Fut = Pin>>>; - fn call(&mut self, event: NextEvent) -> Self::Fut { - match event { + fn call(&mut self, event: LambdaEvent) -> Self::Fut { + match event.next { NextEvent::Shutdown(e) => { info!("[extension] Shutdown event received: {:?}", e); }