AWS Lambda Extensions

In the beginning of October AWS released Lambda Extensions. A new way to customize your Lambda functions. There are two different types of Extensions, external and internal. Internal Extensions run in the same process as the Lambda function while external Extensions run in their own process.

Extensions are still in preview so their might be changes to the API and the way to use them in the future.

In this post I will take internal Extensions for a ride and share my experience. In a later post I will do the same for External. All code used can be found in my GitHub repo

Overview

Extensions, internal and external, share storage, memory, cpu, and other resources with the Lambda function. This is really really important to remember since the performance can be impacted when starting to use an extension. Also the total unzipped package size, function + all extensions can not exceed 250mb. So choose the extensions carefully and be prepared to increase the Lambda memory to increase the performance.
It's possible to have up to 10 extensions per function. I would however make sure to keep it to a minimum, two maybe three max.

What is really good with external extensions are that they start prior to the function and keep running after the function shutdown. This creates possibility to communicate with e.g external monitoring services after shutdown. You will be billed for the entire execution, including the time before function start and after shutdown.

Extensions use the same IAM permissions as the Lambda function, so if your extension will communicate with other AWS services the Lambda function need to have those IAM permissions as well. This was my first gotcha, I would like the possibility to assign IAM permissions directly to the extension. At the same time I can see the problems that can come with that. But still would have been nice.

Extensions can also access the same environment variables as the function, or almost the same, there are some function specific variables that can't be accessed. However all environment variables you have assigned the function your self can be accessed.

Internal Extensions

With internal extensions it's possible to customize the Lambda runtime. The extension run in the same process as the Lambda function and starts and shutdown with the function. There are two ways of doing this. For Java, .Net, and Node there are environment variables that can be used to modify the runtime.
The second method is to use wrapper scripts to modify the runtime. Wrapper scripts are not available for all runtimes, so make sure to check the documentation before you introduce wrapper scripts. Wrapper scripts can be included in the function package or be in a Lambda Layer. Wrapper scripts need to be executable and you use AWS_LAMBDA_EXEC_WRAPPER environment variable to point to it. Here I had my second gotcha and I didn't like how this was done, more on that later.

The secrets extensions

I decided as my first test of extensions take an old helper script and turn that into an internal extension. I created this helper in Python many many years ago and I have been using it a lot since then. The helper is a way to read parameters and secrets from Parameter Store and present them as environment variables. This blog post was the inspiration to it.
The problem with it was that I needed to include it in every Lambda function that wanted to use it, and I needed to call the function to load the secrets. Not as smooth as I hoped for. When Lambda layers was introduced I converted it to a layer and kept using it. However I still needed to call the function in the layer to load the secrets from SSM.
Now with extensions I finally saw the chance not having to do that extra call anymore, just include it as a layer and set a couple of environment variables.

The wrapper

The actual wrapper is a Python script made executable. The wrapper part is not that complicated and basically is just a couple of lines.

#!/usr/bin/env python3

def main():
    load_secrets()

    # Start the function
    args = os.sys.argv[1:]
    os.system(" ".join(args))

if __name__ == "__main__":
    main()

What is needed to be done is to make sure the actual Lambda function is started when the wrapper is done.

Loading secrets and parameters

To load the secrets I call SSM and load all parameters under the path /[environment]/[application]. Environment and application is set as environment variables to the function, remember we can access them from the extension. The code use boto3 when communicating with SSM. Here was one more gotcha. Even though extension run in the same process as the function boto3 was not available for the extension to use. I had to bundle that with the extension.

def load_secrets():

    if not 'APP_PARAMS_PATH' in os.environ:
        return

    app_config_path = os.environ['APP_PARAMS_PATH']
    if 'APP_ENV' in os.environ:
        env = os.environ['APP_ENV']
        full_config_path = '/' + env + '/' + app_config_path
    else:
        full_config_path = '/' + app_config_path

    try:
        client = boto3.client('ssm')
        # Get all parameters for this app
        param_details = client.get_parameters_by_path(
            Path=full_config_path,
            Recursive=False,
            WithDecryption=True
        )

        # Loop through the returned parameters and populate environment variables
        if 'Parameters' in param_details:
            for param in param_details.get('Parameters'):
                param_name = param.get('Name').split("/")[-1]
                param_value = param.get('Value')
                os.environ[param_name] = param_value
    except:
        raise Exception(f"{LAMBDA_WRAPPER_NAME} could not load from SSM")

When the Lambda function start up the internal extension will be invoked and will load all parameters and secrets from parameter store. All of them will then be available as environment variables. That is one true smooth experience.

Point to the wrapper script

As mentioned before I was going to revisit how to point to the wrapper script. At the moment you have to set the full path to the script in the environment variable AWS_LAMBDA_EXEC_WRAPPER. Instead I would like really like the ability to have it in a layer and set the ARN to the layer in AWS_LAMBDA_EXEC_WRAPPER. Then I could create the layer in CloudFormation, get the ARN and be happy. I'm ok to have to have just one executable file at a special path in the layer. Just don't make me specify the full path.
Please AWS give me this possibility!

Conclusion

To wrap everything up. Extensions are a good addition to Lambda. Internal extensions gives the possibility to customize the runtime. Not all runtimes are supported at the moment hopefully there are more to come. I can see myself using this for so many things in the future.

Part 2

Don't miss part 2 that is coming soon, that will go through my experience working external extensions!

Code

All code in this blog can be found on GitHub