Skip to content

Commit 094c3d9

Browse files
subhash686phipag
subhash686
and
phipag
authored
feat: Support CRaC priming of powertools metrics and idempotency-dynamodb (#1861)
* Automatic priming of powertools-metrics * Fixed resource name and added unit tests and javadoc comments * Invoke prime Dynamo persistent store * Fixed Sonar issues * Update classesloaded.txt replaced build time class list with run time class list * Moved priming to MetricsFactory and DynamoDBPersistenceStore, added Priming.md * Update classesloaded.txt classesloaded file generated using the new method * Update LambdaMetricsAspect.java Removed unused import * Update DynamoDBPersistenceStore.java Fixed variable naming Sonar issue * Improved priming.md documentation and static initialization of classes using for priming. Also fixed an unrelated flaky unit test * Fixed a Sonarqube finding and added more unit tests to ensure all available classes are loaded --------- Co-authored-by: Philipp Page <[email protected]>
1 parent 0305642 commit 094c3d9

File tree

11 files changed

+3927
-4
lines changed

11 files changed

+3927
-4
lines changed

‎Priming.md‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Automatic Priming for AWS Lambda Powertools Java
2+
3+
## Table of Contents
4+
- [Overview](#overview)
5+
- [Implementation Steps](#general-implementation-steps)
6+
- [Known Issues](#known-issues-and-solutions)
7+
- [Reference Implementation](#reference-implementation)
8+
9+
## Overview
10+
Priming is the process of preloading dependencies and initializing resources during the INIT phase, rather than during the INVOKE phase to further optimize startup performance with SnapStart.
11+
This is required because Java frameworks that use dependency injection load classes into memory when these classes are explicitly invoked, which typically happens during Lambda’s INVOKE phase.
12+
13+
This documentation provides guidance for automatic class priming in Powertools for AWS Lambda Java modules.
14+
15+
16+
## Implementation Steps
17+
Classes are proactively loaded using Java runtime hooks which are part of the open source [CRaC (Coordinated Restore at Checkpoint) project](https://openjdk.org/projects/crac/).
18+
Implementations across the project use the `beforeCheckpoint()` hook, to prime Snapstart-enabled Java functions via Class Priming.
19+
In order to generate the `classloaded.txt` file for a Java module in this project, follow these general steps.
20+
21+
1. **Add Maven Profile**
22+
- Add maven test profile with the following VM argument for generating classes loaded files.
23+
```shell
24+
-Xlog:class+load=info:classesloaded.txt
25+
```
26+
- You can find an example of this in `generate-classesloaded-file` profile in this [pom.xml](powertools-metrics/pom.xml).
27+
28+
2. **Generate classes loaded file**
29+
- Run tests with `-Pgenerate-classesloaded-file` profile.
30+
```shell
31+
mvn -Pgenerate-classesloaded-file clean test
32+
```
33+
- This will generate a file named `classesloaded.txt` in the target directory of the module.
34+
35+
3. **Cleanup the file**
36+
- The classes loaded file generated in Step 2 has the format
37+
`[0.054s][info][class,load] java.lang.Object source: shared objects file`
38+
but we are only interested in `java.lang.Object` - the fully qualified class name.
39+
- To strip the lines to include only the fully qualified class name,
40+
Use the following regex to replace with empty string.
41+
- `^\[[\[\]0-9.a-z,]+ ` (to replace the left part)
42+
- `( source: )[0-9a-z :/._$-]+` (to replace the right part)
43+
44+
4. **Add file to resources**
45+
- Move the cleaned-up file to the corresponding `src/main/resources` directory of the module. See [example](powertools-metrics/src/main/resources/classesloaded.txt).
46+
47+
5. **Register and checkpoint**
48+
- A class, usually the entry point of the module, should register the CRaC resource in the constructor. [Example](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java)
49+
- Note that AspectJ aspect is not suitable for this purpose, as it does not work with CRaC.
50+
- Add the `beforeCheckpoint()` hook in the same class to invoke `ClassPreLoader.preloadClasses()`. The `ClassPreLoader` class is implemented in `powertools-common` module.
51+
- This will ensure that the classes are already pre-loaded by the Snapstart RESTORE operation leading to a shorter INIT duration.
52+
53+
54+
## Known Issues
55+
- This is a manual process at the moment, but it can be automated in the future.
56+
- `classesloaded.txt` file includes test classes as well because the file is generated while running tests. This is not a problem because all the classes that are not found are ignored by `ClassPreLoader.preloadClasses()`. Also `beforeCheckpoint()` hook is not time-sensitive, it only runs once when a new Lambda version gets published.
57+
58+
## Reference Implementation
59+
Working example is available in the [powertools-metrics](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java).

‎pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@
117117
<mockito.version>5.18.0</mockito.version>
118118
<mockito-junit-jupiter.version>5.18.0</mockito-junit-jupiter.version>
119119
<junit-pioneer.version>2.3.0</junit-pioneer.version>
120+
<crac.version>1.4.0</crac.version>
120121

121122
<!-- As we have a .mvn directory at the root of the project, this will evaluate to the root directory
122123
regardless of where maven is run - sub-module, or root. -->
@@ -264,6 +265,11 @@
264265
<artifactId>logback-ecs-encoder</artifactId>
265266
<version>${elastic.version}</version>
266267
</dependency>
268+
<dependency>
269+
<groupId>org.crac</groupId>
270+
<artifactId>crac</artifactId>
271+
<version>${crac.version}</version>
272+
</dependency>
267273
<dependency>
268274
<groupId>org.slf4j</groupId>
269275
<artifactId>slf4j-api</artifactId>

‎powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/ClassPreLoader.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.common.internal;
15+
16+
import java.io.BufferedReader;
17+
import java.io.IOException;
18+
import java.io.InputStream;
19+
import java.io.InputStreamReader;
20+
import java.nio.charset.StandardCharsets;
21+
import java.net.URL;
22+
import java.net.URLConnection;
23+
import java.util.Enumeration;
24+
25+
/**
26+
* Used to preload classes to support automatic priming for SnapStart
27+
*/
28+
public final class ClassPreLoader {
29+
public static final String CLASSES_FILE = "classesloaded.txt";
30+
31+
private ClassPreLoader() {
32+
// Hide default constructor
33+
}
34+
35+
/**
36+
* Initializes the classes listed in the classesloaded resource
37+
*/
38+
public static void preloadClasses() {
39+
try {
40+
Enumeration<URL> files = ClassPreLoader.class.getClassLoader().getResources(CLASSES_FILE);
41+
// If there are multiple files, preload classes from all of them
42+
while (files.hasMoreElements()) {
43+
URL url = files.nextElement();
44+
URLConnection conn = url.openConnection();
45+
conn.setUseCaches(false);
46+
InputStream is = conn.getInputStream();
47+
preloadClassesFromStream(is);
48+
}
49+
} catch (IOException ignored) {
50+
// No action is required if preloading fails for any reason
51+
}
52+
}
53+
54+
/**
55+
* Loads the list of classes passed as a stream
56+
*
57+
* @param is
58+
*/
59+
private static void preloadClassesFromStream(InputStream is) {
60+
try (is;
61+
InputStreamReader isr = new InputStreamReader(is, StandardCharsets.UTF_8);
62+
BufferedReader reader = new BufferedReader(isr)) {
63+
String line;
64+
while ((line = reader.readLine()) != null) {
65+
int idx = line.indexOf('#');
66+
if (idx != -1) {
67+
line = line.substring(0, idx);
68+
}
69+
final String className = line.stripTrailing();
70+
if (!className.isBlank()) {
71+
loadClassIfFound(className);
72+
}
73+
}
74+
} catch (Exception ignored) {
75+
// No action is required if preloading fails for any reason
76+
}
77+
}
78+
79+
/**
80+
* Initializes the class with given name if found, ignores otherwise
81+
*
82+
* @param className
83+
*/
84+
private static void loadClassIfFound(String className) {
85+
try {
86+
Class.forName(className, true, ClassPreLoader.class.getClassLoader());
87+
} catch (ClassNotFoundException e) {
88+
// No action is required if the class with given name cannot be found
89+
}
90+
}
91+
}

‎powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/ClassPreLoaderTest.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package software.amazon.lambda.powertools.common.internal;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import org.junit.jupiter.api.Test;
6+
7+
class ClassPreLoaderTest {
8+
9+
// Making this volatile so the Thread Context doesn't need any special handling
10+
static volatile boolean dummyClassLoaded = false;
11+
12+
/**
13+
* Dummy class to be loaded by ClassPreLoader in test.
14+
* <b>The class name is referenced in <i>powertools-common/src/test/resources/classesloaded.txt</i></b>
15+
* This class is used to verify that the ClassPreLoader can load valid classes.
16+
* The static block sets a flag to indicate that the class has been loaded.
17+
*/
18+
static class DummyClass {
19+
static {
20+
dummyClassLoaded = true;
21+
}
22+
}
23+
@Test
24+
void preloadClasses_shouldIgnoreInvalidClassesAndLoadValidClasses() {
25+
26+
dummyClassLoaded = false;
27+
// powertools-common/src/test/resources/classesloaded.txt has a class that does not exist
28+
// Verify that the missing class did not throw any exception
29+
assertDoesNotThrow(ClassPreLoader::preloadClasses);
30+
31+
// When the classloaded.txt is a mixed bag of valid and invalid classes, Valid class must load
32+
assertTrue(dummyClassLoaded, "DummyClass should be loaded");
33+
}
34+
}

‎powertools-common/src/test/resources/classesloaded.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
software.amazon.lambda.powertools.common.internal.NonExistingClass
2+
software.amazon.lambda.powertools.common.internal.ClassPreLoaderTest$DummyClass

‎powertools-idempotency/powertools-idempotency-dynamodb/pom.xml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,10 @@
5757
</exclusion>
5858
</exclusions>
5959
</dependency>
60+
<dependency>
61+
<groupId>org.crac</groupId>
62+
<artifactId>crac</artifactId>
63+
</dependency>
6064

6165
<!-- Test dependencies -->
6266
<dependency>

‎powertools-idempotency/powertools-idempotency-dynamodb/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/DynamoDBPersistenceStore.java‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414

1515
package software.amazon.lambda.powertools.idempotency.persistence.dynamodb;
1616

17+
import org.crac.Core;
18+
import org.crac.Resource;
1719
import org.slf4j.Logger;
1820
import org.slf4j.LoggerFactory;
1921
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
@@ -37,6 +39,7 @@
3739
import software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore;
3840

3941
import java.time.Instant;
42+
import java.time.temporal.ChronoUnit;
4043
import java.util.AbstractMap;
4144
import java.util.HashMap;
4245
import java.util.Map;
@@ -52,7 +55,7 @@
5255
* DynamoDB version of the {@link PersistenceStore}. Will store idempotency data in DynamoDB.<br>
5356
* Use the {@link Builder} to create a new instance.
5457
*/
55-
public class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore {
58+
public final class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore, Resource {
5659

5760
public static final String IDEMPOTENCY = "idempotency";
5861
private static final Logger LOG = LoggerFactory.getLogger(DynamoDBPersistenceStore.class);
@@ -109,6 +112,39 @@ private DynamoDBPersistenceStore(String tableName,
109112
this.dynamoDbClient = null;
110113
}
111114
}
115+
Core.getGlobalContext().register(this);
116+
}
117+
118+
/**
119+
* Primes the persistent store by invoking the get record method with a key that doesn't exist.
120+
*
121+
* @param context
122+
* @throws Exception
123+
*/
124+
@Override
125+
public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
126+
try {
127+
String primingRecordKey = "__invoke_prime__";
128+
Instant now = Instant.now();
129+
long expiry = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
130+
DataRecord primingDataRecord = new DataRecord(
131+
primingRecordKey,
132+
DataRecord.Status.COMPLETED,
133+
expiry,
134+
null, // no data
135+
null // no validation
136+
);
137+
putRecord(primingDataRecord, Instant.now());
138+
getRecord(primingRecordKey);
139+
deleteRecord(primingRecordKey);
140+
} catch (Exception unknown) {
141+
// This is unexpected but we must continue without any interruption
142+
}
143+
}
144+
145+
@Override
146+
public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
147+
// This is a no-op, as we don't need to do anything after restore
112148
}
113149

114150
public static Builder builder() {

‎powertools-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaJsonEncoderTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,7 @@ void shouldLogThreadInfo() {
368368
String result = new String(encoded, StandardCharsets.UTF_8);
369369

370370
// THEN
371-
assertThat(result).contains("\"thread\":\"main\",\"thread_id\":1,\"thread_priority\":5");
371+
assertThat(result).contains("\"thread\":\"main\",\"thread_id\":"+ Thread.currentThread().getId() +",\"thread_priority\":5");
372372
}
373373

374374
@Test

‎powertools-metrics/pom.xml‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@
3939
<artifactId>aspectjrt</artifactId>
4040
<scope>provided</scope>
4141
</dependency>
42+
<dependency>
43+
<groupId>org.crac</groupId>
44+
<artifactId>crac</artifactId>
45+
</dependency>
4246
<dependency>
4347
<groupId>software.amazon.lambda</groupId>
4448
<artifactId>powertools-common</artifactId>
@@ -114,6 +118,23 @@
114118
</dependencies>
115119

116120
<profiles>
121+
<profile>
122+
<id>generate-classesloaded-file</id>
123+
<build>
124+
<plugins>
125+
<plugin>
126+
<groupId>org.apache.maven.plugins</groupId>
127+
<artifactId>maven-surefire-plugin</artifactId>
128+
<configuration>
129+
<argLine>-Xlog:class+load=info:classesloaded.txt
130+
--add-opens java.base/java.util=ALL-UNNAMED
131+
--add-opens java.base/java.lang=ALL-UNNAMED
132+
</argLine>
133+
</configuration>
134+
</plugin>
135+
</plugins>
136+
</build>
137+
</profile>
117138
<profile>
118139
<id>generate-graalvm-files</id>
119140
<dependencies>

‎powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414

1515
package software.amazon.lambda.powertools.metrics;
1616

17+
import org.crac.Core;
18+
import org.crac.Resource;
19+
import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
1720
import software.amazon.lambda.powertools.common.internal.LambdaConstants;
1821
import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
1922
import software.amazon.lambda.powertools.metrics.model.DimensionSet;
@@ -23,11 +26,16 @@
2326
/**
2427
* Factory for accessing the singleton Metrics instance
2528
*/
26-
public final class MetricsFactory {
29+
public final class MetricsFactory implements Resource {
2730
private static MetricsProvider provider = new EmfMetricsProvider();
2831
private static Metrics metrics;
2932

30-
private MetricsFactory() {
33+
// Dummy instance to register MetricsFactory with CRaC
34+
private static final MetricsFactory INSTANCE = new MetricsFactory();
35+
36+
// Static block to ensure CRaC registration happens at class loading time
37+
static {
38+
Core.getGlobalContext().register(INSTANCE);
3139
}
3240

3341
/**
@@ -68,4 +76,15 @@ public static synchronized void setMetricsProvider(MetricsProvider metricsProvid
6876
// Reset the metrics instance so it will be recreated with the new provider
6977
metrics = null;
7078
}
79+
80+
@Override
81+
public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
82+
MetricsFactory.getMetricsInstance();
83+
ClassPreLoader.preloadClasses();
84+
}
85+
86+
@Override
87+
public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
88+
// No action needed after restore
89+
}
7190
}

0 commit comments

Comments
 (0)