Managing Connection Resources in AWS Lambda : A Comprehensive Overview

If you’re reading this post, chances are you’ve heard about serverless compute services. For those who are unfamiliar with the concept, let’s start by defining what serverless computing is. At first glance, the term ‘serverless’ may imply that there are no servers involved in executing your code — it almost seems magical! However, the magic lies in the fact that you don’t need to worry about managing servers or clusters.

Some of the most widely-used serverless compute services include AWS Lambda, GCP Cloud Functions, and Azure Functions. These platforms offer a variety of benefits, such as the ability to avoid managing underlying infrastructure, infinite(attention: magic 🧙) scalability, pay-as-you-go pricing, and easy integration with other services.

AWS Lambda

The upcoming sections will primarily address the problem and proposed solutions related to AWS Lambda, although they may also apply to the other serverless compute services (albeit to a lesser extent). In the examples provided, we will be using Python and the PostgreSQL library Psycopg2.

What is the problem?

Although avoiding the management of underlying infrastructure is a very good thing %99 of the time, it can also mean relinquishing control over the runtime environment.

When a Lambda function is invoked for the first time, or after a period of inactivity, it needs to set up its execution environment before processing the event. This initial setup is known as a cold start. During the cold start, AWS provisions a new container to run the function code and initializes the execution environment. However, after the cold start, subsequent function invocations occur in a warm Lambda execution environment, where the container is already running and initialized.

import ...

# The code block here is only executed ONCE 
# during the cold start of the Lambda function.

def lambda_handler(event, context):
  
  # The code block inside the "lambda_handler" function is executed for 
  # each subsequent invocation of the Lambda function.

Even if you can determine the number of events that trigger a Lambda function, it’s still impossible to precisely determine how many new containers (Lambda function executors) are created, how many times each warm Lambda instance is used, and how long each of them lives. After a Lambda function has completed its processing, the container that executed it may be terminated immediately, or it may remain active to handle subsequent events. Therefore, a container’s lifespan can range from a single invocation to multiple invocations over an extended period.

Considering the information above, imagine a scenario in which you need to establish a database connection, Redis connection, or any other resource with limited availability that requires cleanup at the end of use within your Lambda function. In an environment where scaling is infinite and shutdown is out of your control, how can you allocate, utilize, and close resources properly? Managing them incorrectly can lead to longer runtimes(higher expenses) or excessive consumption of compute resources.

Solution Proposals

First and foremost, it’s important to note that none of the methods presented below are universally the best solution for every case. Ultimately, it’s up to you to determine which method is most suitable for your specific problem. Although ways of managing connections in AWS Lambda may vary, there are some important points to keep in mind independent from the solution.

• If applicable, it’s recommended to use connection pooling as it is a more scalable approach compared to a single connection approach, and it simplifies maintenance.
• Additionally, if your use case suits, utilizing the AWS RDS Proxy can also be a good option as it allows you to manage database connections at scale. RDS Proxy mitigates the risk of idle connections and potential resource leaks since it can automatically close idle connections and release unused resources. However, it’s not a silver bullet solution as it has some performance overhead and may not be suitable for all types of workloads or database engines. Additionally, it’s worth noting that AWS RDS Proxy is not a free service.

AWS RDS Proxy

1️⃣ Open a new connection when needed/promptly close after completing task

The simplest approach is to create a connection when it is needed and close it immediately after the relevant code block. While this method may be easy to implement, it goes against best practices for Lambda functions, as it does not reuse the resource for subsequent invocations processed by the same instance. This can increase both the runtime and costs.

Take advantage of execution environment reuse to improve the performance of your function. Initialize SDK clients and database connections outside of the function handler …

import ...

def lambda_handler(event, context):
    try:
        conn = psycopg2.connect(
            host="Host",
            port="Port",
            database="DB_NAME",
            user="USER",
            password="PASSWORD"
        )
        cur = conn.cursor()
        # Execute your queries here
        conn.commit()
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        if conn is not None:
            conn.close()

2️⃣ Only one connection initialized globally, not closed

Another way to manage database connections in AWS Lambda is by initializing a single global connection during the cold start and not closing it anywhere in the code. Instead, the connection should be closed from the other side of the connection, such as the database.

import ...

conn = psycopg2.connect(
            host="Host",
            port="Port",
            database="DB_NAME",
            user="USER",
            password="PASSWORD"
        )

def lambda_handler(event, context):
    global conn
    if conn is None:
        try:
            cur = conn.cursor()
            # Execute your queries here
            conn.commit()
            cur.close()
        except (Exception, psycopg2.DatabaseError) as error:
            print(error)
        # Connection is not going to be closed here
        # finally:
        #   if conn is not None:
        #       conn.close()

However, having a large number of idle connections can create issues for you. It’s important to set a maximum lifetime of 900 seconds for the connection during initialization(if applicable), since a Lambda function can run for at most 15 minutes(hard limit of AWS). After 15 minutes, the entire execution environment, including the container, is terminated by AWS. Therefore, there’s no need to keep the connection open for more than 15 minutes, since subsequent invocations of Lambda definitely will be handled by a new container. It’s also worth noting that the maximum runtime for a Lambda function can be configured individually.

3️⃣ Client side solutions

There are also client-side solutions available, such as Serverless MySQL and PostgreSQL, which adds a connection management component into the most commonly used database libraries. However, if there is no client-side library available for your specific language and database combination, you may not be able to use these solutions. Additionally, the maintenance of these client-side solutions may become a problem in the future, which is one of the downsides of this approach.

4️⃣ Catch shutdown events with Lambda Extensions

The final solution proposal is to add an extension layer to AWS Lambda, which can be used to catch the SIGTERM signal sent by Lambda to indicate a shutdown. By doing so, you can perform graceful shutdown steps within the 300ms time frame provided. It is important to note that if the extension or runtime does not respond within the time limit, Lambda will be killed with a SIGKILL signal. While adding an extension can increase both the runtime and costs, it provides a clean and proper way to manage connection resources.

GitHub - aws-samples/graceful-shutdown-with-aws-lambda

5️⃣ Take remaining time and schedule an event

I haven’t tested this solution, so I’m not sure if it works.

When initializing your function, you can schedule a task to regularly check how much time is left until your Lambda function reaches the timeout. This check can be done by using the getRemainingTimeInMillis() method in the context. If the remaining time is less than a certain value, you can call a callback function to handle a graceful shutdown.

This method of scheduling a task to regularly check the remaining time until the Lambda function reaches its timeout can be challenging to develop, and deciding on an appropriate threshold for handling shutdown can be tricky. Additionally, determining the optimal period for the checks to ensure smooth function operation requires careful consideration. It’s important to note that implementing periodic tasks in this manner can result in increased runtime and cost.

In this post, my goal was to provide solutions and compare them to help you choose the best approach for managing connections in AWS Lambda.

Thanks for reading. If you have questions or comments regarding this article, please feel free to leave a comment below. Also, don’t hesitate to add any other potential solutions that may come into your mind about this topic 👐