feat(dynamic-sampling): enhancing and extracting `get_sliding_window_…

…org_sample_rate` (#79593)

src/sentry/api/serializers/models/organization.py

@@ -54,7 +54,7 @@
 )
 from sentry.db.models.fields.slug import DEFAULT_SLUG_MAX_LENGTH
 from sentry.dynamic_sampling.tasks.common import get_organization_volume
-from sentry.dynamic_sampling.tasks.helpers.sliding_window import get_sliding_window_org_sample_rate
+from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate
 from sentry.killswitches import killswitch_matches_context
 from sentry.lang.native.utils import convert_crashreport_count
 from sentry.models.avatars.organization_avatar import OrganizationAvatar
@@ -644,9 +644,7 @@ def serialize(  # type: ignore[explicit-override, override]
         if sample_rate is not None:
             context["planSampleRate"] = sample_rate
 
-        desired_sample_rate, _ = get_sliding_window_org_sample_rate(
-            org_id=obj.id, default_sample_rate=sample_rate
-        )
+        desired_sample_rate, _ = get_org_sample_rate(org_id=obj.id, default_sample_rate=sample_rate)
         if desired_sample_rate is not None:
             context["desiredSampleRate"] = desired_sample_rate
 

src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py

@@ -49,7 +49,7 @@
 from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import (
     generate_boost_low_volume_projects_cache_key,
 )
-from sentry.dynamic_sampling.tasks.helpers.sliding_window import get_sliding_window_org_sample_rate
+from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate
 from sentry.dynamic_sampling.tasks.logging import log_sample_rate_source
 from sentry.dynamic_sampling.tasks.task_context import TaskContext
 from sentry.dynamic_sampling.tasks.utils import (
@@ -351,7 +351,7 @@ def adjust_sample_rates_of_projects(
 
     # If we have the sliding window org sample rate, we use that or fall back to the blended sample rate in case of
     # issues.
-    sample_rate, success = get_sliding_window_org_sample_rate(
+    sample_rate, success = get_org_sample_rate(
         org_id=org_id,
         default_sample_rate=quotas.backend.get_blended_sample_rate(organization_id=org_id),
     )

src/sentry/dynamic_sampling/tasks/helpers/sample_rate.py

@@ -0,0 +1,63 @@
+from sentry import features
+from sentry.constants import TARGET_SAMPLE_RATE_DEFAULT
+from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds
+from sentry.dynamic_sampling.tasks.helpers.sliding_window import (
+    generate_sliding_window_org_cache_key,
+)
+from sentry.models.options.organization_option import OrganizationOption
+from sentry.models.organization import Organization
+
+__all__ = ["get_org_sample_rate"]
+
+
+def get_org_sample_rate(
+    org_id: int, default_sample_rate: float | None
+) -> tuple[float | None, bool]:
+    """
+    Returns the organization sample rate for dynamic sampling. This returns either the
+    target_sample_rate organization option if custom dynamic sampling is enabled. Otherwise
+    it will fall back on retrieving the sample rate from the sliding window calculations.
+    """
+
+    # check if `organizations:dynamic-sampling-custom` feature flag is enabled for the
+    # organization, if yes, return the `sentry:target_sample_rate` option
+    try:
+        org = Organization.objects.get_from_cache(id=org_id)
+    except Organization.DoesNotExist:
+        org = None
+
+    has_dynamic_sampling_custom = features.has("organizations:dynamic-sampling-custom", org)
+    if has_dynamic_sampling_custom:
+        try:
+            option_inst = OrganizationOption.objects.get(
+                organization=org, key="sentry:target_sample_rate"
+            )
+        except OrganizationOption.DoesNotExist:
+            option_inst = None
+
+        if option_inst is None or option_inst.value is None:
+            if default_sample_rate is not None:
+                return default_sample_rate, False
+            return TARGET_SAMPLE_RATE_DEFAULT, False
+
+        return float(option_inst.value), True
+
+    # fallback to sliding window calculation
+    return _get_sliding_window_org_sample_rate(org_id, default_sample_rate)
+
+
+def _get_sliding_window_org_sample_rate(
+    org_id: int, default_sample_rate: float | None
+) -> tuple[float | None, bool]:
+    redis_client = get_redis_client_for_ds()
+    cache_key = generate_sliding_window_org_cache_key(org_id)
+
+    try:
+        value = redis_client.get(cache_key)
+
+        if value is not None:
+            return float(value), True
+
+        return default_sample_rate, False
+    except (TypeError, ValueError):
+        return default_sample_rate, False

src/sentry/dynamic_sampling/tasks/helpers/sliding_window.py

@@ -37,23 +37,6 @@ def generate_sliding_window_org_cache_key(org_id: int) -> str:
     return f"ds::o:{org_id}:sliding_window_org_sample_rate"
 
 
-def get_sliding_window_org_sample_rate(
-    org_id: int, default_sample_rate: float | None
-) -> tuple[float | None, bool]:
-    redis_client = get_redis_client_for_ds()
-    cache_key = generate_sliding_window_org_cache_key(org_id)
-
-    try:
-        value = redis_client.get(cache_key)
-
-        if value is not None:
-            return float(value), True
-
-        return default_sample_rate, False
-    except (TypeError, ValueError):
-        return default_sample_rate, False
-
-
 def get_sliding_window_size() -> int | None:
     try:
         size = options.get("dynamic-sampling:sliding_window.size")

src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py

@@ -15,7 +15,7 @@
     get_adjusted_factor,
     set_guarded_adjusted_factor,
 )
-from sentry.dynamic_sampling.tasks.helpers.sliding_window import get_sliding_window_org_sample_rate
+from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate
 from sentry.dynamic_sampling.tasks.logging import log_sample_rate_source
 from sentry.dynamic_sampling.tasks.task_context import TaskContext
 from sentry.dynamic_sampling.tasks.utils import (
@@ -88,7 +88,7 @@ def recalibrate_org(org_id: int, total: int, indexed: int) -> None:
 
     # If we have the sliding window org sample rate, we use that or fall back to the blended sample rate in case of
     # issues.
-    target_sample_rate, success = get_sliding_window_org_sample_rate(
+    target_sample_rate, success = get_org_sample_rate(
         org_id=org_id,
         default_sample_rate=quotas.backend.get_blended_sample_rate(organization_id=org_id),
     )

tests/sentry/dynamic_sampling/tasks/helpers/test_sample_rate.py

@@ -0,0 +1,34 @@
+from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate
+from sentry.models.options.organization_option import OrganizationOption
+from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase
+from sentry.testutils.helpers.features import with_feature
+
+
+class PrioritiseProjectsSnubaQueryTest(BaseMetricsLayerTestCase, TestCase, SnubaTestCase):
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_get_org_sample_rate_from_target_sample_rate(self):
+        org1 = self.create_organization("test-org")
+
+        OrganizationOption.objects.create(
+            organization=org1, key="sentry:target_sample_rate", value=0.5
+        )
+
+        sample_rate, success = get_org_sample_rate(org1.id, None)
+        assert success
+        assert sample_rate == 0.5
+
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_get_org_sample_rate_from_target_sample_rate_missing(self):
+        org1 = self.create_organization("test-org")
+
+        sample_rate, success = get_org_sample_rate(org1.id, None)
+        assert not success
+        assert sample_rate == 1.0
+
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_get_org_sample_rate_from_target_sample_rate_missing_default(self):
+        org1 = self.create_organization("test-org")
+
+        sample_rate, success = get_org_sample_rate(org1.id, 0.7)
+        assert not success
+        assert sample_rate == 0.7

tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py

@@ -16,6 +16,7 @@
     generate_sliding_window_org_cache_key,
 )
 from sentry.dynamic_sampling.tasks.task_context import TaskContext
+from sentry.models.options.organization_option import OrganizationOption
 from sentry.models.organization import Organization
 from sentry.models.project import Project
 from sentry.snuba.metrics.naming_layer.mri import TransactionMRI
@@ -62,7 +63,7 @@ def test_simple_one_org_one_project(self):
         assert results[org1.id] == [(p1.id, 4.0, 1, 3)]
 
     @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
-    def test_simple_one_org_one_project_task(self):
+    def test_simple_one_org_one_project_task_sliding_window_sample_rate(self):
         org1 = self.create_organization("test-org")
         p1 = self.create_project(organization=org1)
 
@@ -99,6 +100,43 @@ def test_simple_one_org_one_project_task(self):
         assert got_value
         assert sample_rate == 1.0
 
+    @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"])
+    def test_simple_one_org_one_project_task_target_sample_rate(self):
+        org1 = self.create_organization("test-org")
+        p1 = self.create_project(organization=org1)
+
+        OrganizationOption.objects.create(
+            organization=org1, key="sentry:target_sample_rate", value=0.5
+        )
+
+        self.store_performance_metric(
+            name=TransactionMRI.COUNT_PER_ROOT_PROJECT.value,
+            tags={"transaction": "foo_transaction", "decision": "keep"},
+            minutes_before_now=30,
+            value=1,
+            project_id=p1.id,
+            org_id=org1.id,
+        )
+
+        self.store_performance_metric(
+            name=TransactionMRI.COUNT_PER_ROOT_PROJECT.value,
+            tags={"transaction": "foo_transaction", "decision": "drop"},
+            minutes_before_now=30,
+            value=3,
+            project_id=p1.id,
+            org_id=org1.id,
+        )
+
+        with self.tasks():
+            boost_low_volume_projects_of_org_with_query.delay(org1.id)
+
+        sample_rate, got_value = get_boost_low_volume_projects_sample_rate(
+            org1.id, p1.id, error_sample_rate_fallback=None
+        )
+
+        assert got_value
+        assert sample_rate == 0.5
+
     def test_complex(self):
         context = TaskContext("rebalancing", 20)
         org1 = self.create_organization("test-org1")