ref(dynamic-sampling): simplify bias combination logic (#104119)

Author: shellmayr

src/sentry/dynamic_sampling/rules/base.py

@@ -1,5 +1,4 @@
 import logging
-from collections import OrderedDict
 from datetime import datetime, timedelta, timezone
 
 import sentry_sdk
@@ -8,7 +7,7 @@
 from sentry.constants import TARGET_SAMPLE_RATE_DEFAULT
 from sentry.db.models import Model
 from sentry.dynamic_sampling.rules.biases.base import Bias
-from sentry.dynamic_sampling.rules.combine import get_relay_biases_combinator
+from sentry.dynamic_sampling.rules.combine import get_relay_biases
 from sentry.dynamic_sampling.rules.utils import PolymorphicRule, RuleType, get_enabled_user_biases
 from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import (
     get_boost_low_volume_projects_sample_rate,
@@ -94,7 +93,7 @@ def _get_rules_of_enabled_biases(
     project: Project,
     base_sample_rate: float,
     enabled_biases: set[str],
-    combined_biases: OrderedDict[RuleType, Bias],
+    combined_biases: dict[RuleType, Bias],
 ) -> list[PolymorphicRule]:
     rules = []
 
@@ -124,7 +123,7 @@ def generate_rules(project: Project) -> list[PolymorphicRule]:
         enabled_user_biases = get_enabled_user_biases(
             project.get_option("sentry:dynamic_sampling_biases", None)
         )
-        combined_biases = get_relay_biases_combinator(organization).get_combined_biases()
+        combined_biases = get_relay_biases(organization)
 
         rules = _get_rules_of_enabled_biases(
             project, base_sample_rate, enabled_user_biases, combined_biases

src/sentry/dynamic_sampling/rules/biases/bias_combinator.py

@@ -0,0 +1,16 @@
+from collections.abc import Callable
+
+from sentry.dynamic_sampling.rules.biases.base import Bias
+from sentry.dynamic_sampling.rules.utils import RuleType
+
+
+class OrderedBiasesCombinator:
+    def __init__(self) -> None:
+        self.biases: dict[RuleType, Bias] = {}
+
+    def add_if(self, rule_type: RuleType, bias: Bias, block: Callable[[], bool]) -> None:
+        if block():
+            self.add(rule_type, bias)
+
+    def add(self, rule_type: RuleType, bias: Bias) -> None:
+        self.biases[rule_type] = bias

src/sentry/dynamic_sampling/rules/combinators/__init__.py

@@ -1,4 +1,6 @@
 from sentry import features
+from sentry.dynamic_sampling.rules.biases.base import Bias
+from sentry.dynamic_sampling.rules.biases.bias_combinator import OrderedBiasesCombinator
 from sentry.dynamic_sampling.rules.biases.boost_environments_bias import BoostEnvironmentsBias
 from sentry.dynamic_sampling.rules.biases.boost_latest_releases_bias import BoostLatestReleasesBias
 from sentry.dynamic_sampling.rules.biases.boost_low_volume_projects_bias import (
@@ -15,13 +17,11 @@
 )
 from sentry.dynamic_sampling.rules.biases.minimum_sample_rate_bias import MinimumSampleRateBias
 from sentry.dynamic_sampling.rules.biases.recalibration_bias import RecalibrationBias
-from sentry.dynamic_sampling.rules.combinators.base import BiasesCombinator
-from sentry.dynamic_sampling.rules.combinators.ordered_combinator import OrderedBiasesCombinator
 from sentry.dynamic_sampling.rules.utils import RuleType
 from sentry.models.organization import Organization
 
 
-def get_relay_biases_combinator(organization: Organization) -> BiasesCombinator:
+def get_relay_biases(organization: Organization) -> dict[RuleType, Bias]:
     is_health_checks_trace_based = features.has(
         "organizations:ds-health-checks-trace-based", organization, actor=None
     )
@@ -55,4 +55,4 @@ def get_relay_biases_combinator(organization: Organization) -> BiasesCombinator:
     )
     default_combinator.add(RuleType.BOOST_LOW_VOLUME_PROJECTS_RULE, BoostLowVolumeProjectsBias())
 
-    return default_combinator
+    return default_combinator.biases