Lambda Extension

ShutdownExtension

ShutdownExtension(handler_module: str = 'lambda_function', app_name: str = 'app')

AWS Lambda Extension that executes shutdown handlers on container termination.

This extension registers with the Lambda Runtime Extensions API, waits for SHUTDOWN events, and calls the RestMachine application's shutdown_sync() method.

The extension runs as a separate process from your Lambda handler, monitoring the Lambda lifecycle and ensuring cleanup code runs before the container terminates.

Example

from restmachine_aws.extension import ShutdownExtension

extension = ShutdownExtension(
    handler_module="lambda_function",
    app_name="app"
)
extension.run()

Parameters:

Name	Type	Description	Default
handler_module	str	Python module containing the Lambda handler (default: "lambda_function")	'lambda_function'
app_name	str	Name of the RestApplication variable in the handler module (default: "app")	'app'

Initialize the Lambda Extension.

Parameters:

Name	Type	Description	Default
handler_module	str	Module name where the RestApplication is defined	'lambda_function'
app_name	str	Variable name of the RestApplication instance	'app'

Source code in packages/restmachine-aws/src/restmachine_aws/extension.py

def __init__(self, handler_module: str = "lambda_function", app_name: str = "app"):
    """
    Initialize the Lambda Extension.

    Args:
        handler_module: Module name where the RestApplication is defined
        app_name: Variable name of the RestApplication instance
    """
    self.handler_module = handler_module
    self.app_name = app_name
    self.extension_id: Optional[str] = None

    # Get Lambda Runtime API endpoint from environment
    self.runtime_api = os.environ.get("AWS_LAMBDA_RUNTIME_API")
    if not self.runtime_api:
        raise RuntimeError(
            "AWS_LAMBDA_RUNTIME_API environment variable not set. "
            "This extension must be run within an AWS Lambda environment."
        )

Functions

register

register() -> str

Register this extension with the Lambda Runtime Extensions API.

Sends a POST request to the Extensions API to register for SHUTDOWN events. The Lambda runtime will send a SHUTDOWN event when the container is about to terminate, allowing cleanup handlers to execute.

Returns:

Type	Description
str	Extension ID assigned by the Lambda runtime

Raises:

Type	Description
RuntimeError	If registration fails

Source code in packages/restmachine-aws/src/restmachine_aws/extension.py

def register(self) -> str:
    """
    Register this extension with the Lambda Runtime Extensions API.

    Sends a POST request to the Extensions API to register for SHUTDOWN events.
    The Lambda runtime will send a SHUTDOWN event when the container is about
    to terminate, allowing cleanup handlers to execute.

    Returns:
        Extension ID assigned by the Lambda runtime

    Raises:
        RuntimeError: If registration fails
    """
    url = f"http://{self.runtime_api}/2020-01-01/extension/register"

    payload = {"events": ["SHUTDOWN"]}
    data = json.dumps(payload).encode("utf-8")

    req = request.Request(
        url,
        data=data,
        headers={
            "Lambda-Extension-Name": "restmachine-shutdown",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        # nosec B310: Lambda Extensions API is only accessible over localhost HTTP
        with request.urlopen(req) as response:  # nosec B310
            extension_id = response.headers.get("Lambda-Extension-Identifier")
            if not extension_id:
                raise RuntimeError("Lambda runtime did not return an Extension ID")
            self.extension_id = extension_id
            logger.info(f"Extension registered with ID: {self.extension_id}")
            return self.extension_id
    except Exception as e:
        raise RuntimeError(f"Failed to register extension: {e}") from e

wait_for_event

wait_for_event() -> dict

Wait for the next lifecycle event from Lambda.

This is a blocking call that waits for the Lambda runtime to send the next event. For SHUTDOWN-only extensions, this will block until the container is terminating.

Returns:

Type	Description
dict	Event dictionary containing eventType and other event details

Raises:

Type	Description
RuntimeError	If event retrieval fails or extension is not registered

Source code in packages/restmachine-aws/src/restmachine_aws/extension.py

def wait_for_event(self) -> dict:
    """
    Wait for the next lifecycle event from Lambda.

    This is a blocking call that waits for the Lambda runtime to send the next
    event. For SHUTDOWN-only extensions, this will block until the container
    is terminating.

    Returns:
        Event dictionary containing eventType and other event details

    Raises:
        RuntimeError: If event retrieval fails or extension is not registered
    """
    if not self.extension_id:
        raise RuntimeError("Extension must be registered before waiting for events")

    url = f"http://{self.runtime_api}/2020-01-01/extension/event/next"

    req = request.Request(
        url,
        headers={"Lambda-Extension-Identifier": self.extension_id},
        method="GET",
    )

    try:
        # nosec B310: Lambda Extensions API is only accessible over localhost HTTP
        with request.urlopen(req) as response:  # nosec B310
            event = cast(dict[Any, Any], json.loads(response.read()))
            return event
    except Exception as e:
        raise RuntimeError(f"Failed to get next event: {e}") from e

load_app

load_app() -> Any

Load the RestMachine application from the handler module.

Imports the handler module and retrieves the RestApplication instance. The Lambda task root is added to sys.path to ensure imports work correctly.

Returns:

Type	Description
Any	The RestMachine application instance

Raises:

Type	Description
ImportError	If the handler module or app variable cannot be found

Source code in packages/restmachine-aws/src/restmachine_aws/extension.py

def load_app(self) -> Any:
    """
    Load the RestMachine application from the handler module.

    Imports the handler module and retrieves the RestApplication instance.
    The Lambda task root is added to sys.path to ensure imports work correctly.

    Returns:
        The RestMachine application instance

    Raises:
        ImportError: If the handler module or app variable cannot be found
    """
    # Add Lambda task root to Python path
    task_root = os.environ.get("LAMBDA_TASK_ROOT", ".")
    if task_root not in sys.path:
        sys.path.insert(0, task_root)

    try:
        # Import the handler module
        handler_module = __import__(self.handler_module)
        logger.info(f"Imported module: {self.handler_module}")

        # Get the app instance
        if not hasattr(handler_module, self.app_name):
            raise AttributeError(
                f"Module '{self.handler_module}' has no attribute '{self.app_name}'. "
                f"Available attributes: {dir(handler_module)}"
            )

        app = getattr(handler_module, self.app_name)
        logger.info(f"Loaded app from {self.handler_module}.{self.app_name}")
        return app

    except ImportError as e:
        raise ImportError(
            f"Could not import handler module '{self.handler_module}': {e}"
        ) from e

run

run()

Main extension loop.

Performs the following steps: 1. Register with Lambda Extensions API 2. Load the RestMachine application 3. Wait for lifecycle events 4. On SHUTDOWN event, call app.shutdown_sync()

This method blocks until a SHUTDOWN event is received or an error occurs.

Raises:

Type	Description
Exception	If any step in the extension lifecycle fails

Source code in packages/restmachine-aws/src/restmachine_aws/extension.py

def run(self):
    """
    Main extension loop.

    Performs the following steps:
    1. Register with Lambda Extensions API
    2. Load the RestMachine application
    3. Wait for lifecycle events
    4. On SHUTDOWN event, call app.shutdown_sync()

    This method blocks until a SHUTDOWN event is received or an error occurs.

    Raises:
        Exception: If any step in the extension lifecycle fails
    """
    try:
        # Step 1: Register extension
        logger.info("Registering extension...")
        self.register()

        # Step 2: Load application
        logger.info("Loading application...")
        app = self.load_app()

        # Verify app has shutdown_sync method
        if not hasattr(app, "shutdown_sync"):
            logger.warning(
                f"Application {self.app_name} does not have shutdown_sync() method. "
                "No shutdown handlers will be called."
            )
            # Still wait for SHUTDOWN to prevent extension from exiting early
            app = None

        # Step 3: Wait for events
        logger.info("Extension ready, waiting for events...")
        while True:
            event = self.wait_for_event()
            event_type = event.get("eventType")

            logger.info(f"Received event: {event_type}")

            if event_type == "SHUTDOWN":
                # Step 4: Execute shutdown handlers
                if app and hasattr(app, "shutdown_sync"):
                    logger.info("Executing shutdown handlers...")
                    try:
                        app.shutdown_sync()
                        logger.info("Shutdown handlers completed successfully")
                    except Exception as e:
                        logger.error(f"Error in shutdown handlers: {e}", exc_info=True)
                else:
                    logger.info("No shutdown handlers to execute")

                # Exit after shutdown
                logger.info("Extension shutting down")
                break

    except Exception as e:
        logger.error(f"Extension error: {e}", exc_info=True)
        raise

Overview

The Lambda Extension enables shutdown handlers to run when your Lambda container terminates. This is essential for cleaning up resources like database connections, file handles, and pending operations.

Why Use an Extension?

Lambda functions don't have a traditional shutdown phase - they're frozen after execution. The extension solves this by:

1. Running as a separate process alongside your Lambda function
2. Listening for SHUTDOWN signals from AWS
3. Calling your @app.on_shutdown handlers before termination

Installation

1. Create Extension Script

Create extensions/restmachine-shutdown in your Lambda deployment package:

#!/usr/bin/env python3
from restmachine_aws.extension import main

if __name__ == "__main__":
    main()

2. Make It Executable

chmod +x extensions/restmachine-shutdown

3. Deploy with Lambda

The extension is automatically discovered when included in your deployment package:

lambda_function/
├── lambda_function.py      # Your Lambda handler
├── extensions/
│   └── restmachine-shutdown  # Extension (executable)
└── ...

Using Shutdown Handlers

Define shutdown handlers in your application:

from restmachine import RestApplication
from restmachine_aws import AwsApiGatewayAdapter

app = RestApplication()

@app.on_startup
def database():
    """Open database connection on cold start."""
    print("Opening database connection...")
    return create_db_connection()

@app.on_shutdown
def close_database(database):
    """Close database connection on shutdown."""
    print("Closing database connection...")
    database.close()
    print("Database closed successfully")

# AWS Lambda adapter
adapter = AwsApiGatewayAdapter(app)

def lambda_handler(event, context):
    return adapter.handle_event(event, context)

When AWS terminates the Lambda container, the extension calls close_database() automatically.

How It Works

sequenceDiagram participant AWS as AWS Lambda participant Ext as Extension participant Handler as Lambda Handler participant App as RestMachine AWS->>Ext: INIT Ext->>Ext: Register for SHUTDOWN AWS->>Handler: Cold Start Handler->>App: Run @app.on_startup loop Requests AWS->>Handler: Invoke Handler->>App: Process request App-->>Handler: Response Handler-->>AWS: Return end AWS->>Ext: SHUTDOWN signal Ext->>Handler: Call shutdown Handler->>App: Run @app.on_shutdown App-->>Ext: Complete Ext-->>AWS: Extension stopped

Lifecycle Phases

1. Initialization

When Lambda starts:

# Extension registers with Lambda Runtime API
# Declares interest in SHUTDOWN events

2. Request Processing

Your Lambda function handles requests normally. Shutdown handlers are not called between invocations.

3. Shutdown

When AWS terminates the container:

# Extension receives SHUTDOWN event
# Extension imports your app and calls app.shutdown_sync()
# Your @app.on_shutdown handlers execute
# Extension reports completion to AWS

Configuration

Environment Variables

Configure extension behavior:

# SAM template.yaml
Environment:
  Variables:
    RESTMACHINE_HANDLER_MODULE: "lambda_function"  # Module with your app
    RESTMACHINE_APP_NAME: "app"                    # Variable name

Defaults: - RESTMACHINE_HANDLER_MODULE: "lambda_function" - RESTMACHINE_APP_NAME: "app"

Custom App Location

If your app is in a different module:

# my_api/application.py
from restmachine import RestApplication

my_app = RestApplication()

@my_app.get("/hello")
def hello():
    return {"message": "Hello"}

Set environment variables:

Environment:
  Variables:
    RESTMACHINE_HANDLER_MODULE: "my_api.application"
    RESTMACHINE_APP_NAME: "my_app"

Resource Injection

Shutdown handlers support dependency injection:

@app.on_startup
def database():
    return create_db_connection()

@app.on_startup
def cache():
    return redis.Redis()

@app.on_shutdown
def cleanup(database, cache):
    """Both dependencies are injected automatically."""
    database.close()
    cache.close()

Example: Database Connection Pool

from restmachine import RestApplication
from restmachine_aws import AwsApiGatewayAdapter
import psycopg2.pool

app = RestApplication()

@app.on_startup
def db_pool():
    """Create connection pool on cold start."""
    print("Creating database pool...")
    return psycopg2.pool.SimpleConnectionPool(
        minconn=1,
        maxconn=10,
        host="db.example.com",
        database="myapp",
        user="user",
        password="password"
    )

@app.get('/users/{user_id}')
def get_user(db_pool, request):
    """Use connection from pool."""
    conn = db_pool.getconn()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT * FROM users WHERE id = %s",
                       (request.path_params['user_id'],))
            user = cur.fetchone()
            return {"user": user}
    finally:
        db_pool.putconn(conn)

@app.on_shutdown
def close_pool(db_pool):
    """Close all connections on shutdown."""
    print("Closing database pool...")
    db_pool.closeall()
    print("All connections closed")

adapter = AwsApiGatewayAdapter(app)

def lambda_handler(event, context):
    return adapter.handle_event(event, context)

Monitoring

CloudWatch Logs

Shutdown handlers write to CloudWatch:

[Extension] Shutdown signal received
[Function] Closing database pool...
[Function] All connections closed
[Extension] Shutdown handlers completed successfully

Testing Locally

Test shutdown without deploying:

from restmachine_aws import AwsApiGatewayAdapter

app = RestApplication()

# ... define handlers ...

adapter = AwsApiGatewayAdapter(app)

# Manually trigger shutdown for testing
if hasattr(app, 'shutdown_sync'):
    app.shutdown_sync()

Limitations

• 5-Second Timeout: Extensions have limited time during shutdown
• No Guarantees: AWS may forcibly terminate if timeout exceeded
• Cold Start Only: Shutdown doesn't run between warm invocations
• No Return Values: Shutdown handlers shouldn't return data

Best Practices

1. Keep Shutdown Fast - Close connections quickly, avoid complex cleanup
2. Log Everything - Use logging to track shutdown execution
3. Handle Failures - Use try/except to ensure shutdown completes
4. Test Thoroughly - Test shutdown logic locally before deploying

Troubleshooting

Shutdown Handlers Not Running

• Verify extension is included in deployment package
• Check extensions/ directory exists in Lambda
• Review CloudWatch logs for extension errors

Timeout Errors

• Reduce cleanup operations
• Remove blocking I/O from shutdown handlers

Extension Causing Cold Start Delay

• Extension adds ~50-100ms to cold start
• This is normal and acceptable for most use cases

CLI Tool

Generate extension script automatically:

python -m restmachine_aws create-extension

This creates extensions/restmachine-shutdown with proper permissions.

See Also

• Lambda Deployment Guide - Complete deployment guide
• Lambda Extensions Guide - Detailed usage guide
• Lifecycle Handlers - General lifecycle patterns