Wrap incoming extension events in a struct

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
         }

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
         }

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<Box<dyn Future<Output = Result<(), Error>>>>;
-    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();
             }

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<Output = Result<(), Error>>;
     /// 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<F> {
 
 impl<F, Fut> Extension for ExtensionFn<F>
 where
-    F: Fn(NextEvent) -> Fut,
+    F: Fn(LambdaEvent) -> Fut,
     Fut: Future<Output = Result<(), Error>>,
 {
     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 {

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);
         }

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<Box<dyn Future<Output = Result<(), Error>>>>;
 
-    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);
             }