feat(dynamic-sampling): Add new bias for dev envs [TET-491] (#40382)

This PR add new  bias for dev envs.
Also add common approach to adding new rules like: releases or health
checks to `generate_rules()` function.

Also enable mypy for `src/sentry/dynamic_sampling/`

TODO (fix mypy issues after merge conflicts in) :
- [x] src/sentry/dynamic_sampling/feature_multiplexer.py
- [x] src/sentry/dynamic_sampling/utils.py

mypy.ini

@@ -44,6 +44,7 @@ files = fixtures/mypy-stubs,
         src/sentry/db/models/query.py,
         src/sentry/db/models/utils.py,
         src/sentry/digests/,
+        src/sentry/dynamic_sampling/,
         src/sentry/eventstream/base.py,
         src/sentry/eventstream/snuba.py,
         src/sentry/eventstream/kafka/consumer_strategy.py,

src/sentry/dynamic_sampling/__init__.py

@@ -0,0 +1,32 @@
+from typing import List
+
+import sentry_sdk
+
+from sentry import quotas
+from sentry.dynamic_sampling.feature_multiplexer import DynamicSamplingFeatureMultiplexer
+from sentry.dynamic_sampling.utils import BaseRule, generate_environment_rule, generate_uniform_rule
+from sentry.models import Project
+
+
+def generate_rules(project: Project) -> List[BaseRule]:
+    """
+    This function handles generate rules logic or fallback empty list of rules
+    """
+    rules = []
+
+    sample_rate = quotas.get_blended_sample_rate(project)
+
+    boost_environments = DynamicSamplingFeatureMultiplexer.get_user_bias_by_id(
+        "boostEnvironments", project.get_option("sentry:dynamic_sampling_biases", None)
+    )
+    if sample_rate is None:
+        try:
+            raise Exception("get_blended_sample_rate returns none")
+        except Exception:
+            sentry_sdk.capture_exception()
+    else:
+        if boost_environments["active"] and sample_rate < 1.0:
+            rules.append(generate_environment_rule())
+        rules.append(generate_uniform_rule(sample_rate))
+
+    return rules

src/sentry/dynamic_sampling/feature_multiplexer.py

@@ -1,5 +1,9 @@
+from typing import List, Optional, Set
+
 from sentry import features
-from sentry.dynamic_sampling.utils import DEFAULT_BIASES
+from sentry.dynamic_sampling.utils import DEFAULT_BIASES, Bias
+from sentry.models import Project
+from sentry.models.user import User
 
 
 class DynamicSamplingFeatureMultiplexer:
@@ -13,7 +17,7 @@ class DynamicSamplingFeatureMultiplexer:
     - The `organizations:dynamic-sampling` feature flag is the flag that enables the new adaptive sampling
     """
 
-    def __init__(self, project, user):
+    def __init__(self, project: Project, user: User):
         # Feature flag that informs us that relay is handling DS rules
         self.allow_dynamic_sampling = features.has(
             "organizations:server-side-sampling", project.organization, actor=user
@@ -28,7 +32,7 @@ class DynamicSamplingFeatureMultiplexer:
         )
 
     @property
-    def is_on_dynamic_sampling_deprecated(self):
+    def is_on_dynamic_sampling_deprecated(self) -> bool:
         return (
             self.allow_dynamic_sampling
             and self.deprecated_dynamic_sampling
@@ -36,11 +40,11 @@ class DynamicSamplingFeatureMultiplexer:
         )
 
     @property
-    def is_on_dynamic_sampling(self):
+    def is_on_dynamic_sampling(self) -> bool:
         return self.allow_dynamic_sampling and self.current_dynamic_sampling
 
     @staticmethod
-    def get_user_biases(user_set_biases):
+    def get_user_biases(user_set_biases: Optional[List[Bias]]) -> List[Bias]:
         if user_set_biases is None:
             return DEFAULT_BIASES
 
@@ -54,5 +58,12 @@ class DynamicSamplingFeatureMultiplexer:
         return returned_biases
 
     @staticmethod
-    def get_supported_biases_ids():
+    def get_supported_biases_ids() -> Set[str]:
         return {bias["id"] for bias in DEFAULT_BIASES}
+
+    @classmethod
+    def get_user_bias_by_id(cls, bias_id: str, user_set_biases: Optional[List[Bias]]) -> Bias:
+        for bias in cls.get_user_biases(user_set_biases):
+            if bias["id"] == bias_id:
+                return bias
+        raise ValueError(f"{bias_id} is not in supported biases")

src/sentry/dynamic_sampling/utils.py

@@ -1,15 +1,16 @@
-from typing import Any, List, TypedDict
+from typing import Dict, List, Optional, TypedDict
 
-import sentry_sdk
+UNIFORM_RULE_RESERVED_ID = 0
 
-from sentry import quotas
-from sentry.models import Project
 
-UNIFORM_RULE_RESERVED_ID = 0
+class Bias(TypedDict):
+    id: str
+    active: bool
+
 
 # These represent the biases that are applied to user by default as part of the adaptive dynamic sampling experience.
 # These can be overridden by the project details endpoint
-DEFAULT_BIASES = [
+DEFAULT_BIASES: List[Bias] = [
     {"id": "boostEnvironments", "active": True},
     {
         "id": "boostLatestRelease",
@@ -19,31 +20,27 @@ DEFAULT_BIASES = [
 ]
 
 
-class NoneSampleRateException(Exception):
-    ...
+class Inner(TypedDict):
+    op: str
+    name: str
+    value: List[str]
+    options: Dict[str, bool]
 
 
 class Condition(TypedDict):
     op: str
-    inner: List[Any]
+    inner: List[Optional[Inner]]
 
 
-class UniformRule(TypedDict):
-    sampleRate: float
+class BaseRule(TypedDict):
+    sampleRate: Optional[float]
     type: str
     active: bool
     condition: Condition
     id: int
 
 
-def generate_uniform_rule(project: Project) -> UniformRule:
-    sample_rate = quotas.get_blended_sample_rate(project)
-    if sample_rate is None:
-        try:
-            raise Exception("get_blended_sample_rate returns none")
-        except Exception:
-            sentry_sdk.capture_exception()
-        raise NoneSampleRateException
+def generate_uniform_rule(sample_rate: Optional[float]) -> BaseRule:
     return {
         "sampleRate": sample_rate,
         "type": "trace",
@@ -54,3 +51,23 @@ def generate_uniform_rule(project: Project) -> UniformRule:
         },
         "id": UNIFORM_RULE_RESERVED_ID,
     }
+
+
+def generate_environment_rule() -> BaseRule:
+    return {
+        "sampleRate": 1,
+        "type": "trace",
+        "condition": {
+            "op": "or",
+            "inner": [
+                {
+                    "op": "glob",
+                    "name": "trace.environment",
+                    "value": ["*dev*", "*test*"],
+                    "options": {"ignoreCase": True},
+                }
+            ],
+        },
+        "active": True,
+        "id": 1,
+    }

src/sentry/relay/config/__init__.py

@@ -20,7 +20,7 @@ from sentry_sdk import Hub, capture_exception
 from sentry import features, killswitches, quotas, utils
 from sentry.constants import ObjectStatus
 from sentry.datascrubbing import get_datascrubbing_settings, get_pii_config
-from sentry.dynamic_sampling.utils import NoneSampleRateException, generate_uniform_rule
+from sentry.dynamic_sampling import generate_rules
 from sentry.grouping.api import get_grouping_config_dict_for_project
 from sentry.ingest.inbound_filters import (
     FilterStatKeys,
@@ -136,7 +136,6 @@ def get_quotas(project, keys=None):
 
 def get_project_config(project, full_config=True, project_keys=None):
     """Constructs the ProjectConfig information.
-
     :param project: The project to load configuration for. Ensure that
         organization is bound on this object; otherwise it will be loaded from
         the database.
@@ -147,7 +146,6 @@ def get_project_config(project, full_config=True, project_keys=None):
         no project keys are provided it is assumed that the config does not
         need to contain auth information (this is the case when used in
         python's StoreView)
-
     :return: a ProjectConfig object for the given project
     """
     with sentry_sdk.push_scope() as scope:
@@ -169,12 +167,7 @@ def get_dynamic_sampling_config(project) -> Optional[Mapping[str, Any]]:
     # In this case we should override old conditionnal rules if they exists
     # or just return uniform rule
     if allow_dynamic_sampling:
-        try:
-            return {"rules": [generate_uniform_rule(project)]}
-        except NoneSampleRateException:
-            # just to be consistent with old code, where if there is no active active_rules
-            # we return empty list
-            return {"rules": []}
+        return {"rules": generate_rules(project)}
     elif allow_server_side_sampling:
         dynamic_sampling = project.get_option("sentry:dynamic_sampling")
         if dynamic_sampling is not None:
@@ -208,10 +201,8 @@ def add_experimental_config(
     **kwargs: Any,
 ) -> None:
     """Try to set `config[key] = function(*args, **kwargs)`.
-
     If the result of the function call is None, the key is not set.
     If the function call raises an exception, we log it to sentry and the key remains unset.
-
     NOTE: Only use this function if you expect Relay to behave reasonably
     if ``key`` is missing from the config.
     """
@@ -309,9 +300,7 @@ def _get_project_config(project, full_config=True, project_keys=None):
 class _ConfigBase:
     """
     Base class for configuration objects
-
     Offers a readonly configuration class that can be serialized to json and viewed as a simple dictionary
-
     >>> x = _ConfigBase( a= 1, b="The b", c= _ConfigBase(x=33, y = _ConfigBase(m=3.14159 , w=[1,2,3], z={'t':1})))
     >>> x.a
     1
@@ -321,7 +310,6 @@ class _ConfigBase:
     True
     >>> x.c.y.w
     [1, 2, 3]
-
     """
 
     def __init__(self, **kwargs):
@@ -341,9 +329,7 @@ class _ConfigBase:
     def to_dict(self):
         """
         Converts the config object into a dictionary
-
         :return: A dictionary containing the object properties, with config properties also converted in dictionaries
-
         >>> x = _ConfigBase( a= 1, b="The b", c= _ConfigBase(x=33, y = _ConfigBase(m=3.14159 , w=[1,2,3], z={'t':1})))
         >>> x.to_dict() == {'a': 1, 'c': {'y': {'m': 3.14159, 'w': [1, 2, 3], 'z':{'t': 1}}, 'x': 33}, 'b': 'The b'}
         True
@@ -359,7 +345,6 @@ class _ConfigBase:
         >>> x = _ConfigBase( a = _ConfigBase(b = _ConfigBase( w=[1,2,3])))
         >>> x.to_json_string()
         '{"a": {"b": {"w": [1, 2, 3]}}}'
-
         :return:
         """
         data = self.to_dict()
@@ -368,10 +353,8 @@ class _ConfigBase:
     def get_at_path(self, *args):
         """
         Gets an element at the specified path returning None if the element or the path doesn't exists
-
         :param args: the path to follow ( a list of strings)
         :return: the element if present at specified path or None otherwise)
-
         >>> x = _ConfigBase( a= 1, b="The b", c= _ConfigBase(x=33, y = _ConfigBase(m=3.14159 , w=[1,2,3], z={'t':1})))
         >>> x.get_at_path('c','y','m')
         3.14159
@@ -383,7 +366,6 @@ class _ConfigBase:
         {'t': 1}
         >>> x.get_at_path('c','y','z','t') is None # only navigates in ConfigBase does not try to go into normal dicts.
         True
-
         """
         if len(args) == 0:
             return self
@@ -426,7 +408,6 @@ class ProjectConfig(_ConfigBase):
 def _load_filter_settings(flt, project):
     """
     Returns the filter settings for the specified project
-
     :param flt: the filter function
     :param project: the project for which we want to retrieve the options
     :return: a dictionary with the filter options.
@@ -443,7 +424,6 @@ def _load_filter_settings(flt, project):
 def _filter_option_to_config_setting(flt, setting):
     """
     Encapsulates the logic for associating a filter database option with the filter setting from project_config
-
     :param flt: the filter
     :param setting: the option deserialized from the database
     :return: the option as viewed from project_config

tests/sentry/dynamic_sampling/test_generate_rules.py

@@ -0,0 +1,97 @@
+from unittest.mock import MagicMock, patch
+
+from sentry.dynamic_sampling import generate_rules
+
+
+@patch("sentry.dynamic_sampling.sentry_sdk")
+@patch("sentry.dynamic_sampling.quotas.get_blended_sample_rate")
+def test_generate_rules_capture_exception(get_blended_sample_rate, sentry_sdk):
+    get_blended_sample_rate.return_value = None
+    # since we mock get_blended_sample_rate function
+    # no need to create real project in DB
+    fake_project = MagicMock()
+    # if blended rate is None that means no dynamic sampling behavior should happen.
+    # Therefore no rules should be set.
+    assert generate_rules(fake_project) == []
+    get_blended_sample_rate.assert_called_with(fake_project)
+    sentry_sdk.capture_exception.assert_called_with()
+
+
+@patch(
+    "sentry.dynamic_sampling.feature_multiplexer.DynamicSamplingFeatureMultiplexer.get_user_bias_by_id"
+)
+@patch("sentry.dynamic_sampling.quotas.get_blended_sample_rate")
+def test_generate_rules_return_uniform_rules_with_rate(get_blended_sample_rate, get_user_bias):
+    get_user_bias.return_value = {"id": "boostEnvironments", "active": False}
+    get_blended_sample_rate.return_value = 0.1
+    # since we mock get_blended_sample_rate function
+    # no need to create real project in DB
+    fake_project = MagicMock()
+    assert generate_rules(fake_project) == [
+        {
+            "active": True,
+            "condition": {"inner": [], "op": "and"},
+            "id": 0,
+            "sampleRate": 0.1,
+            "type": "trace",
+        }
+    ]
+    get_blended_sample_rate.assert_called_with(fake_project)
+    get_user_bias.assert_called_with(
+        "boostEnvironments", fake_project.get_option("sentry:dynamic_sampling_biases", None)
+    )
+
+
+@patch("sentry.dynamic_sampling.quotas.get_blended_sample_rate")
+def test_generate_rules_return_uniform_rules_and_env_rule(get_blended_sample_rate):
+    get_blended_sample_rate.return_value = 0.1
+    # since we mock get_blended_sample_rate function
+    # no need to create real project in DB
+    fake_project = MagicMock()
+    assert generate_rules(fake_project) == [
+        {
+            "sampleRate": 1,
+            "type": "trace",
+            "condition": {
+                "op": "or",
+                "inner": [
+                    {
+                        "op": "glob",
+                        "name": "trace.environment",
+                        "value": ["*dev*", "*test*"],
+                        "options": {"ignoreCase": True},
+                    }
+                ],
+            },
+            "active": True,
+            "id": 1,
+        },
+        {
+            "active": True,
+            "condition": {"inner": [], "op": "and"},
+            "id": 0,
+            "sampleRate": 0.1,
+            "type": "trace",
+        },
+    ]
+    get_blended_sample_rate.assert_called_with(fake_project)
+
+
+@patch("sentry.dynamic_sampling.quotas.get_blended_sample_rate")
+def test_generate_rules_return_uniform_rule_with_100_rate_and_without_env_rule(
+    get_blended_sample_rate,
+):
+    get_blended_sample_rate.return_value = 1.0
+    # since we mock get_blended_sample_rate function
+    # no need to create real project in DB
+    fake_project = MagicMock()
+    assert generate_rules(fake_project) == [
+        {
+            "active": True,
+            "condition": {"inner": [], "op": "and"},
+            "id": 0,
+            "sampleRate": 1.0,
+            "type": "trace",
+        },
+    ]
+    get_blended_sample_rate.assert_called_with(fake_project)

tests/sentry/dynamic_sampling/test_utils.py

@@ -1,32 +1,23 @@
-from unittest.mock import MagicMock, patch
+from sentry.dynamic_sampling.utils import generate_environment_rule, generate_uniform_rule
 
-import pytest
 
-from sentry.dynamic_sampling.utils import NoneSampleRateException, generate_uniform_rule
-
-
-@patch("sentry.dynamic_sampling.utils.quotas.get_blended_sample_rate")
-def test_generate_uniform_rule_return_rate(get_blended_sample_rate):
-    get_blended_sample_rate.return_value = 0.1
-    # since we mock get_blended_sample_rate function
-    # no need to create real project in DB
-    fake_project = MagicMock()
-    assert generate_uniform_rule(fake_project) == {
+def test_generate_uniform_rule_return_rate():
+    sample_rate = 0.1
+    assert generate_uniform_rule(sample_rate) == {
         "active": True,
         "condition": {"inner": [], "op": "and"},
         "id": 0,
-        "sampleRate": 0.1,
+        "sampleRate": sample_rate,
         "type": "trace",
     }
-    get_blended_sample_rate.assert_called_with(fake_project)
 
 
-@patch("sentry.dynamic_sampling.utils.quotas.get_blended_sample_rate")
-def test_generate_uniform_rule_raise_exception(get_blended_sample_rate):
-    get_blended_sample_rate.return_value = None
-    # since we mock get_blended_sample_rate function
-    # no need to create real project in DB
-    fake_project = MagicMock()
-    with pytest.raises(NoneSampleRateException):
-        generate_uniform_rule(fake_project)
-    get_blended_sample_rate.assert_called_with(fake_project)
+def test_generate_environment_rule():
+    bias_env_rule = generate_environment_rule()
+    assert bias_env_rule["id"] == 1
+    assert bias_env_rule["condition"]["inner"][0] == {
+        "op": "glob",
+        "name": "trace.environment",
+        "value": ["*dev*", "*test*"],
+        "options": {"ignoreCase": True},
+    }

tests/sentry/relay/test_config.py

@@ -32,6 +32,25 @@ PII_CONFIG = """
 """
 
 
+DEFAULT_ENVIRONMENT_RULE = {
+    "sampleRate": 1,
+    "type": "trace",
+    "condition": {
+        "op": "or",
+        "inner": [
+            {
+                "op": "glob",
+                "name": "trace.environment",
+                "value": ["*dev*", "*test*"],
+                "options": {"ignoreCase": True},
+            }
+        ],
+    },
+    "active": True,
+    "id": 1,
+}
+
+
 @pytest.mark.django_db
 @pytest.mark.parametrize("full", [False, True], ids=["slim_config", "full_config"])
 def test_get_project_config(default_project, insta_snapshot, django_cache, full):
@@ -60,7 +79,7 @@ SOME_EXCEPTION = RuntimeError("foo")
 
 
 @pytest.mark.django_db
-@mock.patch("sentry.relay.config.generate_uniform_rule", side_effect=SOME_EXCEPTION)
+@mock.patch("sentry.relay.config.generate_rules", side_effect=SOME_EXCEPTION)
 @mock.patch("sentry.relay.config.sentry_sdk")
 def test_get_experimental_config(mock_sentry_sdk, _, default_project):
     keys = ProjectKey.objects.filter(project=default_project)
@@ -214,13 +233,14 @@ def test_project_config_with_latest_release_in_dynamic_sampling_rules(default_pr
             {"rules": []},
             {
                 "rules": [
+                    DEFAULT_ENVIRONMENT_RULE,
                     {
                         "sampleRate": 0.1,
                         "type": "trace",
                         "active": True,
                         "condition": {"op": "and", "inner": []},
                         "id": 0,
-                    }
+                    },
                 ]
             },
         ),
@@ -230,7 +250,7 @@ def test_project_config_with_latest_release_in_dynamic_sampling_rules(default_pr
             {
                 "rules": [
                     {
-                        "sampleRate": 0.5,
+                        "sampleRate": 0.1,
                         "type": "trace",
                         "active": True,
                         "condition": {"op": "and", "inner": []},
@@ -240,13 +260,14 @@ def test_project_config_with_latest_release_in_dynamic_sampling_rules(default_pr
             },
             {
                 "rules": [
+                    DEFAULT_ENVIRONMENT_RULE,
                     {
                         "sampleRate": 0.1,
                         "type": "trace",
                         "active": True,
                         "condition": {"op": "and", "inner": []},
                         "id": 0,
-                    }
+                    },
                 ]
             },
         ),
@@ -262,13 +283,14 @@ def test_project_config_with_latest_release_in_dynamic_sampling_rules(default_pr
             {"rules": []},
             {
                 "rules": [
+                    DEFAULT_ENVIRONMENT_RULE,
                     {
                         "sampleRate": 0.1,
                         "type": "trace",
                         "active": True,
                         "condition": {"op": "and", "inner": []},
                         "id": 0,
-                    }
+                    },
                 ]
             },
         ),
@@ -288,9 +310,7 @@ def test_project_config_with_uniform_rules_based_on_plan_in_dynamic_sampling_rul
             "organizations:dynamic-sampling": ds_basic,
         }
     ):
-        with mock.patch(
-            "sentry.dynamic_sampling.utils.quotas.get_blended_sample_rate", return_value=0.1
-        ):
+        with mock.patch("sentry.dynamic_sampling.quotas.get_blended_sample_rate", return_value=0.1):
             cfg = get_project_config(default_project)
 
     cfg = cfg.to_dict()