feat(lambda-extension): Logs API processor

Makefile

@@ -2,7 +2,7 @@ INTEG_STACK_NAME ?= rust-lambda-integration-tests
 INTEG_FUNCTIONS_BUILD := runtime-fn runtime-trait http-fn http-trait
 INTEG_FUNCTIONS_INVOKE := RuntimeFn RuntimeFnAl2 RuntimeTrait RuntimeTraitAl2 Python PythonAl2
 INTEG_API_INVOKE := RestApiUrl HttpApiUrl
-INTEG_EXTENSIONS := extension-fn extension-trait
+INTEG_EXTENSIONS := extension-fn extension-trait logs-trait
 # Using musl to run extensions on both AL1 and AL2
 INTEG_ARCH := x86_64-unknown-linux-musl
 

lambda-extension/Cargo.toml

@@ -11,17 +11,18 @@ keywords = ["AWS", "Lambda", "API"]
 readme = "README.md"
 
 [dependencies]
-tokio = { version = "1.0", features = ["macros", "io-util", "sync", "rt-multi-thread"] }
+async-stream = "0.3"
+bytes = "1.0"
+chrono = { version = "0.4", features = ["serde"] }
+http = "0.2"
 hyper = { version = "0.14", features = ["http1", "client", "server", "stream", "runtime"] }
+lambda_runtime_api_client = { version = "0.4", path = "../lambda-runtime-api-client" }
 serde = { version = "1", features = ["derive"] }
 serde_json = "^1"
-bytes = "1.0"
-http = "0.2"
-async-stream = "0.3"
 tracing = { version = "0.1", features = ["log"] }
-tower = { version = "0.4", features = ["util"] }
+tokio = { version = "1.0", features = ["macros", "io-util", "sync", "rt-multi-thread"] }
 tokio-stream = "0.1.2"
-lambda_runtime_api_client = { version = "0.4", path = "../lambda-runtime-api-client" }
+tower = { version = "0.4", features = ["make", "util"] }
 
 [dev-dependencies]
 simple-error = "0.2"

lambda-extension/README.md

@@ -2,9 +2,11 @@
 
 [![Docs](https://docs.rs/lambda_extension/badge.svg)](https://docs.rs/lambda_extension)
 
-**`lambda-extension`** is a library that makes it easy to write [AWS Lambda Runtime Extensions](https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html) in Rust.
+**`lambda-extension`** is a library that makes it easy to write [AWS Lambda Runtime Extensions](https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html) in Rust. It also helps with using [Lambda Logs API](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-logs-api.html).
 
-## Example extension
+## Example extensions
+
+### Simple extension
 
 The code below creates a simple extension that's registered to every `INVOKE` and `SHUTDOWN` events, and logs them in CloudWatch.
 
@@ -37,6 +39,40 @@ async fn main() -> Result<(), Error> {
 
 ```
 
+### Log processor extension
+
+```rust,no_run
+use lambda_extension::{service_fn, Error, Extension, LambdaLog, LambdaLogRecord, SharedService};
+use tracing::info;
+
+async fn handler(logs: Vec<LambdaLog>) -> Result<(), Error> {
+    for log in logs {
+        match log.record {
+            LambdaLogRecord::Function(_record) => {
+                // do something with the function log record
+            },
+            LambdaLogRecord::Extension(_record) => {
+                // do something with the extension log record
+            },
+            },
+            _ => (),
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    let logs_processor = SharedService::new(service_fn(handler));
+
+    Extension::new().with_logs_processor(logs_processor).run().await?;
+
+    Ok(())
+}
+
+```
+
 ## Deployment
 
 Lambda extensions can be added to your functions either using [Lambda layers](https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html#using-extensions-config), or adding them to [containers images](https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html#invocation-extensions-images).

lambda-extension/examples/combined.rs

@@ -0,0 +1,51 @@
+use lambda_extension::{
+    service_fn, Error, Extension, LambdaEvent, LambdaLog, LambdaLogRecord, NextEvent, SharedService,
+};
+use tracing::info;
+
+async fn my_extension(event: LambdaEvent) -> Result<(), Error> {
+    match event.next {
+        NextEvent::Shutdown(_e) => {
+            // do something with the shutdown event
+        }
+        NextEvent::Invoke(_e) => {
+            // do something with the invoke event
+        }
+    }
+    Ok(())
+}
+
+async fn my_log_processor(logs: Vec<LambdaLog>) -> Result<(), Error> {
+    for log in logs {
+        match log.record {
+            LambdaLogRecord::Function(record) => info!("[logs] [function] {}", record),
+            LambdaLogRecord::Extension(record) => info!("[logs] [extension] {}", record),
+            _ => (),
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // The runtime logging can be enabled here by initializing `tracing` with `tracing-subscriber`
+    // While `tracing` is used internally, `log` can be used as well if preferred.
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        // this needs to be set to false, otherwise ANSI color codes will
+        // show up in a confusing manner in CloudWatch logs.
+        .with_ansi(false)
+        // disabling time is handy because CloudWatch will add the ingestion time.
+        .without_time()
+        .init();
+
+    let func = service_fn(my_extension);
+    let logs_processor = SharedService::new(service_fn(my_log_processor));
+
+    Extension::new()
+        .with_events_processor(func)
+        .with_logs_processor(logs_processor)
+        .run()
+        .await
+}

lambda-extension/examples/custom_events.rs

@@ -1,4 +1,4 @@
-use lambda_extension::{service_fn, Error, LambdaEvent, NextEvent, Runtime};
+use lambda_extension::{service_fn, Error, Extension, LambdaEvent, NextEvent};
 
 async fn my_extension(event: LambdaEvent) -> Result<(), Error> {
     match event.next {
@@ -27,9 +27,9 @@ async fn main() -> Result<(), Error> {
         .without_time()
         .init();
 
-    let func = service_fn(my_extension);
-
-    let runtime = Runtime::builder().with_events(&["SHUTDOWN"]).register().await?;
-
-    runtime.run(func).await
+    Extension::new()
+        .with_events(&["SHUTDOWN"])
+        .with_events_processor(service_fn(my_extension))
+        .run()
+        .await
 }

lambda-extension/examples/custom_logs_service.rs

@@ -0,0 +1,64 @@
+use lambda_extension::{Error, Extension, LambdaLog, LambdaLogRecord, Service, SharedService};
+use std::{
+    future::{ready, Future},
+    pin::Pin,
+    sync::{
+        atomic::{AtomicUsize, Ordering::SeqCst},
+        Arc,
+    },
+    task::Poll,
+};
+use tracing::info;
+
+/// Custom log processor that increments a counter for each log record.
+///
+/// This is a simple example of a custom log processor that can be used to
+/// count the number of log records that are processed.
+///
+/// This needs to derive Clone (and store the counter in an Arc) as the runtime
+/// could need multiple `Service`s to process the logs.
+#[derive(Clone, Default)]
+struct MyLogsProcessor {
+    counter: Arc<AtomicUsize>,
+}
+
+impl MyLogsProcessor {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+/// Implementation of the actual log processor
+///
+/// This receives a `Vec<LambdaLog>` whenever there are new log entries available.
+impl Service<Vec<LambdaLog>> for MyLogsProcessor {
+    type Response = ();
+    type Error = Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+    fn poll_ready(&mut self, _cx: &mut core::task::Context<'_>) -> core::task::Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, logs: Vec<LambdaLog>) -> Self::Future {
+        let counter = self.counter.fetch_add(1, SeqCst);
+        for log in logs {
+            match log.record {
+                LambdaLogRecord::Function(record) => info!("[logs] [function] {}: {}", counter, record),
+                LambdaLogRecord::Extension(record) => info!("[logs] [extension] {}: {}", counter, record),
+                _ => (),
+            }
+        }
+
+        Box::pin(ready(Ok(())))
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    let logs_processor = SharedService::new(MyLogsProcessor::new());
+
+    Extension::new().with_logs_processor(logs_processor).run().await?;
+
+    Ok(())
+}

lambda-extension/examples/logs.rs

@@ -0,0 +1,23 @@
+use lambda_extension::{service_fn, Error, Extension, LambdaLog, LambdaLogRecord, SharedService};
+use tracing::info;
+
+async fn handler(logs: Vec<LambdaLog>) -> Result<(), Error> {
+    for log in logs {
+        match log.record {
+            LambdaLogRecord::Function(record) => info!("[logs] [function] {}", record),
+            LambdaLogRecord::Extension(record) => info!("[logs] [extension] {}", record),
+            _ => (),
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    let logs_processor = SharedService::new(service_fn(handler));
+
+    Extension::new().with_logs_processor(logs_processor).run().await?;
+
+    Ok(())
+}

lambda-extension/src/error.rs

@@ -0,0 +1,23 @@
+/// Error type that extensions may result in
+pub type Error = lambda_runtime_api_client::Error;
+
+/// Simple error that encapsulates human readable descriptions
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ExtensionError {
+    err: String,
+}
+
+impl ExtensionError {
+    pub(crate) fn boxed<T: Into<String>>(str: T) -> Box<ExtensionError> {
+        Box::new(ExtensionError { err: str.into() })
+    }
+}
+
+impl std::fmt::Display for ExtensionError {
+    #[inline]
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.err.fmt(f)
+    }
+}
+
+impl std::error::Error for ExtensionError {}

lambda-extension/src/events.rs

@@ -0,0 +1,71 @@
+use serde::Deserialize;
+
+/// Request tracing information
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Tracing {
+    /// The type of tracing exposed to the extension
+    pub r#type: String,
+    /// The span value
+    pub value: String,
+}
+/// Event received when there is a new Lambda invocation.
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct InvokeEvent {
+    /// The time that the function times out
+    pub deadline_ms: u64,
+    /// The ID assigned to the Lambda request
+    pub request_id: String,
+    /// The function's Amazon Resource Name
+    pub invoked_function_arn: String,
+    /// The request tracing information
+    pub tracing: Tracing,
+}
+
+/// Event received when a Lambda function shuts down.
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ShutdownEvent {
+    /// The reason why the function terminates
+    /// It can be SPINDOWN, TIMEOUT, or FAILURE
+    pub shutdown_reason: String,
+    /// The time that the function times out
+    pub deadline_ms: u64,
+}
+
+/// Event that the extension receives in
+/// either the INVOKE or SHUTDOWN phase
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "UPPERCASE", tag = "eventType")]
+pub enum NextEvent {
+    /// Payload when the event happens in the INVOKE phase
+    Invoke(InvokeEvent),
+    /// Payload when the event happens in the SHUTDOWN phase
+    Shutdown(ShutdownEvent),
+}
+
+impl NextEvent {
+    /// Return whether the event is a [`NextEvent::Invoke`] event or not
+    pub fn is_invoke(&self) -> bool {
+        matches!(self, NextEvent::Invoke(_))
+    }
+}
+
+/// 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,
+}
+
+impl LambdaEvent {
+    pub(crate) fn new(ex_id: &str, next: NextEvent) -> LambdaEvent {
+        LambdaEvent {
+            extension_id: ex_id.into(),
+            next,
+        }
+    }
+}