Add Lambda Extension crate

This new crate encapsulates the logic to create Lambda Extensions.
It includes reference examples.

Signed-off-by: David Calavera <david.calavera@gmail.com>

Author: calavera

Cargo.toml

@@ -3,4 +3,5 @@ members = [
     "lambda-http",
     "lambda-runtime-client",
     "lambda-runtime",
+    "lambda-extension"
 ]

lambda-extension/Cargo.toml

@@ -0,0 +1,35 @@
+[package]
+name = "lambda_extension"
+version = "0.1.0"
+edition = "2021"
+authors = ["David Calavera <david.calavera@gmail.com>"]
+description = "AWS Lambda Extension API"
+license = "Apache-2.0"
+repository = "https://github.com/awslabs/aws-lambda-rust-runtime"
+categories = ["web-programming::http-server"]
+keywords = ["AWS", "Lambda", "API"]
+readme = "../README.md"
+
+[dependencies]
+anyhow = "1.0.48"
+async-trait = "0.1.51"
+tokio = { version = "1.0", features = ["macros", "io-util", "sync", "rt-multi-thread"] }
+hyper = { version = "0.14", features = ["http1", "client", "server", "stream", "runtime"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "^1"
+bytes = "1.0"
+http = "0.2"
+async-stream = "0.3"
+futures = "0.3"
+tracing-error = "0.2"
+tracing = { version = "0.1", features = ["log"] }
+tower-service = "0.3"
+tokio-stream = "0.1.2"
+lambda_runtime_client = { version = "*", path = "../lambda-runtime-client" }
+
+[dev-dependencies]
+tracing-subscriber = "0.3"
+once_cell = "1.4.0"
+simple_logger = "1.6.0"
+log = "^0.4"
+simple-error = "0.2"

lambda-extension/examples/basic.rs

@@ -0,0 +1,26 @@
+use lambda_extension::{run, Error, Extension, InvokeEvent, ShutdownEvent};
+use log::LevelFilter;
+use simple_logger::SimpleLogger;
+
+struct BasicExtension {}
+
+#[async_trait::async_trait]
+impl Extension for BasicExtension {
+    async fn on_invoke(&self, _extension_id: &str, _event: InvokeEvent) -> Result<(), Error> {
+        Ok(())
+    }
+
+    async fn on_shutdown(&self, _extension_id: &str, _event: ShutdownEvent) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // required to enable CloudWatch error logging by the runtime
+    // can be replaced with any other method of initializing `log`
+    SimpleLogger::new().with_level(LevelFilter::Info).init().unwrap();
+
+    run(BasicExtension {}).await?;
+    Ok(())
+}

lambda-extension/examples/error_handling.rs

@@ -0,0 +1,55 @@
+use lambda_extension::{run, Error, Extension, InvokeEvent, ShutdownEvent, requests::exit_error, requests::init_error};
+use lambda_runtime_client::Client;
+use log::LevelFilter;
+use simple_logger::SimpleLogger;
+
+#[derive(Debug)]
+enum ErrorExample {
+    OnInvokeError,
+    OnShutdownError,
+}
+
+impl std::error::Error for ErrorExample {}
+
+impl std::fmt::Display for ErrorExample {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+          ErrorExample::OnInvokeError => write!(f, "error processing invocation call"),
+          ErrorExample::OnShutdownError => write!(f, "error processing shutdown call"),
+        }
+      }
+}
+
+struct ErrorHandlingExtension {
+    client: Client
+}
+
+#[async_trait::async_trait]
+impl Extension for ErrorHandlingExtension {
+    async fn on_invoke(&self, extension_id: &str, _event: InvokeEvent) -> Result<(), Error> {
+        let err = ErrorExample::OnInvokeError;
+        let req = init_error(extension_id, &format!("{}", err), None)?;
+        self.client.call(req).await?;
+        Err(Box::new(err))
+    }
+
+    async fn on_shutdown(&self, extension_id: &str, _event: ShutdownEvent) -> Result<(), Error> {
+        let err = ErrorExample::OnShutdownError;
+        let req = exit_error(extension_id, &format!("{}", err), None)?;
+        self.client.call(req).await?;
+        Err(Box::new(err))
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // required to enable CloudWatch error logging by the runtime
+    // can be replaced with any other method of initializing `log`
+    SimpleLogger::new().with_level(LevelFilter::Info).init().unwrap();
+
+    let client = Client::builder().build()?;
+    let extension = ErrorHandlingExtension { client };
+
+    run(extension).await?;
+    Ok(())
+}

lambda-extension/src/lib.rs

@@ -0,0 +1,129 @@
+// #![deny(clippy::all, clippy::cargo)]
+// #![warn(missing_docs,? nonstandard_style, rust_2018_idioms)]
+
+use async_trait::async_trait;
+use hyper::client::{connect::Connection, HttpConnector};
+use lambda_runtime_client::Client;
+use serde::Deserialize;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_stream::{StreamExt};
+use tower_service::Service;
+use tracing::trace;
+
+pub mod requests;
+
+pub type Error = lambda_runtime_client::Error;
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Tracing {
+    pub r#type: String,
+    pub value: String,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct InvokeEvent {
+    deadline_ms: u64,
+    request_id: String,
+    invoked_function_arn: String,
+    tracing: Tracing,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ShutdownEvent {
+    shutdown_reason: String,
+    deadline_ms: u64,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "UPPERCASE", tag = "eventType")]
+pub enum NextEvent {
+    Invoke(InvokeEvent),
+    Shutdown(ShutdownEvent),
+}
+
+/// A trait describing an asynchronous extension.
+#[async_trait]
+pub trait Extension {
+    async fn on_invoke(&self, extension_id: &str, event: InvokeEvent) -> Result<(), Error>;
+    async fn on_shutdown(&self, extension_id: &str, event: ShutdownEvent) -> Result<(), Error>;
+}
+
+struct Runtime<'a, C: Service<http::Uri> = HttpConnector> {
+    extension_id: &'a str,
+    client: Client<C>,
+}
+
+impl<'a, C> Runtime<'a, C>
+where
+    C: Service<http::Uri> + Clone + Send + Sync + Unpin + 'static,
+    <C as Service<http::Uri>>::Future: Unpin + Send,
+    <C as Service<http::Uri>>::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
+    <C as Service<http::Uri>>::Response: AsyncRead + AsyncWrite + Connection + Unpin + Send + 'static,
+{
+    pub async fn run(&self, extension: impl Extension) -> Result<(), Error> {
+        let client = &self.client;
+        let extension_id = self.extension_id;
+
+        let incoming = async_stream::stream! {
+            loop {
+                trace!("Waiting for next event (incoming loop)");
+                let req = requests::next_event_request(extension_id)?;
+                let res = client.call(req).await;
+                yield res;
+            }
+        };
+
+        tokio::pin!(incoming);
+        while let Some(event) = incoming.next().await {
+            trace!("New event arrived (run loop)");
+            let event = event?;
+            let (_parts, body) = event.into_parts();
+
+            let body = hyper::body::to_bytes(body).await?;
+            trace!("{}", std::str::from_utf8(&body)?); // this may be very verbose
+            let event: NextEvent = serde_json::from_slice(&body)?;
+
+            match event {
+                NextEvent::Invoke(event) => {
+                    extension.on_invoke(extension_id, event).await?;
+                }
+                NextEvent::Shutdown(event) => {
+                    extension.on_shutdown(extension_id, event).await?;
+                }
+            };
+        }
+
+        Ok(())
+    }
+}
+
+async fn register<C>(client: &Client<C>, extension_name: &str) -> Result<String, Error>
+where
+    C: Service<http::Uri> + Clone + Send + Sync + Unpin + 'static,
+    <C as Service<http::Uri>>::Future: Unpin + Send,
+    <C as Service<http::Uri>>::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
+    <C as Service<http::Uri>>::Response: AsyncRead + AsyncWrite + Connection + Unpin + Send + 'static,
+{
+    let req = requests::register_request(extension_name)?;
+    let res = client.call(req).await?;
+    // ensure!(res.status() == http::StatusCode::OK, "Unable to register extension",);
+
+    let ext_id = res.headers().get(requests::EXTENSION_ID_HEADER).unwrap().to_str()?;
+    Ok(ext_id.into())
+}
+
+pub async fn run(extension: impl Extension) -> Result<(), Error> {
+    let args: Vec<String> = std::env::args().collect();
+
+    let client = Client::builder().build().expect("Unable to create a runtime client");
+    let extension_id = register(&client, &args[0]).await?;
+    let runtime = Runtime {
+        extension_id: &extension_id,
+        client,
+    };
+
+    runtime.run(extension).await
+}

lambda-extension/src/requests.rs

@@ -0,0 +1,78 @@
+use crate::Error;
+use http::{Method, Request};
+use hyper::Body;
+use lambda_runtime_client::build_request;
+use serde::Serialize;
+
+const EXTENSION_NAME_HEADER: &str = "Lambda-Extension-Name";
+pub(crate) const EXTENSION_ID_HEADER: &str = "Lambda-Extension-Identifier";
+const EXTENSION_ERROR_TYPE_HEADER: &str = "Lambda-Extension-Function-Error-Type";
+
+pub(crate) fn next_event_request(extension_id: &str) -> Result<Request<Body>, Error> {
+    let req = build_request()
+        .method(Method::GET)
+        .header(EXTENSION_ID_HEADER, extension_id)
+        .uri("/2020-01-01/extension/event/next")
+        .body(Body::empty())?;
+    Ok(req)
+}
+
+pub(crate) fn register_request(extension_name: &str) -> Result<Request<Body>, Error> {
+    let events = serde_json::json!({
+        "events": ["INVOKE", "SHUTDOWN"]
+    });
+
+    let req = build_request()
+        .method(Method::POST)
+        .uri("/2020-01-01/extension/register")
+        .header(EXTENSION_NAME_HEADER, extension_name)
+        .body(Body::from(serde_json::to_string(&events)?))?;
+
+    Ok(req)
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ErrorRequest<'a> {
+    error_message: &'a str,
+    error_type: &'a str,
+    stack_trace: Vec<&'a str>,
+}
+
+pub fn init_error<'a>(
+    extension_id: &str,
+    error_type: &str,
+    request: Option<ErrorRequest<'a>>,
+) -> Result<Request<Body>, Error> {
+    error_request("init", extension_id, error_type, request)
+}
+
+pub fn exit_error<'a>(
+    extension_id: &str,
+    error_type: &str,
+    request: Option<ErrorRequest<'a>>,
+) -> Result<Request<Body>, Error> {
+    error_request("exit", extension_id, error_type, request)
+}
+
+fn error_request<'a>(
+    error_type: &str,
+    extension_id: &str,
+    error_str: &str,
+    request: Option<ErrorRequest<'a>>,
+) -> Result<Request<Body>, Error> {
+    let uri = format!("/2020-01-01/extension/{}/error", error_type);
+
+    let body = match request {
+        None => Body::empty(),
+        Some(err) => Body::from(serde_json::to_string(&err)?),
+    };
+
+    let req = build_request()
+        .method(Method::POST)
+        .uri(uri)
+        .header(EXTENSION_ID_HEADER, extension_id)
+        .header(EXTENSION_ERROR_TYPE_HEADER, error_str)
+        .body(body)?;
+    Ok(req)
+}