Automatic priming of powertools-metrics

Author: subhash686

pom.xml

@@ -91,6 +91,7 @@
         <aws-embedded-metrics.version>4.1.2</aws-embedded-metrics.version>
         <jmespath.version>0.6.0</jmespath.version>
         <elastic.version>1.6.0</elastic.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. -->
@@ -238,6 +239,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>

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

@@ -0,0 +1,71 @@
+/*
+ * 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;
+
+public class ClassPreLoader {
+    public static final String CLASSES_FILE = "classesloaded.txt";
+    public static void preloadClass(String classname, boolean initialize) {
+        try {
+            Class.forName(classname, initialize, ClassPreLoader.class.getClassLoader());
+        } catch (Throwable ignored) {
+        }
+    }
+    public static void preloadClasses(boolean initialize) {
+        try {
+            Enumeration<URL> files = ClassPreLoader.class.getClassLoader().getResources(CLASSES_FILE);
+            while (files.hasMoreElements()) {
+                URL url = files.nextElement();
+                URLConnection conn = url.openConnection();
+                conn.setUseCaches(false);
+                InputStream is = conn.getInputStream();
+                preloadClassesFromStream(is, initialize);
+            }
+        } catch (IOException ignored) {
+        }
+        InputStream is = ClassPreLoader.class
+                .getResourceAsStream( CLASSES_FILE);
+        if (is != null)
+            preloadClassesFromStream(is, initialize);
+    }
+    public static void preloadClassesFromStream(InputStream is, boolean initialize) {
+        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()) {
+                    preloadClass(className, initialize);
+                }
+            }
+        } catch (Exception ignored) {
+        }
+    }
+    public void invokePreloadClasses(boolean initialize) {
+        preloadClasses(initialize);
+    }
+}

powertools-metrics/pom.xml

@@ -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>

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspect.java

@@ -30,21 +30,28 @@
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
+import org.crac.Core;
+import org.crac.Resource;
 import software.amazon.cloudwatchlogs.emf.logger.MetricsLogger;
 import software.amazon.cloudwatchlogs.emf.model.DimensionSet;
 import software.amazon.cloudwatchlogs.emf.model.MetricsContext;
 import software.amazon.cloudwatchlogs.emf.model.Unit;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
 import software.amazon.lambda.powertools.metrics.Metrics;
 import software.amazon.lambda.powertools.metrics.MetricsUtils;
 import software.amazon.lambda.powertools.metrics.ValidationException;
 
 @Aspect
-public class LambdaMetricsAspect {
+public class LambdaMetricsAspect implements Resource {
     public static final String TRACE_ID_PROPERTY = "xray_trace_id";
     public static final String REQUEST_ID_PROPERTY = "function_request_id";
     private static final String NAMESPACE = System.getenv("POWERTOOLS_METRICS_NAMESPACE");
 
+    public LambdaMetricsAspect() {
+        Core.getGlobalContext().register(this);
+    }
+
     private static String service(Metrics metrics) {
         return !"".equals(metrics.service()) ? metrics.service() : serviceName();
     }
@@ -148,4 +155,14 @@ private void validateMetricsAndRefreshOnFailure(Metrics metrics) {
             throw e;
         }
     }
+
+    @Override
+    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
+        System.out.println("before preloading");
+        ClassPreLoader.preloadClasses(true);
+        System.out.println("after preloading");
+    }
+    @Override
+    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
+    }
 }