feat: Support CRaC and priming of powertools logging

docs/core/logging.md

@@ -1645,3 +1645,56 @@ Use the `LambdaEcsEncoder` rather than the `LambdaJsonEncoder` when configuring
         </root>
     </configuration>
     ```
+
+### 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"}.
+
+If you don't set a custom `PowertoolsLogging` in your code yet, make sure to reference `PowertoolsLogging` 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="7"
+    import software.amazon.lambda.powertools.validation.Validation;
+    import software.amazon.lambda.powertools.validation.ValidationConfig;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        public MyFunctionHandler() {
+            PowertoolsLogging.init(); // Ensure PowertoolsLogging is loaded for SnapStart
+        }
+
+        @Override
+        @Validation(inboundSchema = "classpath:/schema_in.json", outboundSchema = "classpath:/schema_out.json")
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```
+
+=== "Static Initializer"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.validation.Validation;
+    import software.amazon.lambda.powertools.validation.ValidationConfig;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        static {
+            PowertoolsLogging.init(); // Ensure PowertoolsLogging is loaded for SnapStart
+        }
+
+        @Override
+        @Validation(inboundSchema = "classpath:/schema_in.json", outboundSchema = "classpath:/schema_out.json")
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```

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;
@@ -64,17 +64,23 @@
  *   <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() {
@@ -366,4 +372,15 @@ public static <T> T withLogging(Context context, double samplingRate, String cor
             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
+    }
 }