ref(dyanmic_sampling): Add more intstrumentation to sliding window org (#53450)

Author: RaduW

src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py

@@ -114,7 +114,10 @@ def boost_low_volume_projects_of_org(
         Tuple[ProjectId, int, DecisionKeepCount, DecisionDropCount]
     ],
 ) -> None:
-    adjust_sample_rates_of_projects(org_id, projects_with_tx_count_and_rates)
+    # secondary tasks should not log the context, I need the context only for calling
+    # `adjust_sample_rates_of_projects`, the accumulated info will be ignored.
+    context = TaskContext("not_used", MAX_TASK_SECONDS)
+    adjust_sample_rates_of_projects(org_id, projects_with_tx_count_and_rates, context)
 
 
 def fetch_projects_with_total_root_transaction_count_and_rates(
@@ -239,6 +242,7 @@ def fetch_projects_with_total_root_transaction_count_and_rates(
 def adjust_sample_rates_of_projects(
     org_id: int,
     projects_with_tx_count: Sequence[Tuple[ProjectId, int, DecisionKeepCount, DecisionDropCount]],
+    context: TaskContext,
 ) -> None:
     """
     Adjusts the sample rates of projects belonging to a specific org.
@@ -253,7 +257,7 @@ def adjust_sample_rates_of_projects(
 
     # We get the sample rate either directly from quotas or from the new sliding window org mechanism.
     if organization is not None and is_sliding_window_org_enabled(organization):
-        sample_rate = get_adjusted_base_rate_from_cache_or_compute(org_id)
+        sample_rate = get_adjusted_base_rate_from_cache_or_compute(org_id, context)
         log_sample_rate_source(
             org_id, None, "boost_low_volume_projects", "sliding_window_org", sample_rate
         )

src/sentry/dynamic_sampling/tasks/common.py

@@ -583,7 +583,11 @@ def are_equal_with_epsilon(a: Optional[float], b: Optional[float]) -> bool:
 
 
 def compute_guarded_sliding_window_sample_rate(
-    org_id: int, project_id: Optional[int], total_root_count: int, window_size: int
+    org_id: int,
+    project_id: Optional[int],
+    total_root_count: int,
+    window_size: int,
+    context: TaskContext,
 ) -> Optional[float]:
     """
     Computes the actual sliding window sample rate by guarding any exceptions and returning None in case
@@ -592,14 +596,16 @@ def compute_guarded_sliding_window_sample_rate(
     try:
         # We want to compute the sliding window sample rate by considering a window of time.
         # This piece of code is very delicate, thus we want to guard it properly and capture any errors.
-        return compute_sliding_window_sample_rate(org_id, project_id, total_root_count, window_size)
+        return compute_sliding_window_sample_rate(
+            org_id, project_id, total_root_count, window_size, context
+        )
     except Exception as e:
         sentry_sdk.capture_exception(e)
         return None
 
 
 def compute_sliding_window_sample_rate(
-    org_id: int, project_id: Optional[int], total_root_count: int, window_size: int
+    org_id: int, project_id: Optional[int], total_root_count: int, window_size: int, context
 ) -> Optional[float]:
     """
     Computes the actual sample rate for the sliding window given the total root count and the size of the
@@ -608,7 +614,8 @@ def compute_sliding_window_sample_rate(
     The org_id is used only because it is required on the quotas side to determine whether dynamic sampling is
     enabled in the first place for that project.
     """
-    extrapolated_volume = extrapolate_monthly_volume(volume=total_root_count, hours=window_size)
+    with context.get_timer("extrapolate_monthly_volume"):
+        extrapolated_volume = extrapolate_monthly_volume(volume=total_root_count, hours=window_size)
     if extrapolated_volume is None:
         with sentry_sdk.push_scope() as scope:
             scope.set_extra("org_id", org_id)
@@ -622,9 +629,10 @@ def compute_sliding_window_sample_rate(
         org_id, project_id, total_root_count, extrapolated_volume, window_size
     )
 
-    sampling_tier = quotas.get_transaction_sampling_tier_for_volume(  # type:ignore
-        org_id, extrapolated_volume
-    )
+    with context.get_timer("get_transaction_sampling_tier_for_volume"):
+        sampling_tier = quotas.get_transaction_sampling_tier_for_volume(  # type:ignore
+            org_id, extrapolated_volume
+        )
     if sampling_tier is None:
         return None
 
@@ -636,7 +644,9 @@ def compute_sliding_window_sample_rate(
     return float(sample_rate)
 
 
-def get_adjusted_base_rate_from_cache_or_compute(org_id: int) -> Optional[float]:
+def get_adjusted_base_rate_from_cache_or_compute(
+    org_id: int, context: TaskContext
+) -> Optional[float]:
     """
     Gets the adjusted base sample rate from the sliding window directly from the Redis cache or tries to compute
     it synchronously.
@@ -656,7 +666,11 @@ def get_adjusted_base_rate_from_cache_or_compute(org_id: int) -> Optional[float]
         )
         if (org_total_root_count := orgs_with_counts.get(org_id)) is not None:
             return compute_guarded_sliding_window_sample_rate(
-                org_id, None, org_total_root_count, window_size
+                org_id,
+                None,
+                org_total_root_count,
+                window_size,
+                context,
             )
 
     return None

src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py

@@ -114,7 +114,9 @@ def recalibrate_org(org_volume: OrganizationDataVolume, context: TaskContext) ->
             "ready_for_recalibration", {"org_id": org_volume.org_id}, orgs_to_check(org_volume)
         )
 
-        target_sample_rate = get_adjusted_base_rate_from_cache_or_compute(org_volume.org_id)
+        target_sample_rate = get_adjusted_base_rate_from_cache_or_compute(
+            org_volume.org_id, context
+        )
         log_sample_rate_source(
             org_volume.org_id, None, "recalibrate_orgs", "sliding_window_org", target_sample_rate
         )

src/sentry/dynamic_sampling/tasks/sliding_window.py

@@ -133,7 +133,11 @@ def adjust_base_sample_rates_of_projects(
 
         for project_id, total_root_count in projects_with_total_root_count:
             sample_rate = compute_guarded_sliding_window_sample_rate(
-                org_id, project_id, total_root_count, window_size
+                org_id,
+                project_id,
+                total_root_count,
+                window_size,
+                context,
             )
 
             # If the sample rate is None, we want to add a sentinel value into Redis, the goal being that when generating

src/sentry/dynamic_sampling/tasks/sliding_window_org.py

@@ -83,25 +83,31 @@ def adjust_base_sample_rate_of_org(
     if time.monotonic() > context.expiration_time:
         raise TimeoutException(context)
 
-    name = adjust_base_sample_rate_of_org.__name__
-    timer = context.get_timer(name)
+    func_name = adjust_base_sample_rate_of_org.__name__
+    timer = context.get_timer(func_name)
     with timer:
-        sample_rate = compute_guarded_sliding_window_sample_rate(
-            org_id, None, total_root_count, window_size
-        )
+        with context.get_timer(compute_guarded_sliding_window_sample_rate.__name__):
+            sample_rate = compute_guarded_sliding_window_sample_rate(
+                org_id,
+                None,
+                total_root_count,
+                window_size,
+                context,
+            )
         # If the sample rate is None, we don't want to store a value into Redis, but we prefer to keep the system
         # with the old value.
         if sample_rate is None:
             return
 
-        redis_client = get_redis_client_for_ds()
-        with redis_client.pipeline(transaction=False) as pipeline:
-            cache_key = generate_sliding_window_org_cache_key(org_id=org_id)
-            pipeline.set(cache_key, sample_rate)
-            pipeline.pexpire(cache_key, DEFAULT_REDIS_CACHE_KEY_TTL)
-            pipeline.execute()
+        with context.get_timer("redis_updates"):
+            redis_client = get_redis_client_for_ds()
+            with redis_client.pipeline(transaction=False) as pipeline:
+                cache_key = generate_sliding_window_org_cache_key(org_id=org_id)
+                pipeline.set(cache_key, sample_rate)
+                pipeline.pexpire(cache_key, DEFAULT_REDIS_CACHE_KEY_TTL)
+                pipeline.execute()
 
-    state = context.get_function_state(name)
+    state = context.get_function_state(func_name)
     state.num_orgs += 1
     state.num_iterations += 1
-    context.set_function_state(name, state)
+    context.set_function_state(func_name, state)