Writing Go Middleware for AWS Lambda

This post assumes a base level of knowledge in AWS Lambda functions written in Go. The Official Documentation is a great place to start if you are unfamiliar with the basics of Lambda functions written in Go.

Introduction

With the addition of native Go support for AWS Lambda, writing Lambda functions in Go has become much simpler since using a shim is no longer needed. The AWS Lambda Go library provides a framework for writing Lambda functions and working with a variety of event sources. The lambda.Start(handler) pattern lets you write a handler function that takes an event and an optional context. The lambda.Start function will unmarshal the event payload into the event struct that your handler accepts and then pass it to your handler. The handler event should be compatible with the encoding/json standard library and the handler has to be one of the supported function signatures (more on that in the official docs).

A standard implementation would look like this:

// Event is a generic Event that the Lambda is invoked with.
type Event struct {
	GenericField string `json:"genericField"`
}

func main() {
	lambda.Start(Process)
}

// Process handles the event and performs business logic.
func Process(ctx *context.Context, event *Event) (interface{}, error) {
	log.Println("processing event")
	// Perform business logic  . . .

	return nil, nil
}

This is a great way to deal with the majority of use cases, however, I’ve recently run into a case where I wanted to have more granular control over the invocation of my business logic handler. In particular, I wanted to be able to bail out early in some situations and perform some basic validation before passing control on to the business logic handler, without cluttering up the logic of my handler. Middleware is an awesome way to do exactly that - it’s a way to isolate repetitive logic that needs to be performed on every request without having it lumped into business logic code.

Writing Middleware

Layering middleware into a Lambda function is similar to how you might implement it in a regular HTTP server, but requires a little bit of set up since we’re not using the net/http library here. The first thing needed is a generic middleware function type that can be used to chain middleware functions together.

// HandlerFunc is a generic JSON Lambda handler used to chain middleware.
type HandlerFunc func(context.Context, json.RawMessage) (interface{}, error)

// MiddlewareFunc is a generic middleware example that takes in a HandlerFunc
// and calls the next middleware in the chain.
func MiddlewareFunc(next HandlerFunc) HandlerFunc {
	return HandlerFunc(func(ctx context.Context, data json.RawMessage) (interface{}, error) {
		return next(ctx, data)
	})
}

This function takes in context and a json.RawMessage and returns interface{} and error. This is one of the compatible types for the aws-lambda-go library and could be modified slightly to handle the other handler types - more on that in a bit. The MiddlewareFunc is a contrived middleware function that doesn’t do anything except call the next middleware in the stack. This function could also do some payload validation, add some logging, handle events differently depending on what’s in them, etc.

We lose some utility from the aws-lambda-go library at this point - HandlerFunc takes in a json.RawMessage, not our final event type. This is so that we can treat the message completely differently within different middleware, and transform it as necessary. This means that we also need a way to unmarshal into our final event type. While this could be implemented in our final business logic handler, we’re in the middleware mindset right now, so it’s much nicer to handle this as a middleware function and let our business logic handler take in the type it needs. It also makes the function signature of the business logic handler more explicit and makes testing much more straightforward.

// ParseEvent is a middleware that unmarshals an Event before passing control.
func ParseEvent(h EventHandler) HandlerFunc {
	return HandlerFunc(func(ctx context.Context, data json.RawMessage) (interface{}, error) {
		var event Event

		if err := json.Unmarshal(data, &event); err != nil {
			return nil, err
		}

		if event.GenericField == "" {
			return nil, errors.New("GenericField must be populated")
		}

		return h(&ctx, &event)
	})
}

// EventHandler is the function signature to process the Event.
type EventHandler func(*context.Context, *Event) (interface{}, error)

// Process satisfies EventHandler and processes the Event.
func Process(ctx *context.Context, event *Event) (interface{}, error) {
	log.Println("processing event")
	// Perform business logic  . . .

	return nil, nil
}

Here we’ve defined a new function type EventHandler - it looks nearly the same as our middleware handler, but now takes in the final event type that we want. Our business logic handler looks exactly the same as it did in the standard implementation - great! That means if we already have code and tests written for the standard implementation, nothing needs to change to layer in middleware.

The ParseEvent middleware is the final middleware that will be called before the business logic handler. It unmarshals the event into our final event type, does some basic field validation, and calls the final business logic handler that was passed to it.

An implementation could look like this:

func main() {
	lambda.Start(
		MiddlewareFunction( // Generic middleware function of the HandlerFunc type.
			ParseEvent( // Parse the standard event - return on error.
				Process, // Process the standard event.
			),
		),
	)
}

So that’s how to set up a middleware pattern for AWS Lambda functions in Go. What would a real world use case be for something like this?

Keeping a Lambda Function Warm

While AWS Lambda functions have a lot of amazing benefits, they also come with some drawbacks. Most of the time, they execute exactly when you want to. But sometimes you get “cold starts” (a delayed function invocation), which happens when AWS has to provision an environment for your Lambda to run in. This can happen in a few different situations:

• A Lambda function hasn’t been ran in awhile and AWS has decided to use the container it was running in for something else.
  
• AWS just decides to recycle the container for whatever other reason.
  
• Another concurrent execution of your Lambda function kicks off.

This means a few things. You cannot rely on regular traffic to keep your Lambda functions “warm”. Regular traffic or even worse, bursts of traffic that kick off additional concurrent Lambda executions, still means your customers are experiencing cold starts and delayed response times. Some have seen cold starts as high as 2.8 seconds!

In many cases this is detrimental to the user experience. One way to solve this is to keep your Lambda functions “warm” by scheduling CloudWatch events to trigger your Lambda functions and handling those events differently. If you need to handle bursts of traffic, you could set up a process to fire off multiple events to keep multiple functions warm so that when your traffic bursts hit, all of your concurrent Lambda functions are ready to go.

You can implement the middleware handler like this:

// StayToasty is a middleware that checks for a ping and returns if it was a ping
// otherwise it will call the next middleware in the chain.
// Can be used to keep a Lambda function warm.
func StayToasty(next HandlerFunc) HandlerFunc {
	type ping struct {
		Ping string `json:"ping"`
	}

	return HandlerFunc(func(ctx context.Context, data json.RawMessage) (interface{}, error) {
		var p ping
		// If unmarshal into the ping struct is successful and there was a value in ping, return out.
		if err := json.Unmarshal(data, &p); err == nil && p.Ping != "" {
			log.Println("ping")
			return "pong", nil
		}
		// Otherwise it's a regular request, call the next middleware.
		return next(ctx, data)
	})
}

This middleware checks for a ping - if the event it received was a ping event (rather than our GenericEvent type), then it will simply return out and end the request. Otherwise it calls the next middleware in the chain.

Drawbacks

There are a couple drawbacks to utilizing a pattern like this in AWS Lambda:

• You lose the built-in payload unmarshaling that the AWS Lambda Go library provides. This generally isn’t a huge deal but that code does have to be implemented by us now.
• There are multiple types of function signatures that the AWS Lambda Go library supports. In this example, we only implemented one of them - if you wanted to implement this pattern in a few different Lambda functions that all have different function signatures you would need to write middleware to handle those different function signatures.

I hope that you found this pattern helpful and can put it to use to improve the performance of your Lambda functions.
You can find the full example source code on GitHub