Add finer grain rollouts to experiments

service/src/main/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicExperimentEnrollmentConfiguration.java

@@ -6,6 +6,7 @@
 package org.whispersystems.textsecuregcm.configuration.dynamic;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
 
 import javax.validation.Valid;
 import javax.validation.constraints.Max;
@@ -16,21 +17,52 @@ import java.util.UUID;
 
 public class DynamicExperimentEnrollmentConfiguration {
 
+  public static class UuidSelector {
+
     @JsonProperty
     @Valid
-    private Set<UUID> enrolledUuids = Collections.emptySet();
+    private Set<UUID> uuids = Collections.emptySet();
 
+    /**
+     * What percentage of enrolled UUIDs should the experiment be enabled for.
+     * <p>
+     * Unlike {@link this#enrollmentPercentage}, this is not stable by UUID. The same UUID may be
+     * enrolled/unenrolled across calls.
+     */
     @JsonProperty
     @Valid
     @Min(0)
     @Max(100)
-    private int enrollmentPercentage = 0;
+    private int uuidEnrollmentPercentage = 100;
 
-    public Set<UUID> getEnrolledUuids() {
+    public Set<UUID> getUuids() {
-        return enrolledUuids;
+      return uuids;
     }
 
-    public int getEnrollmentPercentage() {
+    public int getUuidEnrollmentPercentage() {
-        return enrollmentPercentage;
+      return uuidEnrollmentPercentage;
     }
+
+  }
+
+  private UuidSelector uuidSelector = new UuidSelector();
+
+  /**
+   * If the UUID is not enrolled via {@link UuidSelector#uuids}, what is the percentage chance it should be enrolled.
+   * <p>
+   * This is stable by UUID, for a given configuration if a UUID is enrolled it will always be enrolled on every call.
+   */
+  @JsonProperty
+  @Valid
+  @Min(0)
+  @Max(100)
+  private int enrollmentPercentage = 0;
+
+  public int getEnrollmentPercentage() {
+    return enrollmentPercentage;
+  }
+
+  public UuidSelector getUuidSelector() {
+    return uuidSelector;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/experiment/ExperimentEnrollmentManager.java

@@ -6,7 +6,10 @@
 package org.whispersystems.textsecuregcm.experiment;
 
 import java.util.Optional;
+import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
+import com.google.common.annotations.VisibleForTesting;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicExperimentEnrollmentConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicPreRegistrationExperimentEnrollmentConfiguration;
@@ -16,9 +19,20 @@ import org.whispersystems.textsecuregcm.util.Util;
 public class ExperimentEnrollmentManager {
 
   private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
+  private final Random random;
 
-  public ExperimentEnrollmentManager(final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager) {
+
+  public ExperimentEnrollmentManager(
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager) {
+    this(dynamicConfigurationManager, ThreadLocalRandom.current());
+  }
+
+  @VisibleForTesting
+  ExperimentEnrollmentManager(
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
+      final Random random) {
     this.dynamicConfigurationManager = dynamicConfigurationManager;
+    this.random = random;
   }
 
   public boolean isEnrolled(final UUID accountUuid, final String experimentName) {
@@ -28,8 +42,9 @@ public class ExperimentEnrollmentManager {
 
     return maybeConfiguration.map(config -> {
 
-      if (config.getEnrolledUuids().contains(accountUuid)) {
+      if (config.getUuidSelector().getUuids().contains(accountUuid)) {
-        return true;
+        final int r = random.nextInt(100);
+        return r < config.getUuidSelector().getUuidEnrollmentPercentage();
       }
 
       return isEnrolled(accountUuid, config.getEnrollmentPercentage(), experimentName);

service/src/test/java/org/whispersystems/textsecuregcm/configuration/dynamic/DynamicConfigurationTest.java

@@ -50,16 +50,25 @@ class DynamicConfigurationTest {
             percentageOnly:
               enrollmentPercentage: 12
             uuidsAndPercentage:
-              enrolledUuids:
+              uuidSelector:
-                - 717b1c09-ed0b-4120-bb0e-f4697534b8e1
+                uuids:
-                - 279f264c-56d7-4bbf-b9da-de718ff90903
+                  - 717b1c09-ed0b-4120-bb0e-f4697534b8e1
+                  - 279f264c-56d7-4bbf-b9da-de718ff90903
               enrollmentPercentage: 77
             uuidsOnly:
-              enrolledUuids:
+              uuidSelector:
+                uuids:
                 - 71618739-114c-4b1f-bb0d-6478a44eb600
             uuids-with-dash:
-              enrolledUuids:
+              uuidSelector:
-                - 71618739-114c-4b1f-bb0d-6478ffffffff
+                uuids:
+                  - 71618739-114c-4b1f-bb0d-6478ffffffff
+            uuidsAndSubSelection:
+              uuidSelector:
+                uuids:
+                  - 6664224c-20cc-45a0-829b-95059e8a04f5
+                uuidEnrollmentPercentage: 91
+              enrollmentPercentage: 71
           """);
 
       final DynamicConfiguration config =
@@ -67,27 +76,35 @@ class DynamicConfigurationTest {
 
       assertFalse(config.getExperimentEnrollmentConfiguration("unconfigured").isPresent());
 
-      assertTrue(config.getExperimentEnrollmentConfiguration("percentageOnly").isPresent());
+      final DynamicExperimentEnrollmentConfiguration percentageOnly = config.getExperimentEnrollmentConfiguration("percentageOnly").orElseThrow();
-      assertEquals(12, config.getExperimentEnrollmentConfiguration("percentageOnly").get().getEnrollmentPercentage());
+      assertEquals(12, percentageOnly.getEnrollmentPercentage());
-      assertEquals(Collections.emptySet(),
+      assertEquals(Collections.emptySet(), percentageOnly.getUuidSelector().getUuids());
-          config.getExperimentEnrollmentConfiguration("percentageOnly").get().getEnrolledUuids());
+      assertEquals(100, percentageOnly.getUuidSelector().getUuidEnrollmentPercentage());
 
-      assertTrue(config.getExperimentEnrollmentConfiguration("uuidsAndPercentage").isPresent());
+      final DynamicExperimentEnrollmentConfiguration uuidsAndPercentage = config.getExperimentEnrollmentConfiguration("uuidsAndPercentage").orElseThrow();
-      assertEquals(77,
+      assertEquals(77, uuidsAndPercentage.getEnrollmentPercentage());
-          config.getExperimentEnrollmentConfiguration("uuidsAndPercentage").get().getEnrollmentPercentage());
       assertEquals(Set.of(UUID.fromString("717b1c09-ed0b-4120-bb0e-f4697534b8e1"),
           UUID.fromString("279f264c-56d7-4bbf-b9da-de718ff90903")),
-          config.getExperimentEnrollmentConfiguration("uuidsAndPercentage").get().getEnrolledUuids());
+          uuidsAndPercentage.getUuidSelector().getUuids());
+      assertEquals(100, uuidsAndPercentage.getUuidSelector().getUuidEnrollmentPercentage());
 
-      assertTrue(config.getExperimentEnrollmentConfiguration("uuidsOnly").isPresent());
+      final DynamicExperimentEnrollmentConfiguration uuidsOnly = config.getExperimentEnrollmentConfiguration("uuidsOnly").orElseThrow();
-      assertEquals(0, config.getExperimentEnrollmentConfiguration("uuidsOnly").get().getEnrollmentPercentage());
+      assertEquals(0, uuidsOnly.getEnrollmentPercentage());
       assertEquals(Set.of(UUID.fromString("71618739-114c-4b1f-bb0d-6478a44eb600")),
-          config.getExperimentEnrollmentConfiguration("uuidsOnly").get().getEnrolledUuids());
+          uuidsOnly.getUuidSelector().getUuids());
+      assertEquals(100, uuidsOnly.getUuidSelector().getUuidEnrollmentPercentage());
 
-      assertTrue(config.getExperimentEnrollmentConfiguration("uuids-with-dash").isPresent());
+      final DynamicExperimentEnrollmentConfiguration uuidsWithDash = config.getExperimentEnrollmentConfiguration("uuids-with-dash").orElseThrow();
-      assertEquals(0, config.getExperimentEnrollmentConfiguration("uuids-with-dash").get().getEnrollmentPercentage());
+      assertEquals(0, uuidsWithDash.getEnrollmentPercentage());
       assertEquals(Set.of(UUID.fromString("71618739-114c-4b1f-bb0d-6478ffffffff")),
-          config.getExperimentEnrollmentConfiguration("uuids-with-dash").get().getEnrolledUuids());
+          uuidsWithDash.getUuidSelector().getUuids());
+      assertEquals(100, uuidsWithDash.getUuidSelector().getUuidEnrollmentPercentage());
+
+      final DynamicExperimentEnrollmentConfiguration uuidsAndSubSelection = config.getExperimentEnrollmentConfiguration("uuidsAndSubSelection").orElseThrow();
+      assertEquals(71, uuidsAndSubSelection.getEnrollmentPercentage());
+      assertEquals(Set.of(UUID.fromString("6664224c-20cc-45a0-829b-95059e8a04f5")),
+          uuidsAndSubSelection.getUuidSelector().getUuids());
+      assertEquals(91, uuidsAndSubSelection.getUuidSelector().getUuidEnrollmentPercentage());
     }
   }
 

service/src/test/java/org/whispersystems/textsecuregcm/experiment/ExperimentEnrollmentManagerTest.java

@@ -9,12 +9,18 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 import java.util.Collections;
+import java.util.Map;
 import java.util.Optional;
+import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -29,12 +35,14 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 
 class ExperimentEnrollmentManagerTest {
 
+  private DynamicExperimentEnrollmentConfiguration.UuidSelector uuidSelector;
   private DynamicExperimentEnrollmentConfiguration experimentEnrollmentConfiguration;
   private DynamicPreRegistrationExperimentEnrollmentConfiguration preRegistrationExperimentEnrollmentConfiguration;
 
   private ExperimentEnrollmentManager experimentEnrollmentManager;
 
   private Account account;
+  private Random random;
 
   private static final UUID ACCOUNT_UUID = UUID.randomUUID();
   private static final String UUID_EXPERIMENT_NAME = "uuid_test";
@@ -47,10 +55,14 @@ class ExperimentEnrollmentManagerTest {
   void setUp() {
     final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
     final DynamicConfiguration dynamicConfiguration = mock(DynamicConfiguration.class);
+    random = spy(new Random());
+    experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager, random);
 
-    experimentEnrollmentManager = new ExperimentEnrollmentManager(dynamicConfigurationManager);
+    uuidSelector = mock(DynamicExperimentEnrollmentConfiguration.UuidSelector.class);
+    when(uuidSelector.getUuidEnrollmentPercentage()).thenReturn(100);
 
     experimentEnrollmentConfiguration = mock(DynamicExperimentEnrollmentConfiguration.class);
+    when(experimentEnrollmentConfiguration.getUuidSelector()).thenReturn(uuidSelector);
     preRegistrationExperimentEnrollmentConfiguration = mock(
         DynamicPreRegistrationExperimentEnrollmentConfiguration.class);
 
@@ -70,10 +82,10 @@ class ExperimentEnrollmentManagerTest {
     assertFalse(
         experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME + "-unrelated-experiment"));
 
-    when(experimentEnrollmentConfiguration.getEnrolledUuids()).thenReturn(Set.of(ACCOUNT_UUID));
+    when(uuidSelector.getUuids()).thenReturn(Set.of(ACCOUNT_UUID));
     assertTrue(experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME));
 
-    when(experimentEnrollmentConfiguration.getEnrolledUuids()).thenReturn(Collections.emptySet());
+    when(uuidSelector.getUuids()).thenReturn(Collections.emptySet());
     when(experimentEnrollmentConfiguration.getEnrollmentPercentage()).thenReturn(0);
 
     assertFalse(experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME));
@@ -82,6 +94,24 @@ class ExperimentEnrollmentManagerTest {
     assertTrue(experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME));
   }
 
+  @Test
+  void testIsEnrolled_UuidExperimentPercentage() {
+    when(uuidSelector.getUuids()).thenReturn(Set.of(ACCOUNT_UUID));
+    when(uuidSelector.getUuidEnrollmentPercentage()).thenReturn(0);
+    assertFalse(experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME));
+    when(uuidSelector.getUuidEnrollmentPercentage()).thenReturn(100);
+    assertTrue(experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME));
+
+    when(uuidSelector.getUuidEnrollmentPercentage()).thenReturn(75);
+    final Map<Boolean, Long> counts = IntStream.range(0, 100).mapToObj(i -> {
+          when(random.nextInt(100)).thenReturn(i);
+          return experimentEnrollmentManager.isEnrolled(account.getUuid(), UUID_EXPERIMENT_NAME);
+        })
+        .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
+    assertEquals(25, counts.get(false));
+    assertEquals(75, counts.get(true));
+  }
+
   @ParameterizedTest
   @MethodSource
   void testIsEnrolled_PreRegistrationExperiment(final String e164, final String experimentName,

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/ExperimentHelper.java

@@ -27,8 +27,13 @@ public class ExperimentHelper {
     when(dcm.getConfiguration()).thenReturn(dc);
     final DynamicExperimentEnrollmentConfiguration exp = mock(DynamicExperimentEnrollmentConfiguration.class);
     when(dc.getExperimentEnrollmentConfiguration(experimentName)).thenReturn(Optional.of(exp));
-    when(exp.getEnrolledUuids()).thenReturn(enrolledUuids);
+    final DynamicExperimentEnrollmentConfiguration.UuidSelector uuidSelector =
+        mock(DynamicExperimentEnrollmentConfiguration.UuidSelector.class);
+    when(exp.getUuidSelector()).thenReturn(uuidSelector);
+
     when(exp.getEnrollmentPercentage()).thenReturn(enrollmentPercentage);
+    when(uuidSelector.getUuids()).thenReturn(enrolledUuids);
+    when(uuidSelector.getUuidEnrollmentPercentage()).thenReturn(100);
     return dcm;
   }