Shutdown lambda if container ID in table (#367)

* Shutdown lambda if container ID in table

* log severe error

* rename method

Author: sanchitmalhotra126

cicd/3-app/javabuilder/template.yml.erb

@@ -459,6 +459,7 @@ Resources:
           API_ENDPOINT: !Sub
             - "https://${ApiId}.execute-api.${AWS::Region}.amazonaws.com/${StageName}"
             - ApiId: !Ref WebSocketApi
+          UNHEALTHY_CONTAINERS_TABLE_NAME: !Ref UnhealthyContainersTable
 <%end -%>
 
   ContentBucket:

org-code-javabuilder/lib/build.gradle

@@ -29,6 +29,8 @@ dependencies {
     implementation group: 'com.amazonaws', name: 'aws-java-sdk-s3', version: '1.12.20'
     // https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-cloudwatch
     implementation group: 'com.amazonaws', name: 'aws-java-sdk-cloudwatch', version: '1.12.218'
+    // https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-dynamodb
+    implementation group: 'com.amazonaws', name: 'aws-java-sdk-dynamodb', version: '1.12.285'
     // https://mvnrepository.com/artifact/org.json/json
     implementation group: 'org.json', name: 'json', version: '20210307'
     // jaxb-api is needed to get rid of warning on run

org-code-javabuilder/lib/src/main/java/org/code/javabuilder/LambdaErrorCodes.java

@@ -9,4 +9,5 @@ private LambdaErrorCodes() {
   public static final int LOW_DISK_SPACE_ERROR_CODE = 50;
   public static final int OUT_OF_MEMORY_ERROR_CODE = 60;
   public static final int CONNECTION_POOL_SHUT_DOWN_ERROR_CODE = 70;
+  public static final int UNHEALTHY_CONTAINER_ERROR_CODE = 80;
 }

org-code-javabuilder/lib/src/main/java/org/code/javabuilder/LambdaRequestHandler.java

@@ -1,7 +1,6 @@
 package org.code.javabuilder;
 
 import static org.code.javabuilder.InternalFacingExceptionTypes.INVALID_INPUT;
-import static org.code.protocol.InternalExceptionKey.*;
 import static org.code.protocol.LoggerNames.MAIN_LOGGER;
 
 import com.amazonaws.client.builder.AwsClientBuilder;
@@ -10,6 +9,8 @@
 import com.amazonaws.services.apigatewaymanagementapi.model.DeleteConnectionRequest;
 import com.amazonaws.services.apigatewaymanagementapi.model.GetConnectionRequest;
 import com.amazonaws.services.apigatewaymanagementapi.model.GoneException;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.amazonaws.services.s3.AmazonS3;
@@ -26,6 +27,7 @@
 import java.util.UUID;
 import java.util.logging.Handler;
 import java.util.logging.Logger;
+import org.code.javabuilder.UnhealthyContainerChecker.ShutdownTrigger;
 import org.code.javabuilder.util.LambdaUtils;
 import org.code.protocol.*;
 import org.code.validation.support.UserTestOutputAdapter;
@@ -47,11 +49,18 @@ public class LambdaRequestHandler implements RequestHandler<Map<String, String>,
   private static final String CONTENT_BUCKET_NAME = System.getenv("CONTENT_BUCKET_NAME");
   private static final String CONTENT_BUCKET_URL = System.getenv("CONTENT_BUCKET_URL");
   private static final String API_ENDPOINT = System.getenv("API_ENDPOINT");
+  private static final String UNHEALTHY_CONTAINERS_TABLE_NAME =
+      System.getenv("UNHEALTHY_CONTAINERS_TABLE_NAME");
 
   // Creating these clients here rather than in the request handler method allows us to use
   // provisioned concurrency to decrease cold boot time by 3-10 seconds, depending on the lambda
   private static final AmazonSQS SQS_CLIENT = AmazonSQSClientBuilder.defaultClient();
   private static final AmazonS3 S3_CLIENT = AmazonS3ClientBuilder.standard().build();
+  private static final AmazonDynamoDB DYNAMO_DB_CLIENT =
+      AmazonDynamoDBClientBuilder.defaultClient();
+
+  // Used to check whether this container has been marked unhealthy.
+  private final UnhealthyContainerChecker unhealthyContainerChecker;
 
   // Controls whether the current invocation session has been initialized. This should be reset on
   // every invocation.
@@ -71,6 +80,8 @@ public LambdaRequestHandler() {
             .withEndpointConfiguration(
                 new AwsClientBuilder.EndpointConfiguration(API_ENDPOINT, "us-east-1"))
             .build();
+    this.unhealthyContainerChecker =
+        new UnhealthyContainerChecker(DYNAMO_DB_CLIENT, UNHEALTHY_CONTAINERS_TABLE_NAME);
   }
 
   /**
@@ -177,6 +188,9 @@ public String handleRequest(Map<String, String> lambdaInput, Context context) {
    * creating global objects
    */
   private void initialize(Map<String, String> lambdaInput, String connectionId, Context context) {
+    // Check container health status and exit early if container has been marked unhealthy.
+    this.shutdownContainerIfUnhealthy(ShutdownTrigger.START);
+
     final boolean canAccessDashboardAssets =
         Boolean.parseBoolean(lambdaInput.get("canAccessDashboardAssets"));
 
@@ -320,6 +334,9 @@ private void shutDown(
       System.exit(LambdaErrorCodes.LOW_DISK_SPACE_ERROR_CODE);
     }
 
+    // Check container health status and exit if the container has been marked unhealthy.
+    this.shutdownContainerIfUnhealthy(ShutdownTrigger.END);
+
     this.isSessionInitialized = false;
   }
 
@@ -401,4 +418,14 @@ private void verifyApiClient(String connectionId) {
               .build();
     }
   }
+
+  /**
+   * Checks if this container has been marked unhealthy and if so, forces a shutdown via
+   * System.exit().
+   */
+  private void shutdownContainerIfUnhealthy(ShutdownTrigger trigger) {
+    if (this.unhealthyContainerChecker.shouldForceShutdownContainer(LAMBDA_ID, trigger)) {
+      System.exit(LambdaErrorCodes.UNHEALTHY_CONTAINER_ERROR_CODE);
+    }
+  }
 }

org-code-javabuilder/lib/src/main/java/org/code/javabuilder/UnhealthyContainerChecker.java

@@ -0,0 +1,58 @@
+package org.code.javabuilder;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.model.AttributeValue;
+import java.util.Map;
+import org.code.protocol.LoggerUtils;
+
+/**
+ * Checks if the current container has been marked unhealthy, so that we can shut it down if needed.
+ */
+public class UnhealthyContainerChecker {
+  static final String CONTAINER_ID_KEY_NAME = "container_id";
+
+  /**
+   * When the health status is being checked. This allows us to choose whether to trigger a shutdown
+   * at the beginning or end of the session.
+   */
+  public enum ShutdownTrigger {
+    START("start"),
+    END("end");
+
+    private final String name;
+
+    ShutdownTrigger(String name) {
+      this.name = name;
+    }
+
+    public String getName() {
+      return this.name;
+    }
+  }
+
+  private final AmazonDynamoDB dynamoDBClient;
+  private final String tableName;
+
+  public UnhealthyContainerChecker(AmazonDynamoDB dynamoDBClient, String tableName) {
+    this.dynamoDBClient = dynamoDBClient;
+    this.tableName = tableName;
+  }
+
+  public boolean shouldForceShutdownContainer(String containerId, ShutdownTrigger trigger) {
+    // The container ID value is a concatenation of the ID and the shutdown trigger type
+    final String containerIdCompositeValue = containerId + "#" + trigger.getName();
+    final Map<String, AttributeValue> key =
+        Map.of(CONTAINER_ID_KEY_NAME, new AttributeValue(containerIdCompositeValue));
+    final Map<String, AttributeValue> entry;
+    try {
+      entry = this.dynamoDBClient.getItem(this.tableName, key).getItem();
+    } catch (Exception e) {
+      // Indicates an unexpected error (missing entries should return null); log error and return
+      // false silently to be safe.
+      LoggerUtils.logSevereException(e);
+      return false;
+    }
+
+    return entry != null;
+  }
+}

org-code-javabuilder/lib/src/test/java/org/code/javabuilder/UnhealthyContainerCheckerTest.java

@@ -0,0 +1,76 @@
+package org.code.javabuilder;
+
+import static org.code.javabuilder.UnhealthyContainerChecker.CONTAINER_ID_KEY_NAME;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.model.AttributeValue;
+import com.amazonaws.services.dynamodbv2.model.GetItemResult;
+import java.util.Map;
+import org.code.javabuilder.UnhealthyContainerChecker.ShutdownTrigger;
+import org.code.protocol.JavabuilderContext;
+import org.code.protocol.MetricClient;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+class UnhealthyContainerCheckerTest {
+  private static final String TABLE_NAME = "tableName";
+
+  private GetItemResult getItemResult;
+  private ArgumentCaptor<Map<String, AttributeValue>> keyCaptor;
+  private UnhealthyContainerChecker unitUnderTest;
+
+  @BeforeEach
+  public void setUp() {
+    JavabuilderContext.getInstance().register(MetricClient.class, mock(AWSMetricClient.class));
+    final AmazonDynamoDB dynamoDBClient = mock(AmazonDynamoDB.class);
+    getItemResult = mock(GetItemResult.class);
+    keyCaptor = ArgumentCaptor.forClass(Map.class);
+    when(dynamoDBClient.getItem(anyString(), keyCaptor.capture())).thenReturn(getItemResult);
+
+    unitUnderTest = new UnhealthyContainerChecker(dynamoDBClient, TABLE_NAME);
+  }
+
+  @Test
+  public void testReturnsTrueIfContainerIDFound() {
+    final String containerId = "containerId1234";
+    final ShutdownTrigger trigger = ShutdownTrigger.END;
+    when(getItemResult.getItem()).thenReturn(Map.of());
+
+    assertTrue(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
+
+    this.verifyKey(containerId, trigger);
+  }
+
+  @Test
+  public void testReturnsFalseIfContainerIDNotFound() {
+    final String containerId = "containerId5678";
+    final ShutdownTrigger trigger = ShutdownTrigger.START;
+    when(getItemResult.getItem()).thenReturn(null);
+
+    assertFalse(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
+
+    this.verifyKey(containerId, trigger);
+  }
+
+  @Test
+  public void testReturnsFalseIfClientThrowsException() {
+    final String containerId = "containerId9090";
+    final ShutdownTrigger trigger = ShutdownTrigger.END;
+    when(getItemResult.getItem()).thenThrow(new RuntimeException("exception"));
+
+    assertFalse(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
+
+    this.verifyKey(containerId, trigger);
+  }
+
+  private void verifyKey(String containerId, ShutdownTrigger trigger) {
+    final Map<String, AttributeValue> key = keyCaptor.getValue();
+    final AttributeValue value = key.get(CONTAINER_ID_KEY_NAME);
+    assertEquals(containerId + "#" + trigger.getName(), value.getS());
+  }
+}