Merge pull request #5587 from nickgros/SWC-7044

Sage-Bionetworks · Nov 22, 2024 · c6566c8 · c6566c8
2 parents 1b68750 + 23c243d
commit c6566c8
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 20 deletions.
diff --git a/src/main/java/org/sagebionetworks/web/server/servlet/AppConfigServlet.java b/src/main/java/org/sagebionetworks/web/server/servlet/AppConfigServlet.java
@@ -1,8 +1,10 @@
 package org.sagebionetworks.web.server.servlet;
 
 import com.amazonaws.services.appconfigdata.AWSAppConfigData;
+import com.amazonaws.services.appconfigdata.model.BadRequestException;
 import com.amazonaws.services.appconfigdata.model.GetLatestConfigurationRequest;
 import com.amazonaws.services.appconfigdata.model.GetLatestConfigurationResult;
+import com.amazonaws.services.appconfigdata.model.InternalServerException;
 import com.amazonaws.services.appconfigdata.model.StartConfigurationSessionRequest;
 import com.amazonaws.services.appconfigdata.model.StartConfigurationSessionResult;
 import com.google.gwt.thirdparty.guava.common.base.Supplier;
@@ -95,11 +97,31 @@ public void initializeConfigSupplier() {
       );
   }
 
+  public JSONObjectAdapter getLastConfigValueOrDefault() {
+    if (lastConfigValue != null) {
+      return lastConfigValue;
+    }
+    try {
+      return new JSONObjectAdapterImpl(DEFAULT_CONFIG_VALUE);
+    } catch (JSONObjectAdapterException e) {
+      logger.log(
+        Level.SEVERE,
+        "JSONObjectAdapterException occurred in default configuration",
+        e
+      );
+      throw new RuntimeException(e);
+    }
+  }
+
   public JSONObjectAdapter getLatestConfiguration() {
     try {
       if (configurationToken == null) {
-        logger.log(Level.SEVERE, "returning default config");
-        return new JSONObjectAdapterImpl(DEFAULT_CONFIG_VALUE);
+        logger.log(
+          Level.SEVERE,
+          "The configuration token is null, the last config value will be returned." +
+          " This usually means that the initial call to initialize the configuration session failed."
+        );
+        return getLastConfigValueOrDefault();
       }
       GetLatestConfigurationRequest latestConfigRequest =
         new GetLatestConfigurationRequest()
@@ -116,21 +138,22 @@ public JSONObjectAdapter getLatestConfiguration() {
       if (!newConfigString.isEmpty()) {
         lastConfigValue = new JSONObjectAdapterImpl(newConfigString);
       }
+    } catch (BadRequestException | InternalServerException e) {
+      // Invalid token or server error, re-initialize the session to try to recover.
+      logger.log(
+        Level.SEVERE,
+        "Failed to get latest configuration, returning last or default configuration and attempting to re-initialize the session.",
+        e
+      );
+      initializeAppConfigClient();
+      return getLastConfigValueOrDefault();
     } catch (Exception e) {
-      try {
-        logger.log(
-          Level.SEVERE,
-          "Failed to get or parse latest configuration, returning default configuration",
-          e
-        );
-        return new JSONObjectAdapterImpl(DEFAULT_CONFIG_VALUE);
-      } catch (JSONObjectAdapterException exception) {
-        logger.log(
-          Level.SEVERE,
-          "JSONObjectAdapterException occurred in default configuration",
-          exception
-        );
-      }
+      logger.log(
+        Level.SEVERE,
+        "Failed to get or parse latest configuration, returning last or default configuration.",
+        e
+      );
+      return getLastConfigValueOrDefault();
     }
     return lastConfigValue;
   }

diff --git a/src/test/java/org/sagebionetworks/web/unitserver/servlet/AppConfigServletTest.java b/src/test/java/org/sagebionetworks/web/unitserver/servlet/AppConfigServletTest.java
@@ -2,9 +2,14 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.amazonaws.services.appconfigdata.AWSAppConfigData;
+import com.amazonaws.services.appconfigdata.model.BadRequestException;
 import com.amazonaws.services.appconfigdata.model.GetLatestConfigurationRequest;
 import com.amazonaws.services.appconfigdata.model.GetLatestConfigurationResult;
 import com.amazonaws.services.appconfigdata.model.StartConfigurationSessionRequest;
@@ -20,7 +25,6 @@
 import org.mockito.MockitoAnnotations;
 import org.sagebionetworks.StackConfiguration;
 import org.sagebionetworks.schema.adapter.JSONObjectAdapter;
-import org.sagebionetworks.schema.adapter.JSONObjectAdapterException;
 import org.sagebionetworks.schema.adapter.org.json.JSONObjectAdapterImpl;
 import org.sagebionetworks.web.server.servlet.AppConfigServlet;
 
@@ -134,7 +138,7 @@ public void testGetLatestConfiguration_Success() {
   }
 
   @Test
-  public void testGetLatestConfiguration_Failure() {
+  public void testGetLatestConfiguration_Failure_ReturnDefaultValueWhenNoLastValue() {
     String DEFAULT_CONFIG_VALUE = "{}";
     when(
       mockAppConfigDataClient.getLatestConfiguration(
@@ -143,12 +147,87 @@ public void testGetLatestConfiguration_Failure() {
     )
       .thenThrow(new RuntimeException("Failed to retrieve configuration"));
 
-    servlet.configurationToken = "mockToken"; // Setting the initial configuration token
     JSONObjectAdapter configValue = servlet.getLatestConfiguration();
 
     assertEquals(DEFAULT_CONFIG_VALUE, configValue.toString());
   }
 
+  @Test
+  public void testGetLatestConfiguration_Failure_ReturnLastValue() {
+    // Set the initial configuration token
+    servlet.configurationToken = "mockToken";
+
+    // Set up a successful response
+    ByteBuffer mockByteBuffer = ByteBuffer
+      .wrap("{\"test configuration\":true}".getBytes())
+      .asReadOnlyBuffer();
+    GetLatestConfigurationResult mockConfigResult =
+      new GetLatestConfigurationResult()
+        .withConfiguration(mockByteBuffer)
+        .withNextPollConfigurationToken("new-mock-token");
+
+    when(
+      mockAppConfigDataClient.getLatestConfiguration(
+        any(GetLatestConfigurationRequest.class)
+      )
+    )
+      .thenReturn(mockConfigResult);
+
+    // Initial call succeeds
+    servlet.getLatestConfiguration();
+
+    // Next call will fail, but should return the old value
+    when(
+      mockAppConfigDataClient.getLatestConfiguration(
+        any(GetLatestConfigurationRequest.class)
+      )
+    )
+      .thenThrow(new RuntimeException("Failed to retrieve configuration"));
+
+    JSONObjectAdapter configValue = servlet.getLatestConfiguration();
+
+    assertEquals(mockConfiguration.toString(), configValue.toString());
+
+    // Do not attempt to re-initialize the client
+    verify(mockAppConfigDataClient, never()).startConfigurationSession(any());
+
+    // Try again, but with a recoverable exception
+    when(
+      mockAppConfigDataClient.getLatestConfiguration(
+        any(GetLatestConfigurationRequest.class)
+      )
+    )
+      .thenThrow(new BadRequestException("Token is invalid"));
+
+    configValue = servlet.getLatestConfiguration();
+
+    assertEquals(mockConfiguration.toString(), configValue.toString());
+
+    // We should attempt to re-initialize the client
+    verify(mockAppConfigDataClient).startConfigurationSession(any());
+  }
+
+  @Test
+  public void testGetLatestConfiguration_Failure_NullToken() {
+    String DEFAULT_CONFIG_VALUE = "{}";
+    servlet.configurationToken = null;
+
+    when(
+      mockAppConfigDataClient.getLatestConfiguration(
+        any(GetLatestConfigurationRequest.class)
+      )
+    )
+      .thenThrow(new RuntimeException("Failed to retrieve configuration"));
+
+    JSONObjectAdapter configValue = servlet.getLatestConfiguration();
+
+    assertEquals(DEFAULT_CONFIG_VALUE, configValue.toString());
+    // Verify we never called getLatestConfiguration
+    verify(mockAppConfigDataClient, never()).getLatestConfiguration(any());
+    // Verify we do not try to reinitialize the client
+    verify(mockAppConfigDataClient, never()).startConfigurationSession(any());
+  }
+
   @Test
   public void testInitializeAppConfigClient() {
     servlet.appConfigDataClient = null; // Simulate the client not being injected