Split traffic between the Postgres and Dynamo pre-key stores.

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -420,7 +420,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     AttachmentControllerV1 attachmentControllerV1    = new AttachmentControllerV1(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getBucket());
     AttachmentControllerV2 attachmentControllerV2    = new AttachmentControllerV2(rateLimiters, config.getAwsAttachmentsConfiguration().getAccessKey(), config.getAwsAttachmentsConfiguration().getAccessSecret(), config.getAwsAttachmentsConfiguration().getRegion(), config.getAwsAttachmentsConfiguration().getBucket());
     AttachmentControllerV3 attachmentControllerV3    = new AttachmentControllerV3(rateLimiters, config.getGcpAttachmentsConfiguration().getDomain(), config.getGcpAttachmentsConfiguration().getEmail(), config.getGcpAttachmentsConfiguration().getMaxSizeInBytes(), config.getGcpAttachmentsConfiguration().getPathPrefix(), config.getGcpAttachmentsConfiguration().getRsaSigningKey());
-    KeysController         keysController            = new KeysController(rateLimiters, keys, accountsManager, directoryQueue);
+    KeysController         keysController            = new KeysController(rateLimiters, keys, keysDynamoDb, accountsManager, directoryQueue, experimentEnrollmentManager);
     MessageController      messageController         = new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, messagesManager, apnFallbackManager, featureFlagsManager);
     ProfileController      profileController         = new ProfileController(rateLimiters, accountsManager, profilesManager, usernamesManager, cdnS3Client, profileCdnPolicyGenerator, profileCdnPolicySigner, config.getCdnConfiguration().getBucket(), zkProfileOperations, isZkEnabled);
     StickerController      stickerController         = new StickerController(rateLimiters, config.getCdnConfiguration().getAccessKey(), config.getCdnConfiguration().getAccessSecret(), config.getCdnConfiguration().getRegion(), config.getCdnConfiguration().getBucket());

service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

@@ -6,8 +6,6 @@ package org.whispersystems.textsecuregcm.controllers;
 
 import com.codahale.metrics.annotation.Timed;
 import io.dropwizard.auth.Auth;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AmbiguousIdentifier;
 import org.whispersystems.textsecuregcm.auth.Anonymous;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
@@ -18,6 +16,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
 import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
 import org.whispersystems.textsecuregcm.entities.PreKeyState;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
 import org.whispersystems.textsecuregcm.storage.Account;
@@ -25,6 +24,8 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.KeyRecord;
 import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
+import org.whispersystems.textsecuregcm.storage.PreKeyStore;
 
 import javax.validation.Valid;
 import javax.ws.rs.Consumes;
@@ -45,24 +46,29 @@ import java.util.Optional;
 @Path("/v2/keys")
 public class KeysController {
 
-  private static final Logger logger = LoggerFactory.getLogger(KeysController.class);
+  private final RateLimiters                rateLimiters;
+  private final Keys                        keys;
+  private final KeysDynamoDb                keysDynamoDb;
+  private final AccountsManager             accounts;
+  private final DirectoryQueue              directoryQueue;
+  private final ExperimentEnrollmentManager experimentEnrollmentManager;
 
-  private final RateLimiters    rateLimiters;
+  private static final String DYNAMODB_CONSUMER_EXPERIMENT = "keys_dynamodb_consumer";
-  private final Keys            keys;
+  private static final String DYNAMODB_PRODUCER_EXPERIMENT = "keys_dynamodb_producer";
-  private final AccountsManager accounts;
-  private final DirectoryQueue  directoryQueue;
 
-  public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts, DirectoryQueue directoryQueue) {
+  public KeysController(RateLimiters rateLimiters, Keys keys, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue, final ExperimentEnrollmentManager experimentEnrollmentManager) {
-    this.rateLimiters   = rateLimiters;
+    this.rateLimiters                = rateLimiters;
-    this.keys           = keys;
+    this.keys                        = keys;
-    this.accounts       = accounts;
+    this.keysDynamoDb                = keysDynamoDb;
-    this.directoryQueue = directoryQueue;
+    this.accounts                    = accounts;
+    this.directoryQueue              = directoryQueue;
+    this.experimentEnrollmentManager = experimentEnrollmentManager;
   }
 
   @GET
   @Produces(MediaType.APPLICATION_JSON)
   public PreKeyCount getStatus(@Auth Account account) {
-    int count = keys.getCount(account, account.getAuthenticatedDevice().get().getId());
+    int count = getPreKeyStoreForProducer(account).getCount(account, account.getAuthenticatedDevice().get().getId());
 
     if (count > 0) {
       count = count - 1;
@@ -98,7 +104,7 @@ public class KeysController {
       }
     }
 
-    keys.store(account, device.getId(), preKeys.getPreKeys());
+    getPreKeyStoreForProducer(account).store(account, device.getId(), preKeys.getPreKeys());
   }
 
   @Timed
@@ -179,14 +185,22 @@ public class KeysController {
   private List<KeyRecord> getLocalKeys(Account destination, String deviceIdSelector) {
     try {
       if (deviceIdSelector.equals("*")) {
-        return keys.take(destination);
+        return getPreKeyStoreForConsumer(destination).take(destination);
       }
 
       long deviceId = Long.parseLong(deviceIdSelector);
 
-      return keys.take(destination, deviceId);
+      return getPreKeyStoreForConsumer(destination).take(destination, deviceId);
     } catch (NumberFormatException e) {
       throw new WebApplicationException(Response.status(422).build());
     }
   }
+
+  private PreKeyStore getPreKeyStoreForProducer(final Account account) {
+    return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_PRODUCER_EXPERIMENT) ? keysDynamoDb : keys;
+  }
+
+  private PreKeyStore getPreKeyStoreForConsumer(final Account account) {
+    return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_CONSUMER_EXPERIMENT) ? keysDynamoDb : keys;
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java

@@ -21,6 +21,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyCount;
 import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
 import org.whispersystems.textsecuregcm.entities.PreKeyState;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
@@ -29,6 +30,7 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.KeyRecord;
 import org.whispersystems.textsecuregcm.storage.Keys;
+import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 
 import javax.ws.rs.client.Entity;
@@ -69,10 +71,12 @@ public class KeysControllerTest {
   private final SignedPreKey SAMPLE_SIGNED_KEY3      = new SignedPreKey( 3333, "barfoo", "sig33"    );
   private final SignedPreKey VALID_DEVICE_SIGNED_KEY = new SignedPreKey(89898, "zoofarb", "sigvalid");
 
-  private final Keys            keys           = mock(Keys.class           );
+  private final Keys                        keys                        = mock(Keys.class                       );
-  private final AccountsManager accounts       = mock(AccountsManager.class);
+  private final KeysDynamoDb                keysDynamoDb                = mock(KeysDynamoDb.class               );
-  private final DirectoryQueue  directoryQueue = mock(DirectoryQueue.class);
+  private final AccountsManager             accounts                    = mock(AccountsManager.class            );
-  private final Account         existsAccount  = mock(Account.class        );
+  private final DirectoryQueue              directoryQueue              = mock(DirectoryQueue.class             );
+  private final Account                     existsAccount               = mock(Account.class                    );
+  private final ExperimentEnrollmentManager experimentEnrollmentManager = mock(ExperimentEnrollmentManager.class);
 
   private RateLimiters          rateLimiters  = mock(RateLimiters.class);
   private RateLimiter           rateLimiter   = mock(RateLimiter.class );
@@ -82,7 +86,7 @@ public class KeysControllerTest {
                                                             .addProvider(AuthHelper.getAuthFilter())
                                                             .addProvider(new PolymorphicAuthValueFactoryProvider.Binder<>(ImmutableSet.of(Account.class, DisabledPermittedAccount.class)))
                                                             .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
-                                                            .addResource(new KeysController(rateLimiters, keys, accounts, directoryQueue))
+                                                            .addResource(new KeysController(rateLimiters, keys, keysDynamoDb, accounts, directoryQueue, experimentEnrollmentManager))
                                                             .build();
 
   @Before