Skip to content

Commit b377a56

Browse files
hrodmnhrodmn
authored
feat: add option to enable SnapStart for Lambda functions (#198)
Enable SnapStart to reduce cold start latency. SnapStart creates a snapshot of the initialized Lambda function, allowing new instances to start from this pre-initialized state instead of starting from scratch. Benefits: - Significantly reduces cold start times (typically 10x faster) - Improves API response time for infrequent requests Considerations: - Additional cost: charges for snapshot storage and restore operations - Requires Lambda versioning (automatically configured by this construct) - Database connections are recreated on restore using snapshot lifecycle hooks see https://docs.aws.amazon.com/lambda/latest/dg/snapstart.html
1 parent cd7a842 commit b377a56

File tree

7 files changed

+312
-33
lines changed

7 files changed

+312
-33
lines changed

‎lib/lambda-api-gateway/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ export interface LambdaApiGatewayProps {
1010
/**
1111
* Lambda function to integrate with the API Gateway.
1212
*/
13-
readonly lambdaFunction: lambda.Function;
13+
readonly lambdaFunction: lambda.Function | lambda.Version;
1414

1515
/**
1616
* Custom Domain Name for the API. If defined, will create the

‎lib/stac-api/index.ts‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ export class PgStacApiLambdaRuntime extends Construct {
6969

7070
const enabledExtensions = props.enabledExtensions || defaultExtensions;
7171

72-
const { code: userCode, ...otherLambdaOptions } = props.lambdaFunctionOptions || {};
72+
const { code: userCode, ...otherLambdaOptions } =
73+
props.lambdaFunctionOptions || {};
7374

7475
this.lambdaFunction = new lambda.Function(this, "lambda", {
7576
// defaults
@@ -78,14 +79,10 @@ export class PgStacApiLambdaRuntime extends Construct {
7879
memorySize: 8192,
7980
logRetention: aws_logs.RetentionDays.ONE_WEEK,
8081
timeout: Duration.seconds(30),
81-
code: resolveLambdaCode(
82-
userCode,
83-
path.join(__dirname, ".."),
84-
{
85-
file: "stac-api/runtime/Dockerfile",
86-
buildArgs: { PYTHON_VERSION: "3.12" },
87-
}
88-
),
82+
code: resolveLambdaCode(userCode, path.join(__dirname, ".."), {
83+
file: "stac-api/runtime/Dockerfile",
84+
buildArgs: { PYTHON_VERSION: "3.12" },
85+
}),
8986
vpc: props.vpc,
9087
vpcSubnets: props.subnetSelection,
9188
allowPublicSubnet: true,
@@ -96,6 +93,9 @@ export class PgStacApiLambdaRuntime extends Construct {
9693
ENABLED_EXTENSIONS: enabledExtensions.join(","),
9794
...props.apiEnv,
9895
},
96+
snapStart: props.enableSnapStart
97+
? lambda.SnapStartConf.ON_PUBLISHED_VERSIONS
98+
: undefined,
9999
// overwrites defaults with user-provided configurable properties (excluding code)
100100
...otherLambdaOptions,
101101
});
@@ -145,6 +145,26 @@ export interface PgStacApiLambdaRuntimeProps {
145145
*/
146146
readonly enabledExtensions?: ExtensionType[];
147147

148+
/**
149+
* Enable SnapStart to reduce cold start latency.
150+
*
151+
* SnapStart creates a snapshot of the initialized Lambda function, allowing new instances
152+
* to start from this pre-initialized state instead of starting from scratch.
153+
*
154+
* Benefits:
155+
* - Significantly reduces cold start times (typically 10x faster)
156+
* - Improves API response time for infrequent requests
157+
*
158+
* Considerations:
159+
* - Additional cost: charges for snapshot storage and restore operations
160+
* - Requires Lambda versioning (automatically configured by this construct)
161+
* - Database connections are recreated on restore using snapshot lifecycle hooks
162+
*
163+
* @see https://docs.aws.amazon.com/lambda/latest/dg/snapstart.html
164+
* @default false
165+
*/
166+
readonly enableSnapStart?: boolean;
167+
148168
/**
149169
* Can be used to override the default lambda function properties.
150170
*
@@ -179,12 +199,15 @@ export class PgStacApiLambda extends Construct {
179199
dbSecret: props.dbSecret,
180200
enabledExtensions: props.enabledExtensions,
181201
apiEnv: props.apiEnv,
202+
enableSnapStart: props.enableSnapStart,
182203
lambdaFunctionOptions: props.lambdaFunctionOptions,
183204
});
184205
this.stacApiLambdaFunction = this.lambdaFunction = runtime.lambdaFunction;
185206

186207
const { api } = new LambdaApiGateway(this, "stac-api", {
187-
lambdaFunction: runtime.lambdaFunction,
208+
lambdaFunction: props.enableSnapStart!
209+
? runtime.lambdaFunction.currentVersion
210+
: runtime.lambdaFunction,
188211
domainName: props.domainName ?? props.stacApiDomainName,
189212
});
190213

‎lib/stac-api/runtime/src/handler.py‎

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
import os
77

88
from mangum import Mangum
9-
from stac_fastapi.pgstac.app import app
9+
from snapshot_restore_py import register_after_restore, register_before_snapshot
10+
from stac_fastapi.pgstac.app import app, with_transactions
1011
from stac_fastapi.pgstac.config import PostgresSettings
1112
from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
1213
from utils import get_secret_dict
@@ -21,12 +22,92 @@
2122
postgres_port=int(secret["port"]),
2223
)
2324

25+
_connection_initialized = False
26+
27+
28+
@register_before_snapshot
29+
def on_snapshot():
30+
"""
31+
Runtime hook called by Lambda before taking a snapshot.
32+
We close database connections that shouldn't be in the snapshot.
33+
"""
34+
35+
# Close any existing database connections before the snapshot is taken
36+
if hasattr(app, "state") and hasattr(app.state, "readpool") and app.state.readpool:
37+
try:
38+
app.state.readpool.close()
39+
app.state.readpool = None
40+
except Exception as e:
41+
print(f"SnapStart: Error closing database readpool: {e}")
42+
43+
if hasattr(app, "state") and hasattr(app.state, "writepool") and app.state.writepool:
44+
try:
45+
app.state.writepool.close()
46+
app.state.writepool = None
47+
except Exception as e:
48+
print(f"SnapStart: Error closing database writepool: {e}")
49+
50+
return {"statusCode": 200}
51+
52+
53+
@register_after_restore
54+
def on_snap_restore():
55+
"""
56+
Runtime hook called by Lambda after restoring from a snapshot.
57+
We recreate database connections that were closed before the snapshot.
58+
"""
59+
global _connection_initialized
60+
61+
try:
62+
# Get the event loop or create a new one
63+
try:
64+
loop = asyncio.get_running_loop()
65+
except RuntimeError:
66+
loop = asyncio.new_event_loop()
67+
asyncio.set_event_loop(loop)
68+
69+
# Close any existing pool (from snapshot)
70+
if hasattr(app.state, "readpool") and app.state.readpool:
71+
try:
72+
app.state.readpool.close()
73+
except Exception as e:
74+
print(f"SnapStart: Error closing stale readpool: {e}")
75+
app.state.readpool = None
76+
77+
if hasattr(app.state, "writepool") and app.state.writepool:
78+
try:
79+
app.state.writepool.close()
80+
except Exception as e:
81+
print(f"SnapStart: Error closing stale writepool: {e}")
82+
app.state.writepool = None
83+
84+
# Create fresh connection pool
85+
loop.run_until_complete(
86+
connect_to_db(
87+
app,
88+
postgres_settings=postgres_settings,
89+
add_write_connection_pool=with_transactions,
90+
)
91+
)
92+
93+
_connection_initialized = True
94+
95+
except Exception as e:
96+
print(f"SnapStart: Failed to initialize database connection: {e}")
97+
raise
98+
99+
return {"statusCode": 200}
100+
24101

25102
@app.on_event("startup")
26103
async def startup_event():
27104
"""Connect to database on startup."""
28105
print("Setting up DB connection...")
29-
await connect_to_db(app, postgres_settings=postgres_settings)
106+
await connect_to_db(
107+
app,
108+
postgres_settings=postgres_settings,
109+
add_write_connection_pool=with_transactions,
110+
)
30111
print("DB connection setup.")
31112

32113

‎lib/tipg-api/index.ts‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ export class TiPgApiLambdaRuntime extends Construct {
2020
constructor(scope: Construct, id: string, props: TiPgApiLambdaRuntimeProps) {
2121
super(scope, id);
2222

23-
const { code: userCode, ...otherLambdaOptions } = props.lambdaFunctionOptions || {};
23+
const { code: userCode, ...otherLambdaOptions } =
24+
props.lambdaFunctionOptions || {};
2425

2526
this.lambdaFunction = new lambda.Function(this, "lambda", {
2627
// defaults
@@ -29,14 +30,10 @@ export class TiPgApiLambdaRuntime extends Construct {
2930
memorySize: 1024,
3031
logRetention: logs.RetentionDays.ONE_WEEK,
3132
timeout: Duration.seconds(30),
32-
code: resolveLambdaCode(
33-
userCode,
34-
path.join(__dirname, ".."),
35-
{
36-
file: "tipg-api/runtime/Dockerfile",
37-
buildArgs: { PYTHON_VERSION: "3.12" },
38-
}
39-
),
33+
code: resolveLambdaCode(userCode, path.join(__dirname, ".."), {
34+
file: "tipg-api/runtime/Dockerfile",
35+
buildArgs: { PYTHON_VERSION: "3.12" },
36+
}),
4037
vpc: props.vpc,
4138
vpcSubnets: props.subnetSelection,
4239
allowPublicSubnet: true,
@@ -46,6 +43,9 @@ export class TiPgApiLambdaRuntime extends Construct {
4643
DB_MAX_CONN_SIZE: "1",
4744
...props.apiEnv,
4845
},
46+
snapStart: props.enableSnapStart
47+
? lambda.SnapStartConf.ON_PUBLISHED_VERSIONS
48+
: undefined,
4949
// overwrites defaults with user-provided configurable properties (excluding code)
5050
...otherLambdaOptions,
5151
});
@@ -88,6 +88,26 @@ export interface TiPgApiLambdaRuntimeProps {
8888
*/
8989
readonly apiEnv?: Record<string, string>;
9090

91+
/**
92+
* Enable SnapStart to reduce cold start latency.
93+
*
94+
* SnapStart creates a snapshot of the initialized Lambda function, allowing new instances
95+
* to start from this pre-initialized state instead of starting from scratch.
96+
*
97+
* Benefits:
98+
* - Significantly reduces cold start times (typically 10x faster)
99+
* - Improves API response time for infrequent requests
100+
*
101+
* Considerations:
102+
* - Additional cost: charges for snapshot storage and restore operations
103+
* - Requires Lambda versioning (automatically configured by this construct)
104+
* - Database connections are recreated on restore using snapshot lifecycle hooks
105+
*
106+
* @see https://docs.aws.amazon.com/lambda/latest/dg/snapstart.html
107+
* @default false
108+
*/
109+
readonly enableSnapStart?: boolean;
110+
91111
/**
92112
* Can be used to override the default lambda function properties.
93113
*
@@ -121,12 +141,15 @@ export class TiPgApiLambda extends Construct {
121141
db: props.db,
122142
dbSecret: props.dbSecret,
123143
apiEnv: props.apiEnv,
144+
enableSnapStart: props.enableSnapStart,
124145
lambdaFunctionOptions: props.lambdaFunctionOptions,
125146
});
126147
this.tiPgLambdaFunction = this.lambdaFunction = runtime.lambdaFunction;
127148

128149
const { api } = new LambdaApiGateway(this, "api", {
129-
lambdaFunction: runtime.lambdaFunction,
150+
lambdaFunction: props.enableSnapStart!
151+
? runtime.lambdaFunction.currentVersion
152+
: runtime.lambdaFunction,
130153
domainName: props.domainName ?? props.tipgApiDomainName,
131154
});
132155

‎lib/tipg-api/runtime/src/handler.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import os
77

88
from mangum import Mangum
9+
from snapshot_restore_py import register_after_restore, register_before_snapshot
910
from tipg.collections import register_collection_catalog
1011
from tipg.database import connect_to_db
1112
from tipg.main import app
@@ -27,6 +28,77 @@
2728
db_settings = DatabaseSettings()
2829
custom_sql_settings = CustomSQLSettings()
2930

31+
_connection_initialized = False
32+
33+
34+
@register_before_snapshot
35+
def on_snapshot():
36+
"""
37+
Runtime hook called by Lambda before taking a snapshot.
38+
We close database connections that shouldn't be in the snapshot.
39+
"""
40+
41+
# Close any existing database connections before the snapshot is taken
42+
if hasattr(app, "state") and hasattr(app.state, "pool") and app.state.pool:
43+
try:
44+
app.state.pool.close()
45+
app.state.pool = None
46+
except Exception as e:
47+
print(f"SnapStart: Error closing database pool: {e}")
48+
49+
return {"statusCode": 200}
50+
51+
52+
@register_after_restore
53+
def on_snap_restore():
54+
"""
55+
Runtime hook called by Lambda after restoring from a snapshot.
56+
We recreate database connections that were closed before the snapshot.
57+
"""
58+
global _connection_initialized
59+
60+
try:
61+
# Get the event loop or create a new one
62+
try:
63+
loop = asyncio.get_running_loop()
64+
except RuntimeError:
65+
loop = asyncio.new_event_loop()
66+
asyncio.set_event_loop(loop)
67+
68+
# Close any existing pool (from snapshot)
69+
if hasattr(app.state, "pool") and app.state.pool:
70+
try:
71+
app.state.pool.close()
72+
except Exception as e:
73+
print(f"SnapStart: Error closing stale pool: {e}")
74+
app.state.pool = None
75+
76+
# Create fresh connection pool
77+
loop.run_until_complete(
78+
connect_to_db(
79+
app,
80+
schemas=db_settings.schemas,
81+
tipg_schema=db_settings.tipg_schema,
82+
user_sql_files=custom_sql_settings.sql_files,
83+
settings=postgres_settings,
84+
)
85+
)
86+
87+
loop.run_until_complete(
88+
register_collection_catalog(
89+
app,
90+
db_settings=db_settings,
91+
)
92+
)
93+
94+
_connection_initialized = True
95+
96+
except Exception as e:
97+
print(f"SnapStart: Failed to initialize database connection: {e}")
98+
raise
99+
100+
return {"statusCode": 200}
101+
30102

31103
@app.on_event("startup")
32104
async def startup_event() -> None:

0 commit comments

Comments
 (0)