feat: Support CRaC and priming of powertools logging

docs/core/logging.md

@@ -1339,6 +1339,70 @@ sequenceDiagram
 8. **What happens if Lambda times out without flushing the buffer?** Logs that are still in the buffer will be lost.
 9. **How does the `BufferingAppender` work with different appenders?** The `BufferingAppender` is designed to wrap arbitrary appenders, providing maximum flexibility. You can wrap console appenders, file appenders, or any custom appenders with buffering functionality.
 
+### Lambda SnapStart priming
+
+The PowertoolsLogging class integrates with AWS Lambda SnapStart to improve restore durations. To make sure the SnapStart
+priming logic of this class runs correctly, you need an explicit reference to `PowertoolsLogging` in your code to allow the
+library to register before SnapStart takes a memory snapshot. Learn more about what priming is in
+this [blog post](https://aws.amazon.com/blogs/compute/optimizing-cold-start-performance-of-aws-lambda-using-advanced-priming-strategies-with-snapstart/)
+{target="_blank"}.
+
+Make sure to call `PowertoolsLogging.init()` in your Lambda handler initialization code. This can be done by adding one of the following lines to your handler class:
+
+=== "Constructor"
+
+    ```java hl_lines="12"
+    import org.slf4j.Logger;
+    import org.slf4j.LoggerFactory;
+    import software.amazon.lambda.powertools.logging.Logging;
+    import software.amazon.lambda.powertools.logging.PowertoolsLogging;
+    // ... other imports
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        private static final Logger LOGGER = LoggerFactory.getLogger(MyFunctionHandler.class);
+
+        public MyFunctionHandler() {
+            PowertoolsLogging.init(); // Ensure PowertoolsLogging is loaded for SnapStart
+        }
+
+        @Override
+        @Logging
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            LOGGER.info("Example logging");
+            // ...
+            return something;
+        }
+    }
+    ```
+
+=== "Static Initializer"
+
+    ```java hl_lines="12"
+    import org.slf4j.Logger;
+    import org.slf4j.LoggerFactory;
+    import software.amazon.lambda.powertools.logging.Logging;
+    import software.amazon.lambda.powertools.logging.PowertoolsLogging;
+    // ... other imports
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        private static final Logger LOGGER = LoggerFactory.getLogger(MyFunctionHandler.class);
+
+        static {
+            PowertoolsLogging.init(); // Ensure PowertoolsLogging is loaded for SnapStart
+        }
+
+        @Override
+        @Logging
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            LOGGER.info("Example logging");
+            // ...
+            return something;
+        }
+    }
+    ```
+
 ## Sampling debug logs
 
 You can dynamically set a percentage of your logs to`DEBUG` level to be included in the logger output, regardless of configured log level, using the`POWERTOOLS_LOGGER_SAMPLE_RATE` environment variable,

powertools-logging/pom.xml

@@ -55,6 +55,10 @@
             <artifactId>aspectjrt</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -116,6 +120,24 @@
         </dependency>
     </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>native</id>
             <build>

powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/PowertoolsLogging.java

@@ -26,21 +26,21 @@
 import static software.amazon.lambda.powertools.logging.internal.PowertoolsLoggedFields.SAMPLING_RATE;
 import static software.amazon.lambda.powertools.logging.internal.PowertoolsLoggedFields.SERVICE;
 
+import com.amazonaws.services.lambda.runtime.Context;
+import com.fasterxml.jackson.databind.JsonNode;
+import io.burt.jmespath.Expression;
 import java.util.Arrays;
 import java.util.Locale;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
-
+import org.crac.Core;
+import org.crac.Resource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.slf4j.MDC;
 import org.slf4j.event.Level;
-
-import com.amazonaws.services.lambda.runtime.Context;
-import com.fasterxml.jackson.databind.JsonNode;
-
-import io.burt.jmespath.Expression;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.logging.internal.BufferManager;
 import software.amazon.lambda.powertools.logging.internal.LoggingManager;
 import software.amazon.lambda.powertools.logging.internal.LoggingManagerRegistry;
@@ -52,10 +52,10 @@
  * This class abstracts away the underlying logging framework (Log4j2, Logback) and provides a unified
  * interface for Lambda context extraction, correlation ID handling, sampling rate configuration,
  * log buffering operations, and other Lambda-specific logging features.
- * 
+ *
  * <p>This class serves as a programmatic alternative to AspectJ-based {@code @Logging} annotation,
  * allowing developers to integrate Powertools logging capabilities without AspectJ dependencies.</p>
- * 
+ * <p>
  * Key features:
  * <ul>
  *   <li>Lambda context initialization with function metadata, trace ID, and service name</li>
@@ -64,27 +64,38 @@
  *   <li>MDC state management for structured logging</li>
  * </ul>
  */
-public final class PowertoolsLogging {
+public final class PowertoolsLogging implements Resource {
     private static final Logger LOG = LoggerFactory.getLogger(PowertoolsLogging.class);
     private static final ThreadLocal<Random> SAMPLER = ThreadLocal.withInitial(Random::new);
     private static AtomicBoolean hasBeenInitialized = new AtomicBoolean(false);
 
+    // Dummy instance to register PowertoolsLogging with CRaC
+    private static final PowertoolsLogging INSTANCE = new PowertoolsLogging();
+
     static {
         initializeLogLevel();
+        Core.getGlobalContext().register(INSTANCE);
     }
 
     private PowertoolsLogging() {
         // Utility class
     }
 
+    public static void init() {
+        // Placeholder method used to enable SnapStart priming. Users need a direct reference to this class in order
+        // for the CRaC hooks to execute.
+        new PowertoolsLogging();
+    }
+
     private static void initializeLogLevel() {
         if (POWERTOOLS_LOG_LEVEL != null) {
             Level powertoolsLevel = getLevelFromString(POWERTOOLS_LOG_LEVEL);
             if (LAMBDA_LOG_LEVEL != null) {
                 Level lambdaLevel = getLevelFromString(LAMBDA_LOG_LEVEL);
                 if (powertoolsLevel.toInt() < lambdaLevel.toInt()) {
                     LOG.warn(
-                            "Current log level ({}) does not match AWS Lambda Advanced Logging Controls minimum log level ({}). This can lead to data loss, consider adjusting them.",
+                            "Current log level ({}) does not match AWS Lambda Advanced Logging Controls minimum log " +
+                                    "level ({}). This can lead to data loss, consider adjusting them.",
                             POWERTOOLS_LOG_LEVEL, LAMBDA_LOG_LEVEL);
                 }
             }
@@ -140,10 +151,10 @@ public static void clearBuffer() {
      * Initializes Lambda logging context with standard Powertools fields.
      * This method should be called at the beginning of your Lambda handler to set up
      * logging context with Lambda function information, trace ID, and service name.
-     * 
+     *
      * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
      * {@link #withLogging(Context, Supplier)} to handle cleanup automatically.</p>
-     * 
+     *
      * @param context the Lambda context provided by AWS Lambda runtime
      */
     public static void initializeLogging(Context context) {
@@ -154,11 +165,11 @@ public static void initializeLogging(Context context) {
      * Initializes Lambda logging context with sampling rate configuration.
      * This method sets up logging context and optionally enables DEBUG logging
      * based on the provided sampling rate.
-     * 
+     *
      * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
      * {@link #withLogging(Context, double, Supplier)} to handle cleanup automatically.</p>
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
+     *
+     * @param context      the Lambda context provided by AWS Lambda runtime
      * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
      */
     public static void initializeLogging(Context context, double samplingRate) {
@@ -169,13 +180,13 @@ public static void initializeLogging(Context context, double samplingRate) {
      * Initializes Lambda logging context with correlation ID extraction.
      * This method sets up logging context and extracts correlation ID from the event
      * using the provided JSON path.
-     * 
+     *
      * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
      * {@link #withLogging(Context, String, Object, Supplier)} to handle cleanup automatically.</p>
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
+     *
+     * @param context           the Lambda context provided by AWS Lambda runtime
      * @param correlationIdPath JSON path to extract correlation ID from event
-     * @param event the Lambda event object
+     * @param event             the Lambda event object
      */
     public static void initializeLogging(Context context, String correlationIdPath, Object event) {
         initializeLogging(context, 0.0, correlationIdPath, event);
@@ -186,16 +197,16 @@ public static void initializeLogging(Context context, String correlationIdPath,
      * This method sets up logging context with Lambda function information,
      * configures sampling rate for DEBUG logging, and optionally extracts
      * correlation ID from the event.
-     * 
+     *
      * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
      * {@link #withLogging(Context, double, String, Object, Supplier)} to handle cleanup automatically.</p>
-     * 
+     *
      * <p>This method is thread-safe.</p>
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
-     * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
+     *
+     * @param context           the Lambda context provided by AWS Lambda runtime
+     * @param samplingRate      sampling rate for DEBUG logging (0.0 to 1.0)
      * @param correlationIdPath JSON path to extract correlation ID from event (can be null)
-     * @param event the Lambda event object (required if correlationIdPath is provided)
+     * @param event             the Lambda event object (required if correlationIdPath is provided)
      */
     public static void initializeLogging(Context context, double samplingRate, String correlationIdPath, Object event) {
 
@@ -254,7 +265,8 @@ private static double getEffectiveSamplingRate(double samplingRate) {
                 return Double.parseDouble(envSampleRate);
             } catch (NumberFormatException e) {
                 LOG.warn(
-                        "Skipping sampling rate on environment variable configuration because of invalid value. Sampling rate: {}",
+                        "Skipping sampling rate on environment variable configuration because of invalid value. " +
+                                "Sampling rate: {}",
                         envSampleRate);
             }
         }
@@ -281,7 +293,7 @@ private static void captureCorrelationId(String correlationIdPath, Object event)
 
     /**
      * Clears MDC state and log buffer.
-     * 
+     *
      * @param clearMdcState whether to clear MDC state
      */
     public static void clearState(boolean clearMdcState) {
@@ -294,10 +306,10 @@ public static void clearState(boolean clearMdcState) {
 
     /**
      * Executes code with logging context initialized and automatically clears state.
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
+     *
+     * @param context  the Lambda context provided by AWS Lambda runtime
      * @param supplier the code to execute with logging context
-     * @param <T> the return type
+     * @param <T>      the return type
      * @return the result of the supplier execution
      */
     public static <T> T withLogging(Context context, Supplier<T> supplier) {
@@ -311,11 +323,11 @@ public static <T> T withLogging(Context context, Supplier<T> supplier) {
 
     /**
      * Executes code with logging context initialized with sampling rate and automatically clears state.
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
+     *
+     * @param context      the Lambda context provided by AWS Lambda runtime
      * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
-     * @param supplier the code to execute with logging context
-     * @param <T> the return type
+     * @param supplier     the code to execute with logging context
+     * @param <T>          the return type
      * @return the result of the supplier execution
      */
     public static <T> T withLogging(Context context, double samplingRate, Supplier<T> supplier) {
@@ -329,12 +341,12 @@ public static <T> T withLogging(Context context, double samplingRate, Supplier<T
 
     /**
      * Executes code with logging context initialized with correlation ID extraction and automatically clears state.
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
+     *
+     * @param context           the Lambda context provided by AWS Lambda runtime
      * @param correlationIdPath JSON path to extract correlation ID from event
-     * @param event the Lambda event object
-     * @param supplier the code to execute with logging context
-     * @param <T> the return type
+     * @param event             the Lambda event object
+     * @param supplier          the code to execute with logging context
+     * @param <T>               the return type
      * @return the result of the supplier execution
      */
     public static <T> T withLogging(Context context, String correlationIdPath, Object event, Supplier<T> supplier) {
@@ -348,22 +360,33 @@ public static <T> T withLogging(Context context, String correlationIdPath, Objec
 
     /**
      * Executes code with logging context initialized with full configuration and automatically clears state.
-     * 
-     * @param context the Lambda context provided by AWS Lambda runtime
-     * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
+     *
+     * @param context           the Lambda context provided by AWS Lambda runtime
+     * @param samplingRate      sampling rate for DEBUG logging (0.0 to 1.0)
      * @param correlationIdPath JSON path to extract correlation ID from event (can be null)
-     * @param event the Lambda event object (required if correlationIdPath is provided)
-     * @param supplier the code to execute with logging context
-     * @param <T> the return type
+     * @param event             the Lambda event object (required if correlationIdPath is provided)
+     * @param supplier          the code to execute with logging context
+     * @param <T>               the return type
      * @return the result of the supplier execution
      */
     public static <T> T withLogging(Context context, double samplingRate, String correlationIdPath, Object event,
-            Supplier<T> supplier) {
+                                    Supplier<T> supplier) {
         initializeLogging(context, samplingRate, correlationIdPath, event);
         try {
             return supplier.get();
         } finally {
             clearState(true);
         }
     }
+
+    @Override
+    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
+        init();
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }