fix(dynamic-sampling): add option for switching bw spans and transactions (v2)

src/sentry/core/endpoints/organization_details.py

@@ -1152,10 +1152,10 @@ def _compute_project_target_sample_rates(self, request: Request, organization: O
         #       so we need to refactor this into an async task we can run and observe
         org_id = organization.id
         measure = SamplingMeasure.TRANSACTIONS
-        if options.get("dynamic-sampling.check_span_feature_flag") and features.has(
-            "organizations:dynamic-sampling-spans", organization
-        ):
-            measure = SamplingMeasure.SPANS
+        if options.get("dynamic-sampling.check_span_feature_flag"):
+            span_org_ids = options.get("dynamic-sampling.measure.spans") or []
+            if org_id in span_org_ids:
+                measure = SamplingMeasure.SPANS
 
         projects_with_tx_count_and_rates = []
         for chunk in query_project_counts_by_org(

src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py

@@ -133,38 +133,33 @@ def partition_by_measure(
 
     # Exclude orgs with project-mode sampling from the start. We know the
     # default is DynamicSamplingMode.ORGANIZATION.
-    orgs = [org for org, mode in modes.items() if mode != DynamicSamplingMode.PROJECT]
+    filtered_org_ids = {
+        org.id for org, mode in modes.items() if mode != DynamicSamplingMode.PROJECT
+    }
 
     if not options.get("dynamic-sampling.check_span_feature_flag"):
-        metrics.incr("dynamic_sampling.partition_by_measure.transactions", amount=len(orgs))
-        return {SamplingMeasure.TRANSACTIONS: [org.id for org in orgs]}
-
-    spans = []
-    transactions = []
+        metrics.incr(
+            "dynamic_sampling.partition_by_measure.transactions", amount=len(filtered_org_ids)
+        )
+        return {SamplingMeasure.TRANSACTIONS: sorted(filtered_org_ids)}
 
-    # Use batch feature flag check to avoid N+1 queries.
-    feature_results = features.batch_has_for_organizations(
-        "organizations:dynamic-sampling-spans", orgs
-    )
-    if feature_results is None:
-        metrics.incr("dynamic_sampling.partition_by_measure.transactions", amount=len(orgs))
-        logger.error("dynamic_sampling.partition_by_measure.features_none", extra={"orgs": orgs})
-        return {SamplingMeasure.TRANSACTIONS: [org.id for org in orgs]}
+    span_org_ids = set(options.get("dynamic-sampling.measure.spans") or [])
+    span_org_ids = span_org_ids & filtered_org_ids
+    transactions_org_ids = filtered_org_ids - span_org_ids
 
     logger.info(
-        "dynamic_sampling.partition_by_measure.batched_feature_check",
-        extra={"feature_results": feature_results},
+        "dynamic_sampling.partition_by_measure.options_check",
+        extra={"span_org_ids": span_org_ids},
     )
 
-    for org in orgs:
-        if feature_results.get(f"organization:{org.id}"):
-            spans.append(org.id)
-        else:
-            transactions.append(org.id)
-
-    metrics.incr("dynamic_sampling.partition_by_measure.spans", amount=len(spans))
-    metrics.incr("dynamic_sampling.partition_by_measure.transactions", amount=len(transactions))
-    return {SamplingMeasure.SPANS: spans, SamplingMeasure.TRANSACTIONS: transactions}
+    metrics.incr("dynamic_sampling.partition_by_measure.spans", amount=len(span_org_ids))
+    metrics.incr(
+        "dynamic_sampling.partition_by_measure.transactions", amount=len(transactions_org_ids)
+    )
+    return {
+        SamplingMeasure.SPANS: sorted(span_org_ids),
+        SamplingMeasure.TRANSACTIONS: sorted(transactions_org_ids),
+    }
 
 
 @instrumented_task(
@@ -190,10 +185,10 @@ def boost_low_volume_projects_of_org_with_query(org_id: OrganizationId) -> None:
         return
 
     measure = SamplingMeasure.TRANSACTIONS
-    if options.get("dynamic-sampling.check_span_feature_flag") and features.has(
-        "organizations:dynamic-sampling-spans", org
-    ):
-        measure = SamplingMeasure.SPANS
+    if options.get("dynamic-sampling.check_span_feature_flag"):
+        span_org_ids = options.get("dynamic-sampling.measure.spans") or []
+        if org_id in span_org_ids:
+            measure = SamplingMeasure.SPANS
 
     projects_with_tx_count_and_rates = fetch_projects_with_total_root_transaction_count_and_rates(
         org_ids=[org_id],

src/sentry/options/defaults.py

@@ -2181,6 +2181,14 @@
     flags=FLAG_AUTOMATOR_MODIFIABLE | FLAG_MODIFIABLE_RATE,
 )
 
+# List of organization IDs that should be using spans for rebalancing in dynamic sampling.
+register(
+    "dynamic-sampling.measure.spans",
+    default=[],
+    type=Sequence,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 # === Hybrid cloud subsystem options ===
 # UI rollout
 register(

tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py

@@ -287,9 +287,11 @@ def test_complex(self) -> None:
 class TestPartitionByMeasure(TestCase):
     def test_partition_by_measure_with_spans_feature(self) -> None:
         org = self.create_organization("test-org1")
-        with (
-            self.options({"dynamic-sampling.check_span_feature_flag": True}),
-            self.feature({"organizations:dynamic-sampling-spans": True}),
+        with self.options(
+            {
+                "dynamic-sampling.check_span_feature_flag": True,
+                "dynamic-sampling.measure.spans": [org.id],
+            }
         ):
             result = partition_by_measure([org.id])
             assert SamplingMeasure.SPANS in result
@@ -299,9 +301,11 @@ def test_partition_by_measure_with_spans_feature(self) -> None:
 
     def test_partition_by_measure_without_spans_feature(self) -> None:
         org = self.create_organization("test-org1")
-        with (
-            self.options({"dynamic-sampling.check_span_feature_flag": True}),
-            self.feature({"organizations:dynamic-sampling-spans": False}),
+        with self.options(
+            {
+                "dynamic-sampling.check_span_feature_flag": True,
+                "dynamic-sampling.measure.spans": [],
+            }
         ):
             result = partition_by_measure([org.id])
             assert SamplingMeasure.SPANS in result
@@ -311,11 +315,48 @@ def test_partition_by_measure_without_spans_feature(self) -> None:
 
     def test_partition_by_measure_with_span_feature_flag_disabled(self) -> None:
         org = self.create_organization("test-org1")
-        with (
-            self.options({"dynamic-sampling.check_span_feature_flag": False}),
-            self.feature({"organizations:dynamic-sampling-spans": True}),
+        with self.options(
+            {
+                "dynamic-sampling.check_span_feature_flag": False,
+                "dynamic-sampling.measure.spans": [org.id],
+            }
         ):
             result = partition_by_measure([org.id])
             assert SamplingMeasure.TRANSACTIONS in result
             assert SamplingMeasure.SPANS not in result
             assert result[SamplingMeasure.TRANSACTIONS] == [org.id]
+
+    def test_partition_by_measure_returns_sorted_output_multiple_orgs(self) -> None:
+        orgs = [self.create_organization(f"test-org{i}") for i in range(10)]
+        org_ids = [org.id for org in reversed(orgs)]
+
+        with self.options(
+            {
+                "dynamic-sampling.check_span_feature_flag": True,
+                "dynamic-sampling.measure.spans": [orgs[2].id, orgs[7].id, orgs[5].id],
+            }
+        ):
+            result = partition_by_measure(org_ids)
+
+            assert result[SamplingMeasure.SPANS] == sorted([orgs[2].id, orgs[7].id, orgs[5].id])
+            expected_transaction_orgs = sorted(
+                [org.id for org in orgs if org.id not in [orgs[2].id, orgs[7].id, orgs[5].id]]
+            )
+            assert result[SamplingMeasure.TRANSACTIONS] == expected_transaction_orgs
+
+    def test_partition_by_measure_returns_sorted_when_feature_disabled(self) -> None:
+        org1 = self.create_organization("test-org1")
+        org2 = self.create_organization("test-org2")
+        org3 = self.create_organization("test-org3")
+
+        org_ids = [org3.id, org1.id, org2.id]
+
+        with self.options(
+            {
+                "dynamic-sampling.check_span_feature_flag": False,
+            }
+        ):
+            result = partition_by_measure(org_ids)
+
+            assert result[SamplingMeasure.TRANSACTIONS] == sorted(org_ids)
+            assert SamplingMeasure.SPANS not in result