Skip to content

feat: Support CRaC and priming of powertools metrics and idempotency-dynamodb #1861

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 17 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions Priming.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# Automatic Priming for AWS Lambda Powertools Java

## Table of Contents
- [Overview](#overview)
- [Implementation Steps](#general-implementation-steps)
- [Known Issues](#known-issues-and-solutions)
- [Reference Implementation](#reference-implementation)

## Overview
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.
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.

This documentation provides guidance for automatic class priming in AWS Lambda Powertools Java modules.


## Implementation Steps
Classes are proactively loaded using Java runtime hooks which are part of the open source CRaC (Coordinated Restore at Checkpoint) project.
This implementation uses this hook, called `beforeCheckpoint()`, to prime SnapStart-enabled Java functions via Class Priming.
In order to generate the `classloaded` files for Powertools java module, follow these general steps.

1. **Add Maven Profile**
- Add maven test profile with the following VM argument for generating classes loaded files.
```shell
-Xlog:class+load=info:classesloaded.txt
```
- You can find an example of this in profile `generate-classesloaded-file` in this [pom.xml](powertools-metrics/pom.xml).

2. **Generate classes loaded file**
- Run tests with `-Pgenerate-classesloaded-file` profile.
```shell
mvn -Pgenerate-classesloaded-file clean test
```
- This will generate a file named `classesloaded.txt` in the target directory of the module.

3. **Cleanup the file**
- The classes loaded file generated in Step 2 has the format
`[0.054s][info][class,load] java.lang.Object source: shared objects file`
but we are only interested in `java.lang.Object` the fully qualified class name.
- To strip the lines to include only the fully qualified class name,
Use the following regex to replace with empty string.
- `^\[[\[\]0-9.a-z,]+ ` (to replace the left part)
- `( source: )[0-9a-z :/._$-]+` (to replace the right part)

4. **Place the file**
- Move the cleanedup file containing to the corresponding resources directory of the module. See [example](powertools-metrics/src/resources/classesloaded.txt).

5. **Register and checkpoint**
- Register the CRaC Resource in the constructor of the main class.
- Add the `beforeCheckpoint()` hook in the same class to invoke `ClassPreLoader.preloadClasses()`.
- This will ensure that the classes are loaded during the INIT phase of the Lambda function. See [example](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java)


## Known Issues
- This is a manual process at the moment, but it can be automated in the future.
- Because the file is generated while running tests, it includes test classes as well.

## Reference Implementation
Working example is available in the [powertools-metrics](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java).
6 changes: 6 additions & 0 deletions pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@
<mockito.version>5.18.0</mockito.version>
<mockito-junit-jupiter.version>5.18.0</mockito-junit-jupiter.version>
<junit-pioneer.version>2.3.0</junit-pioneer.version>
<crac.version>1.4.0</crac.version>

<!-- As we have a .mvn directory at the root of the project, this will evaluate to the root directory
regardless of where maven is run - sub-module, or root. -->
Expand Down Expand Up @@ -264,6 +265,11 @@
<artifactId>logback-ecs-encoder</artifactId>
<version>${elastic.version}</version>
</dependency>
<dependency>
<groupId>org.crac</groupId>
<artifactId>crac</artifactId>
<version>${crac.version}</version>
</dependency>
<dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-api</artifactId>
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
/*
* Copyright 2023 Amazon.com, Inc. or its affiliates.
* Licensed under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
* http://www.apache.org/licenses/LICENSE-2.0
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package software.amazon.lambda.powertools.common.internal;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.net.URL;
import java.net.URLConnection;
import java.util.Enumeration;

/**
* Used to preload classes to support automatic priming for SnapStart
*/
public class ClassPreLoader {
public static final String CLASSES_FILE = "classesloaded.txt";

private ClassPreLoader() {
// Hide default constructor
}
/**
* Initializes the classes listed in the classesloaded resource
*/
public static void preloadClasses() {
try {
Enumeration<URL> files = ClassPreLoader.class.getClassLoader().getResources(CLASSES_FILE);
// If there are multiple files, preload classes from all of them
while (files.hasMoreElements()) {
URL url = files.nextElement();
URLConnection conn = url.openConnection();
conn.setUseCaches(false);
InputStream is = conn.getInputStream();
preloadClassesFromStream(is);
}
} catch (IOException ignored) {
// No action is required if preloading fails for any reason
}
}

/**
* Loads the list of classes passed as a stream
*
* @param is
*/
private static void preloadClassesFromStream(InputStream is) {
try (is;
InputStreamReader isr = new InputStreamReader(is, StandardCharsets.UTF_8);
BufferedReader reader = new BufferedReader(isr)) {
String line;
while ((line = reader.readLine()) != null) {
int idx = line.indexOf('#');
if (idx != -1) {
line = line.substring(0, idx);
}
final String className = line.stripTrailing();
if (!className.isBlank()) {
Class.forName(className, true, ClassPreLoader.class.getClassLoader());
}
}
} catch (Exception ignored) {
// No action is required if preloading fails for any reason
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
package software.amazon.lambda.powertools.common.internal;

import static org.junit.jupiter.api.Assertions.*;

import org.junit.jupiter.api.Test;

class ClassPreLoaderTest {

@Test
void preloadClasses_shouldIgnoreInvalidClassesAndLoadValidClasses() {
// Verify that the missing class did not throw any exception
assertDoesNotThrow(ClassPreLoader::preloadClasses);
}
}
1 change: 1 addition & 0 deletions powertools-common/src/test/resources/classesloaded.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
software.amazon.lambda.powertools.common.internal.NonExistingClass
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>org.crac</groupId>
<artifactId>crac</artifactId>
</dependency>

<!-- Test dependencies -->
<dependency>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@

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

import org.crac.Core;
import org.crac.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
Expand All @@ -37,6 +39,7 @@
import software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.AbstractMap;
import java.util.HashMap;
import java.util.Map;
Expand All @@ -52,7 +55,7 @@
* DynamoDB version of the {@link PersistenceStore}. Will store idempotency data in DynamoDB.<br>
* Use the {@link Builder} to create a new instance.
*/
public class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore {
public class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore, Resource {

public static final String IDEMPOTENCY = "idempotency";
private static final Logger LOG = LoggerFactory.getLogger(DynamoDBPersistenceStore.class);
Expand Down Expand Up @@ -109,6 +112,39 @@ private DynamoDBPersistenceStore(String tableName,
this.dynamoDbClient = null;
}
}
Core.getGlobalContext().register(this);
}

/**
* Primes the persistent store by invoking the get record method with a key that doesn't exist.
*
* @param context
* @throws Exception
*/
@Override
public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
try {
String primingRecordKey = "__invoke_prime__";
Instant now = Instant.now();
long expiry = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
DataRecord primingDataRecord = new DataRecord(
primingRecordKey,
DataRecord.Status.COMPLETED,
expiry,
null, // no data
null // no validation
);
putRecord(primingDataRecord, Instant.now());
getRecord(primingRecordKey);
deleteRecord(primingRecordKey);
} catch (Exception unknown) {
// This is unexpected but we must continue without any interruption
}
}

@Override
public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
// This is a no-op, as we don't need to do anything after restore
}

public static Builder builder() {
Expand Down
21 changes: 21 additions & 0 deletions powertools-metrics/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@
<artifactId>aspectjrt</artifactId>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.crac</groupId>
<artifactId>crac</artifactId>
</dependency>
<dependency>
<groupId>software.amazon.lambda</groupId>
<artifactId>powertools-common</artifactId>
Expand Down Expand Up @@ -114,6 +118,23 @@
</dependencies>

<profiles>
<profile>
<id>generate-classesloaded-file</id>
<build>
<plugins>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-surefire-plugin</artifactId>
<configuration>
<argLine>-Xlog:class+load=info:classesloaded.txt
--add-opens java.base/java.util=ALL-UNNAMED
--add-opens java.base/java.lang=ALL-UNNAMED
</argLine>
</configuration>
</plugin>
</plugins>
</build>
</profile>
<profile>
<id>generate-graalvm-files</id>
<dependencies>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@

package software.amazon.lambda.powertools.metrics;

import org.crac.Core;
import org.crac.Resource;
import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
import software.amazon.lambda.powertools.common.internal.LambdaConstants;
import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
import software.amazon.lambda.powertools.metrics.model.DimensionSet;
Expand All @@ -23,11 +26,12 @@
/**
* Factory for accessing the singleton Metrics instance
*/
public final class MetricsFactory {
public final class MetricsFactory implements Resource {
private static MetricsProvider provider = new EmfMetricsProvider();
private static Metrics metrics;

private MetricsFactory() {
Core.getGlobalContext().register(this);
}

/**
Expand Down Expand Up @@ -68,4 +72,14 @@ public static synchronized void setMetricsProvider(MetricsProvider metricsProvid
// Reset the metrics instance so it will be recreated with the new provider
metrics = null;
}

@Override
public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
ClassPreLoader.preloadClasses();
}

@Override
public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
// No action needed after restore
}
}
Loading