From 49d0726017f3b57a85c94eebf0a351f72e8a20bf Mon Sep 17 00:00:00 2001 From: David Calavera Date: Fri, 7 Jan 2022 16:09:17 -0800 Subject: [PATCH] Wrap incoming extension events in a struct 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 --- lambda-extension/examples/basic.rs | 6 +++--- lambda-extension/examples/custom_events.rs | 6 +++--- .../examples/custom_trait_implementation.rs | 6 +++--- lambda-extension/src/lib.rs | 20 ++++++++++++++++--- .../src/bin/extension-fn.rs | 6 +++--- .../src/bin/extension-trait.rs | 6 +++--- 6 files changed, 32 insertions(+), 18 deletions(-) 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); }