Repository: awslabs/aws-lambda-rust-runtime (language: rust)

Issue:
Support for internal Lambda extensions
First off, thanks for these useful crates! It's nice to see official Rust support from AWS.

The `lambda-extension` crate currently only supports **external** [Lambda extensions](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html) (i.e., those running in a separate binary from the main runtime). AWS Lambda also supports **internal** Lambda extensions that run in the same process as the runtime, and it would be great if this crate also supported those.

The main requirement for supporting internal extensions is that each internal extension must register itself with the Lambda Runtime API [*before* the runtime thread calls `Next`](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html#runtimes-extensions-api-reg):
> The `Init` phase completes after the runtime and each registered extension indicate completion by sending a `Next` API request.

Once the main runtime calls `Next`, it's too late to register any additional extensions. The Lambda Runtime API returns an error if an extension tries to register itself at that point.

The current implementation of `Extension::run()` [registers the extension](https://github.com/awslabs/aws-lambda-rust-runtime/blob/3c53476206482535f6b73f118e09258aeeb1d6f8/lambda-extension/src/extension.rs#L222) with the Lambda Runtime API and then proceeds to [call `Next`](https://github.com/awslabs/aws-lambda-rust-runtime/blob/3c53476206482535f6b73f118e09258aeeb1d6f8/lambda-extension/src/extension.rs#L314). Since `Extension::run()` returns a single future that includes both registration and invocation of the extension, there's no way for the caller to determine when registration is complete, and that it's safe to call `Runtime::run()` (which immediately [calls `Next`](https://github.com/awslabs/aws-lambda-rust-runtime/blob/3c53476206482535f6b73f118e09258aeeb1d6f8/lambda-runtime/src/lib.rs#L121) and ends the `Init` phase). In other words, calling `Extension::run()` and `Runtime::run()` concurrently creates a race condition between registering the extensions and the runtime calling `Next`.

A straightforward solution would be to replace (or augment) the `Extension::run()` function with an API that looks like the following (note the new `Extension::register()` function and a new type representing a *registered* extension that's ready to be run):
```rust
let extension = Extension::new()
  .with_events(&["INVOKE"])
  .with_events_processor(...);
let registered_extension = extension.register().await;

futures::future::try_join(
  registered_extension.run(),
  Runtime::run(...),
).await
```

A potential downside to this API is that `.with_events` will allow a user to attempt to register the `SHUTDOWN` event, which will fail at runtime since it's not permitted for internal extensions. Also, I'm not sure whether the telemetry and logs processing functionality works for internal extensions or not. These issues could be solved by having separate types for internal and external extensions, along with an enum to represent the supported events instead of accepting raw strings. However, that's a more involved API redesign that may not be worth it at this point.

I'm happy to contribute a PR once there's consensus on the desired API!


Discussion / hints:
I don't know much about internal extensions, so my suggestion would be to start with the most simple thing we can do to address the immediate problem.

Let's create that `register` function that only registers the extension. If a user sends events that are not supported by an internal extension, or tries unsupported apis, they'll receive an error. If the error is not clear, we can improve the api.

Fix patch:
diff --git a/examples/extension-internal-flush/Cargo.toml b/examples/extension-internal-flush/Cargo.toml
new file mode 100644
index 00000000..daadd0eb
--- /dev/null
+++ b/examples/extension-internal-flush/Cargo.toml
@@ -0,0 +1,12 @@
+[package]
+name = "extension-internal-flush"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+anyhow = "1"
+aws_lambda_events = { path = "../../lambda-events" }
+lambda-extension = { path = "../../lambda-extension" }
+lambda_runtime = { path = "../../lambda-runtime" }
+serde = "1.0.136"
+tokio = { version = "1", features = ["macros", "sync"] }
diff --git a/examples/extension-internal-flush/README.md b/examples/extension-internal-flush/README.md
new file mode 100644
index 00000000..553f7a3d
--- /dev/null
+++ b/examples/extension-internal-flush/README.md
@@ -0,0 +1,30 @@
+# AWS Lambda runtime + internal extension example
+
+This example demonstrates how to build an AWS Lambda function that includes a
+[Lambda internal extension](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html).
+Unlike external extensions that run as separate processes, an internal extension runs within the
+main runtime process.
+
+One use case for internal extensions is to flush logs or telemetry data after the Lambda runtime
+handler has finished processing an event but before the execution environment is frozen awaiting the
+arrival of the next event. Without an explicit flush, telemetry data may never be sent since the
+execution environment will remain frozen and eventually be terminated if no additional events arrive.
+
+Note that for
+[synchronous](https://docs.aws.amazon.com/lambda/latest/dg/invocation-sync.html) Lambda invocations
+(e.g., via
+[Amazon API Gateway](https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-integrations.html)),
+the Lambda service returns the response to the caller immediately. Extensions may continue to run
+without introducing an observable delay.
+
+## Build & Deploy
+
+1. Install [cargo-lambda](https://github.com/cargo-lambda/cargo-lambda#installation)
+2. Build the extension with `cargo lambda build --release`
+3. Deploy the function to AWS Lambda with `cargo lambda deploy --iam-role YOUR_ROLE`
+
+The last command will give you an ARN for the extension layer that you can use in your functions.
+
+## Build for ARM 64
+
+Build the extension with `cargo lambda build --release --arm64`
diff --git a/examples/extension-internal-flush/src/main.rs b/examples/extension-internal-flush/src/main.rs
new file mode 100644
index 00000000..3706809d
--- /dev/null
+++ b/examples/extension-internal-flush/src/main.rs
@@ -0,0 +1,112 @@
+use anyhow::anyhow;
+use aws_lambda_events::sqs::{SqsBatchResponse, SqsEventObj};
+use lambda_extension::{service_fn, Error, Extension, NextEvent};
+use serde::{Deserialize, Serialize};
+use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
+use tokio::sync::Mutex;
+
+use std::sync::Arc;
+
+/// Implements an internal Lambda extension to flush logs/telemetry after each request.
+struct FlushExtension {
+    request_done_receiver: Mutex<UnboundedReceiver<()>>,
+}
+
+impl FlushExtension {
+    pub fn new(request_done_receiver: UnboundedReceiver<()>) -> Self {
+        Self {
+            request_done_receiver: Mutex::new(request_done_receiver),
+        }
+    }
+
+    pub async fn invoke(&self, event: lambda_extension::LambdaEvent) -> Result<(), Error> {
+        match event.next {
+            // NB: Internal extensions only support the INVOKE event.
+            NextEvent::Shutdown(shutdown) => {
+                return Err(anyhow!("extension received unexpected SHUTDOWN event: {:?}", shutdown).into());
+            }
+            NextEvent::Invoke(_e) => {}
+        }
+
+        eprintln!("[extension] waiting for event to be processed");
+
+        // Wait for runtime to finish processing event.
+        self.request_done_receiver
+            .lock()
+            .await
+            .recv()
+            .await
+            .ok_or_else(|| anyhow!("channel is closed"))?;
+
+        eprintln!("[extension] flushing logs and telemetry");
+
+        // <flush logs and telemetry here>
+
+        Ok(())
+    }
+}
+
+/// Object that you send to SQS and plan to process with the function.
+#[derive(Debug, Deserialize, Serialize)]
+struct Data {
+    a: String,
+    b: i64,
+}
+
+/// Implements the main event handler for processing events from an SQS queue.
+struct EventHandler {
+    request_done_sender: UnboundedSender<()>,
+}
+
+impl EventHandler {
+    pub fn new(request_done_sender: UnboundedSender<()>) -> Self {
+        Self { request_done_sender }
+    }
+
+    pub async fn invoke(
+        &self,
+        event: lambda_runtime::LambdaEvent<SqsEventObj<Data>>,
+    ) -> Result<SqsBatchResponse, Error> {
+        let data = &event.payload.records[0].body;
+        eprintln!("[runtime] received event {data:?}");
+
+        // <process event here>
+
+        // Notify the extension to flush traces.
+        self.request_done_sender.send(()).map_err(Box::new)?;
+
+        Ok(SqsBatchResponse::default())
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    let (request_done_sender, request_done_receiver) = unbounded_channel::<()>();
+
+    let flush_extension = Arc::new(FlushExtension::new(request_done_receiver));
+    let extension = Extension::new()
+        // Internal extensions only support INVOKE events.
+        .with_events(&["INVOKE"])
+        .with_events_processor(service_fn(|event| {
+            let flush_extension = flush_extension.clone();
+            async move { flush_extension.invoke(event).await }
+        }))
+        // Internal extension names MUST be unique within a given Lambda function.
+        .with_extension_name("internal-flush")
+        // Extensions MUST be registered before calling lambda_runtime::run(), which ends the Init
+        // phase and begins the Invoke phase.
+        .register()
+        .await?;
+
+    let handler = Arc::new(EventHandler::new(request_done_sender));
+
+    tokio::try_join!(
+        lambda_runtime::run(service_fn(|event| {
+            let handler = handler.clone();
+            async move { handler.invoke(event).await }
+        })),
+        extension.run(),
+    )?;
+
+    Ok(())
+}
diff --git a/lambda-extension/src/extension.rs b/lambda-extension/src/extension.rs
index 4747b041..d653e0dc 100644
--- a/lambda-extension/src/extension.rs
+++ b/lambda-extension/src/extension.rs
@@ -215,14 +215,21 @@ where
         }
     }
 
-    /// Execute the given extension
-    pub async fn run(self) -> Result<(), Error> {
+    /// Register the extension.
+    ///
+    /// Performs the
+    /// [init phase](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-ib)
+    /// Lambda lifecycle operations to register the extension. When implementing an internal Lambda
+    /// extension, it is safe to call `lambda_runtime::run` once the future returned by this
+    /// function resolves.
+    pub async fn register(self) -> Result<RegisteredExtension<E>, Error> {
         let client = &Client::builder().build()?;
 
         let extension_id = register(client, self.extension_name, self.events).await?;
         let extension_id = extension_id.to_str()?;
-        let mut ep = self.events_processor;
 
+        // Logs API subscriptions must be requested during the Lambda init phase (see
+        // https://docs.aws.amazon.com/lambda/latest/dg/runtimes-logs-api.html#runtimes-logs-api-subscribing).
         if let Some(mut log_processor) = self.logs_processor {
             trace!("Log processor found");
 
@@ -262,6 +269,8 @@ where
             trace!("Registered extension with Logs API");
         }
 
+        // Telemetry API subscriptions must be requested during the Lambda init phase (see
+        // https://docs.aws.amazon.com/lambda/latest/dg/telemetry-api.html#telemetry-api-registration
         if let Some(mut telemetry_processor) = self.telemetry_processor {
             trace!("Telemetry processor found");
 
@@ -301,6 +310,42 @@ where
             trace!("Registered extension with Telemetry API");
         }
 
+        Ok(RegisteredExtension {
+            extension_id: extension_id.to_string(),
+            events_processor: self.events_processor,
+        })
+    }
+
+    /// Execute the given extension.
+    pub async fn run(self) -> Result<(), Error> {
+        self.register().await?.run().await
+    }
+}
+
+/// An extension registered by calling [`Extension::register`].
+pub struct RegisteredExtension<E> {
+    extension_id: String,
+    events_processor: E,
+}
+
+impl<E> RegisteredExtension<E>
+where
+    E: Service<LambdaEvent>,
+    E::Future: Future<Output = Result<(), E::Error>>,
+    E::Error: Into<Box<dyn std::error::Error + Send + Sync>> + fmt::Display + fmt::Debug,
+{
+    /// Execute the extension's run loop.
+    ///
+    /// Performs the
+    /// [invoke](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-invoke)
+    /// and, for external Lambda extensions registered to receive the `SHUTDOWN` event, the
+    /// [shutdown](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-shutdown)
+    /// Lambda lifecycle phases.
+    pub async fn run(self) -> Result<(), Error> {
+        let client = &Client::builder().build()?;
+        let mut ep = self.events_processor;
+        let extension_id = &self.extension_id;
+
         let incoming = async_stream::stream! {
             loop {
                 trace!("Waiting for next event (incoming loop)");
@@ -351,6 +396,8 @@ where
                 return Err(err.into());
             }
         }
+
+        // Unreachable.
         Ok(())
     }
 }


Test patch: