[ab] IJPL-186222: Introduce a simpler API for running AB experiments

(cherry picked from commit d37a06ebb60306d38f9e5a04d98df6669f07946a)


(cherry picked from commit f9ef6adbd996974ebedfe1f46e5d28e959ee8c01)

IJ-CR-176469

GitOrigin-RevId: c98c80f169ec9f508971dd84a45a63050b0c101e

platform/experiment/BUILD.bazel

@@ -23,4 +23,32 @@ jvm_library(
   ],
   runtime_deps = [":experiment_resources"]
 )
-### auto-generated section `build intellij.platform.experiment` end
+
+jvm_library(
+  name = "experiment_test_lib",
+  visibility = ["//visibility:public"],
+  srcs = glob(["test/**/*.kt", "test/**/*.java", "test/**/*.form"], allow_empty = True),
+  associates = [":experiment"],
+  deps = [
+    "@lib//:kotlin-stdlib",
+    "//platform/core-api:core",
+    "//platform/statistics",
+    "//platform/statistics:statistics_test_lib",
+    "//platform/editor-ui-api:editor-ui",
+    "//platform/projectModel-api:projectModel",
+    "//platform/projectModel-impl",
+    "//platform/platform-impl:ide-impl",
+    "@lib//:junit5",
+  ],
+  runtime_deps = [":experiment_resources"]
+)
+### auto-generated section `build intellij.platform.experiment` end
+
+### auto-generated section `test intellij.platform.experiment` start
+load("@community//build:tests-options.bzl", "jps_test")
+
+jps_test(
+  name = "experiment_test",
+  runtime_deps = [":experiment_test_lib"]
+)
+### auto-generated section `test intellij.platform.experiment` end

platform/experiment/intellij.platform.experiment.iml

@@ -5,6 +5,7 @@
     <content url="file://$MODULE_DIR$">
       <sourceFolder url="file://$MODULE_DIR$/resources" type="java-resource" />
       <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/test" isTestSource="true" />
     </content>
     <orderEntry type="inheritedJdk" />
     <orderEntry type="sourceFolder" forTests="false" />
@@ -15,5 +16,6 @@
     <orderEntry type="module" module-name="intellij.platform.projectModel" />
     <orderEntry type="module" module-name="intellij.platform.projectModel.impl" />
     <orderEntry type="module" module-name="intellij.platform.ide.impl" />
+    <orderEntry type="library" scope="TEST" name="JUnit5" level="project" />
   </component>
 </module>

platform/experiment/src/com/intellij/platform/experiment/ab/demo/ABExperimentDemoAction.kt

@@ -4,18 +4,16 @@ package com.intellij.platform.experiment.ab.demo
 import com.intellij.openapi.actionSystem.ActionUpdateThread
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
-import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.components.service
-import com.intellij.platform.experiment.ab.impl.experiment.ABExperiment
-import com.intellij.platform.experiment.ab.impl.experiment.ABExperimentImpl
+import com.intellij.platform.experiment.ab.impl.ABExperimentOption
+import com.intellij.platform.experiment.ab.impl.getUserBucketNumber
+import com.intellij.platform.experiment.ab.impl.reportableName
 
 internal class ABExperimentDemoAction : AnAction() {
   override fun actionPerformed(e: AnActionEvent) {
-    val service = ApplicationManager.getApplication().service<ABExperiment>() as ABExperimentImpl
-
-    println("User experiment option is: " + service.getUserExperimentOption())
-    println("User experiment option id is: " + service.getUserExperimentOptionId())
-    println("Is control experiment option enabled: " + service.isControlExperimentOptionEnabled())
+    val experimentValues = ABExperimentOption.entries.filter { it != ABExperimentOption.UNASSIGNED }
+    for (experimentValue in experimentValues) {
+      println("Experiment value: $experimentValue; reportable name: ${experimentValue.reportableName()}; user bucket: ${getUserBucketNumber()} is enabled: ${experimentValue.isEnabled()}")
+    }
   }
 
   override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT

platform/experiment/src/com/intellij/platform/experiment/ab/impl/experiments.kt

@@ -0,0 +1,102 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.experiment.ab.impl
+
+import com.intellij.internal.statistic.eventLog.fus.MachineIdManager
+import com.intellij.openapi.application.ApplicationInfo.getInstance
+import com.intellij.openapi.diagnostic.debug
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.diagnostic.runAndLogException
+import com.intellij.platform.experiment.ab.impl.ABExperimentOption.*
+import com.intellij.platform.experiment.ab.impl.statistic.ABExperimentCountCollector
+import org.jetbrains.annotations.VisibleForTesting
+import kotlin.math.absoluteValue
+
+/**
+ * Complete list of all available AB experiments.
+ *
+ * The plugins are welcome to use [ABExperimentOption.isEnabled] to check whether the experiment is enabled on the user's machine
+ */
+enum class ABExperimentOption {
+  EXPERIMENT_1,
+  EXPERIMENT_2,
+  EXPERIMENT_3,
+
+  /**
+   * A group for users which are not assigned to any experiment.
+   */
+  UNASSIGNED;
+
+  fun isEnabled(): Boolean {
+    require(this != UNASSIGNED) {
+      "UNASSIGNED experiment option is not supposed to be used in the isEnabled() method"
+    }
+    val decision = thisUserDecision
+    ABExperimentCountCollector.logABExperimentOptionUsed(decision)
+    return isAllowed(decision.option) && decision.option == this && !decision.isControlGroup
+  }
+}
+
+/**
+ * Total number of "containers" where users are distributed. Each user belongs exactly to one bucket.
+ * Each bucket has at most one experimental option enabled, which are controlled by [experimentsPartition].
+ */
+internal const val NUMBER_OF_BUCKETS: Int = 1024
+
+/**
+ * Mapping of buckets to experiments.
+ * Each experiment is assigned a non-overlapping range of buckets in [0..[NUMBER_OF_BUCKETS])
+ * The buckets that are not assigned to any experiment are considered to be in the UNASSIGNED experiment (i.e., no experiments are enabled for them).
+ */
+@VisibleForTesting
+internal val experimentsPartition: List<ExperimentAssignment> = listOf(
+  ExperimentAssignment(experiment = EXPERIMENT_1, experimentBuckets = (0 until 180).toSet(), controlBuckets = (180 until 256).toSet()),
+  ExperimentAssignment(experiment = EXPERIMENT_2, experimentBuckets = (256 until 384).toSet(), controlBuckets = (384 until 512).toSet()),
+  ExperimentAssignment(experiment = EXPERIMENT_3, experimentBuckets = (512 until 640).toSet(), controlBuckets = (640 until 768).toSet()),
+  // the rest belongs to the unassigned experiment
+)
+
+/**
+ * This method can be configured to allow options only in particular IDEs.
+ */
+fun isAllowed(option: ABExperimentOption): Boolean {
+  return true
+}
+
+// ================= IMPLEMENTATION ====================
+
+internal data class ABExperimentDecision(val option: ABExperimentOption, val isControlGroup: Boolean, val bucketNumber: Int)
+
+internal fun ABExperimentOption.reportableName(): String {
+  return toString().lowercase().replace('_', '.')
+}
+
+private val thisUserDecision: ABExperimentDecision by lazy {
+  val currentBucket = getUserBucketNumber()
+  val option = experimentsPartition.find {
+    it.experimentBuckets.contains(currentBucket) || it.controlBuckets.contains(currentBucket)
+  } ?: return@lazy ABExperimentDecision(option = UNASSIGNED, isControlGroup = true, bucketNumber = currentBucket)
+  ABExperimentDecision(option = option.experiment, isControlGroup = option.controlBuckets.contains(currentBucket), bucketNumber = currentBucket)
+}
+
+data class ExperimentAssignment(val experiment: ABExperimentOption, val experimentBuckets: Set<Int>, val controlBuckets: Set<Int>)
+
+internal fun getUserBucketNumber(): Int {
+  val overridingBucket = Integer.getInteger("ide.ab.test.overriding.bucket")
+  if (overridingBucket != null) {
+    LOG.info("Overriding bucket number: $overridingBucket")
+    return overridingBucket
+  }
+  val deviceId = LOG.runAndLogException {
+    MachineIdManager.getAnonymizedMachineId(getDeviceIdPurpose())
+  }
+
+  val bucketNumber = deviceId.hashCode().absoluteValue % NUMBER_OF_BUCKETS
+  LOG.debug { "User bucket number is: $bucketNumber." }
+  return bucketNumber
+}
+
+private fun getDeviceIdPurpose(): String {
+  return "A/B Experiment" + getInstance().shortVersion
+}
+
+private val LOG = logger<ABExperimentOption>()

platform/experiment/src/com/intellij/platform/experiment/ab/impl/statistic/ABExperimentCountCollector.kt

@@ -4,47 +4,22 @@ package com.intellij.platform.experiment.ab.impl.statistic
 import com.intellij.internal.statistic.eventLog.EventLogGroup
 import com.intellij.internal.statistic.eventLog.events.EventFields
 import com.intellij.internal.statistic.service.fus.collectors.CounterUsagesCollector
-import com.intellij.platform.experiment.ab.impl.experiment.ABExperimentOptionId
-import com.intellij.platform.experiment.ab.impl.experiment.OPTION_ID_FREE_GROUP
-import com.intellij.platform.experiment.ab.impl.experiment.getJbABExperimentOptionList
+import com.intellij.platform.experiment.ab.impl.ABExperimentDecision
+import com.intellij.platform.experiment.ab.impl.reportableName
 
 internal object ABExperimentCountCollector : CounterUsagesCollector() {
-  private val GROUP = EventLogGroup("experiment.ab", 3)
-
-  /**
-   * For the case when user enables a plugin and then disables it.
-   *
-   * When the plugin is disabled, then a corresponding option is missing and the validation rule will reject such an option id,
-   * because the option is not present at runtime.
-   *
-   * To overcome such cases, an original option id is replaced with an artificial one for statistics.
-   */
-  internal val OPTION_ID_MISSING = ABExperimentOptionId("missing.option")
+  private val GROUP = EventLogGroup("experiment.ab", 4)
 
   private val AB_EXPERIMENT_OPTION_USED = GROUP.registerEvent(
     "option.used",
     EventFields.StringValidatedByCustomRule("id", ABExperimentOptionIdValidationRule::class.java),
-    EventFields.Int("group"),
+    EventFields.Boolean("is_control_group"),
     EventFields.Int("bucket")
   )
 
-  fun logABExperimentOptionUsed(userOptionId: ABExperimentOptionId?, userGroupNumber: Int, userBucketNumber: Int) {
-    if (userOptionId == null) {
-      return
-    }
-
-    if (userOptionId == OPTION_ID_FREE_GROUP) {
-      AB_EXPERIMENT_OPTION_USED.log(userOptionId.value, userGroupNumber, userBucketNumber)
-      return
-    }
-
-    val option = getJbABExperimentOptionList().find { it.id.value == userOptionId.value }
-    if (option != null) {
-      AB_EXPERIMENT_OPTION_USED.log(option.id.value, userGroupNumber, userBucketNumber)
-      return
-    }
-
-    AB_EXPERIMENT_OPTION_USED.log(OPTION_ID_MISSING.value, userGroupNumber, userBucketNumber)
+  fun logABExperimentOptionUsed(decision: ABExperimentDecision) {
+    val optionName = decision.option.reportableName()
+    AB_EXPERIMENT_OPTION_USED.log(optionName, decision.isControlGroup, decision.bucketNumber)
   }
 
   override fun getGroup(): EventLogGroup = GROUP

platform/experiment/src/com/intellij/platform/experiment/ab/impl/statistic/ABExperimentOptionIdValidationRule.kt

@@ -4,17 +4,17 @@ package com.intellij.platform.experiment.ab.impl.statistic
 import com.intellij.internal.statistic.eventLog.validator.ValidationResultType
 import com.intellij.internal.statistic.eventLog.validator.rules.EventContext
 import com.intellij.internal.statistic.eventLog.validator.rules.impl.CustomValidationRule
-import com.intellij.platform.experiment.ab.impl.experiment.OPTION_ID_FREE_GROUP
-import com.intellij.platform.experiment.ab.impl.experiment.getJbABExperimentOptionList
-import com.intellij.platform.experiment.ab.impl.statistic.ABExperimentCountCollector.OPTION_ID_MISSING
+import com.intellij.platform.experiment.ab.impl.ABExperimentOption
+import com.intellij.platform.experiment.ab.impl.reportableName
 
 internal class ABExperimentOptionIdValidationRule : CustomValidationRule() {
   override fun getRuleId(): String = "ab_experiment_option_id"
 
   override fun doValidate(data: String, context: EventContext): ValidationResultType {
-    return if (getJbABExperimentOptionList().any { it.id.value == data } ||
-               data == OPTION_ID_FREE_GROUP.value ||
-               data == OPTION_ID_MISSING.value) {
+    val optionNames = ABExperimentOption.entries
+      .filter { it != ABExperimentOption.UNASSIGNED }
+      .map { it.reportableName() }
+    return if (data in optionNames) {
       ValidationResultType.ACCEPTED
     }
     else {

platform/experiment/test/com/intellij/platform/experiment/ab/test.kt

@@ -0,0 +1,41 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.experiment.ab
+
+import com.intellij.platform.experiment.ab.impl.NUMBER_OF_BUCKETS
+import com.intellij.platform.experiment.ab.impl.experimentsPartition
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+
+class ABExperimentSanityTest {
+
+  @Test
+  fun `no intersections in experiment partitions`() {
+    for (experiment1 in experimentsPartition) {
+      for (experiment2 in experimentsPartition) {
+        if (experiment1.experiment == experiment2.experiment) {
+          assertEmptyIntersection(experiment1.controlBuckets, experiment2.experimentBuckets)
+        }
+        else {
+          assertEmptyIntersection(experiment1.experimentBuckets, experiment2.controlBuckets)
+          assertEmptyIntersection(experiment1.controlBuckets, experiment2.controlBuckets)
+          assertEmptyIntersection(experiment1.experimentBuckets, experiment2.experimentBuckets)
+          assertEmptyIntersection(experiment1.controlBuckets, experiment2.experimentBuckets)
+        }
+      }
+    }
+  }
+
+  @Test
+  fun `all ranges are within bounds`() {
+    for (experiment in experimentsPartition) {
+      assertWithinBounds(experiment.controlBuckets)
+      assertWithinBounds(experiment.experimentBuckets)
+    }
+  }
+
+  fun assertEmptyIntersection(range1: Set<*>, range2: Set<*>) {
+    Assertions.assertTrue(range1.intersect(range2).isEmpty())
+  }
+
+  fun assertWithinBounds(range: Set<Int>) = Assertions.assertTrue(range.subtract(0 until NUMBER_OF_BUCKETS).isEmpty())
+}

python/ide/impl/BUILD.bazel

@@ -97,6 +97,7 @@ jvm_library(
     "//platform/whatsNew:whatsNew_test_lib",
     "//platform/util/coroutines",
     "//platform/experiment",
+    "//platform/experiment:experiment_test_lib",
     "@lib//:jna",
     "//python/services/system-python",
     "//python/services/system-python:system-python_test_lib",