From f664b7a8b68391dc0faf350527dbb1d94196ee90 Mon Sep 17 00:00:00 2001
From: Tony Xiao <txiao@sentry.io>
Date: Wed, 3 Nov 2021 11:48:19 -0400
Subject: [PATCH] feat(suspect-spans): Allow enabling ingestion by percentage
 (#2181)

Previously, suspect span ingestion can only be enabled on a per project basis.
This change allows us to enable it for a percentage of projects at a time. See
getsentry/getsentry#6531 and getsentry/sentry#29692  for the sentry PRs.
---
 snuba/datasets/transactions_processor.py     | 10 +++++++--
 snuba/state/__init__.py                      | 23 +++++++++++++++++++-
 tests/datasets/test_transaction_processor.py |  2 +-
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/snuba/datasets/transactions_processor.py b/snuba/datasets/transactions_processor.py
index 3ac9632aca1..871ca24ca38 100644
--- a/snuba/datasets/transactions_processor.py
+++ b/snuba/datasets/transactions_processor.py
@@ -27,7 +27,11 @@
     _ensure_valid_ip,
     _unicodify,
 )
-from snuba.state import get_config, is_project_in_rollout_group
+from snuba.state import (
+    get_config,
+    is_project_in_rollout_group,
+    is_project_in_rollout_list,
+)
 from snuba.utils.metrics.wrapper import MetricsWrapper
 
 logger = logging.getLogger(__name__)
@@ -299,8 +303,10 @@ def _process_spans(
             max_spans_per_transaction = 2000
 
         try:
-            if not is_project_in_rollout_group(
+            if not is_project_in_rollout_list(
                 "write_span_columns_projects", processed["project_id"]
+            ) and not is_project_in_rollout_group(
+                "write_span_columns_rollout_percentage", processed["project_id"]
             ):
                 return
 
diff --git a/snuba/state/__init__.py b/snuba/state/__init__.py
index 9223e769d3d..23541098847 100644
--- a/snuba/state/__init__.py
+++ b/snuba/state/__init__.py
@@ -258,8 +258,15 @@ def get_queries() -> Sequence[Mapping[str, Optional[Any]]]:
     return queries
 
 
-def is_project_in_rollout_group(rollout_key: str, project_id: int) -> bool:
+def is_project_in_rollout_list(rollout_key: str, project_id: int) -> bool:
+    """
+    Helper function for doing selective rollouts by project ids. The config
+    value is assumed to be a string of the form `[project,project,...]`.
+
+    Returns `True` if `project_id` is present in the config.
+    """
     project_rollout_setting = get_config(rollout_key, "")
+    assert isinstance(project_rollout_setting, str)
     if project_rollout_setting:
         # The expected format is [project,project,...]
         project_rollout_setting = project_rollout_setting[1:-1]
@@ -268,3 +275,17 @@ def is_project_in_rollout_group(rollout_key: str, project_id: int) -> bool:
                 if int(p.strip()) == project_id:
                     return True
     return False
+
+
+def is_project_in_rollout_group(rollout_key: str, project_id: int) -> bool:
+    """
+    Helper function for doing consistent incremental rollouts by project ids.
+    The config value is assumed to be an integer between 0 and 100.
+
+    Returns `True` if `project_id` falls within the rollout group.
+    """
+    project_rollout_percentage = get_config(rollout_key, 0)
+    assert isinstance(project_rollout_percentage, (int, float))
+    if project_rollout_percentage:
+        return project_id % 100 < project_rollout_percentage
+    return False
diff --git a/tests/datasets/test_transaction_processor.py b/tests/datasets/test_transaction_processor.py
index 5ae23201b5b..d5ea4c5752f 100644
--- a/tests/datasets/test_transaction_processor.py
+++ b/tests/datasets/test_transaction_processor.py
@@ -318,7 +318,7 @@ def test_missing_trace_context(self) -> None:
     def test_base_process(self) -> None:
         old_skip_context = settings.TRANSACT_SKIP_CONTEXT_STORE
         settings.TRANSACT_SKIP_CONTEXT_STORE = {1: {"experiments"}}
-        set_config("write_span_columns_projects", "[1]")
+        set_config("write_span_columns_rollout_percentage", 100)
 
         start, finish = self.__get_timestamps()
         message = TransactionEvent(