Revert "fix(dynamic-sampling): use options instead of feature flags for switching b/w span and transaction based rebalancing (#102576)"

This reverts commit 8ecceb3.

Co-authored-by: shellmayr <[email protected]>

Author: 2 people

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"):
-            span_org_ids = options.get("dynamic-sampling.measure.spans") or []
-            if org_id in span_org_ids:
-                measure = SamplingMeasure.SPANS
+        if options.get("dynamic-sampling.check_span_feature_flag") and features.has(
+            "organizations:dynamic-sampling-spans", organization
+        ):
+            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,33 +133,38 @@ def partition_by_measure(
 
     # Exclude orgs with project-mode sampling from the start. We know the
     # default is DynamicSamplingMode.ORGANIZATION.
-    filtered_org_ids = {
-        org.id for org, mode in modes.items() if mode != DynamicSamplingMode.PROJECT
-    }
+    orgs = [org 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(filtered_org_ids)
-        )
-        return {SamplingMeasure.TRANSACTIONS: sorted(filtered_org_ids)}
+        metrics.incr("dynamic_sampling.partition_by_measure.transactions", amount=len(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
+    spans = []
+    transactions = []
 
-    logger.info(
-        "dynamic_sampling.partition_by_measure.options_check",
-        extra={"span_org_ids": span_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]}
 
-    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)
+    logger.info(
+        "dynamic_sampling.partition_by_measure.batched_feature_check",
+        extra={"feature_results": feature_results},
     )
-    return {
-        SamplingMeasure.SPANS: sorted(span_org_ids),
-        SamplingMeasure.TRANSACTIONS: sorted(transactions_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}
 
 
 @instrumented_task(
@@ -185,10 +190,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"):
-        span_org_ids = options.get("dynamic-sampling.measure.spans") or []
-        if org_id in span_org_ids:
-            measure = SamplingMeasure.SPANS
+    if options.get("dynamic-sampling.check_span_feature_flag") and features.has(
+        "organizations:dynamic-sampling-spans", org
+    ):
+        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,14 +2181,6 @@
     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,11 +287,9 @@ 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,
-                "dynamic-sampling.measure.spans": [org.id],
-            }
+        with (
+            self.options({"dynamic-sampling.check_span_feature_flag": True}),
+            self.feature({"organizations:dynamic-sampling-spans": True}),
         ):
             result = partition_by_measure([org.id])
             assert SamplingMeasure.SPANS in result
@@ -301,11 +299,9 @@ 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,
-                "dynamic-sampling.measure.spans": [],
-            }
+        with (
+            self.options({"dynamic-sampling.check_span_feature_flag": True}),
+            self.feature({"organizations:dynamic-sampling-spans": False}),
         ):
             result = partition_by_measure([org.id])
             assert SamplingMeasure.SPANS in result
@@ -315,48 +311,11 @@ 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,
-                "dynamic-sampling.measure.spans": [org.id],
-            }
+        with (
+            self.options({"dynamic-sampling.check_span_feature_flag": False}),
+            self.feature({"organizations:dynamic-sampling-spans": True}),
         ):
             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