Adaptive Batching Example Extension (#45)

Co-authored-by: Michael Lanthier <[email protected]>

Author: mmaall

README.md

@@ -27,6 +27,7 @@ In this repository you'll find a number of different sample projects and demos t
 * [Logs API extension in Python](python-example-logs-api-extension/)
 * [Logs API extension in Python for Elasticsearch](python-example-elasticsearch-extension/)
 * [Logs API extension in Node.js](nodejs-example-logs-api-extension/)
+* [Adaptive Batching extension in Go](go-example-adaptive-batching-extension/)
 * [Inter-process communication extension in Go](go-example-ipc-extension/)
 * [Crash uploader extension in Go](go-example-crash-uploader-extension/)
 * [Lambda layer extension using SAM](go-example-extension-sam-layer/)

go-example-adaptive-batching-extension/Makefile

@@ -0,0 +1,9 @@
+build:
+	GOOS=linux GOARCH=amd64 go build -o bin/extensions/go-example-adaptive-batching main.go
+
+build-GoExampleAdaptiveBatchingExtensionLayer:
+	GOOS=linux GOARCH=amd64 go build -o $(ARTIFACTS_DIR)/extensions/go-example-adaptive-batching main.go
+	chmod +x $(ARTIFACTS_DIR)/extensions/go-example-adaptive-batching
+
+run-GoExampleAdaptiveBatchingExtensionLayer:
+	go run go-example-adaptive-batching/main.go

go-example-adaptive-batching-extension/README.md

@@ -0,0 +1,98 @@
+# Adaptive Batching Extension in Go
+
+The provided code sample demonstrates how to use the Logs API extension written in Go to adaptively batch log data to a destination (S3). 
+
+> This is a simple example extension to help you start investigating the Lambda Runtime Logs API. This example code is not production ready. Use it with your own discretion after testing thoroughly.
+
+This sample extension: 
+* Subscribes to receive platform and function logs
+* Runs with a main and a helper goroutine: The main goroutine registers to ExtensionAPI and process its invoke and shutdown events (see nextEvent call). The helper goroutine:
+    - starts a local HTTP server at the provided port (default 1234) that receives requests from Logs API
+    - puts the logs in a synchronized queue (Producer) to be processed by the main goroutine (Consumer)
+* The main go routine tracks the number of invokes, last time logs were shipped, and the size of the logs that have been accumulated since the last ship using a structure. Once one of these fields exceeds a set value the log shipping process begins and a new file with the logs is created in S3.
+
+## Requirements
+
+* [AWS SAM CLI ](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-install.html) - **minimum version 0.48**.
+
+## Installation instructions
+
+1. [Create an AWS account](https://portal.aws.amazon.com/gp/aws/developer/registration/index.html) if you do not already have one and login.
+
+2. Clone the repository onto your local development machine:
+```bash
+git clone https://github.com/aws-samples/aws-lambda-extensions.git
+```
+3. Enter the directory 
+```bash
+cd go-example-adaptive-batching/
+```
+
+4. Run the following command for AWS SAM to deploy the components as specified in the `template.yaml` file:
+```bash
+sam build
+# If you don't have 'Python' or 'make' installed, you can use the option to build using a container which uses a python3.8 Docker container image
+# sam build --use-container
+sam deploy --stack-name adaptive-batching-extension --guided
+```
+
+Following the prompts:
+
+* Accept the default Stack Name `adaptive-batching-extension`.
+* Enter your preferred Region
+* Accept the defaults for the remaining questions.
+
+AWS SAM deploys the application stack which includes the Lambda function and an IAM Role. AWS SAM creates a layer for the runtime, a layer for the extension, and adds them to the function.
+
+Note the outputted S3 bucket name. This is where the logs will be outputted. 
+
+## Invoking the Lambda function
+You can now invoke the Lambda function. Amend the Region and use the following command:
+```bash
+aws lambda invoke \
+ --function-name "adaptive-batching-extension-demo-function" \
+ --payload '{"payload": "hello"}' /tmp/invoke-result \
+ --cli-binary-format raw-in-base64-out \
+ --log-type Tail \
+ --region <use your Region>
+```
+The function should return `"StatusCode": 200`
+
+Browse to the [Amazon CloudWatch Console](https://console.aws.amazon.com/cloudwatch). Navigate to *Logs\Log Groups*. Select the log group **/aws/lambda/adaptive-batching-extension-demo-function**.
+
+View the log stream to see the platform, function, and extensions each logging while they are processing.
+
+The logging extension also receives the log stream directly from Lambda, and copies the logs to S3.
+
+Browse to the [Amazon S3 Console](https://console.aws.amazon.com/S3). Navigate to the S3 bucket created as part of the SAM deployment. 
+
+Downloading the file object containing the copied log stream. The log contains the same platform and function logs, but not the extension logs, as specified during the subscription.
+
+
+## Environment Variables
+
+This section details environment variables that can be used to modify the functionality of the extension. Some are required and are marked as such. 
+
+* ADAPTIVE_BATCHING_EXTENSION_S3_BUCKET : (REQUIRED) The value of this variable will be used to create a bucket or use an existing bucket if it is created previously. The logs received from Logs API will be written in a file inside that bucket. For S3 bucket naming rules, see [AWS docs](https://docs.aws.amazon.com/AmazonS3/latest/dev/BucketRestrictions.html). The SAM template will set this by default. 
+* ADAPTIVE_BATCHING_EXTENSION_SHIP_RATE_BYTES : Logs are shipped to S3 once log size reaches the number of bytes defined here. For example a value of 1024 here would result in logs being shipped once the log size exceeds 1 kilobyte in size. Default value of 4096 bytes (4 kilobytes).
+* ADAPTIVE_BATCHING_EXTENSION_SHIP_RATE_INVOKES : Logs are shipped to S3 once the number of invokes reaches the number defined here. For example a value of 10  would result in logs being shipped at least once every 10 invokes since the last time logs were shipped. Default value is 10 invokes. 
+* ADAPTIVE_BATCHING_EXTENSION_SHIP_RATE_MILLISECONDS : Logs are shipped to S3 once the amount of time elapsed since the last time logs were shipped is exceeded. For example a value of 60,000 here would result in logs being shipped once every 60 seconds has passed. The default value is 10,0000 milliseconds. 
+* ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES: This is a JSON array the log types that can be requested from the Logs API. These are the supported log types `["platform", "function", "extension"]`. If not included or parsing errors occur, the log types default to `["platform", "function"]`. 
+
+## Performance, maximums, and environment shutdown 
+
+Logs are shipped to S3 based off of the rates defined in the environment variables or the default values provided. If any of the conditions are met, the logs queued will be shipped to S3. The rates defined are only checked when an invoke to lambda occurs. This means for ADAPTIVE_BATCHING_EXTENSION_SHIP_RATE_MILLISECONDS, the time elapsed may be exceeded and the metric will only be checked once an invocation occurs. So for a rate of 100 ms, if there are 200 ms gaps between lambda invocations, logs will be shipped every 200 ms, once invocations occur. 
+
+Lambda extensions share Lambda function resources, like CPU, memory, and storage, with your function code. The extension here has a limit written into `agent/metrics.go` with MAX_SHIP_RATE_BYTES to prevent users from using too much memory. It is currently set to 50 megabytes, meaning the maximum rate that can be set for ADAPTIVE_BATCHING_EXTENSION_SHIP_RATE_BYTES is 50 megabytes. Any value set above will default to that maximum value of 50 megabytes. That maximum can be increased by modifying the constant in `agent/metrics/go`.
+
+In the case of the Lambda environment shutting down, either from error or stagnation the extension will flush the log queue and upload the final log file to S3. For information about the Lambda environment shutdown phase, see [AWS docs](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html#runtimes-lifecycle-shutdown)
+
+## Outputs 
+Log files shipped to S3 are put into prefixes depending on where the execution is occurring. For every lambda execution environment that is started, the extension will generate an environment uuid. A prefix for that execution environment is then created using this format `<year>-<month>-<day>-<environment-uuid>/`. Within a given prefix files will show up in the S3 bucket depending on how many logs are generated and the thresholds defined. Files are named as follows, `<function-name>-<timestamp>-<file-uuid>.log`. 
+
+This results in the following full format for a log file placed in S3.
+`<year>-<month>-<day>-<environment-uuid>/<function-name>-<timestamp>-<file-uuid>.log`
+
+
+
+

go-example-adaptive-batching-extension/agent/http.go

@@ -0,0 +1,220 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"time"
+
+	"aws-lambda-extensions/go-example-adaptive-batching-extension/logsapi"
+	"aws-lambda-extensions/go-example-adaptive-batching-extension/queuewrapper"
+	log "github.com/sirupsen/logrus"
+)
+
+var httpLogger = log.WithFields(log.Fields{"agent": "httpAgent"})
+
+// DefaultHttpListenerPort is used to set the URL where the logs will be sent by Logs API
+const DefaultHttpListenerPort = "1234"
+
+// LogsApiHttpListener is used to listen to the Logs API using HTTP
+type LogsApiHttpListener struct {
+	httpServer *http.Server
+	// logQueue is a synchronous queue and is used to put the received logs to be consumed later (see main)
+	logQueue *queuewrapper.QueueWrapper
+}
+
+// NewLogsApiHttpListener returns a LogsApiHttpListener with the given log queue
+func NewLogsApiHttpListener(lq *queuewrapper.QueueWrapper) (*LogsApiHttpListener, error) {
+
+	return &LogsApiHttpListener{
+		httpServer: nil,
+		logQueue:   lq,
+	}, nil
+}
+
+func ListenOnAddress() string {
+	env_aws_local, ok := os.LookupEnv("AWS_SAM_LOCAL")
+	if ok && "true" == env_aws_local {
+		return ":" + DefaultHttpListenerPort
+	}
+
+	return "sandbox.localdomain:" + DefaultHttpListenerPort
+}
+
+// Start initiates the server in a goroutine where the logs will be sent
+func (s *LogsApiHttpListener) Start() (bool, error) {
+	address := ListenOnAddress()
+	s.httpServer = &http.Server{Addr: address}
+	http.HandleFunc("/", s.http_handler)
+	go func() {
+		logger.Infof("Serving agent on %s", address)
+		err := s.httpServer.ListenAndServe()
+		if err != http.ErrServerClosed {
+			logger.Errorf("Unexpected stop on Http Server: %v", err)
+			s.Shutdown()
+		} else {
+			logger.Errorf("Http Server closed %v", err)
+		}
+	}()
+	return true, nil
+}
+
+// http_handler handles the requests coming from the Logs API.
+// Everytime Logs API sends logs, this function will read the logs from the response body
+// and put them into a synchronous queue to be read by the main goroutine.
+// Logging or printing besides the error cases below is not recommended if you have subscribed to receive extension logs.
+// Otherwise, logging here will cause Logs API to send new logs for the printed lines which will create an infinite loop.
+func (h *LogsApiHttpListener) http_handler(w http.ResponseWriter, r *http.Request) {
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		logger.Errorf("Error reading body: %+v", err)
+		return
+	}
+
+	//fmt.Println("Logs API event received:", string(body))
+
+	// Puts the log message into the queue
+	err = h.logQueue.Put(string(body))
+	if err != nil {
+		logger.Errorf("Can't push logs to destination: %v", err)
+	}
+}
+
+// Shutdown terminates the HTTP server listening for logs
+func (s *LogsApiHttpListener) Shutdown() {
+	if s.httpServer != nil {
+		ctx, _ := context.WithTimeout(context.Background(), 1*time.Second)
+		err := s.httpServer.Shutdown(ctx)
+		if err != nil {
+			logger.Errorf("Failed to shutdown http server gracefully %s", err)
+		} else {
+			s.httpServer = nil
+		}
+	}
+}
+
+// HttpAgent has the listener that receives the logs and the logger that handles the received logs
+type HttpAgent struct {
+	listener *LogsApiHttpListener
+	logger   *S3Logger
+}
+
+// NewHttpAgent returns an agent to listen and handle logs coming from Logs API for HTTP
+// Make sure the agent is initialized by calling Init(agentId) before subscription for the Logs API.
+func NewHttpAgent(s3Logger *S3Logger, jq *queuewrapper.QueueWrapper) (*HttpAgent, error) {
+
+	logsApiListener, err := NewLogsApiHttpListener(jq)
+	if err != nil {
+		return nil, err
+	}
+
+	return &HttpAgent{
+		logger:   s3Logger,
+		listener: logsApiListener,
+	}, nil
+}
+
+// Init initializes the configuration for the Logs API and subscribes to the Logs API for HTTP
+func (a HttpAgent) Init(agentID string) error {
+	extensions_api_address, ok := os.LookupEnv("AWS_LAMBDA_RUNTIME_API")
+	if !ok {
+		return errors.New("AWS_LAMBDA_RUNTIME_API is not set")
+	}
+
+	logsApiBaseUrl := fmt.Sprintf("http://%s", extensions_api_address)
+
+	logsApiClient, err := logsapi.NewClient(logsApiBaseUrl)
+	if err != nil {
+		return err
+	}
+
+	_, err = a.listener.Start()
+	if err != nil {
+		return err
+	}
+
+	// Read environment variable ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES
+	inputJson := os.Getenv("ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES")
+	inputJsonBytes := []byte(inputJson)
+
+	var eventTypes []logsapi.EventType
+
+	// No Json included
+	if inputJson == "" {
+		// Hold defaults
+		eventTypes = append(eventTypes, logsapi.Platform, logsapi.Function)
+		httpLogger.Info("ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES not included, subscribing to default log types")
+	} else if !json.Valid(inputJsonBytes) {
+		// Invalid JSON provided
+		eventTypes = append(eventTypes, logsapi.Platform, logsapi.Function)
+		httpLogger.Info("ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES includes invalid JSON, subscribing to default log types")
+	} else {
+
+		// Unmarshal json into structure
+		var jsonArray []logsapi.EventType
+
+		err = json.Unmarshal(inputJsonBytes, &jsonArray)
+		if err != nil {
+			// Error unmarshaling json
+			eventTypes = append(eventTypes, logsapi.Platform, logsapi.Function)
+			httpLogger.Info("Unable to unmarshal json from ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES, subscribing to default log types")
+		}
+
+		// If array is empty, use default values
+		if len(jsonArray) == 0 {
+			eventTypes = append(eventTypes, logsapi.Platform, logsapi.Function)
+			httpLogger.Info("LogTypes in ADAPTIVE_BATCHING_EXTENSION_LOG_TYPES does not include any elements, subscribing to default log types")
+		}
+
+		// loop through elements, and check if required elements are included
+
+		for _, logType := range jsonArray {
+			switch logType {
+			case logsapi.Platform:
+				eventTypes = append(eventTypes, logsapi.Platform)
+			case logsapi.Function:
+				eventTypes = append(eventTypes, logsapi.Function)
+			case logsapi.Extension:
+				eventTypes = append(eventTypes, logsapi.Extension)
+			default:
+				httpLogger.Info("Log type ", logType, " is not valid. Not including")
+			}
+		}
+
+	}
+
+	bufferingCfg := logsapi.BufferingCfg{
+		MaxItems:  1000,
+		MaxBytes:  262144,
+		TimeoutMS: 25,
+	}
+	if err != nil {
+		return err
+	}
+	destination := logsapi.Destination{
+		Protocol:   logsapi.HttpProto,
+		URI:        logsapi.URI(fmt.Sprintf("http://sandbox.localdomain:%s", DefaultHttpListenerPort)),
+		HttpMethod: logsapi.HttpPost,
+		Encoding:   logsapi.JSON,
+	}
+
+	_, err = logsApiClient.Subscribe(eventTypes, bufferingCfg, destination, agentID)
+	return err
+}
+
+// Shutdown finalizes the logging and terminates the listener
+func (a *HttpAgent) Shutdown() {
+	err := a.logger.Shutdown()
+	if err != nil {
+		logger.Errorf("Error when trying to shutdown logger: %v", err)
+	}
+
+	a.listener.Shutdown()
+}