Merge branch 'main' into sanchit/catch-cleanup-exception

Author: sanchitmalhotra126

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

@@ -8,6 +8,7 @@
 import com.amazonaws.services.apigatewaymanagementapi.AmazonApiGatewayManagementApi;
 import com.amazonaws.services.apigatewaymanagementapi.AmazonApiGatewayManagementApiClientBuilder;
 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.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
@@ -49,24 +50,27 @@ public class LambdaRequestHandler implements RequestHandler<Map<String, String>,
 
   // 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 AmazonApiGatewayManagementApi API_CLIENT =
-      AmazonApiGatewayManagementApiClientBuilder.standard()
-          .withEndpointConfiguration(
-              new AwsClientBuilder.EndpointConfiguration(API_ENDPOINT, "us-east-1"))
-          .build();
   private static final AmazonSQS SQS_CLIENT = AmazonSQSClientBuilder.defaultClient();
   private static final AmazonS3 S3_CLIENT = AmazonS3ClientBuilder.standard().build();
 
   // Controls whether the current invocation session has been initialized. This should be reset on
   // every invocation.
   private boolean isSessionInitialized = false;
+  // API Gateway Client. We create this in the constructor so we can recreate it if it goes away for
+  // some reason.
+  private AmazonApiGatewayManagementApi apiClient;
 
   public LambdaRequestHandler() {
     // create CachedResources once for the entire container.
     // This will only be called once in the initial creation of the lambda instance.
     // Documentation: https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html
     CachedResources.create();
     COLD_BOOT_END = Clock.systemUTC().instant();
+    this.apiClient =
+        AmazonApiGatewayManagementApiClientBuilder.standard()
+            .withEndpointConfiguration(
+                new AwsClientBuilder.EndpointConfiguration(API_ENDPOINT, "us-east-1"))
+            .build();
   }
 
   /**
@@ -134,7 +138,7 @@ public String handleRequest(Map<String, String> lambdaInput, Context context) {
     }
 
     final ExceptionHandler exceptionHandler =
-        new ExceptionHandler(outputAdapter, new AWSSystemExitHelper(connectionId, API_CLIENT));
+        new ExceptionHandler(outputAdapter, new AWSSystemExitHelper(connectionId, this.apiClient));
     final TempDirectoryManager tempDirectoryManager = new AWSTempDirectoryManager();
 
     CodeExecutionManager codeExecutionManager = null;
@@ -150,7 +154,7 @@ public String handleRequest(Map<String, String> lambdaInput, Context context) {
       // Create and start thread that that will notify us if we're nearing the timeout limit
       timeoutNotifierThread =
           this.createTimeoutThread(
-              context, outputAdapter, codeExecutionManager, connectionId, API_CLIENT);
+              context, outputAdapter, codeExecutionManager, connectionId, this.apiClient);
       timeoutNotifierThread.start();
 
       // Initialize and start code execution
@@ -162,7 +166,7 @@ public String handleRequest(Map<String, String> lambdaInput, Context context) {
       if (timeoutNotifierThread != null) {
         timeoutNotifierThread.interrupt();
       }
-      this.shutDown(codeExecutionManager, connectionId, API_CLIENT);
+      this.shutDown(codeExecutionManager, connectionId, this.apiClient);
     }
 
     return "done";
@@ -188,6 +192,8 @@ private void initialize(Map<String, String> lambdaInput, String connectionId, Co
     // /opt is the folder all layer files go into.
     props.put("sun.awt.fontconfig", "/opt/fontconfig.properties");
 
+    this.verifyApiClient(connectionId);
+
     this.isSessionInitialized = true;
   }
 
@@ -214,7 +220,7 @@ private OutputAdapter createOutputAdapter(Map<String, String> lambdaInput)
     if (connectionId == null) {
       throw new InternalFacingException(INVALID_INPUT, new Exception("Missing connection ID"));
     }
-    final AWSOutputAdapter awsOutputAdapter = new AWSOutputAdapter(connectionId, API_CLIENT);
+    final AWSOutputAdapter awsOutputAdapter = new AWSOutputAdapter(connectionId, this.apiClient);
 
     try {
       final ExecutionType executionType = ExecutionType.valueOf(lambdaInput.get("executionType"));
@@ -270,7 +276,7 @@ private CodeExecutionManager createExecutionManager(
         compileList,
         tempDirectoryManager,
         contentManager,
-        new AWSSystemExitHelper(connectionId, API_CLIENT));
+        new AWSSystemExitHelper(connectionId, this.apiClient));
   }
 
   /**
@@ -376,4 +382,23 @@ private void cleanUpAWSResources(String connectionId, AmazonApiGatewayManagement
       Logger.getLogger(MAIN_LOGGER).removeHandler(allHandlers[i]);
     }
   }
+
+  private void verifyApiClient(String connectionId) {
+    GetConnectionRequest connectionRequest =
+        new GetConnectionRequest().withConnectionId(connectionId);
+    try {
+      this.apiClient.getConnection(connectionRequest);
+    } catch (IllegalStateException e) {
+      // This can occur if the api client has been shut down, which we have seen happen on occasion.
+      // Recreate the api client in this case. Log a warning so we can track when this happens.
+      Logger.getLogger(MAIN_LOGGER)
+          .warning(
+              "Received illegal state exception when trying to talk to API Gateway. Recreating api client.");
+      this.apiClient =
+          AmazonApiGatewayManagementApiClientBuilder.standard()
+              .withEndpointConfiguration(
+                  new AwsClientBuilder.EndpointConfiguration(API_ENDPOINT, "us-east-1"))
+              .build();
+    }
+  }
 }