diff --git a/src/sentry/core/endpoints/organization_details.py b/src/sentry/core/endpoints/organization_details.py index f57653db6c97d1..9b07d955d6d771 100644 --- a/src/sentry/core/endpoints/organization_details.py +++ b/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( diff --git a/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py b/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py index 4c5a8ee02b5111..89ecc52456f4bd 100644 --- a/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py +++ b/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], diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index 77a0b706ac26bb..8d1936dbf0be85 100644 --- a/src/sentry/options/defaults.py +++ b/src/sentry/options/defaults.py @@ -2182,6 +2182,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( diff --git a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py index be449560c26342..7b562512b29218 100644 --- a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py +++ b/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