Wrap incoming extension events in a struct (#394)

This new structure can be used to pass additional information to the extension in each invocation, like the extension id that the Runtime assigns to the extension when it's initialized.

Signed-off-by: David Calavera <[email protected]>

Author: calavera

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