Runtime hooks (#412)

* Add crac context implementation to RIC
* Add restore endpoint and runtime hooks execution logic

---------

Co-authored-by: Anton Stepanov <110172761+anton-stepanof@users.noreply.github.com>

Author: smirnoal

aws-lambda-java-runtime-interface-client/pom.xml

@@ -0,0 +1,12 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+public class CheckpointException extends Exception {
+    private static final long serialVersionUID = -4956873658083157585L;
+    public CheckpointException() {
+        super();
+    }
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/CheckpointException.java

@@ -0,0 +1,21 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+public abstract class Context<R extends Resource> implements Resource {
+
+    protected Context() {
+    }
+
+    @Override
+    public abstract void beforeCheckpoint(Context<? extends Resource> context)
+            throws CheckpointException;
+
+    @Override
+    public abstract void afterRestore(Context<? extends Resource> context)
+            throws RestoreException;
+
+    public abstract void register(R resource);
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/Context.java

@@ -0,0 +1,90 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.stream.Collectors;
+
+
+/**
+ * Spec reference: https://crac.github.io/openjdk-builds/javadoc/api/java.base/jdk/crac/package-summary.html
+ */
+
+public class ContextImpl extends Context<Resource> {
+
+    private volatile long order = -1L;
+    private final WeakHashMap<Resource, Long> checkpointQueue = new WeakHashMap<>();
+
+    @Override
+    public synchronized void beforeCheckpoint(Context<? extends Resource> context) throws CheckpointException {
+
+        List<Throwable> exceptionsThrown = new ArrayList<>();
+        for (Resource resource : getCheckpointQueueReverseOrderOfRegistration()) {
+            try {
+                resource.beforeCheckpoint(this);
+            } catch (CheckpointException e) {
+                Collections.addAll(exceptionsThrown, e.getSuppressed());
+            } catch (Exception e) {
+                exceptionsThrown.add(e);
+            }
+        }
+
+        if (!exceptionsThrown.isEmpty()) {
+            CheckpointException checkpointException = new CheckpointException();
+            for (Throwable t : exceptionsThrown) {
+                checkpointException.addSuppressed(t);
+            }
+            throw checkpointException;
+        }
+    }
+
+    @Override
+    public synchronized void afterRestore(Context<? extends Resource> context) throws RestoreException {
+
+        List<Throwable> exceptionsThrown = new ArrayList<>();
+        for (Resource resource : getCheckpointQueueForwardOrderOfRegistration()) {
+            try {
+                resource.afterRestore(this);
+            } catch (RestoreException e) {
+                Collections.addAll(exceptionsThrown, e.getSuppressed());
+            } catch (Exception e) {
+                exceptionsThrown.add(e);
+            }
+        }
+
+        if (!exceptionsThrown.isEmpty()) {
+            RestoreException restoreException = new RestoreException();
+            for (Throwable t : exceptionsThrown) {
+                restoreException.addSuppressed(t);
+            }
+            throw restoreException;
+        }
+    }
+
+    @Override
+    public synchronized void register(Resource resource) {
+        checkpointQueue.put(resource, ++order);
+    }
+
+    private List<Resource> getCheckpointQueueReverseOrderOfRegistration() {
+        return checkpointQueue.entrySet()
+                .stream()
+                .sorted((r1, r2) -> (int) (r2.getValue() - r1.getValue()))
+                .map(Map.Entry::getKey)
+                .collect(Collectors.toList());
+    }
+
+    private List<Resource> getCheckpointQueueForwardOrderOfRegistration() {
+        return checkpointQueue.entrySet()
+                .stream()
+                .sorted((r1, r2) -> (int) (r1.getValue() - r2.getValue()))
+                .map(Map.Entry::getKey)
+                .collect(Collectors.toList());
+    }
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/ContextImpl.java

@@ -0,0 +1,28 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+/**
+ * Provides the global context for registering resources.
+ */
+public final class Core {
+
+    private Core() {
+    }
+
+    private static Context<Resource> globalContext = new ContextImpl();
+
+    public static Context<Resource> getGlobalContext() {
+        return globalContext;
+    }
+
+    public static void checkpointRestore() {
+        throw new UnsupportedOperationException();
+    }
+
+    static void resetGlobalContext() {
+        globalContext = new ContextImpl();
+    }
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/Core.java

@@ -0,0 +1,11 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+public interface Resource {
+    void afterRestore(Context<? extends Resource> context) throws Exception;
+
+    void beforeCheckpoint(Context<? extends Resource> context) throws Exception;
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/Resource.java

@@ -0,0 +1,13 @@
+/*
+ *  Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ */
+
+package com.amazonaws.services.lambda.crac;
+
+public class RestoreException extends Exception {
+    private static final long serialVersionUID = -823900409868237860L;
+
+    public RestoreException() {
+        super();
+    }
+}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/crac/RestoreException.java

@@ -5,6 +5,7 @@
 //
 package com.amazonaws.services.lambda.runtime.api.client;
 
+import com.amazonaws.services.lambda.crac.Core;
 import com.amazonaws.services.lambda.runtime.LambdaLogger;
 import com.amazonaws.services.lambda.runtime.api.client.LambdaRequestHandler.UserFaultHandler;
 import com.amazonaws.services.lambda.runtime.api.client.logging.FramedTelemetryLogSink;
@@ -32,6 +33,7 @@
 import java.security.Security;
 import java.util.Properties;
 
+
 /**
  * The entrypoint of this class is {@link AWSLambda#startRuntime}. It performs two main tasks:
  *
@@ -59,6 +61,10 @@ public class AWSLambda {
     // https://github.com/aws/aws-xray-sdk-java/blob/2f467e50db61abb2ed2bd630efc21bddeabd64d9/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java#L39-L40
     private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
 
+    private static final String INIT_TYPE_SNAP_START = "snap-start";
+
+    private static final String AWS_LAMBDA_INITIALIZATION_TYPE = System.getenv(ReservedRuntimeEnvironmentVariables.AWS_LAMBDA_INITIALIZATION_TYPE);
+
     static {
         // Override the disabledAlgorithms setting to match configuration for openjdk8-u181.
         // This is to keep DES ciphers around while we deploying security updates.
@@ -211,14 +217,13 @@ private static void startRuntime(String handler, LambdaLogger lambdaLogger) thro
             requestHandler = findRequestHandler(handler, customerClassLoader);
         } catch (UserFault userFault) {
             lambdaLogger.log(userFault.reportableError());
-            ByteArrayOutputStream payload = new ByteArrayOutputStream(1024);
-            Failure failure = new Failure(userFault);
-            GsonFactory.getInstance().getSerializer(Failure.class).toJson(failure, payload);
-            runtimeClient.postInitError(payload.toByteArray(), failure.getErrorType());
+            reportInitError(new Failure(userFault), runtimeClient);
             System.exit(1);
             return;
         }
-
+        if (INIT_TYPE_SNAP_START.equals(AWS_LAMBDA_INITIALIZATION_TYPE)) {
+            onInitComplete(runtimeClient, lambdaLogger);
+        }
         boolean shouldExit = false;
         while (!shouldExit) {
             UserFault userFault = null;
@@ -260,6 +265,49 @@ private static void startRuntime(String handler, LambdaLogger lambdaLogger) thro
         }
     }
 
+    static void onInitComplete(final LambdaRuntimeClient runtimeClient, final LambdaLogger lambdaLogger) throws IOException {
+        try {
+            Core.getGlobalContext().beforeCheckpoint(null);
+            // Blocking call to RAPID /restore/next API, will return after taking snapshot.
+            // This will also be the 'entrypoint' when resuming from snapshots.
+            runtimeClient.getRestoreNext();
+        } catch (Exception e1) {
+            logExceptionCloudWatch(lambdaLogger, e1);
+            reportInitError(new Failure(e1), runtimeClient);
+            System.exit(64);
+        }
+        try {
+            Core.getGlobalContext().afterRestore(null);
+        } catch (Exception restoreExc) {
+            logExceptionCloudWatch(lambdaLogger, restoreExc);
+            Failure errorPayload = new Failure(restoreExc);
+            reportRestoreError(errorPayload, runtimeClient);
+            System.exit(64);
+        }
+    }
+
+    private static void logExceptionCloudWatch(LambdaLogger lambdaLogger, Exception exc) {
+        UserFault.filterStackTrace(exc);
+        UserFault userFault = UserFault.makeUserFault(exc, true);
+        lambdaLogger.log(userFault.reportableError());
+    }
+
+    static void reportInitError(final Failure failure,
+                                final LambdaRuntimeClient runtimeClient) throws IOException {
+
+        ByteArrayOutputStream payload = new ByteArrayOutputStream(1024);
+        JacksonFactory.getInstance().getSerializer(Failure.class).toJson(failure, payload);
+        runtimeClient.postInitError(payload.toByteArray(), failure.getErrorType());
+    }
+
+    static int reportRestoreError(final Failure failure,
+                                  final LambdaRuntimeClient runtimeClient) throws IOException {
+
+        ByteArrayOutputStream payload = new ByteArrayOutputStream(1024);
+        JacksonFactory.getInstance().getSerializer(Failure.class).toJson(failure, payload);
+        return runtimeClient.postRestoreError(payload.toByteArray(), failure.getErrorType());
+    }
+
     private static PojoSerializer<XRayErrorCause> xRayErrorCauseSerializer;
 
     /**

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/AWSLambda.java

@@ -81,4 +81,8 @@ public int compare(Class o1, Class o2) {
     public String getErrorType() {
         return errorType;
     }
+
+    public String getErrorMessage() {
+        return errorMessage;
+    }
 }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/Failure.java

@@ -87,6 +87,12 @@ public interface ReservedRuntimeEnvironmentVariables {
      */
     String AWS_LAMBDA_RUNTIME_API = "AWS_LAMBDA_RUNTIME_API";
 
+
+    /**
+     * Initialization type
+     */
+    String AWS_LAMBDA_INITIALIZATION_TYPE = "AWS_LAMBDA_INITIALIZATION_TYPE";
+
     /**
      * The path to your Lambda function code.
      */

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/ReservedRuntimeEnvironmentVariables.java

@@ -80,6 +80,12 @@ public static <T extends Throwable> T filterStackTrace(T t) {
         if(cause != null) {
             filterStackTrace(cause);
         }
+
+        Throwable[] suppressedExceptions = t.getSuppressed();
+        for(Throwable suppressed: suppressedExceptions) {
+            filterStackTrace(suppressed);
+        }
+
         return t;
     }