Bug 1829014 - Re-evaluate bucketing when updating rollouts r=emcminn

Differential Revision: https://phabricator.services.mozilla.com/D176177

UltraBlame original commit: 57a5cee8e8143d46a4459b6a65642068177e5cdb

Author: marco-c

toolkit/components/nimbus/lib/ExperimentManager.sys.mjs

@@ -160,7 +160,7 @@ export class _ExperimentManager {
     this.sessions.get(source).add(slug);
 
     if (this.store.has(slug)) {
-      this.updateEnrollment(recipe);
+      await this.updateEnrollment(recipe);
     } else if (isEnrollmentPaused) {
       lazy.log.debug(`Enrollment is paused for "${slug}"`);
     } else if (!(await this.isInBucketAllocation(recipe.bucketConfig))) {
@@ -497,7 +497,7 @@ export class _ExperimentManager {
    * @param {RecipeArgs} recipe
    * @returns {boolean} whether the enrollment is still active
    */
-  updateEnrollment(recipe) {
+  async updateEnrollment(recipe) {
     /** @type Enrollment */
     const enrollment = this.store.get(recipe.slug);
 
@@ -507,6 +507,14 @@ export class _ExperimentManager {
       return false;
     }
 
+    if (
+      recipe.isRollout &&
+      !(await this.isInBucketAllocation(recipe.bucketConfig))
+    ) {
+      this.unenroll(recipe.slug, "bucketing");
+      return false;
+    }
+
     // Stay in the same branch, don't re-sample every time.
     const branch = recipe.branches.find(
       branch => branch.slug === enrollment.branch.slug

toolkit/components/nimbus/test/unit/test_ExperimentManager_lifecycle.js

@@ -215,7 +215,7 @@ add_task(async function test_onRecipe_rollout_update() {
     "should call .updateEnrollment() if the recipe has already been enrolled"
   );
   Assert.ok(
-    manager.updateEnrollment.alwaysReturned(true),
+    manager.updateEnrollment.alwaysReturned(Promise.resolve(true)),
     "updateEnrollment will confirm the enrolled branch still exists in the recipe and exit"
   );
   Assert.ok(

toolkit/components/nimbus/test/unit/test_RemoteSettingsExperimentLoader_updateRecipes.js

@@ -973,3 +973,128 @@ add_task(async function test_updateRecipes_invalidFeature_mismatch() {
 
   targetingSpy.restore();
 });
+
+add_task(async function test_updateRecipes_rollout_bucketing() {
+  TelemetryEvents.init();
+  Services.fog.testResetFOG();
+  Services.telemetry.snapshotEvents(
+    Ci.nsITelemetry.DATASET_PRERELEASE_CHANNELS,
+     true
+  );
+
+  const loader = ExperimentFakes.rsLoader();
+  const manager = loader.manager;
+
+  const experiment = ExperimentFakes.recipe("experiment", {
+    branches: [
+      {
+        slug: "control",
+        ratio: 1,
+        features: [
+          {
+            featureId: "testFeature",
+            value: {},
+          },
+        ],
+      },
+    ],
+    bucketConfig: {
+      namespace: "nimbus-test-utils",
+      randomizationUnit: "normandy_id",
+      start: 0,
+      count: 1000,
+      total: 1000,
+    },
+  });
+  const rollout = ExperimentFakes.recipe("rollout", {
+    isRollout: true,
+    branches: [
+      {
+        slug: "rollout",
+        ratio: 1,
+        features: [
+          {
+            featureId: "testFeature",
+            value: {},
+          },
+        ],
+      },
+    ],
+    bucketConfig: {
+      namespace: "nimbus-test-utils",
+      randomizationUnit: "normandy_id",
+      start: 0,
+      count: 1000,
+      total: 1000,
+    },
+  });
+
+  await loader.init();
+  await manager.onStartup();
+  await manager.store.ready();
+
+  sinon
+    .stub(loader.remoteSettingsClient, "get")
+    .resolves([experiment, rollout]);
+
+  await loader.updateRecipes();
+
+  Assert.equal(
+    manager.store.getExperimentForFeature("testFeature")?.slug,
+    experiment.slug,
+    "Should enroll in experiment"
+  );
+  Assert.equal(
+    manager.store.getRolloutForFeature("testFeature")?.slug,
+    rollout.slug,
+    "Should enroll in rollout"
+  );
+
+  experiment.bucketConfig.count = 0;
+  rollout.bucketConfig.count = 0;
+
+  await loader.updateRecipes();
+
+  Assert.equal(
+    manager.store.getExperimentForFeature("testFeature")?.slug,
+    experiment.slug,
+    "Should stay enrolled in experiment -- experiments cannot be resized"
+  );
+  Assert.ok(
+    !manager.store.getRolloutForFeature("testFeature"),
+    "Should unenroll from rollout"
+  );
+
+  const unenrollmentEvents = Glean.nimbusEvents.unenrollment.testGetValue();
+  Assert.equal(
+    unenrollmentEvents.length,
+    1,
+    "Should be one unenrollment event"
+  );
+  Assert.equal(
+    unenrollmentEvents[0].extra.experiment,
+    rollout.slug,
+    "Experiment slug should match"
+  );
+  Assert.equal(
+    unenrollmentEvents[0].extra.reason,
+    "bucketing",
+    "Reason should match"
+  );
+
+  TelemetryTestUtils.assertEvents(
+    [
+      {
+        value: rollout.slug,
+        extra: {
+          reason: "bucketing",
+        },
+      },
+    ],
+    {
+      category: "normandy",
+      method: "unenroll",
+      object: "nimbus_experiment",
+    }
+  );
+});