Use GetLatestConfiguration in config manager

Use StartConfigurationSession/GetLatestConfiguration instead of
GetConfiguration since the latter has been deprecated

service/pom.xml

@@ -277,6 +277,10 @@
       <groupId>software.amazon.awssdk</groupId>
       <artifactId>appconfig</artifactId>
     </dependency>
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>appconfigdata</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.amazonaws</groupId>
       <artifactId>aws-java-sdk-core</artifactId>

service/src/main/java/org/whispersystems/textsecuregcm/storage/DynamicConfigurationManager.java

@@ -6,24 +6,26 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
 import com.google.common.annotations.VisibleForTesting;
-import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.validation.ConstraintViolation;
 import javax.validation.Validation;
 import javax.validation.Validator;
 import io.micrometer.core.instrument.Metrics;
-import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.util.Util;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
-import software.amazon.awssdk.services.appconfig.AppConfigClient;
+import software.amazon.awssdk.services.appconfigdata.AppConfigDataClient;
-import software.amazon.awssdk.services.appconfig.model.GetConfigurationRequest;
+import software.amazon.awssdk.services.appconfigdata.model.GetLatestConfigurationRequest;
-import software.amazon.awssdk.services.appconfig.model.GetConfigurationResponse;
+import software.amazon.awssdk.services.appconfigdata.model.GetLatestConfigurationResponse;
+import software.amazon.awssdk.services.appconfigdata.model.StartConfigurationSessionRequest;
+import software.amazon.awssdk.services.appconfigdata.model.StartConfigurationSessionResponse;
 
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
@@ -32,17 +34,15 @@ public class DynamicConfigurationManager<T> {
   private final String application;
   private final String environment;
   private final String configurationName;
-  private final String clientId;
+  private final AppConfigDataClient appConfigClient;
-  private final AppConfigClient appConfigClient;
-
   private final Class<T> configurationClass;
 
+  // Set on initial config fetch
   private final AtomicReference<T> configuration = new AtomicReference<>();
-
+  private String configurationToken = null;
-  private GetConfigurationResponse lastConfigResult;
-
   private boolean initialized = false;
 
+
   private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(new YAMLFactory())
       .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
       .registerModule(new JavaTimeModule());
@@ -57,22 +57,22 @@ public class DynamicConfigurationManager<T> {
 
   public DynamicConfigurationManager(String application, String environment, String configurationName,
       Class<T> configurationClass) {
-    this(AppConfigClient.builder()
+    this(AppConfigDataClient
+            .builder()
             .overrideConfiguration(ClientOverrideConfiguration.builder()
-                .apiCallTimeout(Duration.ofMillis(10000))
+                .apiCallTimeout(Duration.ofSeconds(10))
-                .apiCallAttemptTimeout(Duration.ofMillis(10000)).build())
+                .apiCallAttemptTimeout(Duration.ofSeconds(10)).build())
             .build(),
-        application, environment, configurationName, UUID.randomUUID().toString(), configurationClass);
+        application, environment, configurationName, configurationClass);
   }
 
   @VisibleForTesting
-  DynamicConfigurationManager(AppConfigClient appConfigClient, String application, String environment,
+  DynamicConfigurationManager(AppConfigDataClient appConfigClient, String application, String environment,
-      String configurationName, String clientId, Class<T> configurationClass) {
+      String configurationName, Class<T> configurationClass) {
     this.appConfigClient = appConfigClient;
     this.application = application;
     this.environment = environment;
     this.configurationName = configurationName;
-    this.clientId = clientId;
     this.configurationClass = configurationClass;
   }
 
@@ -82,13 +82,11 @@ public class DynamicConfigurationManager<T> {
         Util.wait(this);
       }
     }
-
     return configuration.get();
   }
 
   public void start() {
     configuration.set(retrieveInitialDynamicConfiguration());
-
     synchronized (this) {
       this.initialized = true;
       this.notifyAll();
@@ -98,8 +96,8 @@ public class DynamicConfigurationManager<T> {
       while (true) {
         try {
           retrieveDynamicConfiguration().ifPresent(configuration::set);
-        } catch (Throwable t) {
+        } catch (Exception e) {
-          logger.warn("Error retrieving dynamic configuration", t);
+          logger.warn("Error retrieving dynamic configuration", e);
         }
 
         Util.sleep(5000);
@@ -111,44 +109,38 @@ public class DynamicConfigurationManager<T> {
   }
 
   private Optional<T> retrieveDynamicConfiguration() throws JsonProcessingException {
-    final String previousVersion = lastConfigResult != null ? lastConfigResult.configurationVersion() : null;
+    if (configurationToken == null) {
-
+        logger.error("Invalid configuration token, will not be able to fetch configuration updates");
+    }
+    GetLatestConfigurationResponse latestConfiguration;
     try {
-      lastConfigResult = appConfigClient.getConfiguration(GetConfigurationRequest.builder()
+      latestConfiguration = appConfigClient.getLatestConfiguration(GetLatestConfigurationRequest.builder()
-          .application(application)
+          .configurationToken(configurationToken)
-          .environment(environment)
-          .configuration(configurationName)
-          .clientId(clientId)
-          .clientConfigurationVersion(previousVersion)
           .build());
+      // token to use in the next fetch
+      configurationToken = latestConfiguration.nextPollConfigurationToken();
+      logger.debug("next token: {}", configurationToken);
     } catch (final RuntimeException e) {
       Metrics.counter(ERROR_COUNTER_NAME, ERROR_TYPE_TAG_NAME, "fetch").increment();
       throw e;
     }
 
-    final Optional<T> maybeDynamicConfiguration;
+    if (!latestConfiguration.configuration().asByteBuffer().hasRemaining()) {
-
+      // empty configuration means nothing has changed
-    if (!StringUtils.equals(lastConfigResult.configurationVersion(), previousVersion)) {
+      return Optional.empty();
-      logger.info("Received new config version: {}", lastConfigResult.configurationVersion());
-
-      try {
-        maybeDynamicConfiguration =
-            parseConfiguration(
-                StandardCharsets.UTF_8.decode(lastConfigResult.content().asByteBuffer().asReadOnlyBuffer()).toString(),
-                configurationClass);
-      } catch (final JsonProcessingException e) {
-        Metrics.counter(ERROR_COUNTER_NAME,
-            ERROR_TYPE_TAG_NAME, "parse",
-            CONFIG_CLASS_TAG_NAME, configurationClass.getName()).increment();
-
-        throw e;
-      }
-    } else {
-      // No change since last version
-      maybeDynamicConfiguration = Optional.empty();
     }
+    logger.info("Received new config of length {}, next configuration token: {}",
+        latestConfiguration.configuration().asByteBuffer().remaining(),
+        configurationToken);
 
-    return maybeDynamicConfiguration;
+    try {
+      return parseConfiguration(latestConfiguration.configuration().asUtf8String(), configurationClass);
+    } catch (final JsonProcessingException e) {
+      Metrics.counter(ERROR_COUNTER_NAME,
+          ERROR_TYPE_TAG_NAME, "parse",
+          CONFIG_CLASS_TAG_NAME, configurationClass.getName()).increment();
+      throw e;
+    }
   }
 
   @VisibleForTesting
@@ -172,15 +164,18 @@ public class DynamicConfigurationManager<T> {
   private T retrieveInitialDynamicConfiguration() {
     for (;;) {
       try {
-        final Optional<T> maybeDynamicConfiguration = retrieveDynamicConfiguration();
+        if (configurationToken == null) {
-
+          // first time around, start the configuration session
-        if (maybeDynamicConfiguration.isPresent()) {
+          final StartConfigurationSessionResponse startResponse = appConfigClient
-          return maybeDynamicConfiguration.get();
+              .startConfigurationSession(StartConfigurationSessionRequest.builder()
-        } else {
+                  .applicationIdentifier(application)
-          throw new IllegalStateException("No initial configuration available");
+                  .environmentIdentifier(environment)
+                  .configurationProfileIdentifier(configurationName).build());
+          configurationToken = startResponse.initialConfigurationToken();
         }
-      } catch (Throwable t) {
+        return retrieveDynamicConfiguration().orElseThrow(() -> new IllegalStateException("No initial configuration available"));
-        logger.warn("Error retrieving initial dynamic configuration", t);
+      } catch (Exception e) {
+        logger.warn("Error retrieving initial dynamic configuration", e);
         Util.sleep(1000);
       }
     }

service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamicConfigurationManagerTest.java

@@ -6,37 +6,127 @@ import static org.mockito.Mockito.when;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.mockito.ArgumentCaptor;
+import org.junit.jupiter.api.Timeout;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import software.amazon.awssdk.core.SdkBytes;
-import software.amazon.awssdk.services.appconfig.AppConfigClient;
+import software.amazon.awssdk.services.appconfigdata.AppConfigDataClient;
-import software.amazon.awssdk.services.appconfig.model.GetConfigurationRequest;
+import software.amazon.awssdk.services.appconfigdata.model.GetLatestConfigurationRequest;
-import software.amazon.awssdk.services.appconfig.model.GetConfigurationResponse;
+import software.amazon.awssdk.services.appconfigdata.model.GetLatestConfigurationResponse;
+import software.amazon.awssdk.services.appconfigdata.model.StartConfigurationSessionRequest;
+import software.amazon.awssdk.services.appconfigdata.model.StartConfigurationSessionResponse;
+import java.util.concurrent.TimeUnit;
 
 class DynamicConfigurationManagerTest {
 
   private DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
-  private AppConfigClient             appConfig;
+  private AppConfigDataClient appConfig;
+  private StartConfigurationSessionRequest startConfigurationSession;
 
   @BeforeEach
   void setup() {
-    this.appConfig                   = mock(AppConfigClient.class);
+    this.appConfig = mock(AppConfigDataClient.class);
-    this.dynamicConfigurationManager = new DynamicConfigurationManager<>(appConfig, "foo", "bar", "baz", "poof", DynamicConfiguration.class);
+    this.dynamicConfigurationManager = new DynamicConfigurationManager<>(
+        appConfig, "foo", "bar", "baz", DynamicConfiguration.class);
+    this.startConfigurationSession = StartConfigurationSessionRequest.builder()
+        .applicationIdentifier("foo")
+        .environmentIdentifier("bar")
+        .configurationProfileIdentifier("baz")
+        .build();
   }
 
   @Test
-  void testGetConfig() {
+  void testGetInitalConfig() {
-    ArgumentCaptor<GetConfigurationRequest> captor = ArgumentCaptor.forClass(GetConfigurationRequest.class);
+    when(appConfig.startConfigurationSession(startConfigurationSession))
-    when(appConfig.getConfiguration(captor.capture())).thenReturn(
+        .thenReturn(StartConfigurationSessionResponse.builder()
-        GetConfigurationResponse.builder().content(SdkBytes.fromByteArray("test: true".getBytes())).configurationVersion("1").build());
+            .initialConfigurationToken("initial")
+            .build());
+
+    // call with initial token will return a real config
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder()
+        .configurationToken("initial").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String("test: true"))
+            .nextPollConfigurationToken("next").build());
+
+    // subsequent config calls will return empty (no update)
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("next").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String(""))
+            .nextPollConfigurationToken("next").build());
 
     dynamicConfigurationManager.start();
-
-    assertThat(captor.getValue().application()).isEqualTo("foo");
-    assertThat(captor.getValue().environment()).isEqualTo("bar");
-    assertThat(captor.getValue().configuration()).isEqualTo("baz");
-    assertThat(captor.getValue().clientId()).isEqualTo("poof");
-
     assertThat(dynamicConfigurationManager.getConfiguration()).isNotNull();
   }
+
+  @Test
+  void testBadConfig() {
+    when(appConfig.startConfigurationSession(startConfigurationSession))
+        .thenReturn(StartConfigurationSessionResponse.builder()
+            .initialConfigurationToken("initial")
+            .build());
+
+    // call with initial token will return a bad config
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder()
+        .configurationToken("initial").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String("zzz"))
+            .nextPollConfigurationToken("goodconfig").build());
+
+    // next config call will return a good config
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("goodconfig").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String("test: true"))
+            .nextPollConfigurationToken("next").build());
+
+    // all subsequent config calls will return an empty config (no update)
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("next").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String(""))
+            .nextPollConfigurationToken("next").build());
+    dynamicConfigurationManager.start();
+    assertThat(dynamicConfigurationManager.getConfiguration()).isNotNull();
+  }
+
+  @Test
+  @Timeout(value=5, unit= TimeUnit.SECONDS)
+  void testGetConfigMultiple() throws InterruptedException {
+    when(appConfig.startConfigurationSession(startConfigurationSession))
+        .thenReturn(StartConfigurationSessionResponse.builder()
+            .initialConfigurationToken("0")
+            .build());
+
+    // initial config
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("0").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String("test: true"))
+            .nextPollConfigurationToken("1").build());
+
+    // config update with a real config
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("1").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String("""
+                featureFlags:
+                  - testFlag
+                """))
+            .nextPollConfigurationToken("2").build());
+
+    // all subsequent are no update
+    when(appConfig.getLatestConfiguration(GetLatestConfigurationRequest.builder().
+        configurationToken("2").build()))
+        .thenReturn(GetLatestConfigurationResponse.builder()
+            .configuration(SdkBytes.fromUtf8String(""))
+            .nextPollConfigurationToken("2").build());
+
+    // we should eventually get the updated config (or the test will timeout)
+    dynamicConfigurationManager.start();
+    while (dynamicConfigurationManager.getConfiguration().getActiveFeatureFlags().isEmpty()) {
+      Thread.sleep(100);
+    }
+    assertThat(dynamicConfigurationManager.getConfiguration().getActiveFeatureFlags()).containsExactly("testFlag");
+  }
 }