AWS Lambda SnapStart Python Runtime Hooks

Warnings

• gotcha Runtime hooks (@register_before_snapshot, @register_after_restore) must be imported and defined at the top level of your Python code, outside of the `lambda_handler` function. Importing or defining them within the handler will cause them to be ignored and not executed by SnapStart.
  Fix: Ensure `from snapshot_restore import ...` and the decorated functions are placed globally in your module.
• gotcha The `snapshot-restore-py` library is pre-installed in AWS Lambda's Python 3.9+ managed runtimes. You should *not* include this library in your deployment package when deploying to Lambda, as it can lead to conflicts or increase your package size unnecessarily. For local development and testing, however, you *do* need to install it (`pip install snapshot-restore-py`).
  Fix: Exclude `snapshot-restore-py` from your Lambda deployment package (e.g., using `.dockerignore` or `exclude` in `serverless.yml`). Install it only for local development.
• breaking Global state (e.g., timestamps, external service tokens, credentials, or open network/database connections) captured in the snapshot will be stale or invalid upon restoration. You *must* implement logic within `register_after_restore` hooks to re-initialize or refresh any such time-sensitive or ephemeral data.
  Fix: Identify all volatile global state. In your `@register_after_restore` hook, add code to re-establish connections, re-fetch credentials, or re-generate random seeds.
• gotcha The total execution time for all `@register_after_restore` hooks combined is limited to 10 seconds. If your after-restore logic exceeds this timeout, the Lambda function will fail to restore from the snapshot with a `SnapStartTimeoutException`.
  Fix: Optimize your `after_restore_hook` functions to perform only essential, quick operations. Consider asynchronous patterns or lazy initialization for longer tasks if absolutely necessary.

Install

• pip install snapshot-restore-py Install for local development/testing

Imports

• register_before_snapshot

  from snapshot_restore import register_before_snapshot

• register_after_restore

  from snapshot_restore import register_after_restore

Quickstart

This quickstart demonstrates how to use `register_before_snapshot` and `register_after_restore` decorators for an AWS Lambda function. The `load_config` and `connect_to_db` functions simulate resource initialization that benefits from SnapStart. The `before_snapshot_hook` cleans up resources (like database connections) that cannot be snapshotted, while the `after_restore_hook` re-establishes them or refreshes time-sensitive data after a function is restored from a snapshot. The example includes local testing instructions to simulate the hook execution flow.

import os
from snapshot_restore import register_before_snapshot, register_after_restore

# Global state that will be snapshotted
config = None
db_connection = None

def load_config():
    """Simulate loading configuration from environment or Parameter Store."""
    global config
    if config is None:
        print("Loading config...")
        # In a real scenario, fetch config from a secure source
        config = os.environ.get('MY_APP_CONFIG', 'default_config_value')
        print(f"Config loaded: {config}")
    return config

def connect_to_db():
    """Simulate connecting to a database."""
    global db_connection
    if db_connection is None:
        print("Connecting to database...")
        # In a real scenario, use actual DB connection, e.g., psycopg2.connect
        db_connection = f"DB_Connection_to_{os.environ.get('DB_HOST', 'localhost')}"
        print(f"Database connected: {db_connection}")
    return db_connection

@register_before_snapshot
def before_snapshot_hook():
    """Hook to run before Lambda takes a snapshot.
    Use this to clean up resources that cannot be snapshotted (e.g., open network connections).
    """
    global db_connection
    print("Executing before_snapshot_hook: Closing DB connection.")
    if db_connection:
        # Proper closing of connection in a real app
        db_connection = None # Simulate closing

@register_after_restore
def after_restore_hook():
    """Hook to run after Lambda restores from a snapshot.
    Use this to re-establish connections or refresh volatile data.
    """
    global config, db_connection
    print("Executing after_restore_hook: Re-establishing resources.")
    # Refresh config if it contains time-sensitive data
    load_config()
    # Re-establish DB connection
    connect_to_db()

def lambda_handler(event, context):
    """Main Lambda handler function."""
    load_config()
    db = connect_to_db()
    print(f"Handler executed with config: {config} and db: {db}")
    return {
        'statusCode': 200,
        'body': 'Function executed successfully!'
    }

# Example of local testing (not part of actual Lambda deployment process)
if __name__ == "__main__":
    os.environ['MY_APP_CONFIG'] = 'local_test_config'
    os.environ['DB_HOST'] = 'test_db_instance'
    
    print("--- First 'invocation' (initialization) ---")
    lambda_handler({}, {})
    
    print("\n--- Simulating Snapshot and Restore (local only, no actual snapshot) ---")
    before_snapshot_hook()
    after_restore_hook()
    
    print("\n--- Second 'invocation' (after restore) ---")
    lambda_handler({}, {})