Skip to content

Commit b436d1e

Browse files
sanchitmalhotra126sanchitmalhotra126
authored
Shutdown lambda if container ID in table (#367)
* Shutdown lambda if container ID in table * log severe error * rename method
1 parent 0f8d84d commit b436d1e

File tree

6 files changed

+166
-1
lines changed

6 files changed

+166
-1
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,7 @@ Resources:
459459
API_ENDPOINT: !Sub
460460
- "https://${ApiId}.execute-api.${AWS::Region}.amazonaws.com/${StageName}"
461461
- ApiId: !Ref WebSocketApi
462+
UNHEALTHY_CONTAINERS_TABLE_NAME: !Ref UnhealthyContainersTable
462463
<%end -%>
463464

464465
ContentBucket:

‎org-code-javabuilder/lib/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ dependencies {
2929
implementation group: 'com.amazonaws', name: 'aws-java-sdk-s3', version: '1.12.20'
3030
// https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-cloudwatch
3131
implementation group: 'com.amazonaws', name: 'aws-java-sdk-cloudwatch', version: '1.12.218'
32+
// https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-dynamodb
33+
implementation group: 'com.amazonaws', name: 'aws-java-sdk-dynamodb', version: '1.12.285'
3234
// https://mvnrepository.com/artifact/org.json/json
3335
implementation group: 'org.json', name: 'json', version: '20210307'
3436
// jaxb-api is needed to get rid of warning on run

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@ private LambdaErrorCodes() {
99
public static final int LOW_DISK_SPACE_ERROR_CODE = 50;
1010
public static final int OUT_OF_MEMORY_ERROR_CODE = 60;
1111
public static final int CONNECTION_POOL_SHUT_DOWN_ERROR_CODE = 70;
12+
public static final int UNHEALTHY_CONTAINER_ERROR_CODE = 80;
1213
}

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

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package org.code.javabuilder;
22

33
import static org.code.javabuilder.InternalFacingExceptionTypes.INVALID_INPUT;
4-
import static org.code.protocol.InternalExceptionKey.*;
54
import static org.code.protocol.LoggerNames.MAIN_LOGGER;
65

76
import com.amazonaws.client.builder.AwsClientBuilder;
@@ -10,6 +9,8 @@
109
import com.amazonaws.services.apigatewaymanagementapi.model.DeleteConnectionRequest;
1110
import com.amazonaws.services.apigatewaymanagementapi.model.GetConnectionRequest;
1211
import com.amazonaws.services.apigatewaymanagementapi.model.GoneException;
12+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
13+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
1314
import com.amazonaws.services.lambda.runtime.Context;
1415
import com.amazonaws.services.lambda.runtime.RequestHandler;
1516
import com.amazonaws.services.s3.AmazonS3;
@@ -26,6 +27,7 @@
2627
import java.util.UUID;
2728
import java.util.logging.Handler;
2829
import java.util.logging.Logger;
30+
import org.code.javabuilder.UnhealthyContainerChecker.ShutdownTrigger;
2931
import org.code.javabuilder.util.LambdaUtils;
3032
import org.code.protocol.*;
3133
import org.code.validation.support.UserTestOutputAdapter;
@@ -47,11 +49,18 @@ public class LambdaRequestHandler implements RequestHandler<Map<String, String>,
4749
private static final String CONTENT_BUCKET_NAME = System.getenv("CONTENT_BUCKET_NAME");
4850
private static final String CONTENT_BUCKET_URL = System.getenv("CONTENT_BUCKET_URL");
4951
private static final String API_ENDPOINT = System.getenv("API_ENDPOINT");
52+
private static final String UNHEALTHY_CONTAINERS_TABLE_NAME =
53+
System.getenv("UNHEALTHY_CONTAINERS_TABLE_NAME");
5054

5155
// Creating these clients here rather than in the request handler method allows us to use
5256
// provisioned concurrency to decrease cold boot time by 3-10 seconds, depending on the lambda
5357
private static final AmazonSQS SQS_CLIENT = AmazonSQSClientBuilder.defaultClient();
5458
private static final AmazonS3 S3_CLIENT = AmazonS3ClientBuilder.standard().build();
59+
private static final AmazonDynamoDB DYNAMO_DB_CLIENT =
60+
AmazonDynamoDBClientBuilder.defaultClient();
61+
62+
// Used to check whether this container has been marked unhealthy.
63+
private final UnhealthyContainerChecker unhealthyContainerChecker;
5564

5665
// Controls whether the current invocation session has been initialized. This should be reset on
5766
// every invocation.
@@ -71,6 +80,8 @@ public LambdaRequestHandler() {
7180
.withEndpointConfiguration(
7281
new AwsClientBuilder.EndpointConfiguration(API_ENDPOINT, "us-east-1"))
7382
.build();
83+
this.unhealthyContainerChecker =
84+
new UnhealthyContainerChecker(DYNAMO_DB_CLIENT, UNHEALTHY_CONTAINERS_TABLE_NAME);
7485
}
7586

7687
/**
@@ -177,6 +188,9 @@ public String handleRequest(Map<String, String> lambdaInput, Context context) {
177188
* creating global objects
178189
*/
179190
private void initialize(Map<String, String> lambdaInput, String connectionId, Context context) {
191+
// Check container health status and exit early if container has been marked unhealthy.
192+
this.shutdownContainerIfUnhealthy(ShutdownTrigger.START);
193+
180194
final boolean canAccessDashboardAssets =
181195
Boolean.parseBoolean(lambdaInput.get("canAccessDashboardAssets"));
182196

@@ -320,6 +334,9 @@ private void shutDown(
320334
System.exit(LambdaErrorCodes.LOW_DISK_SPACE_ERROR_CODE);
321335
}
322336

337+
// Check container health status and exit if the container has been marked unhealthy.
338+
this.shutdownContainerIfUnhealthy(ShutdownTrigger.END);
339+
323340
this.isSessionInitialized = false;
324341
}
325342

@@ -401,4 +418,14 @@ private void verifyApiClient(String connectionId) {
401418
.build();
402419
}
403420
}
421+
422+
/**
423+
* Checks if this container has been marked unhealthy and if so, forces a shutdown via
424+
* System.exit().
425+
*/
426+
private void shutdownContainerIfUnhealthy(ShutdownTrigger trigger) {
427+
if (this.unhealthyContainerChecker.shouldForceShutdownContainer(LAMBDA_ID, trigger)) {
428+
System.exit(LambdaErrorCodes.UNHEALTHY_CONTAINER_ERROR_CODE);
429+
}
430+
}
404431
}

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package org.code.javabuilder;
2+
3+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
4+
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
5+
import java.util.Map;
6+
import org.code.protocol.LoggerUtils;
7+
8+
/**
9+
* Checks if the current container has been marked unhealthy, so that we can shut it down if needed.
10+
*/
11+
public class UnhealthyContainerChecker {
12+
static final String CONTAINER_ID_KEY_NAME = "container_id";
13+
14+
/**
15+
* When the health status is being checked. This allows us to choose whether to trigger a shutdown
16+
* at the beginning or end of the session.
17+
*/
18+
public enum ShutdownTrigger {
19+
START("start"),
20+
END("end");
21+
22+
private final String name;
23+
24+
ShutdownTrigger(String name) {
25+
this.name = name;
26+
}
27+
28+
public String getName() {
29+
return this.name;
30+
}
31+
}
32+
33+
private final AmazonDynamoDB dynamoDBClient;
34+
private final String tableName;
35+
36+
public UnhealthyContainerChecker(AmazonDynamoDB dynamoDBClient, String tableName) {
37+
this.dynamoDBClient = dynamoDBClient;
38+
this.tableName = tableName;
39+
}
40+
41+
public boolean shouldForceShutdownContainer(String containerId, ShutdownTrigger trigger) {
42+
// The container ID value is a concatenation of the ID and the shutdown trigger type
43+
final String containerIdCompositeValue = containerId + "#" + trigger.getName();
44+
final Map<String, AttributeValue> key =
45+
Map.of(CONTAINER_ID_KEY_NAME, new AttributeValue(containerIdCompositeValue));
46+
final Map<String, AttributeValue> entry;
47+
try {
48+
entry = this.dynamoDBClient.getItem(this.tableName, key).getItem();
49+
} catch (Exception e) {
50+
// Indicates an unexpected error (missing entries should return null); log error and return
51+
// false silently to be safe.
52+
LoggerUtils.logSevereException(e);
53+
return false;
54+
}
55+
56+
return entry != null;
57+
}
58+
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package org.code.javabuilder;
2+
3+
import static org.code.javabuilder.UnhealthyContainerChecker.CONTAINER_ID_KEY_NAME;
4+
import static org.junit.jupiter.api.Assertions.*;
5+
import static org.mockito.ArgumentMatchers.anyString;
6+
import static org.mockito.Mockito.mock;
7+
import static org.mockito.Mockito.when;
8+
9+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
10+
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
11+
import com.amazonaws.services.dynamodbv2.model.GetItemResult;
12+
import java.util.Map;
13+
import org.code.javabuilder.UnhealthyContainerChecker.ShutdownTrigger;
14+
import org.code.protocol.JavabuilderContext;
15+
import org.code.protocol.MetricClient;
16+
import org.junit.jupiter.api.BeforeEach;
17+
import org.junit.jupiter.api.Test;
18+
import org.mockito.ArgumentCaptor;
19+
20+
class UnhealthyContainerCheckerTest {
21+
private static final String TABLE_NAME = "tableName";
22+
23+
private GetItemResult getItemResult;
24+
private ArgumentCaptor<Map<String, AttributeValue>> keyCaptor;
25+
private UnhealthyContainerChecker unitUnderTest;
26+
27+
@BeforeEach
28+
public void setUp() {
29+
JavabuilderContext.getInstance().register(MetricClient.class, mock(AWSMetricClient.class));
30+
final AmazonDynamoDB dynamoDBClient = mock(AmazonDynamoDB.class);
31+
getItemResult = mock(GetItemResult.class);
32+
keyCaptor = ArgumentCaptor.forClass(Map.class);
33+
when(dynamoDBClient.getItem(anyString(), keyCaptor.capture())).thenReturn(getItemResult);
34+
35+
unitUnderTest = new UnhealthyContainerChecker(dynamoDBClient, TABLE_NAME);
36+
}
37+
38+
@Test
39+
public void testReturnsTrueIfContainerIDFound() {
40+
final String containerId = "containerId1234";
41+
final ShutdownTrigger trigger = ShutdownTrigger.END;
42+
when(getItemResult.getItem()).thenReturn(Map.of());
43+
44+
assertTrue(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
45+
46+
this.verifyKey(containerId, trigger);
47+
}
48+
49+
@Test
50+
public void testReturnsFalseIfContainerIDNotFound() {
51+
final String containerId = "containerId5678";
52+
final ShutdownTrigger trigger = ShutdownTrigger.START;
53+
when(getItemResult.getItem()).thenReturn(null);
54+
55+
assertFalse(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
56+
57+
this.verifyKey(containerId, trigger);
58+
}
59+
60+
@Test
61+
public void testReturnsFalseIfClientThrowsException() {
62+
final String containerId = "containerId9090";
63+
final ShutdownTrigger trigger = ShutdownTrigger.END;
64+
when(getItemResult.getItem()).thenThrow(new RuntimeException("exception"));
65+
66+
assertFalse(unitUnderTest.shouldForceShutdownContainer(containerId, trigger));
67+
68+
this.verifyKey(containerId, trigger);
69+
}
70+
71+
private void verifyKey(String containerId, ShutdownTrigger trigger) {
72+
final Map<String, AttributeValue> key = keyCaptor.getValue();
73+
final AttributeValue value = key.get(CONTAINER_ID_KEY_NAME);
74+
assertEquals(containerId + "#" + trigger.getName(), value.getS());
75+
}
76+
}

0 commit comments

Comments
 (0)