From 7543a49c713c75277af3efeaf324c53e5a9c74e9 Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 10 Sep 2026 10:53:13 +0200 Subject: [PATCH 1/2] ref(dynamic-sampling): Delete the killswitched legacy dynamic sampling jobs The four scheduled legacy jobs (boost_low_volume_projects, boost_low_volume_transactions, recalibrate_orgs, sliding_window_org) have been killswitched since the per-org pipeline took over serving. Delete them, their fan-out tasks, their schedule entries and their SDK sampling entries, together with the generic-metrics queries only they used. The organization details endpoint was the last caller of the legacy project balancing. On a target sample rate change it now schedules the per-org calculation for that organization, and on a switch to project mode it seeds the project rates from 30 days of EAP volume balanced by the per-org model, so both paths stay off generic metrics. The transaction volume query stays in place for now, since the per-org query module still imports it on master; a follow-up removes the remaining generic-metrics queries. The options the deleted jobs read stay registered until the options automator has unset them. Co-Authored-By: Claude Fable 5.1 --- src/sentry/conf/server.py | 20 - .../core/endpoints/organization_details.py | 60 +- .../dynamic_sampling/per_org/calculations.py | 6 +- .../tasks/boost_low_volume_projects.py | 502 -------- .../tasks/boost_low_volume_transactions.py | 448 +------ src/sentry/dynamic_sampling/tasks/common.py | 159 +-- .../dynamic_sampling/tasks/constants.py | 2 - .../tasks/recalibrate_orgs.py | 215 ---- .../tasks/sliding_window_org.py | 96 -- src/sentry/dynamic_sampling/tasks/utils.py | 55 - src/sentry/options/defaults.py | 15 +- src/sentry/snuba/referrer.py | 6 - src/sentry/utils/sdk.py | 4 - .../endpoints/test_organization_details.py | 52 +- .../tasks/test_boost_low_volume_projects.py | 718 ----------- .../test_boost_low_volume_transactions.py | 264 +---- .../dynamic_sampling/tasks/test_common.py | 158 --- .../dynamic_sampling/tasks/test_tasks.py | 1054 ----------------- 18 files changed, 67 insertions(+), 3767 deletions(-) delete mode 100644 src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py delete mode 100644 src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py delete mode 100644 src/sentry/dynamic_sampling/tasks/sliding_window_org.py delete mode 100644 src/sentry/dynamic_sampling/tasks/utils.py delete mode 100644 tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py delete mode 100644 tests/sentry/dynamic_sampling/tasks/test_tasks.py diff --git a/src/sentry/conf/server.py b/src/sentry/conf/server.py index 90239ed1c88f..fcf74b91a808 100644 --- a/src/sentry/conf/server.py +++ b/src/sentry/conf/server.py @@ -881,10 +881,6 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: "sentry.demo_mode.tasks", "sentry.dynamic_sampling.per_org.feature_cache", "sentry.dynamic_sampling.per_org.scheduler", - "sentry.dynamic_sampling.tasks.boost_low_volume_projects", - "sentry.dynamic_sampling.tasks.boost_low_volume_transactions", - "sentry.dynamic_sampling.tasks.recalibrate_orgs", - "sentry.dynamic_sampling.tasks.sliding_window_org", "sentry.feedback.tasks.update_user_reports", "sentry.hybridcloud.tasks.deliver_from_outbox", "sentry.hybridcloud.tasks.deliver_webhooks", @@ -1144,22 +1140,6 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: "task": "performance:sentry.ingest.transaction_clusterer.tasks.spawn_clusterers", "schedule": crontab("17", "*", "*", "*", "*"), }, - "dynamic-sampling-boost-low-volume-projects": { - "task": "telemetry-experience:sentry.dynamic_sampling.tasks.boost_low_volume_projects", - "schedule": crontab("*/10", "*", "*", "*", "*"), - }, - "dynamic-sampling-boost-low-volume-transactions": { - "task": "telemetry-experience:sentry.dynamic_sampling.tasks.boost_low_volume_transactions", - "schedule": crontab("*/10", "*", "*", "*", "*"), - }, - "dynamic-sampling-recalibrate-orgs": { - "task": "telemetry-experience:sentry.dynamic_sampling.tasks.recalibrate_orgs", - "schedule": crontab("*/10", "*", "*", "*", "*"), - }, - "dynamic-sampling-sliding-window-org": { - "task": "telemetry-experience:sentry.dynamic_sampling.tasks.sliding_window_org", - "schedule": crontab("*/10", "*", "*", "*", "*"), - }, "dynamic-sampling-schedule-per-org-calculations": { "task": "telemetry-experience:sentry.dynamic_sampling.per_org.schedule_per_org_calculations", "schedule": timedelta(seconds=10), diff --git a/src/sentry/core/endpoints/organization_details.py b/src/sentry/core/endpoints/organization_details.py index ec45db4b76fc..4eb32300db64 100644 --- a/src/sentry/core/endpoints/organization_details.py +++ b/src/sentry/core/endpoints/organization_details.py @@ -73,16 +73,16 @@ SEER_AUTOMATED_RUN_STOPPING_POINT_DEFAULT, SEER_DEFAULT_CODING_AGENT_DEFAULT, TARGET_SAMPLE_RATE_DEFAULT, - ObjectStatus, ) from sentry.core.endpoints.project_details import MAX_SENSITIVE_FIELD_CHARS from sentry.deletions.models.scheduleddeletion import CellScheduledDeletion -from sentry.dynamic_sampling.tasks.boost_low_volume_projects import ( - boost_low_volume_projects_of_org_with_query, - calculate_sample_rates_of_projects, - query_project_counts_by_org, +from sentry.dynamic_sampling.per_org.calculations import run_project_balancing +from sentry.dynamic_sampling.per_org.configuration import ( + CustomDynamicSamplingOrganizationConfiguration, ) -from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure +from sentry.dynamic_sampling.per_org.queries import get_eap_project_volumes +from sentry.dynamic_sampling.per_org.scheduler import run_calculations_per_org_task_entry +from sentry.dynamic_sampling.types import DynamicSamplingMode from sentry.dynamic_sampling.utils import ( has_custom_dynamic_sampling, is_organization_mode_sampling, @@ -101,7 +101,6 @@ from sentry.models.options.project_option import ProjectOption from sentry.models.organization import Organization, OrganizationStatus from sentry.models.organizationmember import OrganizationMember -from sentry.models.project import Project from sentry.organizations.services.organization import organization_service from sentry.organizations.services.organization.model import ( RpcOrganization, @@ -1218,7 +1217,7 @@ def put( if "samplingMode" in changed_data: with transaction.atomic(router.db_for_write(ProjectOption)): if is_project_mode_sampling(organization): - self._compute_project_target_sample_rates(request, organization) + self._compute_project_target_sample_rates(organization) organization.delete_option("sentry:target_sample_rate") changed_data["samplingMode"] = "to Advanced Mode" @@ -1246,9 +1245,7 @@ def put( if is_org_mode and ( "samplingMode" in changed_data or "targetSampleRate" in changed_data ): - boost_low_volume_projects_of_org_with_query.delay( - organization.id, - ) + run_calculations_per_org_task_entry.delay(organization.id) if is_org_mode and "defaultAutofixAutomationTuning" in changed_data: organization.update_option( @@ -1314,36 +1311,23 @@ def put( return self.respond(context) return self.respond(as_validation_errors(serializer), status=status.HTTP_400_BAD_REQUEST) - def _compute_project_target_sample_rates(self, request: Request, organization: Organization): + def _compute_project_target_sample_rates(self, organization: Organization) -> None: + """Seed every active project with the rate it had under organization mode. + + Balances the last 30 days of volume at the organization's target rate, which is + still set when this runs, so that the switch to project mode keeps the stored + volume per project unchanged. + """ # TODO: this will take a long time for organizations with a lot of projects # so we need to refactor this into an async task we can run and observe - org_id = organization.id - measure = SamplingMeasure.SEGMENTS - projects_with_tx_count_and_rates = [] - for chunk in query_project_counts_by_org( - [org_id], measure, query_interval=timedelta(days=30) - ): - for row in chunk: - projects_with_tx_count_and_rates.append(row[1:]) - - rebalanced_projects = calculate_sample_rates_of_projects( - org_id, projects_with_tx_count_and_rates - ) - - project_ids = set( - Project.objects.filter(organization_id=org_id, status=ObjectStatus.ACTIVE).values_list( - "id", flat=True + config = CustomDynamicSamplingOrganizationConfiguration(organization) + project_volumes = get_eap_project_volumes(config, time_interval=timedelta(days=30)) + for rebalanced_item in run_project_balancing(config, project_volumes): + ProjectOption.objects.update_or_create( + project_id=rebalanced_item.id, + key="sentry:target_sample_rate", + defaults={"value": round(rebalanced_item.new_sample_rate, 4)}, ) - ) - - if rebalanced_projects is not None: - for rebalanced_item in rebalanced_projects: - if int(rebalanced_item.id) in project_ids: - ProjectOption.objects.update_or_create( - project_id=rebalanced_item.id, - key="sentry:target_sample_rate", - defaults={"value": round(rebalanced_item.new_sample_rate, 4)}, - ) def handle_delete(self, request: Request, organization: Organization): """ diff --git a/src/sentry/dynamic_sampling/per_org/calculations.py b/src/sentry/dynamic_sampling/per_org/calculations.py index 1ed90a09b3e2..94ac349e58c5 100644 --- a/src/sentry/dynamic_sampling/per_org/calculations.py +++ b/src/sentry/dynamic_sampling/per_org/calculations.py @@ -145,10 +145,8 @@ def run_transaction_balancing( "its transactions" ) continue - # Mirror the legacy pipeline (boost_low_volume_transactions_of_project): at a 100% - # project rate every transaction is kept anyway, so the legacy task skips the model - # and writes no cache entry. Skipping here keeps parity and avoids comparison log - # lines that would only ever hit cache misses. + # At a 100% project rate every transaction is kept anyway, so there is nothing to + # balance and no cache entry to write. if sample_rate == 1.0: continue named_rates, implicit_rate = TransactionsRebalancingModel().run( diff --git a/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py b/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py deleted file mode 100644 index babb580014a7..000000000000 --- a/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py +++ /dev/null @@ -1,502 +0,0 @@ -from __future__ import annotations - -import logging -from collections import defaultdict -from collections.abc import Iterator, Mapping, Sequence -from datetime import timedelta - -import sentry_sdk -from snuba_sdk import ( - Column, - Condition, - Direction, - Entity, - Function, - Granularity, - Limit, - LimitBy, - Op, - OrderBy, - Query, - Request, -) -from taskbroker_client.retry import Retry - -from sentry import quotas -from sentry.constants import ObjectStatus -from sentry.dynamic_sampling.models.common import RebalancedItem, guarded_run -from sentry.dynamic_sampling.models.projects_rebalancing import ( - ProjectsRebalancingInput, - ProjectsRebalancingModel, -) -from sentry.dynamic_sampling.rules.utils import ( - DecisionDropCount, - DecisionKeepCount, - OrganizationId, - ProjectId, - get_redis_client_for_ds, -) -from sentry.dynamic_sampling.tasks.common import ( - MEASURE_CONFIGS, - GetActiveOrgs, - are_equal_with_epsilon, - sample_rate_to_float, -) -from sentry.dynamic_sampling.tasks.constants import ( - CHUNK_SIZE, - DEFAULT_REDIS_CACHE_KEY_TTL, - MAX_PROJECTS_PER_QUERY, - MAX_TRANSACTIONS_PER_PROJECT, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( - generate_boost_low_volume_projects_cache_key, -) -from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate -from sentry.dynamic_sampling.tasks.utils import ( - dynamic_sampling_task, - legacy_pipeline_killswitched, -) -from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure -from sentry.dynamic_sampling.utils import has_dynamic_sampling, is_project_mode_sampling -from sentry.models.options import OrganizationOption -from sentry.models.organization import Organization -from sentry.models.project import Project -from sentry.sentry_metrics import indexer -from sentry.silo.base import SiloMode -from sentry.snuba.dataset import Dataset, EntityKey -from sentry.snuba.referrer import Referrer -from sentry.tasks.base import instrumented_task -from sentry.tasks.relay import schedule_invalidate_project_config -from sentry.taskworker.namespaces import telemetry_experience_tasks -from sentry.utils import metrics -from sentry.utils.dates import deprecated_utcnow -from sentry.utils.snuba import raw_snql_query - -# This set contains all the projects for which we want to start extracting the sample rate over time. This is done -# as a temporary solution to dogfood our own product without exploding the cardinality of the project_id tag. -PROJECTS_WITH_METRICS = {1, 11276} # sentry # javascript -logger = logging.getLogger(__name__) - -# a tuple type alias of project_id, root_count, keep_count, drop_count, to be used in extraction of metrics for a specific project -ProjectVolumes = tuple[ProjectId, int, DecisionKeepCount, DecisionDropCount] - -# the same as ProjectVolumes, but with the organization ID added -OrgProjectVolumes = tuple[OrganizationId, ProjectId, int, DecisionKeepCount, DecisionDropCount] - - -@metrics.wraps("dynamic_sampling.filter_project_mode_orgs") -def _without_project_mode_orgs(org_ids: list[int]) -> list[int]: - """ - Drop the organizations that sample per project. Their rates come from - project options, so the rebalancing task must not overwrite them. - """ - modes_per_org = OrganizationOption.objects.get_value_bulk_id(org_ids, "sentry:sampling_mode") - return sorted( - org_id for org_id, mode in modes_per_org.items() if mode != DynamicSamplingMode.PROJECT - ) - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.boost_low_volume_projects", - namespace=telemetry_experience_tasks, - processing_deadline_duration=20 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def boost_low_volume_projects() -> None: - """ - Task to adjusts the sample rates of all projects in all active organizations. - """ - if legacy_pipeline_killswitched("boost_low_volume_projects"): - return - - for orgs in GetActiveOrgs( - max_projects=MAX_PROJECTS_PER_QUERY, - granularity=Granularity(60), - measure=SamplingMeasure.SEGMENTS, - ): - _process_orgs_for_boost(_without_project_mode_orgs(orgs), SamplingMeasure.SEGMENTS) - - -def _process_orgs_for_boost( - org_ids: list[int], - measure: SamplingMeasure, -) -> None: - """ - Process organizations for boost_low_volume_projects. - - Dispatches to the per-org task for each org with project volume data. - """ - if not org_ids: - return - - metrics.incr( - "dynamic_sampling.boost_low_volume_projects.orgs_processed", - amount=len(org_ids), - tags={"measure": str(measure.value)}, - ) - - for org_id, projects in fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=org_ids, measure=measure - ).items(): - boost_low_volume_projects_of_org.apply_async( - kwargs={ - "org_id": org_id, - "projects_with_tx_count_and_rates": projects, - }, - headers={"sentry-propagate-traces": False}, - ) - - -@instrumented_task( - name="sentry.dynamic_sampling.boost_low_volume_projects_of_org_with_query", - namespace=telemetry_experience_tasks, - processing_deadline_duration=3 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def boost_low_volume_projects_of_org_with_query(org_id: OrganizationId) -> None: - """ - Task to adjust the sample rates of the projects of a single organization specified by an - organization ID. Transaction counts and rates are fetched within this task. - """ - logger.info( - "boost_low_volume_projects_of_org_with_query", - extra={"traceparent": sentry_sdk.get_traceparent(), "baggage": sentry_sdk.get_baggage()}, - ) - - org = Organization.objects.get_from_cache(id=org_id) - if is_project_mode_sampling(org): - return - - projects_with_tx_count_and_rates = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org_id], - measure=SamplingMeasure.SEGMENTS, - )[org_id] - rebalanced_projects = calculate_sample_rates_of_projects( - org_id, projects_with_tx_count_and_rates - ) - if rebalanced_projects is not None: - store_rebalanced_projects(org_id, rebalanced_projects) - - -@instrumented_task( - name="sentry.dynamic_sampling.boost_low_volume_projects_of_org", - namespace=telemetry_experience_tasks, - processing_deadline_duration=3 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def boost_low_volume_projects_of_org( - org_id: OrganizationId, - projects_with_tx_count_and_rates: Sequence[ProjectVolumes], -) -> None: - """ - Task to adjust the sample rates of the projects of a single organization specified by an - organization ID. Transaction counts and rates have to be provided. - """ - - try: - rebalanced_projects = calculate_sample_rates_of_projects( - org_id, projects_with_tx_count_and_rates - ) - except Exception as e: - sentry_sdk.capture_exception(e) - raise - - logger.info( - "boost_low_volume_projects_of_org", - extra={ - "traceparent": sentry_sdk.get_traceparent(), - "baggage": sentry_sdk.get_baggage(), - "org_id": org_id, - }, - ) - if rebalanced_projects is not None: - store_rebalanced_projects(org_id, rebalanced_projects) - metrics.incr( - "dynamic_sampling.boost_low_volume_projects_of_org.success", - tags={"type": "rebalanced"}, - sample_rate=1, - ) - else: - metrics.incr( - "dynamic_sampling.boost_low_volume_projects_of_org.success", - tags={"type": "not_rebalanced"}, - sample_rate=1, - ) - - -@metrics.wraps("dynamic_sampling.fetch_projects_with_total_root_transaction_count_and_rates") -def fetch_projects_with_total_root_transaction_count_and_rates( - org_ids: list[int], - measure: SamplingMeasure, - query_interval: timedelta | None = None, -) -> Mapping[OrganizationId, Sequence[ProjectVolumes]]: - """ - Fetches for each org and each project the total root transaction count and how many transactions were kept and - dropped. - """ - aggregated_projects = defaultdict(list) - project_count_query_iter = query_project_counts_by_org(org_ids, measure, query_interval) - for chunk in project_count_query_iter: - for org_id, project_id, root_count_value, keep_count, drop_count in chunk: - aggregated_projects[org_id].append( - ( - project_id, - root_count_value, - keep_count, - drop_count, - ) - ) - - return aggregated_projects - - -@dynamic_sampling_task -def query_project_counts_by_org( - org_ids: list[int], measure: SamplingMeasure, query_interval: timedelta | None = None -) -> Iterator[Sequence[OrgProjectVolumes]]: - """Queries the total root transaction count and how many transactions were kept and dropped - for each project in a given interval (defaults to the last hour). - - Yields chunks of result rows, to allow timeouts to be handled in the caller. - """ - if not org_ids: - return - - if query_interval is None: - query_interval = timedelta(hours=1) - - if query_interval > timedelta(days=1): - granularity = Granularity(24 * 3600) - else: - granularity = Granularity(60) - - metrics.incr( - "dynamic_sampling.query_project_counts_by_org.count", - amount=len(org_ids), - tags={"measure": str(measure.value)}, - ) - - org_ids = list(org_ids) - project_ids = list( - Project.objects.filter(organization_id__in=org_ids, status=ObjectStatus.ACTIVE).values_list( - "id", flat=True - ) - ) - decision_string_id = indexer.resolve_shared_org("decision") - decision_tag = f"tags_raw[{decision_string_id}]" - - config = MEASURE_CONFIGS.get(measure) - if config is None: - raise ValueError(f"Unsupported measure: {measure}") - - metric_id = indexer.resolve_shared_org(str(config["mri"])) - use_case_id = config["use_case_id"] - - where_conditions = [ - Condition(Column("timestamp"), Op.GTE, deprecated_utcnow() - query_interval), - Condition(Column("timestamp"), Op.LT, deprecated_utcnow()), - Condition(Column("metric_id"), Op.EQ, metric_id), - Condition(Column("org_id"), Op.IN, org_ids), - Condition(Column("project_id"), Op.IN, project_ids), - ] - - # Add tag filters from config - for tag_name, tag_value in config["tags"].items(): - tag_string_id = indexer.resolve_shared_org(tag_name) - tag_column = f"tags_raw[{tag_string_id}]" - where_conditions.append(Condition(Column(tag_column), Op.EQ, tag_value)) - - query = Query( - match=Entity(EntityKey.GenericOrgMetricsCounters.value), - select=[ - Function("sum", [Column("value")], "root_count_value"), - Column("org_id"), - Column("project_id"), - Function( - "sumIf", - [ - Column("value"), - Function("equals", [Column(decision_tag), "keep"]), - ], - alias="keep_count", - ), - Function( - "sumIf", - [ - Column("value"), - Function("equals", [Column(decision_tag), "drop"]), - ], - alias="drop_count", - ), - ], - groupby=[Column("org_id"), Column("project_id")], - where=where_conditions, - granularity=granularity, - orderby=[ - OrderBy(Column("org_id"), Direction.ASC), - OrderBy(Column("project_id"), Direction.ASC), - ], - limitby=LimitBy( - columns=[Column("org_id"), Column("project_id")], - count=MAX_TRANSACTIONS_PER_PROJECT, - ), - # we are fetching one more than the chunk size to determine if there are more results - limit=Limit(CHUNK_SIZE + 1), - ) - - offset = 0 - more_results: bool = True - while more_results: - with metrics.timer( - "dynamic_sampling.query_project_counts_by_org.query_time", - tags={"measure": str(measure.value)}, - ): - request = Request( - dataset=Dataset.PerformanceMetrics.value, - app_id="dynamic_sampling", - query=query.set_offset(offset), - tenant_ids={"use_case_id": use_case_id.value, "cross_org_query": 1}, - ) - data = raw_snql_query( - request, - referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value, - )["data"] - - more_results = len(data) > CHUNK_SIZE - offset += CHUNK_SIZE - - # re-adjust, for the extra row we fetched - if more_results: - data = data[:-1] - - yield [ - ( - row["org_id"], - row["project_id"], - row["root_count_value"], - row["keep_count"], - row["drop_count"], - ) - for row in data - ] - - -@dynamic_sampling_task -def calculate_sample_rates_of_projects( - org_id: int, - projects_with_tx_count: Sequence[ProjectVolumes], -) -> list[RebalancedItem] | None: - """ - Calculates the sample rates of projects belonging to a specific org. - """ - try: - # We need the organization object for the feature flag. - organization = Organization.objects.get_from_cache(id=org_id) - except Organization.DoesNotExist: - # In case an org is not found, it might be that it has been deleted in the time between - # the query triggering this job and the actual execution of the job. - organization = None - - # If the org doesn't have dynamic sampling, we want to early return to avoid unnecessary work. - if not has_dynamic_sampling(organization): - return 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. - - default_sample_rate = quotas.backend.get_blended_sample_rate(organization_id=org_id) - sample_rate, success = get_org_sample_rate( - org_id=org_id, - default_sample_rate=default_sample_rate, - ) - - # If we didn't find any sample rate, it doesn't make sense to run the adjustment model. - if sample_rate is None: - sentry_sdk.capture_message( - "Sample rate of org not found when trying to adjust the sample rates of its projects" - ) - return None - - projects_with_counts = { - project_id: count_per_root for project_id, count_per_root, _, _ in projects_with_tx_count - } - - # The rebalancing will not work (or would make sense) when we have only projects with zero-counts. - if not any(projects_with_counts.values()): - return None - - # Since we don't mind about strong consistency, we query a replica of the main database with the possibility of - # having out of date information. This is a trade-off we accept, since we work under the assumption that eventually - # the projects of an org will be replicated consistently across replicas, because no org should continue to create - # new projects. - all_projects_ids = ( - Project.objects.using_replica() - .filter(organization=organization) - .values_list("id", flat=True) - ) - for project_id in all_projects_ids: - # In case a specific project has not been considered in the count query, it means that no metrics were extracted - # for it, thus we consider it as having 0 transactions for the query's time window. - if project_id not in projects_with_counts: - projects_with_counts[project_id] = 0 - - projects = [] - for project_id, count_per_root in projects_with_counts.items(): - projects.append( - RebalancedItem( - id=project_id, - count=count_per_root, - ) - ) - - model = ProjectsRebalancingModel() - rebalanced_projects: list[RebalancedItem] | None = guarded_run( - model, ProjectsRebalancingInput(classes=projects, sample_rate=sample_rate) - ) - - return rebalanced_projects - - -@dynamic_sampling_task -def store_rebalanced_projects(org_id: int, rebalanced_projects: list[RebalancedItem]) -> None: - """Stores the rebalanced projects in the cache and invalidates the project configs.""" - redis_client = get_redis_client_for_ds() - with redis_client.pipeline(transaction=False) as pipeline: - for rebalanced_project in rebalanced_projects: - cache_key = generate_boost_low_volume_projects_cache_key(org_id=org_id) - # We want to get the old sample rate, which will be None in case it was not set. - old_sample_rate = sample_rate_to_float( - redis_client.hget(cache_key, str(rebalanced_project.id)) - ) - - if rebalanced_project.id in PROJECTS_WITH_METRICS: - metrics.gauge( - "dynamic_sampling.project_sample_rate", - rebalanced_project.new_sample_rate * 100, - tags={"project_id": rebalanced_project.id}, - unit="percent", - ) - - # We want to store the new sample rate as a string. - pipeline.hset( - cache_key, - str(rebalanced_project.id), - rebalanced_project.new_sample_rate, # redis stores is as string - ) - pipeline.pexpire(cache_key, DEFAULT_REDIS_CACHE_KEY_TTL) - - # We invalidate the caches only if there was a change in the sample rate. This is to avoid flooding the - # system with project config invalidations, especially for projects with no volume. - if not are_equal_with_epsilon(old_sample_rate, rebalanced_project.new_sample_rate): - schedule_invalidate_project_config( - project_id=rebalanced_project.id, - trigger="dynamic_sampling_boost_low_volume_projects", - ) - - pipeline.execute() diff --git a/src/sentry/dynamic_sampling/tasks/boost_low_volume_transactions.py b/src/sentry/dynamic_sampling/tasks/boost_low_volume_transactions.py index 8f9f722b3221..c35ff12a3b88 100644 --- a/src/sentry/dynamic_sampling/tasks/boost_low_volume_transactions.py +++ b/src/sentry/dynamic_sampling/tasks/boost_low_volume_transactions.py @@ -1,9 +1,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator, Sequence from typing import TypedDict -import sentry_sdk from snuba_sdk import ( AliasedExpression, Column, @@ -18,42 +16,16 @@ Query, Request, ) -from taskbroker_client.retry import Retry -from sentry import options, quotas -from sentry.dynamic_sampling.models.common import RebalancedItem, guarded_run -from sentry.dynamic_sampling.models.transactions_rebalancing import ( - TransactionsRebalancingInput, - TransactionsRebalancingModel, -) -from sentry.dynamic_sampling.tasks.common import MEASURE_CONFIGS, GetActiveOrgs +from sentry.dynamic_sampling.tasks.common import MEASURE_CONFIGS from sentry.dynamic_sampling.tasks.constants import ( BOOST_LOW_VOLUME_TRANSACTIONS_QUERY_INTERVAL, CHUNK_SIZE, - DEFAULT_REDIS_CACHE_KEY_TTL, - MAX_PROJECTS_PER_QUERY, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( - get_boost_low_volume_projects_sample_rate, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_transactions import ( - set_transactions_resampling_rates, -) -from sentry.dynamic_sampling.tasks.utils import ( - dynamic_sampling_task, - legacy_pipeline_killswitched, ) from sentry.dynamic_sampling.types import SamplingMeasure -from sentry.dynamic_sampling.utils import has_dynamic_sampling, is_project_mode_sampling -from sentry.models.options.project_option import ProjectOption -from sentry.models.organization import Organization from sentry.sentry_metrics import indexer -from sentry.silo.base import SiloMode from sentry.snuba.dataset import Dataset, EntityKey from sentry.snuba.referrer import Referrer -from sentry.tasks.base import instrumented_task -from sentry.tasks.relay import schedule_invalidate_project_config -from sentry.taskworker.namespaces import telemetry_experience_tasks from sentry.utils import metrics from sentry.utils.dates import deprecated_utcnow from sentry.utils.snuba import raw_snql_query @@ -79,347 +51,6 @@ class ProjectTransactions(ProjectIdentity, total=True): total_num_classes: int | None -class ProjectTransactionsTotals(ProjectIdentity, total=True): - total_num_transactions: float - total_num_classes: int | float - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.boost_low_volume_transactions", - namespace=telemetry_experience_tasks, - processing_deadline_duration=10 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def boost_low_volume_transactions() -> None: - if legacy_pipeline_killswitched("boost_low_volume_transactions"): - return - - num_big_trans = int( - options.get("dynamic-sampling.prioritise_transactions.num_explicit_large_transactions") - ) - - for orgs in GetActiveOrgs( - max_projects=MAX_PROJECTS_PER_QUERY, - granularity=Granularity(60), - measure=SamplingMeasure.SEGMENTS, - ): - metrics.incr( - "dynamic_sampling.boost_low_volume_transactions.orgs_partitioned", - tags={"metric_type": "segment"}, - amount=len(orgs), - ) - _process_orgs_for_boost_low_volume_transactions( - orgs, num_big_trans, measure=SamplingMeasure.SEGMENTS - ) - - -def _process_orgs_for_boost_low_volume_transactions( - orgs: list[int], - num_big_trans: int, - measure: SamplingMeasure, -) -> None: - """ - Process a batch of organizations for boost low volume transactions. - """ - if not orgs: - return - - totals_it = FetchProjectTransactionTotals(orgs, measure=measure) - big_transactions_it = FetchProjectTransactionVolumes( - orgs, - max_transactions=num_big_trans, - measure=measure, - ) - - for project_transactions in transactions_zip(totals_it, big_transactions_it): - boost_low_volume_transactions_of_project.apply_async( - kwargs={"project_transactions": project_transactions}, - headers={"sentry-propagate-traces": False}, - ) - - -def _factor_bucket(factor: float) -> str: - """ - Buckets a sampling factor (multiplier) into a coarse, low-cardinality range for use as a metric - tag. - """ - # (inclusive lower bound, label), ordered high -> low; the first bound the factor clears wins. - factor_buckets: tuple[tuple[float, str], ...] = ( - (0.1, "1-0.1"), - (0.01, "0.1-0.01"), - (0.001, "0.01-0.001"), - (0.0001, "0.001-0.0001"), - (0.00001, "0.0001-0.00001"), - ) - if factor > 1: - return ">1" - return next((label for bound, label in factor_buckets if factor >= bound), "<0.00001") - - -def _emit_smallest_transaction_factor_bucket( - named_rates: Sequence[RebalancedItem], implicit_rate: float -) -> None: - """ - Emits the bucket of the sampling factor (multiplier) the model assigned to the lowest-volume - (smallest) explicit transaction of a project, once per project run. - - This records the value that actually lands in the project config rather than an absolute sample - rate: Relay multiplies the factor onto the implicit rate, so the factor is - ``new_sample_rate / implicit_rate`` (matching BoostLowVolumeTransactionsBias). - """ - if not options.get( - "dynamic-sampling.boost_low_volume_transactions.emit_smallest_transaction_factor_metric" - ): - return - - if not named_rates: - return - - smallest = min(named_rates, key=lambda item: item.count) - denominator = implicit_rate if implicit_rate != 0.0 else 1.0 - factor = smallest.new_sample_rate / denominator - metrics.incr( - "dynamic_sampling.boost_low_volume_transactions.smallest_transaction_factor", - tags={"factor_bucket": _factor_bucket(factor)}, - sample_rate=1.0, - ) - - -@instrumented_task( - name="sentry.dynamic_sampling.boost_low_volume_transactions_of_project", - namespace=telemetry_experience_tasks, - processing_deadline_duration=4 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def boost_low_volume_transactions_of_project(project_transactions: ProjectTransactions) -> None: - org_id = project_transactions["org_id"] - project_id = project_transactions["project_id"] - total_num_transactions = project_transactions.get("total_num_transactions") - total_num_classes = project_transactions.get("total_num_classes") - transactions = [ - RebalancedItem(id=id, count=count) - for id, count in project_transactions["transaction_counts"] - ] - - try: - organization = Organization.objects.get_from_cache(id=org_id) - except Organization.DoesNotExist: - organization = None - - # If the org doesn't have dynamic sampling, we want to early return to avoid unnecessary work. - if not has_dynamic_sampling(organization): - return - - if is_project_mode_sampling(organization): - sample_rate = ProjectOption.objects.get_value(project_id, "sentry:target_sample_rate") - else: - # We try to use the sample rate that was individually computed for each project, but if we don't find it, we will - # resort to the blended sample rate of the org. - sample_rate, success = get_boost_low_volume_projects_sample_rate( - org_id=org_id, - project_id=project_id, - error_sample_rate_fallback=quotas.backend.get_blended_sample_rate( - organization_id=org_id - ), - ) - if sample_rate is None: - sentry_sdk.capture_message( - "Sample rate of project not found when trying to adjust the sample rates of " - "its transactions" - ) - return - - if sample_rate == 1.0: - return - - # the model fails when we are not having any transactions, thus we can simply return here - if len(transactions) == 0: - return - - intensity = options.get("dynamic-sampling.prioritise_transactions.rebalance_intensity", 1.0) - min_sample_rate = options.get("dynamic-sampling.prioritise_transactions.min_sample_rate") - - model = TransactionsRebalancingModel() - rebalanced_transactions = guarded_run( - model, - TransactionsRebalancingInput( - classes=transactions, - sample_rate=sample_rate, - total_num_classes=total_num_classes, - total=total_num_transactions, - intensity=intensity, - min_sample_rate=min_sample_rate, - ), - ) - # In case the result of the model is None, it means that an error occurred, thus we want to early return. - if rebalanced_transactions is None: - return - - # Only after checking the nullability of rebalanced_transactions, we want to unpack the tuple. - named_rates, implicit_rate = rebalanced_transactions - _emit_smallest_transaction_factor_bucket(named_rates, implicit_rate) - if sample_rate > 0: - implicit_factor = implicit_rate / sample_rate - comparison = ( - "below" if implicit_factor < 1.0 else "above" if implicit_factor > 1.0 else "equal" - ) - metrics.incr( - "dynamic_sampling.boost_low_volume_transactions.implicit_factor", - tags={"comparison": comparison}, - ) - set_transactions_resampling_rates( - org_id=org_id, - proj_id=project_id, - named_rates=named_rates, - default_rate=implicit_rate, - ttl_ms=DEFAULT_REDIS_CACHE_KEY_TTL, - ) - - schedule_invalidate_project_config( - project_id=project_id, trigger="dynamic_sampling_boost_low_volume_transactions" - ) - - -def is_same_project(left: ProjectIdentity | None, right: ProjectIdentity | None) -> bool: - if left is None or right is None: - return False - - return left["project_id"] == right["project_id"] and left["org_id"] == right["org_id"] - - -def is_project_identity_before(left: ProjectIdentity, right: ProjectIdentity) -> bool: - return left["org_id"] < right["org_id"] or ( - left["org_id"] == right["org_id"] and left["project_id"] < right["project_id"] - ) - - -class FetchProjectTransactionTotals: - """ - Fetches the total number of transactions and the number of distinct transaction types for each - project in the given organizations - """ - - def __init__(self, orgs: Sequence[int], measure: SamplingMeasure = SamplingMeasure.SEGMENTS): - transaction_string_id = indexer.resolve_shared_org("transaction") - self.transaction_tag = f"tags_raw[{transaction_string_id}]" - - config = MEASURE_CONFIGS[measure] - self.metric_id = indexer.resolve_shared_org(str(config["mri"])) - self.use_case_id = config["use_case_id"] - self.tag_filters = config["tags"] - self.measure = measure - - self.org_ids = list(orgs) - self.offset = 0 - self.has_more_results = True - self.cache: list[dict[str, int | float]] = [] - self.last_org_id: int | None = None - - def __iter__(self) -> FetchProjectTransactionTotals: - return self - - def __next__(self) -> ProjectTransactionsTotals: - if not self._cache_empty(): - return self._get_from_cache() - - granularity = Granularity(60) - - if self.has_more_results: - where_conditions = [ - Condition( - Column("timestamp"), - Op.GTE, - deprecated_utcnow() - BOOST_LOW_VOLUME_TRANSACTIONS_QUERY_INTERVAL, - ), - Condition(Column("timestamp"), Op.LT, deprecated_utcnow()), - Condition(Column("metric_id"), Op.EQ, self.metric_id), - Condition(Column("org_id"), Op.IN, self.org_ids), - ] - # Add tag filters from config - for tag_name, tag_value in self.tag_filters.items(): - tag_string_id = indexer.resolve_shared_org(tag_name) - tag_column = f"tags_raw[{tag_string_id}]" - where_conditions.append(Condition(Column(tag_column), Op.EQ, tag_value)) - - query = ( - Query( - match=Entity(EntityKey.GenericOrgMetricsCounters.value), - select=[ - Function("sum", [Column("value")], "num_transactions"), - Function("uniq", [Column(self.transaction_tag)], "num_classes"), - Column("org_id"), - Column("project_id"), - ], - groupby=[ - Column("org_id"), - Column("project_id"), - ], - where=where_conditions, - granularity=granularity, - orderby=[ - OrderBy(Column("org_id"), Direction.ASC), - OrderBy(Column("project_id"), Direction.ASC), - ], - ) - .set_limit(CHUNK_SIZE + 1) - .set_offset(self.offset) - ) - request = Request( - dataset=Dataset.PerformanceMetrics.value, - app_id="dynamic_sampling", - query=query, - tenant_ids={"use_case_id": self.use_case_id.value, "cross_org_query": 1}, - ) - data = raw_snql_query( - request, - referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_TRANSACTION_TOTALS.value, - )["data"] - - metric_type = self.measure.value - metrics.incr( - "dynamic_sampling.boost_low_volume_transactions.query", - tags={"query_type": "totals", "metric_type": metric_type}, - sample_rate=1, - ) - - count = len(data) - self.has_more_results = count > CHUNK_SIZE - self.offset += CHUNK_SIZE - - if self.has_more_results: - data = data[:-1] - self.cache.extend(data) - - return self._get_from_cache() - - def _get_from_cache(self) -> ProjectTransactionsTotals: - if self._cache_empty(): - raise StopIteration() - - row = self.cache.pop(0) - proj_id = int(row["project_id"]) - org_id = int(row["org_id"]) - num_transactions = row["num_transactions"] - num_classes = int(row["num_classes"]) - - if self.last_org_id != org_id: - self.last_org_id = org_id - - return { - "project_id": proj_id, - "org_id": org_id, - "total_num_transactions": num_transactions, - "total_num_classes": num_classes, - } - - def _cache_empty(self) -> bool: - return not self.cache - - class FetchProjectTransactionVolumes: """ Fetch the highest-volume transactions for all orgs and all projects with pagination @@ -600,80 +231,3 @@ def _get_from_cache(self) -> ProjectTransactions: raise StopIteration() return self.cache.pop(0) - - -def merge_transactions( - transactions: ProjectTransactions, - totals: ProjectTransactionsTotals | None, -) -> ProjectTransactions: - if totals is not None and not is_same_project(transactions, totals): - raise ValueError( - "mismatched projectTransaction and projectTransactionTotals", - (transactions["org_id"], transactions["project_id"]), - (totals["org_id"], totals["project_id"]), - ) - - total_num_classes = totals.get("total_num_classes") if totals is not None else None - - return { - "org_id": transactions["org_id"], - "project_id": transactions["project_id"], - "transaction_counts": transactions["transaction_counts"], - "total_num_transactions": ( - totals.get("total_num_transactions") if totals is not None else None - ), - "total_num_classes": int(total_num_classes) if total_num_classes is not None else None, - } - - -def next_totals( - totals: Iterator[ProjectTransactionsTotals], -) -> Callable[[ProjectIdentity], ProjectTransactionsTotals | None]: - """ - Advances the total iterator until it reaches the required identity - - Given a match the iterator returns None if it cannot find it ( i.e. it is - already past it) or it is at the end (it never terminates, DO NOT use it - in a for loop). If it finds the match it will return the total for the match. - - """ - current: list[ProjectTransactionsTotals | None] = [None] - # protection for the case when the caller passes a list instead of an iterator - totals = iter(totals) - - def inner(match: ProjectIdentity) -> ProjectTransactionsTotals | None: - if is_same_project(current[0], match): - temp = current[0] - current[0] = None - return temp - - if current[0] is not None and is_project_identity_before(match, current[0]): - # still haven't reach current no point looking further - return None - - for total in totals: - if is_same_project(total, match): - # found it - return total - - if is_project_identity_before(match, total): - # we passed after match, remember were we are no need to go further - current[0] = total - return None - return None - - return inner - - -def transactions_zip( - totals: Iterator[ProjectTransactionsTotals], - transactions: Iterator[ProjectTransactions], -) -> Iterator[ProjectTransactions]: - """ - Consolidates each project's transaction volumes with its totals information, - when a matching totals entry exists. - """ - get_next_total = next_totals(totals) - - for project_transactions in transactions: - yield merge_transactions(project_transactions, get_next_total(project_transactions)) diff --git a/src/sentry/dynamic_sampling/tasks/common.py b/src/sentry/dynamic_sampling/tasks/common.py index 601a42b1249c..93de3445b968 100644 --- a/src/sentry/dynamic_sampling/tasks/common.py +++ b/src/sentry/dynamic_sampling/tasks/common.py @@ -6,18 +6,7 @@ from typing import TypedDict import sentry_sdk -from snuba_sdk import ( - Column, - Condition, - Direction, - Entity, - Function, - Granularity, - Op, - OrderBy, - Query, - Request, -) +from snuba_sdk import Column, Condition, Entity, Function, Granularity, Op, Query, Request from sentry import quotas from sentry.dynamic_sampling.tasks.constants import CHUNK_SIZE, MAX_ORGS_PER_QUERY @@ -31,9 +20,6 @@ from sentry.utils.dates import deprecated_utcnow from sentry.utils.snuba import raw_snql_query -ACTIVE_ORGS_DEFAULT_TIME_INTERVAL = timedelta(hours=1) -ACTIVE_ORGS_DEFAULT_GRANULARITY = Granularity(3600) - ACTIVE_ORGS_VOLUMES_DEFAULT_TIME_INTERVAL = timedelta(minutes=5) ACTIVE_ORGS_VOLUMES_DEFAULT_GRANULARITY = Granularity(60) @@ -63,149 +49,6 @@ class MeasureConfig(TypedDict): } -class GetActiveOrgs: - """ - Fetch organizations in batches. - A batch will return at max max_orgs elements - It will accumulate org ids in the list until either it accumulates max_orgs or the - number of projects in the already accumulated orgs is more than max_projects or there - are no more orgs - """ - - def __init__( - self, - max_orgs: int = MAX_ORGS_PER_QUERY, - max_projects: int | None = None, - time_interval: timedelta = ACTIVE_ORGS_DEFAULT_TIME_INTERVAL, - granularity: Granularity = ACTIVE_ORGS_DEFAULT_GRANULARITY, - measure: SamplingMeasure = SamplingMeasure.SEGMENTS, - ) -> None: - config = MEASURE_CONFIGS[measure] - self.metric_id = indexer.resolve_shared_org(str(config["mri"])) - self.use_case_id = config["use_case_id"] - self.tag_filters = config["tags"] - - self.offset = 0 - self.last_result: list[tuple[int, int]] = [] - self.has_more_results = True - self.max_orgs = max_orgs - self.max_projects = max_projects - self.time_interval = time_interval - self.granularity = granularity - - def __iter__(self) -> GetActiveOrgs: - return self - - def __next__(self) -> list[int]: - if self._enough_results_cached(): - # we have enough in the cache to satisfy the current iteration - return self._get_from_cache() - - if self.has_more_results: - # not enough for the current iteration and data still in the db top it up from db - where_conditions = [ - Condition( - Column("timestamp"), - Op.GTE, - deprecated_utcnow() - self.time_interval, - ), - Condition(Column("timestamp"), Op.LT, deprecated_utcnow()), - Condition(Column("metric_id"), Op.EQ, self.metric_id), - ] - for tag_name, tag_value in self.tag_filters.items(): - tag_string_id = indexer.resolve_shared_org(tag_name) - tag_column = f"tags_raw[{tag_string_id}]" - where_conditions.append(Condition(Column(tag_column), Op.EQ, tag_value)) - - query = ( - Query( - match=Entity(EntityKey.GenericOrgMetricsCounters.value), - select=[ - Function("uniq", [Column("project_id")], "num_projects"), - Column("org_id"), - ], - groupby=[ - Column("org_id"), - ], - where=where_conditions, - orderby=[ - OrderBy(Column("org_id"), Direction.ASC), - ], - granularity=self.granularity, - ) - .set_limit(CHUNK_SIZE + 1) - .set_offset(self.offset) - ) - request = Request( - dataset=Dataset.PerformanceMetrics.value, - app_id="dynamic_sampling", - query=query, - tenant_ids={ - "use_case_id": self.use_case_id.value, - "cross_org_query": 1, - }, - ) - data = raw_snql_query( - request, - referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value, - )["data"] - count = len(data) - - self.has_more_results = count > CHUNK_SIZE - self.offset += CHUNK_SIZE - if self.has_more_results: - data = data[:-1] - for row in data: - self.last_result.append((row["org_id"], row["num_projects"])) - - if len(self.last_result) > 0: - # we have some data left return up to the max amount - return self._get_from_cache() # we still have something left in cache - else: - # nothing left in the DB or cache - raise StopIteration() - - def _enough_results_cached(self) -> bool: - """ - Return true if we have enough data to return a full batch in the cache (i.e. last_result) - """ - if len(self.last_result) >= self.max_orgs: - return True - - if self.max_projects is not None: - total_projects = 0 - for _, num_projects in self.last_result: - total_projects += num_projects - if num_projects >= self.max_projects: - return True - return False - - def _get_orgs(self, orgs_and_counts: list[tuple[int, int]]) -> list[int]: - """ - Extracts the orgs from last_result - """ - return [org for org, _ in orgs_and_counts] - - def _get_from_cache(self) -> list[int]: - """ - Returns a batch from cache and removes the elements returned from the cache - """ - count_projects = 0 - for idx, (org_id, num_projects) in enumerate(self.last_result): - count_projects += num_projects - if idx >= (self.max_orgs - 1) or ( - self.max_projects is not None and count_projects >= self.max_projects - ): - # we got to the number of elements desired - ret_val = self._get_orgs(self.last_result[: idx + 1]) - self.last_result = self.last_result[idx + 1 :] - return ret_val - # if we are here we haven't reached our max limit, return everything - ret_val = self._get_orgs(self.last_result) - self.last_result = [] - return ret_val - - @dataclass(frozen=True) class OrganizationDataVolume: """ diff --git a/src/sentry/dynamic_sampling/tasks/constants.py b/src/sentry/dynamic_sampling/tasks/constants.py index 51d564513d52..933ab1e799b5 100644 --- a/src/sentry/dynamic_sampling/tasks/constants.py +++ b/src/sentry/dynamic_sampling/tasks/constants.py @@ -15,8 +15,6 @@ def adjusted_factor_ttl_ms() -> int: # Parameters to bound the queries run in Snuba. MAX_ORGS_PER_QUERY = 80 -MAX_PROJECTS_PER_QUERY = 4000 -MAX_TRANSACTIONS_PER_PROJECT = 20 # MIN and MAX rebalance factor in order to make sure we don't go crazy when rebalancing orgs. MIN_REBALANCE_FACTOR = 0.1 diff --git a/src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py b/src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py deleted file mode 100644 index f1493b2215a4..000000000000 --- a/src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py +++ /dev/null @@ -1,215 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence - -import sentry_sdk -from taskbroker_client.retry import Retry - -from sentry import quotas -from sentry.constants import SAMPLING_MODE_DEFAULT, TARGET_SAMPLE_RATE_DEFAULT -from sentry.dynamic_sampling.per_org.gate import is_org_in_serving_rollout -from sentry.dynamic_sampling.per_org.serving import get_previous_recalibration_factor -from sentry.dynamic_sampling.rules.utils import DecisionKeepCount, OrganizationId, ProjectId -from sentry.dynamic_sampling.tasks.boost_low_volume_projects import ( - fetch_projects_with_total_root_transaction_count_and_rates, -) -from sentry.dynamic_sampling.tasks.common import GetActiveOrgsVolumes, OrganizationDataVolume -from sentry.dynamic_sampling.tasks.constants import bounded_rebalance_factor -from sentry.dynamic_sampling.tasks.helpers.recalibrate_orgs import ( - compute_adjusted_factor, - delete_adjusted_factor, - delete_adjusted_project_factor, - get_adjusted_project_factor, - set_guarded_adjusted_factor, - set_guarded_adjusted_project_factor, -) -from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate -from sentry.dynamic_sampling.tasks.utils import ( - dynamic_sampling_task, - legacy_pipeline_killswitched, -) -from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure -from sentry.dynamic_sampling.utils import has_dynamic_sampling -from sentry.models.options.organization_option import OrganizationOption -from sentry.models.options.project_option import ProjectOption -from sentry.models.organization import Organization -from sentry.silo.base import SiloMode -from sentry.tasks.base import instrumented_task -from sentry.taskworker.namespaces import telemetry_experience_tasks -from sentry.utils import metrics - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.recalibrate_orgs", - namespace=telemetry_experience_tasks, - processing_deadline_duration=1 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def recalibrate_orgs() -> None: - if legacy_pipeline_killswitched("recalibrate_orgs"): - return - - for segment_volumes in GetActiveOrgsVolumes(measure=SamplingMeasure.SEGMENTS): - _process_orgs_volumes(segment_volumes) - - -def _process_orgs_volumes(org_volumes: Sequence[OrganizationDataVolume]) -> None: - """ - Process organization volumes for recalibration. - - Args: - org_volumes: Volumes to process for recalibration. - """ - if not org_volumes: - return - - modes = OrganizationOption.objects.get_value_bulk_id( - [v.org_id for v in org_volumes], "sentry:sampling_mode", SAMPLING_MODE_DEFAULT - ) - orgs_batch = [] - projects_batch = [] - for org_volume in org_volumes: - if not org_volume.is_valid_for_recalibration(): - continue - if modes[org_volume.org_id] == DynamicSamplingMode.PROJECT: - projects_batch.append(org_volume.org_id) - else: - orgs_batch.append((org_volume.org_id, org_volume.total, org_volume.indexed)) - - if orgs_batch: - recalibrate_orgs_batch.delay(orgs_batch) - if projects_batch: - recalibrate_projects_batch.delay(projects_batch) - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.recalibrate_orgs_batch", - namespace=telemetry_experience_tasks, - processing_deadline_duration=6 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def recalibrate_orgs_batch(orgs: Sequence[tuple[OrganizationId, int, int]]) -> None: - for org_id, total, indexed in orgs: - try: - recalibrate_org(org_id, total, indexed) - except Exception as e: - sentry_sdk.capture_exception(e) - continue - - -def recalibrate_org(org_id: OrganizationId, total: int, indexed: int) -> None: - if is_org_in_serving_rollout(org_id): - metrics.incr("dynamic_sampling.tasks.recalibrate_orgs.skipped_served_per_org") - return - - try: - # We need the organization object for the feature flag. - organization = Organization.objects.get_from_cache(id=org_id) - except Organization.DoesNotExist: - # In case an org is not found, it might be that it has been deleted in the time between - # the query triggering this job and the actual execution of the job. - organization = None - - # If the org doesn't have dynamic sampling, we want to early return to avoid unnecessary work. - if not has_dynamic_sampling(organization): - return - - # 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_org_sample_rate( - org_id=org_id, - default_sample_rate=quotas.backend.get_blended_sample_rate(organization_id=org_id), - ) - - # If we didn't find any sample rate, we can't recalibrate the organization. - if target_sample_rate is None: - sentry_sdk.capture_message("Sample rate of org not found when trying to recalibrate it") - return - - # We compute the effective sample rate that we had in the last considered time window. - effective_sample_rate = indexed / total - # We get the previous factor that was used for the recalibration. - previous_factor = get_previous_recalibration_factor(org_id) - - # We want to compute the new adjusted factor. - adjusted_factor = compute_adjusted_factor( - previous_factor, effective_sample_rate, target_sample_rate - ) - if adjusted_factor is None: - sentry_sdk.capture_message( - "The adjusted factor for org recalibration could not be computed" - ) - return - - bounded_factor = bounded_rebalance_factor(adjusted_factor) - if bounded_factor is None: - # In case the new factor would result into too much recalibration, we want to remove it from cache, - # effectively removing the generated rule. - delete_adjusted_factor(org_id) - return - - # At the end we set the adjusted factor. - set_guarded_adjusted_factor(org_id, bounded_factor) - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.recalibrate_projects_batch", - namespace=telemetry_experience_tasks, - processing_deadline_duration=2 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def recalibrate_projects_batch(orgs: list[OrganizationId]) -> None: - for org_id, projects in fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=orgs, measure=SamplingMeasure.SPANS - ).items(): - sample_rates = ProjectOption.objects.get_value_bulk_id( - [t[0] for t in projects], "sentry:target_sample_rate" - ) - - for project_id, total, keep, _ in projects: - try: - recalibrate_project(org_id, project_id, total, keep, sample_rates[project_id]) - except Exception as e: - sentry_sdk.capture_exception(e) - continue - - -def recalibrate_project( - org_id: OrganizationId, - project_id: ProjectId, - total: int, - indexed: DecisionKeepCount, - target_sample_rate: float | None, -) -> None: - if target_sample_rate is None: - target_sample_rate = TARGET_SAMPLE_RATE_DEFAULT - # We compute the effective sample rate that we had in the last considered time window. - effective_sample_rate = indexed / total - # We get the previous factor that was used for the recalibration. - previous_factor = get_adjusted_project_factor(project_id, source="task") - - # We want to compute the new adjusted factor. - adjusted_factor = compute_adjusted_factor( - previous_factor, effective_sample_rate, target_sample_rate - ) - if adjusted_factor is None: - sentry_sdk.capture_message( - "The adjusted factor for org recalibration could not be computed" - ) - return - - bounded_factor = bounded_rebalance_factor(adjusted_factor) - if bounded_factor is None: - # In case the new factor would result into too much recalibration, we want to remove it from cache, - # effectively removing the generated rule. - delete_adjusted_project_factor(project_id) - return - - # At the end we set the adjusted factor. - set_guarded_adjusted_project_factor(project_id, bounded_factor) diff --git a/src/sentry/dynamic_sampling/tasks/sliding_window_org.py b/src/sentry/dynamic_sampling/tasks/sliding_window_org.py deleted file mode 100644 index 346889abfe76..000000000000 --- a/src/sentry/dynamic_sampling/tasks/sliding_window_org.py +++ /dev/null @@ -1,96 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from datetime import timedelta - -from taskbroker_client.retry import Retry - -from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds -from sentry.dynamic_sampling.tasks.common import ( - GetActiveOrgsVolumes, - OrganizationDataVolume, - compute_guarded_sliding_window_sample_rate, -) -from sentry.dynamic_sampling.tasks.constants import CHUNK_SIZE, DEFAULT_REDIS_CACHE_KEY_TTL -from sentry.dynamic_sampling.tasks.helpers.sliding_window import ( - generate_sliding_window_org_cache_key, - get_sliding_window_size, - mark_sliding_window_org_executed, -) -from sentry.dynamic_sampling.tasks.utils import ( - dynamic_sampling_task, - legacy_pipeline_killswitched, -) -from sentry.dynamic_sampling.types import SamplingMeasure -from sentry.silo.base import SiloMode -from sentry.tasks.base import instrumented_task -from sentry.taskworker.namespaces import telemetry_experience_tasks - - -@instrumented_task( - name="sentry.dynamic_sampling.tasks.sliding_window_org", - namespace=telemetry_experience_tasks, - processing_deadline_duration=15 * 60 + 5, - retry=Retry(times=5, delay=5), - silo_mode=SiloMode.CELL, -) -@dynamic_sampling_task -def sliding_window_org() -> None: - if legacy_pipeline_killswitched("sliding_window_org"): - return - - window_size = get_sliding_window_size() - # In case the size is None it means that we disabled the sliding window entirely. - if window_size is None: - return - - for segment_volumes in GetActiveOrgsVolumes( - max_orgs=CHUNK_SIZE, - time_interval=timedelta(hours=window_size), - include_keep=False, - measure=SamplingMeasure.SEGMENTS, - ): - _process_org_volumes(segment_volumes, window_size) - - # Due to the synchronous nature of the sliding window org, when we arrived here, we can confidently say - # that the execution of the sliding window org was successful. We will keep this state for 1 hour. - mark_sliding_window_org_executed() - - -def _process_org_volumes(org_volumes: Sequence[OrganizationDataVolume], window_size: int) -> None: - """ - Process sliding window calculations for the given organization volumes. - - Args: - org_volumes: The organization volumes to process. - window_size: The sliding window size in hours. - """ - for org_volume in org_volumes: - adjust_base_sample_rate_of_org( - org_id=org_volume.org_id, - total_root_count=org_volume.total, - window_size=window_size, - ) - - -def adjust_base_sample_rate_of_org(org_id: int, total_root_count: int, window_size: int) -> None: - """ - Adjusts the base sample rate per org by considering its volume and how it fits w.r.t. to the sampling tiers. - """ - sample_rate = compute_guarded_sliding_window_sample_rate( - org_id, - None, - total_root_count, - window_size, - ) - # 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() diff --git a/src/sentry/dynamic_sampling/tasks/utils.py b/src/sentry/dynamic_sampling/tasks/utils.py deleted file mode 100644 index 8416f3474de0..000000000000 --- a/src/sentry/dynamic_sampling/tasks/utils.py +++ /dev/null @@ -1,55 +0,0 @@ -from collections.abc import Callable -from functools import wraps -from random import random -from typing import Any - -import sentry_sdk - -from sentry import options -from sentry.utils import metrics - -LEGACY_KILLSWITCH_OPTION = "dynamic-sampling.legacy.killswitch" - - -def sample_function(function: Callable[..., Any], _sample_rate: float = 1.0, **kwargs: Any) -> None: - """ - Calls the supplied function with a uniform probability of `_sample_rate`. - """ - if _sample_rate >= 1.0 or 0.0 <= random() <= _sample_rate: - function(**kwargs) - - -def _compute_task_name(function_name: str) -> str: - return f"sentry.tasks.dynamic_sampling.{function_name}" - - -def dynamic_sampling_task(func: Callable[..., Any]) -> Callable[..., Any]: - """ - Decorator to wrap dynamic sampling related tasks to record metrics for the execution of - the task, durations associated with it as metrics, and capture all exceptions in sentry. - """ - - @wraps(func) - def _wrapper(*args: Any, **kwargs: Any) -> Any: - function_name = func.__name__ - task_name = _compute_task_name(function_name) - metrics.incr(f"{task_name}.start", sample_rate=1.0) - with metrics.timer(task_name, sample_rate=1.0): - try: - return func(*args, **kwargs) - except Exception as e: - sentry_sdk.capture_exception(e) - raise - - return _wrapper - - -def legacy_pipeline_killswitched(task_name: str) -> bool: - """ - Reports whether the legacy dynamic sampling pipeline is switched off. - Scheduled legacy jobs call this first and return before they do any work when it is True. - """ - if not options.get(LEGACY_KILLSWITCH_OPTION): - return False - metrics.incr(f"{_compute_task_name(task_name)}.killswitched", sample_rate=1.0) - return True diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index d41963833df9..00497aaee8b0 100644 --- a/src/sentry/options/defaults.py +++ b/src/sentry/options/defaults.py @@ -2410,7 +2410,8 @@ 30, flags=FLAG_AUTOMATOR_MODIFIABLE, ) -# Toggles emitting the smallest-transaction sampling-factor bucket metric during transaction rebalancing. +# Nothing reads this option any more. It stays registered until the options automator +# has unset it, since the automator can only unset a registered option. register( "dynamic-sampling.boost_low_volume_transactions.emit_smallest_transaction_factor_metric", default=False, @@ -2459,9 +2460,8 @@ flags=FLAG_AUTOMATOR_MODIFIABLE, ) -# Killswitch for the legacy dynamic sampling pipeline. When set to True, the four -# scheduled jobs (sliding_window_org, boost_low_volume_projects, -# boost_low_volume_transactions, recalibrate_orgs) exit before they do any work. +# Nothing reads this option any more. It stays registered until the options automator +# has unset it, since the automator can only unset a registered option. register( "dynamic-sampling.legacy.killswitch", default=False, @@ -2554,11 +2554,8 @@ flags=FLAG_AUTOMATOR_MODIFIABLE, ) -# Controls the intensity of dynamic sampling transaction rebalancing. 0.0 = explict rebalancing -# not performed, 1.0= full rebalancing (tries to bring everything to mean). Note that even at 0.0 -# there will still be some rebalancing between the explicit and implicit transactions ( so setting rebalancing -# to 0.0 is not the same as no rebalancing. To effectively disable rebalancing set the number of explicit -# transactions to be rebalance (both small and large) to 0. +# Nothing reads this option any more. It stays registered until the options automator +# has unset it, since the automator can only unset a registered option. register( "dynamic-sampling.prioritise_transactions.rebalance_intensity", default=0.8, diff --git a/src/sentry/snuba/referrer.py b/src/sentry/snuba/referrer.py index 814440d32d44..16901b96d089 100644 --- a/src/sentry/snuba/referrer.py +++ b/src/sentry/snuba/referrer.py @@ -648,9 +648,6 @@ class Referrer(StrEnum): DYNAMIC_SAMPLING_COUNTERS_GET_ORG_TRANSACTION_VOLUMES = ( "dynamic_sampling.counters.get_org_transaction_volumes" ) - DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT = ( - "dynamic_sampling.distribution.fetch_projects_with_count_per_root_total_volumes" - ) DYNAMIC_SAMPLING_PER_ORG_GET_EAP_ORG_VOLUME = "dynamic_sampling.per_org.get_eap_org_volume" DYNAMIC_SAMPLING_PER_ORG_GET_EAP_PROJECT_VOLUMES = ( "dynamic_sampling.per_org.get_eap_project_volumes" @@ -661,9 +658,6 @@ class Referrer(StrEnum): DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION = ( "dynamic_sampling.counters.fetch_projects_with_count_per_transaction_volumes" ) - DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_TRANSACTION_TOTALS = ( - "dynamic_sampling.counters.fetch_projects_with_transaction_totals" - ) DYNAMIC_SAMPLING_SETTINGS_GET_SPAN_COUNTS = "dynamic_sampling.settings.get_project_span_counts" ESCALATING_GROUPS = "sentry.issues.escalating" EVENTSTORE_GET_EVENT_BY_ID_NODESTORE = "eventstore.backend.get_event_by_id_nodestore" diff --git a/src/sentry/utils/sdk.py b/src/sentry/utils/sdk.py index ae484dbd34bf..e17b297d06e1 100644 --- a/src/sentry/utils/sdk.py +++ b/src/sentry/utils/sdk.py @@ -89,12 +89,8 @@ "sentry.tasks.summaries.weekly_reports.schedule_organizations": 1.0, "sentry.profiles.task.process_profile": 0.1 * settings.SENTRY_BACKEND_APM_SAMPLING, "sentry.monitors.tasks.clock_pulse": 1.0, - "sentry.dynamic_sampling.tasks.boost_low_volume_projects": 1.0, - "sentry.dynamic_sampling.tasks.boost_low_volume_transactions": 1.0, - "sentry.dynamic_sampling.tasks.recalibrate_orgs": 0.2 * settings.SENTRY_BACKEND_APM_SAMPLING, "sentry.dynamic_sampling.per_org.run_calculations_per_org": 1.0, "sentry.dynamic_sampling.per_org.schedule_per_org_calculations": 1.0, - "sentry.dynamic_sampling.tasks.sliding_window_org": 0.2 * settings.SENTRY_BACKEND_APM_SAMPLING, "sentry.tasks.autofix.configure_seer_for_existing_org": 1.0, "sentry.tasks.seer.context_engine_index.schedule_context_engine_indexing_tasks": 1.0, } diff --git a/tests/sentry/core/endpoints/test_organization_details.py b/tests/sentry/core/endpoints/test_organization_details.py index 19a0b6fb5613..7fc357a48f7e 100644 --- a/tests/sentry/core/endpoints/test_organization_details.py +++ b/tests/sentry/core/endpoints/test_organization_details.py @@ -43,8 +43,8 @@ from sentry.replays.models import OrganizationMemberReplayAccess from sentry.signals import project_created from sentry.silo.safety import unguarded_write -from sentry.snuba.metrics import SpanMRI -from sentry.testutils.cases import APITestCase, BaseMetricsLayerTestCase, TwoFactorAPITestCase +from sentry.testutils.cases import APITestCase, SnubaTestCase, SpanTestCase, TwoFactorAPITestCase +from sentry.testutils.helpers.datetime import before_now from sentry.testutils.helpers.features import with_feature from sentry.testutils.outbox import outbox_runner from sentry.testutils.pytest.fixtures import django_db_all @@ -95,7 +95,7 @@ def has_scope(self, scope): @cell_silo_test(cells=cells, include_monolith_run=True) -class OrganizationDetailsTest(OrganizationDetailsTestBase, BaseMetricsLayerTestCase): +class OrganizationDetailsTest(OrganizationDetailsTestBase, SnubaTestCase, SpanTestCase): @property def now(self): return datetime.now().replace(microsecond=0) @@ -568,11 +568,29 @@ def test_sampling_mode_change_requires_write_scope(self) -> None: assert response.status_code == 403 + @django_db_all + def test_change_org_target_sample_rate_schedules_per_org_calculation(self) -> None: + self.organization.update_option( + "sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION.value + ) + + with ( + self.feature("organizations:dynamic-sampling-custom"), + patch( + "sentry.core.endpoints.organization_details.run_calculations_per_org_task_entry" + ) as task, + ): + response = self.get_response(self.organization.slug, method="put", targetSampleRate=0.1) + + assert response.status_code == 200 + task.delay.assert_called_once_with(self.organization.id) + @django_db_all @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_sampling_mode_change_with_deleted_projects_that_had_metrics(self) -> None: + def test_sampling_mode_change_with_deleted_projects_that_had_spans(self) -> None: project_1 = self.create_project(organization=self.organization) project_2 = self.create_project(organization=self.organization) + self.organization.update_option("sentry:target_sample_rate", 0.5) # Create a team member for project_1 only team_1 = self.create_team(organization=self.organization) @@ -583,21 +601,17 @@ def test_sampling_mode_change_with_deleted_projects_that_had_metrics(self) -> No ) self.login_as(user=member_user) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"is_segment": "true", "decision": "keep"}, - minutes_before_now=60 * 24 * 12, - value=1, - project_id=project_1.id, - org_id=self.organization.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"is_segment": "true", "decision": "keep"}, - minutes_before_now=60 * 24 * 12, - value=1, - project_id=project_2.id, - org_id=self.organization.id, + timestamp = before_now(days=12) + self.store_spans( + [ + self.create_span( + {"is_segment": True, "sentry_tags": {"dsc.project_id": str(project.id)}}, + organization=self.organization, + project=project, + start_ts=timestamp, + ) + for project in (project_1, project_2) + ] ) project_2.delete() diff --git a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py deleted file mode 100644 index 3549abc82992..000000000000 --- a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py +++ /dev/null @@ -1,718 +0,0 @@ -from datetime import datetime, timedelta -from typing import cast -from unittest.mock import patch - -from django.utils import timezone - -from sentry.dynamic_sampling.rules.base import get_guarded_project_sample_rate -from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds -from sentry.dynamic_sampling.tasks.boost_low_volume_projects import ( - boost_low_volume_projects, - boost_low_volume_projects_of_org_with_query, - fetch_projects_with_total_root_transaction_count_and_rates, - query_project_counts_by_org, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( - get_boost_low_volume_projects_sample_rate, -) -from sentry.dynamic_sampling.tasks.helpers.sliding_window import ( - generate_sliding_window_org_cache_key, -) -from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure -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 SpanMRI -from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase -from sentry.testutils.helpers.datetime import freeze_time -from sentry.testutils.helpers.features import with_feature -from sentry.testutils.helpers.options import override_options - -MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace( - hour=0, minute=0, second=0, microsecond=0 -) - - -@freeze_time(MOCK_DATETIME) -class PrioritiseProjectsSnubaQueryTest(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - @property - def now(self) -> datetime: - return MOCK_DATETIME - - def test_simple_one_org_one_project(self) -> None: - org1 = self.create_organization("test-org") - p1 = self.create_project(organization=org1) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p1.id, - org_id=org1.id, - ) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=p1.id, - org_id=org1.id, - ) - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org1.id], measure=SamplingMeasure.SEGMENTS - ) - assert results[org1.id] == [(p1.id, 4.0, 1, 3)] - - def test_deleted_projects_are_not_queried(self) -> None: - org1 = self.create_organization("test-org") - p1 = self.create_project(organization=org1) - p2 = self.create_project(organization=org1) - - for p in [p1, p2]: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p.id, - org_id=org1.id, - ) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=p.id, - org_id=org1.id, - ) - p2.delete() - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org1.id], measure=SamplingMeasure.SEGMENTS - ) - 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_sliding_window_sample_rate(self) -> None: - org1 = self.create_organization("test-org") - p1 = self.create_project(organization=org1) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p1.id, - org_id=org1.id, - ) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=p1.id, - org_id=org1.id, - ) - - # simulate having a sliding window sample rate for the org - redis_client = get_redis_client_for_ds() - cache_key = generate_sliding_window_org_cache_key(org1.id) - redis_client.set(cache_key, 1.0) - - 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 == 1.0 - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_simple_one_org_one_project_task_target_sample_rate(self) -> None: - 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=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p1.id, - org_id=org1.id, - ) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - 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 (sample_rate, got_value) == (0.5, True) - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_per_project_sample_rate_override(self) -> None: - # A per-project override configured via options hard-replaces the rate the - # custom dynamic sampling path would otherwise resolve for that project -- - # winning even over the recently-added 100% boost -- and leaves other projects - # untouched. - org1 = self.create_organization("am3-override-org") - org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) - org1.update_option("sentry:target_sample_rate", 0.5) - overridden = self.create_project(organization=org1) - normal = self.create_project(organization=org1) - - # Baseline: freshly-created projects are boosted to 1.0 by the recently-added - # rule, so neither resolves to the org target yet. - assert get_guarded_project_sample_rate(org1, overridden) == 1.0 - - with override_options( - {"dynamic-sampling.sample-rate-override-per-project": {str(overridden.id): 0.9}} - ): - assert get_guarded_project_sample_rate(org1, overridden) == 0.9 - # Not in the override map -> unaffected by the override. - assert get_guarded_project_sample_rate(org1, normal) == 1.0 - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_per_project_sample_rate_override_ignores_out_of_range(self) -> None: - org1 = self.create_organization("am3-override-org-bad") - org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) - org1.update_option("sentry:target_sample_rate", 0.5) - project = self.create_project(organization=org1) - - baseline = get_guarded_project_sample_rate(org1, project) - with override_options( - {"dynamic-sampling.sample-rate-override-per-project": {str(project.id): 2.0}} - ): - # Out-of-range override is ignored; the resolved rate is unchanged. - assert get_guarded_project_sample_rate(org1, project) == baseline - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_project_mode_sampling_with_query(self) -> None: - org1 = self.create_organization("test-org") - p1 = self.create_project(organization=org1) - - org1.update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) - p1.update_option("sentry:target_sample_rate", 0.2) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p1.id, - org_id=org1.id, - ) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=p1.id, - org_id=org1.id, - ) - - # bulk task - with self.tasks(): - boost_low_volume_projects.delay() - - sample_rate, got_value = get_boost_low_volume_projects_sample_rate( - org1.id, p1.id, error_sample_rate_fallback=None - ) - assert (sample_rate, got_value) == (None, False) - - # single-org task - 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 (sample_rate, got_value) == (None, False) - - assert get_guarded_project_sample_rate(org1, p1) == 0.2 - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_project_mode_sampling_with_query_zero_metrics(self) -> None: - organization = self.create_organization("test-org") - project = self.create_project(organization=organization) - - organization.update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) - project.update_option("sentry:target_sample_rate", 0.2) - - # make sure that no rebalancing is actually run - with patch( - "sentry.dynamic_sampling.models.projects_rebalancing.ProjectsRebalancingModel._run" - ) as mock_run: - with self.tasks(): - boost_low_volume_projects.delay() - assert not mock_run.called - - def test_complex(self) -> None: - org1 = self.create_organization("test-org1") - p1_1 = self.create_project(organization=org1, name="p1_1") - p1_2 = self.create_project(organization=org1, name="p1_2") - org2 = self.create_organization("test-org2") - p2_1 = self.create_project(organization=org2, name="p2_1") - p2_2 = self.create_project(organization=org2, name="p2_2") - - proj_orgs = [ - {"org": org1, "projects": [p1_1, p1_2]}, - {"org": org2, "projects": [p2_1, p2_2]}, - ] - - proj_counts = {"p1_1": (1, 2), "p1_2": (3, 4), "p2_1": (5, 6), "p2_2": (7, 8)} # keep,drop - - for org_info in proj_orgs: - org = cast(Organization, org_info.get("org")) - projects = cast(list[Project], org_info.get("projects")) - for project in projects: - keep, drop = proj_counts[project.name] - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={ - "transaction": "foo_transaction", - "decision": "keep", - "is_segment": "true", - }, - minutes_before_now=29, - value=keep, - project_id=project.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={ - "transaction": "foo_transaction", - "decision": "drop", - "is_segment": "true", - }, - minutes_before_now=29, - value=drop, - project_id=project.id, - org_id=org.id, - ) - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org1.id, org2.id], measure=SamplingMeasure.SEGMENTS - ) - - assert len(results) == 2 # two orgs - - org_1_results = results[org1.id] - - assert len(org_1_results) == 2 - - # (p.id, total, keep, drop) == result - assert (p1_1.id, 3, 1, 2) in org_1_results - assert (p1_2.id, 7, 3, 4) in org_1_results - - org_2_results = results[org2.id] - assert len(org_2_results) == 2 - assert (p2_1.id, 11, 5, 6) in org_2_results - assert (p2_2.id, 15, 7, 8) in org_2_results - - -@freeze_time(MOCK_DATETIME) -class TestQueryProjectCountsByOrgEmptyOrgIds(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - """ - Test that query_project_counts_by_org correctly skips Snuba queries - when org_ids is empty, avoiding unnecessary queries. - """ - - @property - def now(self) -> datetime: - return MOCK_DATETIME - - def test_query_skips_for_empty_org_ids_when_option_enabled(self) -> None: - """ - Confirms that query_project_counts_by_org does NOT make a Snuba query - when called with an empty org_ids list. - """ - with patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects.raw_snql_query" - ) as mock_query: - mock_query.return_value = {"data": []} - - list(query_project_counts_by_org([], SamplingMeasure.SEGMENTS)) - - assert mock_query.call_count == 0 - - def test_fetch_projects_only_queries_measures_with_orgs(self) -> None: - """ - Confirms that fetch_projects_with_total_root_transaction_count_and_rates - does NOT make a Snuba query when called with an empty org_ids list. - """ - org = self.create_organization("test-org") - self.create_project(organization=org) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=org.project_set.first().id, - org_id=org.id, - ) - - with patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects.raw_snql_query" - ) as mock_query: - mock_query.return_value = {"data": []} - - # Query with org should make one call - fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org.id], measure=SamplingMeasure.SEGMENTS - ) - assert mock_query.call_count == 1 - - # Query with empty list should not make any additional calls - fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[], measure=SamplingMeasure.SEGMENTS - ) - assert mock_query.call_count == 1 - - -@freeze_time(MOCK_DATETIME) -class TestSpanMetricQuery(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - """ - Tests that verify the span metric query works correctly with is_segment filter. - """ - - @property - def now(self) -> datetime: - return MOCK_DATETIME - - def test_span_metric_with_is_segment_filter(self) -> None: - """ - Test that span metric queries only count spans with is_segment=true. - """ - org = self.create_organization("test-org") - project = self.create_project(organization=org) - - # Store span metrics with is_segment=true (should be counted) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=5, - project_id=project.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=10, - project_id=project.id, - org_id=org.id, - ) - - # Store span metrics without is_segment tag (should NOT be counted) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar_transaction", "decision": "keep"}, - minutes_before_now=30, - value=100, - project_id=project.id, - org_id=org.id, - ) - - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org.id], measure=SamplingMeasure.SEGMENTS - ) - - # Should only count the is_segment=true metrics (5 + 10 = 15) - assert results[org.id] == [(project.id, 15.0, 5, 10)] - - def test_span_metric_multiple_projects(self) -> None: - """ - Test span metric query with multiple projects. - """ - org = self.create_organization("test-org") - p1 = self.create_project(organization=org) - p2 = self.create_project(organization=org) - - # Project 1: 3 keep, 7 drop - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=p1.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=7, - project_id=p1.id, - org_id=org.id, - ) - - # Project 2: 2 keep, 8 drop - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=2, - project_id=p2.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=8, - project_id=p2.id, - org_id=org.id, - ) - - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org.id], measure=SamplingMeasure.SPANS - ) - - assert len(results[org.id]) == 2 - assert (p1.id, 10.0, 3, 7) in results[org.id] - assert (p2.id, 10.0, 2, 8) in results[org.id] - - -@freeze_time(MOCK_DATETIME) -class TestEndToEndMeasureDispatching(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - """ - End-to-end tests verifying that the boost_low_volume_projects task correctly - dispatches orgs to the right measure and that segment, transaction, and span - processing are all executed correctly. - """ - - @property - def now(self) -> datetime: - return MOCK_DATETIME - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_org_uses_segments_measure_in_with_query_task(self) -> None: - """ - boost_low_volume_projects_of_org_with_query should use SEGMENTS measure. - """ - org = self.create_organization("test-org") - p1 = self.create_project(organization=org) - - # Store span metrics with is_segment=true (used by SEGMENTS measure) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=5, - project_id=p1.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=10, - project_id=p1.id, - org_id=org.id, - ) - - redis_client = get_redis_client_for_ds() - cache_key = generate_sliding_window_org_cache_key(org.id) - redis_client.set(cache_key, 0.5) - - with self.tasks(): - boost_low_volume_projects_of_org_with_query.delay(org.id) - - sample_rate, got_value = get_boost_low_volume_projects_sample_rate( - org.id, p1.id, error_sample_rate_fallback=None - ) - assert got_value - assert sample_rate is not None - - def test_main_task_dispatches_correct_measures(self) -> None: - """ - The main boost_low_volume_projects task should call _process_orgs_for_boost - with SEGMENTS measure for all orgs discovered via GetActiveOrgs scan. - """ - org1 = self.create_organization("org-1") - org2 = self.create_organization("org-2") - p1 = self.create_project(organization=org1) - p2 = self.create_project(organization=org2) - - # Both orgs emit SpanMRI with is_segment=true - for p, org in [(p1, org1), (p2, org2)]: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=p.id, - org_id=org.id, - ) - - with patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects._process_orgs_for_boost" - ) as mock_process: - with self.tasks(): - boost_low_volume_projects() - - # Collect all calls: (org_ids, measure) pairs - calls_by_measure: dict[SamplingMeasure, list[int]] = {} - for call in mock_process.call_args_list: - org_ids = call[0][0] - measure = call[0][1] - calls_by_measure.setdefault(measure, []).extend(org_ids) - - assert org1.id in calls_by_measure.get(SamplingMeasure.SEGMENTS, []) - assert org2.id in calls_by_measure.get(SamplingMeasure.SEGMENTS, []) - - def test_segment_only_org_is_discovered_by_main_task(self) -> None: - """ - An org that emits segment metrics (SpanMRI with is_segment=true) - must be discovered and processed by the main boost_low_volume_projects task. - """ - org = self.create_organization("segment-only-org") - project = self.create_project(organization=org) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=5, - project_id=project.id, - org_id=org.id, - ) - - with patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects._process_orgs_for_boost" - ) as mock_process: - with self.tasks(): - boost_low_volume_projects() - - calls_by_measure: dict[SamplingMeasure, list[int]] = {} - for call in mock_process.call_args_list: - org_ids = call[0][0] - measure = call[0][1] - calls_by_measure.setdefault(measure, []).extend(org_ids) - - assert org.id in calls_by_measure.get(SamplingMeasure.SEGMENTS, []) - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_segments_query_uses_span_mri_with_is_segment_tag(self) -> None: - """ - When processing an org with SEGMENTS measure, the Snuba query should use - SpanMRI and filter by is_segment=true, not TransactionMRI. - """ - org = self.create_organization("test-org") - project = self.create_project(organization=org) - - # Store ONLY span metrics with is_segment=true - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=project.id, - org_id=org.id, - ) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, - minutes_before_now=30, - value=7, - project_id=project.id, - org_id=org.id, - ) - - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org.id], measure=SamplingMeasure.SEGMENTS - ) - - # Should only see the span/segment metrics (3 + 7 = 10), not the transaction metrics (100) - assert results[org.id] == [(project.id, 10.0, 3, 7)] - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_spans_query_uses_span_mri_without_is_segment(self) -> None: - """ - When processing an org with SPANS measure, the Snuba query should use - SpanMRI but NOT filter by is_segment (counts all spans). - """ - org = self.create_organization("test-org") - project = self.create_project(organization=org) - - # Store span metrics WITH is_segment=true - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=3, - project_id=project.id, - org_id=org.id, - ) - - # Store span metrics WITHOUT is_segment tag - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar", "decision": "keep"}, - minutes_before_now=30, - value=7, - project_id=project.id, - org_id=org.id, - ) - - results = fetch_projects_with_total_root_transaction_count_and_rates( - org_ids=[org.id], measure=SamplingMeasure.SPANS - ) - - # SPANS measure should count ALL spans (both with and without is_segment) - # Total = 3 + 7 = 10, all keeps - assert results[org.id] == [(project.id, 10.0, 10, 0)] - - @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_with_query_task_skips_project_mode_orgs(self) -> None: - """ - boost_low_volume_projects_of_org_with_query should early-return for - project-mode orgs without storing any rebalanced rates. - """ - org = self.create_organization("test-org") - p1 = self.create_project(organization=org) - org.update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=5, - project_id=p1.id, - org_id=org.id, - ) - - redis_client = get_redis_client_for_ds() - cache_key = generate_sliding_window_org_cache_key(org.id) - redis_client.set(cache_key, 0.5) - - with self.tasks(): - boost_low_volume_projects_of_org_with_query.delay(org.id) - - sample_rate, got_value = get_boost_low_volume_projects_sample_rate( - org.id, p1.id, error_sample_rate_fallback=None - ) - assert not got_value - assert sample_rate is None diff --git a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_transactions.py b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_transactions.py index a28136bc4082..aebddbc4940e 100644 --- a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_transactions.py +++ b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_transactions.py @@ -4,18 +4,9 @@ from django.utils import timezone from sentry.dynamic_sampling.tasks.boost_low_volume_transactions import ( - FetchProjectTransactionTotals, FetchProjectTransactionVolumes, - ProjectIdentity, - ProjectTransactions, - ProjectTransactionsTotals, - is_project_identity_before, - is_same_project, - merge_transactions, - next_totals, - transactions_zip, ) -from sentry.dynamic_sampling.tasks.common import MEASURE_CONFIGS, GetActiveOrgs +from sentry.dynamic_sampling.tasks.common import MEASURE_CONFIGS from sentry.dynamic_sampling.types import SamplingMeasure from sentry.sentry_metrics import indexer from sentry.snuba.metrics.naming_layer.mri import SpanMRI @@ -28,7 +19,7 @@ @freeze_time(MOCK_DATETIME) -class PrioritiseProjectsSnubaQueryTest(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): +class FetchProjectTransactionVolumesTest(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): @property def now(self): return MOCK_DATETIME @@ -73,27 +64,6 @@ def get_count_for_transaction(self, idx: int, name: str): } return idx + counts[name] - def get_total_counts_for_project(self, idx: int): - """ - Get the total number of transactions and the number of transaction classes for a proj_idx - """ - return 1 + 100 + 1000 + 2000 + 3000 + idx * 5, 5 - - def test_get_orgs_with_transactions_respects_max_orgs(self) -> None: - actual = list(GetActiveOrgs(2, 20)) - - orgs = self.org_ids - # we should return groups of 2 orgs at a time - assert actual == [[orgs[0], orgs[1]], [orgs[2]]] - - def test_get_orgs_with_transactions_respects_max_projs(self) -> None: - actual = list(GetActiveOrgs(10, 5)) - - orgs = [org["org_id"] for org in self.orgs_info] - # since each org has 3 projects and we have a limit of 5 proj - # we should return 2 orgs at a time - assert actual == [[orgs[0], orgs[1]], [orgs[2]]] - def test_fetch_transactions_with_total_volumes_large(self) -> None: """ Create some transactions in some orgs and project and verify @@ -111,33 +81,6 @@ def test_fetch_transactions_with_total_volumes_large(self) -> None: assert name in expected_names assert count == self.get_count_for_transaction(idx, name) - def test_fetch_transactions_with_total_volumes(self) -> None: - """ - Create some transactions in some orgs and project and verify - that the total counts and total transaction types per project are - correctly returned - """ - - orgs = self.org_ids - - for idx, totals in enumerate(FetchProjectTransactionTotals(orgs)): - total_counts, num_classes = self.get_total_counts_for_project(idx) - assert totals["total_num_transactions"] == total_counts - assert totals["total_num_classes"] == num_classes - - def test_fetch_project_transaction_totals_uses_segment_metric_by_default(self) -> None: - """ - Verify that FetchProjectTransactionTotals uses the span count per root metric - with is_segment tag by default (measure=SEGMENTS). - """ - orgs = self.org_ids - fetcher = FetchProjectTransactionTotals(orgs) - - expected_metric_id = indexer.resolve_shared_org(str(SpanMRI.COUNT_PER_ROOT_PROJECT.value)) - assert fetcher.metric_id == expected_metric_id - assert fetcher.measure == SamplingMeasure.SEGMENTS - assert fetcher.tag_filters == MEASURE_CONFIGS[SamplingMeasure.SEGMENTS]["tags"] - def test_fetch_project_transaction_volumes_uses_segment_metric_by_default(self) -> None: """ Verify that FetchProjectTransactionVolumes uses the span count per root metric @@ -151,19 +94,6 @@ def test_fetch_project_transaction_volumes_uses_segment_metric_by_default(self) assert fetcher.measure == SamplingMeasure.SEGMENTS assert fetcher.tag_filters == MEASURE_CONFIGS[SamplingMeasure.SEGMENTS]["tags"] - def test_fetch_project_transaction_totals_uses_segment_metric_when_enabled(self) -> None: - """ - Verify that FetchProjectTransactionTotals uses the span count per root metric - with is_segment tag when measure=SEGMENTS. - """ - orgs = self.org_ids - fetcher = FetchProjectTransactionTotals(orgs, measure=SamplingMeasure.SEGMENTS) - - expected_metric_id = indexer.resolve_shared_org(str(SpanMRI.COUNT_PER_ROOT_PROJECT.value)) - assert fetcher.metric_id == expected_metric_id - assert fetcher.measure == SamplingMeasure.SEGMENTS - assert fetcher.tag_filters == MEASURE_CONFIGS[SamplingMeasure.SEGMENTS]["tags"] - def test_fetch_project_transaction_volumes_uses_segment_metric_when_enabled(self) -> None: """ Verify that FetchProjectTransactionVolumes uses the span count per root metric @@ -179,31 +109,6 @@ def test_fetch_project_transaction_volumes_uses_segment_metric_when_enabled(self assert fetcher.measure == SamplingMeasure.SEGMENTS assert fetcher.tag_filters == MEASURE_CONFIGS[SamplingMeasure.SEGMENTS]["tags"] - @patch("sentry.dynamic_sampling.tasks.boost_low_volume_transactions.raw_snql_query") - def test_fetch_project_transaction_totals_query_includes_is_segment_filter_for_segments( - self, mock_raw_snql_query - ) -> None: - """ - Verify that the query sent to Snuba includes the is_segment=true filter for SEGMENTS measure. - """ - mock_raw_snql_query.return_value = {"data": []} - - orgs = self.org_ids - fetcher = FetchProjectTransactionTotals(orgs, measure=SamplingMeasure.SEGMENTS) - try: - next(fetcher) - except StopIteration: - pass - - assert mock_raw_snql_query.called - call_args = mock_raw_snql_query.call_args - request = call_args[0][0] - - query_str = str(request.query) - is_segment_id = indexer.resolve_shared_org("is_segment") - assert f"tags_raw[{is_segment_id}]" in query_str - assert "'true'" in query_str - @patch("sentry.dynamic_sampling.tasks.boost_low_volume_transactions.raw_snql_query") def test_fetch_project_transaction_volumes_query_includes_is_segment_filter_for_segments( self, mock_raw_snql_query @@ -230,168 +135,3 @@ def test_fetch_project_transaction_volumes_query_includes_is_segment_filter_for_ is_segment_id = indexer.resolve_shared_org("is_segment") assert f"tags_raw[{is_segment_id}]" in query_str assert "'true'" in query_str - - -def test_merge_transactions_with_totals() -> None: - t1: ProjectTransactions = { - "project_id": 1, - "org_id": 2, - "transaction_counts": [("ts1", 10), ("tm2", 100)], - "total_num_transactions": None, - "total_num_classes": None, - } - counts: ProjectTransactionsTotals = { - "project_id": 1, - "org_id": 2, - "total_num_transactions": 5555, - "total_num_classes": 20, - } - actual = merge_transactions(t1, counts) - - expected: ProjectTransactions = { - "project_id": 1, - "org_id": 2, - "transaction_counts": [("ts1", 10), ("tm2", 100)], - "total_num_transactions": 5555, - "total_num_classes": 20, - } - - assert actual == expected - - -def test_merge_transactions_missing_totals() -> None: - t1: ProjectTransactions = { - "project_id": 1, - "org_id": 2, - "transaction_counts": [("ts1", 10), ("tm2", 100)], - "total_num_transactions": None, - "total_num_classes": None, - } - - actual = merge_transactions(t1, None) - - expected: ProjectTransactions = { - "project_id": 1, - "org_id": 2, - "transaction_counts": [("ts1", 10), ("tm2", 100)], - "total_num_transactions": None, - "total_num_classes": None, - } - - assert actual == expected - - -def test_transactions_zip() -> None: - def pt(org_id: int, proj_id: int, add_totals: bool = False): - return { - "project_id": proj_id, - "org_id": org_id, - "transaction_counts": [("tm2", 100), ("tl3", 1000)], - "total_num_transactions": 5000 if add_totals else None, - "total_num_classes": 5 if add_totals else None, - } - - def tot(org_id, proj_id): - return { - "project_id": proj_id, - "org_id": org_id, - "total_num_transactions": 5000, - "total_num_classes": 5, - } - - trans = [pt(1, 1), pt(1, 2), pt(2, 1), pt(2, 3), pt(3, 2)] - totals = [tot(1, 0), tot(1, 2), tot(1, 3), tot(2, 1), tot(2, 4), tot(3, 2)] - - expected = [ - pt(1, 1), - pt(1, 2, True), - pt(2, 1, True), - pt(2, 3), - pt(3, 2, True), - ] - - actual = list(transactions_zip((x for x in totals), (x for x in trans))) - - assert actual == expected - - -def test_same_project() -> None: - p1: ProjectIdentity = {"project_id": 1, "org_id": 2} - p1bis: ProjectIdentity = {"project_id": 1, "org_id": 2} - p2: ProjectIdentity = {"project_id": 1, "org_id": 3} - p3: ProjectIdentity = {"project_id": 2, "org_id": 1} - p4: ProjectIdentity = {"project_id": 3, "org_id": 4} - - assert is_same_project(p1, p1bis) - assert not is_same_project(p1, p2) - assert not is_same_project(p1, p3) - assert not is_same_project(p1, p4) - - -def test_project_before() -> None: - p1: ProjectIdentity = {"project_id": 1, "org_id": 2} - p1bis: ProjectIdentity = {"project_id": 1, "org_id": 2} - p2: ProjectIdentity = {"project_id": 1, "org_id": 3} - p3: ProjectIdentity = {"project_id": 2, "org_id": 2} - p4: ProjectIdentity = {"project_id": 2, "org_id": 1} - - # same project - assert not is_project_identity_before(p1, p1bis) - assert not is_project_identity_before(p1bis, p1) - - # different project_id - assert is_project_identity_before(p1, p2) - assert not is_project_identity_before(p2, p1) - - # different org_id - assert is_project_identity_before(p1, p3) - assert not is_project_identity_before(p3, p1) - - # just different - assert is_project_identity_before(p4, p1) - assert not is_project_identity_before(p1, p4) - - -def test_next_totals() -> None: - def ct(org_id: int, project_id: int) -> ProjectTransactionsTotals: - return { - "project_id": project_id, - "org_id": org_id, - "total_num_transactions": 123, - "total_num_classes": 5, - } - - def pi(org_id: int, project_id: int) -> ProjectIdentity: - return { - "project_id": project_id, - "org_id": org_id, - } - - my_totals = iter([ct(1, 2), ct(1, 4), ct(1, 5), ct(1, 6), ct(1, 9), ct(2, 1)]) - - get_totals = next_totals(my_totals) - - # current should be 1,2 - # ask for something before 1,2 - assert get_totals(pi(0, 1)) is None - assert get_totals(pi(0, 2)) is None - assert get_totals(pi(1, 1)) is None - - # ask for 1.2 - assert get_totals(pi(1, 2)) == ct(1, 2) - # ask again - assert get_totals(pi(1, 2)) is None - # jump a few totals - assert get_totals(pi(1, 6)) == ct(1, 6) - # make sure we don't go back - assert get_totals(pi(1, 5)) is None - # forcing it to go forward jumps just enough - assert get_totals(pi(1, 10)) is None - # but not too much - assert get_totals(pi(1, 11)) is None - assert get_totals(pi(1, 12)) is None - assert get_totals(pi(2, 1)) == ct(2, 1) - # and from now on we return None - assert get_totals(pi(3, 1)) is None - assert get_totals(pi(3, 2)) is None - assert get_totals(pi(3, 3)) is None diff --git a/tests/sentry/dynamic_sampling/tasks/test_common.py b/tests/sentry/dynamic_sampling/tasks/test_common.py index b7286a7e2aca..5e2ae97bf303 100644 --- a/tests/sentry/dynamic_sampling/tasks/test_common.py +++ b/tests/sentry/dynamic_sampling/tasks/test_common.py @@ -4,7 +4,6 @@ from django.utils import timezone from sentry.dynamic_sampling.tasks.common import ( - GetActiveOrgs, GetActiveOrgsVolumes, OrganizationDataVolume, get_organization_volume, @@ -19,55 +18,6 @@ ) -@freeze_time(MOCK_DATETIME) -class TestGetActiveOrgs(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - def setUp(self) -> None: - # create 10 orgs each with 10 transactions - for i in range(10): - org = self.create_organization(f"org-{i}") - for i in range(10): - project = self.create_project(organization=org) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={ - "transaction": "foo_transaction", - "decision": "keep", - "is_segment": "true", - }, - minutes_before_now=30, - value=1, - project_id=project.id, - org_id=org.id, - ) - - @property - def now(self): - return MOCK_DATETIME - - def test_get_active_orgs_no_max_projects(self) -> None: - total_orgs = 0 - for idx, orgs in enumerate(GetActiveOrgs(3)): - num_orgs = len(orgs) - total_orgs += num_orgs - if idx in [0, 1, 2]: - assert num_orgs == 3 # first batch should be full - else: - assert num_orgs == 1 # second should contain the remaining 3 - assert total_orgs == 10 - - def test_get_active_orgs_with_max_projects(self) -> None: - total_orgs = 0 - for orgs in GetActiveOrgs(3, 18): - # we ask for max 18 proj (that's 2 org per request since one org has 10 ) - num_orgs = len(orgs) - total_orgs += num_orgs - assert num_orgs == 2 # only 2 orgs since we limit the number of projects - assert total_orgs == 10 - - -NOW_ISH = timezone.now().replace(second=0, microsecond=0) - - @freeze_time(MOCK_DATETIME) class TestGetActiveOrgsVolumes(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): def setUp(self) -> None: @@ -148,114 +98,6 @@ def test_get_organization_volume_missing_org(self) -> None: assert org_volume is None -@freeze_time(MOCK_DATETIME) -class TestGetActiveOrgsMeasureFiltering(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - """ - Tests that SEGMENTS and SPANS measures filter metrics correctly. - """ - - @property - def now(self): - return MOCK_DATETIME - - def test_segments_measure_only_counts_segment_spans(self) -> None: - """ - Test that SEGMENTS measure only counts SpanMRI with is_segment=true. - """ - org1 = self.create_organization("test-org-1") - project1 = self.create_project(organization=org1) - org2 = self.create_organization("test-org-2") - project2 = self.create_project(organization=org2) - - # Store span metric with is_segment=true (should be counted by SEGMENTS measure) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=project1.id, - org_id=org1.id, - ) - - # Store span metric without is_segment (should NOT be counted by SEGMENTS measure) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar", "decision": "keep"}, - minutes_before_now=30, - value=100, - project_id=project2.id, - org_id=org2.id, - ) - - found_orgs = [] - for orgs in GetActiveOrgs(max_orgs=10, measure=SamplingMeasure.SEGMENTS): - found_orgs.extend(orgs) - - assert org1.id in found_orgs - assert org2.id not in found_orgs - - def test_segments_measure_excludes_non_segment_spans(self) -> None: - """ - Test that SEGMENTS measure excludes SpanMRI without is_segment=true. - """ - org1 = self.create_organization("test-org-1") - project1 = self.create_project(organization=org1) - org2 = self.create_organization("test-org-2") - project2 = self.create_project(organization=org2) - - # Store span metric with is_segment=true (should be counted) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=project1.id, - org_id=org1.id, - ) - - # Store span metric without is_segment (should NOT be counted) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "bar", "decision": "keep"}, - minutes_before_now=30, - value=100, - project_id=project2.id, - org_id=org2.id, - ) - - found_orgs = [] - for orgs in GetActiveOrgs(max_orgs=10, measure=SamplingMeasure.SEGMENTS): - found_orgs.extend(orgs) - - assert org1.id in found_orgs - assert org2.id not in found_orgs - - def test_segments_measure_multiple_orgs(self) -> None: - """ - Test GetActiveOrgs with SEGMENTS measure for multiple organizations. - """ - created_org_ids = [] - for i in range(5): - org = self.create_organization(f"segment-org-{i}") - created_org_ids.append(org.id) - project = self.create_project(organization=org) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "tx", "decision": "keep", "is_segment": "true"}, - minutes_before_now=30, - value=1, - project_id=project.id, - org_id=org.id, - ) - - found_orgs = [] - for orgs in GetActiveOrgs(max_orgs=10, measure=SamplingMeasure.SEGMENTS): - found_orgs.extend(orgs) - - for org_id in created_org_ids: - assert org_id in found_orgs - - @freeze_time(MOCK_DATETIME) class TestGetActiveOrgsVolumesMeasureFiltering(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): """ diff --git a/tests/sentry/dynamic_sampling/tasks/test_tasks.py b/tests/sentry/dynamic_sampling/tasks/test_tasks.py deleted file mode 100644 index bd6baa351808..000000000000 --- a/tests/sentry/dynamic_sampling/tasks/test_tasks.py +++ /dev/null @@ -1,1054 +0,0 @@ -from collections.abc import Callable -from datetime import timedelta -from unittest.mock import MagicMock, patch - -import pytest -from django.utils import timezone - -from sentry.dynamic_sampling import RuleType, generate_rules, get_redis_client_for_ds -from sentry.dynamic_sampling.per_org import cache as per_org_cache -from sentry.dynamic_sampling.rules.base import NEW_MODEL_THRESHOLD_IN_MINUTES -from sentry.dynamic_sampling.rules.biases.recalibration_bias import RecalibrationBias -from sentry.dynamic_sampling.tasks.boost_low_volume_projects import boost_low_volume_projects -from sentry.dynamic_sampling.tasks.boost_low_volume_transactions import ( - boost_low_volume_transactions, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( - generate_boost_low_volume_projects_cache_key, -) -from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_transactions import ( - get_transactions_resampling_rates, -) -from sentry.dynamic_sampling.tasks.helpers.recalibrate_orgs import ( - generate_recalibrate_orgs_cache_key, - generate_recalibrate_projects_cache_key, -) -from sentry.dynamic_sampling.tasks.helpers.sliding_window import ( - generate_sliding_window_org_cache_key, - mark_sliding_window_org_executed, -) -from sentry.dynamic_sampling.tasks.recalibrate_orgs import recalibrate_orgs -from sentry.dynamic_sampling.tasks.sliding_window_org import sliding_window_org -from sentry.dynamic_sampling.types import DynamicSamplingMode -from sentry.snuba.metrics.naming_layer.mri import SpanMRI -from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase -from sentry.testutils.helpers import with_feature -from sentry.testutils.helpers.datetime import freeze_time -from sentry.testutils.helpers.options import override_options - -MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace( - hour=0, minute=0, second=0, microsecond=0 -) - - -class TasksTestCase(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): - @staticmethod - def old_date(): - return timezone.now() - timedelta(minutes=NEW_MODEL_THRESHOLD_IN_MINUTES + 1) - - @staticmethod - def disable_all_biases(project): - project.update_option( - "sentry:dynamic_sampling_biases", - [ - {"id": RuleType.BOOST_ENVIRONMENTS_RULE.value, "active": False}, - {"id": RuleType.IGNORE_HEALTH_CHECKS_RULE.value, "active": False}, - {"id": RuleType.BOOST_LATEST_RELEASES_RULE.value, "active": False}, - {"id": RuleType.BOOST_KEY_TRANSACTIONS_RULE.value, "active": False}, - {"id": RuleType.BOOST_LOW_VOLUME_TRANSACTIONS_RULE.value, "active": False}, - {"id": RuleType.BOOST_REPLAY_ID_RULE.value, "active": False}, - ], - ) - - def create_old_organization(self, name): - return self.create_organization(name=name, date_added=self.old_date()) - - def create_old_project(self, name, organization): - return self.create_project(name=name, organization=organization, date_added=self.old_date()) - - def create_project_and_add_metrics(self, name, count, org, tags=None, is_old=True): - if tags is None: - tags = {"transaction": "foo_transaction", "is_segment": "true"} - elif "is_segment" not in tags: - tags = {**tags, "is_segment": "true"} - - if is_old: - proj = self.create_old_project(name=name, organization=org) - else: - proj = self.create_project(name=name, organization=org) - - self.disable_all_biases(project=proj) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags=tags, - minutes_before_now=30, - value=count, - project_id=proj.id, - org_id=org.id, - ) - - return proj - - def create_project_without_metrics(self, name, org, is_old=True): - if is_old: - proj = self.create_old_project(name=name, organization=org) - else: - proj = self.create_project(name=name, organization=org) - - self.disable_all_biases(project=proj) - - return proj - - -@freeze_time(MOCK_DATETIME) -class TestBoostLowVolumeProjectsTasks(TasksTestCase): - @property - def now(self): - return MOCK_DATETIME - - @staticmethod - def add_sample_rate_per_project(org_id: int, project_id: int, sample_rate: float): - redis_client = get_redis_client_for_ds() - redis_client.hset( - name=generate_boost_low_volume_projects_cache_key(org_id), - key=str(project_id), - value=sample_rate, - ) - - @staticmethod - def sampling_tier_side_effect(*args, **kwargs): - volume = args[1] - - if volume == 20: - return 100_000, 0.25 - # We want to also hardcode the error case, to test how the system reacts to errors. - elif volume == 0: - return None - - return volume, 1.0 - - @staticmethod - def forecasted_volume_side_effect(*args, **kwargs): - return kwargs["volume"] - - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_projects_with_no_dynamic_sampling(self, get_blended_sample_rate): - get_blended_sample_rate.return_value = 0.25 - test_org = self.create_old_organization(name="sample-org") - - self.create_project_and_add_metrics("a", 9, test_org) - self.create_project_and_add_metrics("b", 7, test_org) - self.create_project_and_add_metrics("c", 3, test_org) - self.create_project_and_add_metrics("d", 1, test_org) - - with self.tasks(): - sliding_window_org() - boost_low_volume_projects() - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_projects_simple( - self, - get_blended_sample_rate, - ): - get_blended_sample_rate.return_value = 0.25 - # Create a org - test_org = self.create_old_organization(name="sample-org") - - # Create 4 projects - proj_a = self.create_project_and_add_metrics("a", 9, test_org) - proj_b = self.create_project_and_add_metrics("b", 7, test_org) - proj_c = self.create_project_and_add_metrics("c", 3, test_org) - proj_d = self.create_project_and_add_metrics("d", 1, test_org) - - with self.tasks(): - sliding_window_org() - boost_low_volume_projects() - - # we expect only uniform rule - # also we test here that `generate_rules` can handle trough redis long floats - assert generate_rules(proj_a)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.14814814814814817), - } - assert generate_rules(proj_b)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.1904761904761905), - } - assert generate_rules(proj_c)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.4444444444444444), - } - assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0} - - @with_feature("organizations:dynamic-sampling") - @override_options({"dynamic-sampling.legacy.killswitch": True}) - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_projects_killswitch(self, get_blended_sample_rate: MagicMock) -> None: - get_blended_sample_rate.return_value = 0.25 - test_org = self.create_old_organization(name="sample-org") - self.create_project_and_add_metrics("a", 9, test_org) - self.create_project_and_add_metrics("b", 1, test_org) - - with self.tasks(): - boost_low_volume_projects() - - redis_client = get_redis_client_for_ds() - assert redis_client.hgetall(generate_boost_low_volume_projects_cache_key(test_org.id)) == {} - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_projects_simple_with_empty_project( - self, - get_blended_sample_rate, - ): - get_blended_sample_rate.return_value = 0.25 - test_org = self.create_old_organization(name="sample-org") - - proj_a = self.create_project_and_add_metrics("a", 9, test_org) - proj_b = self.create_project_and_add_metrics("b", 7, test_org) - proj_c = self.create_project_and_add_metrics("c", 3, test_org) - proj_d = self.create_project_and_add_metrics("d", 1, test_org) - proj_e = self.create_project_without_metrics("e", test_org) - - with self.tasks(): - sliding_window_org() - boost_low_volume_projects() - - # we expect only uniform rule - # also we test here that `generate_rules` can handle trough redis long floats - assert generate_rules(proj_a)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.14814814814814817), - } - assert generate_rules(proj_b)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.1904761904761905), - } - assert generate_rules(proj_c)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.4444444444444444), - } - assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0} - assert generate_rules(proj_e)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0} - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume") - @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume") - def test_boost_low_volume_projects_simple_with_sliding_window_org_from_cache( - self, - extrapolate_monthly_volume, - get_transaction_sampling_tier_for_volume, - get_blended_sample_rate, - ): - extrapolate_monthly_volume.side_effect = self.forecasted_volume_side_effect - get_transaction_sampling_tier_for_volume.side_effect = self.sampling_tier_side_effect - get_blended_sample_rate.return_value = 0.8 - - test_org = self.create_old_organization(name="sample-org") - - proj_a = self.create_project_and_add_metrics("a", 9, test_org) - proj_b = self.create_project_and_add_metrics("b", 7, test_org) - proj_c = self.create_project_and_add_metrics("c", 3, test_org) - proj_d = self.create_project_and_add_metrics("d", 1, test_org) - - with self.tasks(): - sliding_window_org() - boost_low_volume_projects() - - # we expect only uniform rule - # also we test here that `generate_rules` can handle trough redis long floats - assert generate_rules(proj_a)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.14814814814814817), - } - assert generate_rules(proj_b)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.1904761904761905), - } - assert generate_rules(proj_c)[0]["samplingValue"] == { - "type": "sampleRate", - "value": pytest.approx(0.4444444444444444), - } - assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0} - - @with_feature("organizations:dynamic-sampling") - @patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects.schedule_invalidate_project_config" - ) - @patch("sentry.quotas.backend.get_blended_sample_rate") - @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume") - @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume") - def test_config_invalidation_when_sample_rates_change( - self, - extrapolate_monthly_volume, - get_transaction_sampling_tier_for_volume, - get_blended_sample_rate, - schedule_invalidate_project_config, - ): - extrapolate_monthly_volume.side_effect = self.forecasted_volume_side_effect - get_transaction_sampling_tier_for_volume.side_effect = self.sampling_tier_side_effect - get_blended_sample_rate.return_value = 0.8 - - test_org = self.create_old_organization(name="sample-org") - - proj_a = self.create_project_and_add_metrics("a", 9, test_org) - proj_b = self.create_project_and_add_metrics("b", 7, test_org) - - self.add_sample_rate_per_project(org_id=test_org.id, project_id=proj_a.id, sample_rate=0.1) - self.add_sample_rate_per_project(org_id=test_org.id, project_id=proj_b.id, sample_rate=0.2) - - with self.tasks(): - sliding_window_org() - boost_low_volume_projects() - - assert schedule_invalidate_project_config.call_count == 2 - - @with_feature("organizations:dynamic-sampling") - @patch( - "sentry.dynamic_sampling.tasks.boost_low_volume_projects.schedule_invalidate_project_config" - ) - @patch("sentry.quotas.backend.get_blended_sample_rate") - @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume") - @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume") - def test_config_invalidation_when_sample_rates_do_not_change( - self, - extrapolate_monthly_volume, - get_transaction_sampling_tier_for_volume, - get_blended_sample_rate, - schedule_invalidate_project_config, - ): - extrapolate_monthly_volume.side_effect = self.forecasted_volume_side_effect - get_transaction_sampling_tier_for_volume.side_effect = self.sampling_tier_side_effect - get_blended_sample_rate.return_value = 1.0 - - test_org = self.create_old_organization(name="sample-org") - - proj_a = self.create_project_and_add_metrics("a", 9, test_org) - proj_b = self.create_project_and_add_metrics("b", 7, test_org) - - self.add_sample_rate_per_project(org_id=test_org.id, project_id=proj_a.id, sample_rate=1.0) - self.add_sample_rate_per_project(org_id=test_org.id, project_id=proj_b.id, sample_rate=1.0) - - with self.tasks(): - boost_low_volume_projects() - - schedule_invalidate_project_config.assert_not_called() - - -@freeze_time(MOCK_DATETIME) -class TestBoostLowVolumeTransactionsTasks(TasksTestCase): - @property - def now(self): - return MOCK_DATETIME - - def setUp(self) -> None: - super().setUp() - self.orgs_info = [] - num_orgs = 3 - num_proj_per_org = 3 - for org_idx in range(num_orgs): - org = self.create_old_organization(f"test-org{org_idx}") - org_info = {"org_id": org.id, "project_ids": []} - self.orgs_info.append(org_info) - for proj_idx in range(num_proj_per_org): - p = self.create_old_project(name=f"test-project-{proj_idx}", organization=org) - org_info["project_ids"].append(p.id) - # create 5 transaction types - for name in ["ts1", "ts2", "tm3", "tl4", "tl5"]: - # make up some unique count - idx = org_idx * num_orgs + proj_idx - num_transactions = self.get_count_for_transaction(idx, name) - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": name, "is_segment": "true"}, - minutes_before_now=30, - value=num_transactions, - project_id=p.id, - org_id=org.id, - ) - self.org_ids = [org["org_id"] for org in self.orgs_info] - - def get_count_for_transaction(self, idx: int, name: str): - """ - Create some known count based on transaction name and the order (based on org and project) - """ - counts = { - "ts1": 1, - "ts2": 100, - "tm3": 1000, - "tl4": 2000, - "tl5": 3000, - } - return idx + counts[name] - - @staticmethod - def flush_redis(): - get_redis_client_for_ds().flushdb() - - @staticmethod - def set_boost_low_volume_projects_cache_entry(org_id: int, project_id: int, value: str): - redis = get_redis_client_for_ds() - cache_key = generate_boost_low_volume_projects_cache_key(org_id=org_id) - redis.hset(name=cache_key, key=str(project_id), value=value) - - def set_boost_low_volume_projects_sample_rate( - self, org_id: int, project_id: int, sample_rate: float - ): - self.set_boost_low_volume_projects_cache_entry(org_id, project_id, str(sample_rate)) - - def set_prioritise_by_project_invalid(self, org_id: int, project_id: int): - # We want also to test for this case in order to verify the fallback to the `get_blended_sample_rate`. - self.set_boost_low_volume_projects_cache_entry(org_id, project_id, "invalid") - - def for_all_orgs_and_projects(self, block: Callable[[int, int], None]): - for org in self.orgs_info: - org_id = org["org_id"] - for project_id in org["project_ids"]: - block(org_id, project_id) - - def set_boost_low_volume_projects_invalid_for_all(self): - self.for_all_orgs_and_projects( - lambda org_id, project_id: self.set_prioritise_by_project_invalid(org_id, project_id) - ) - - def set_boost_low_volume_projects_for_all(self, sample_rate: float): - self.for_all_orgs_and_projects( - lambda org_id, project_id: self.set_boost_low_volume_projects_sample_rate( - org_id, project_id, sample_rate - ) - ) - - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_transactions_with_blended_sample_rate_and_no_dynamic_sampling( - self, get_blended_sample_rate - ): - """ - Create orgs projects & transactions and then check that the rebalancing model is not called because dynamic - sampling is disabled - """ - BLENDED_RATE = 0.25 - get_blended_sample_rate.return_value = BLENDED_RATE - - with self.tasks(): - boost_low_volume_transactions() - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_transactions_with_blended_sample_rate( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Create orgs projects & transactions and then check that the task creates rebalancing data - in Redis. - """ - BLENDED_RATE = 0.25 - get_blended_sample_rate.return_value = BLENDED_RATE - - with self.tasks(): - boost_low_volume_transactions() - - # now redis should contain rebalancing data for our projects - for org in self.orgs_info: - org_id = org["org_id"] - for proj_id in org["project_ids"]: - tran_rate, global_rate = get_transactions_resampling_rates( - org_id=org_id, proj_id=proj_id, default_rate=0.1 - ) - for transaction_name in ["ts1", "ts2", "tm3", "tl4", "tl5"]: - assert ( - transaction_name in tran_rate - ) # check we have some rate calculated for each transaction - assert global_rate == BLENDED_RATE - - @with_feature("organizations:dynamic-sampling") - @override_options({"dynamic-sampling.legacy.killswitch": True}) - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_transactions_killswitch( - self, get_blended_sample_rate: MagicMock - ) -> None: - get_blended_sample_rate.return_value = 0.25 - - with self.tasks(): - boost_low_volume_transactions() - - for org in self.orgs_info: - for proj_id in org["project_ids"]: - tran_rate, _ = get_transactions_resampling_rates( - org_id=org["org_id"], proj_id=proj_id, default_rate=0.1 - ) - assert tran_rate == {} - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_transactions_with_sliding_window_org( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Create orgs projects & transactions and then check that the task creates rebalancing data - in Redis with the sliding window per org enabled. - """ - BLENDED_RATE = 0.25 - get_blended_sample_rate.return_value = BLENDED_RATE - - for sliding_window_step, used_sample_rate in ((1, 1.0), (2, BLENDED_RATE), (3, 0.5)): - # We flush redis after each run, to make sure no data persists. - self.flush_redis() - - # No value in cache and sliding window org executed. - if sliding_window_step == 1: - mark_sliding_window_org_executed() - # Invalid value in cache. - elif sliding_window_step == 2: - self.set_boost_low_volume_projects_invalid_for_all() - # Value in cache. - elif sliding_window_step == 3: - self.set_boost_low_volume_projects_for_all(used_sample_rate) - - with self.tasks(): - boost_low_volume_transactions() - - # now redis should contain rebalancing data for our projects - for org in self.orgs_info: - org_id = org["org_id"] - for proj_id in org["project_ids"]: - tran_rate, global_rate = get_transactions_resampling_rates( - org_id=org_id, proj_id=proj_id, default_rate=0.1 - ) - - if sliding_window_step == 1: - # If the sample rate is 100%, we will not find anything in cache, since we don't - # need to run and store the rebalancing. - assert tran_rate == {} - else: - # If the sample rate is < 100%, we want to check that in cache we have a value with - # the correct global rate. - for transaction_name in ["ts1", "ts2", "tm3", "tl4", "tl5"]: - assert ( - transaction_name in tran_rate - ) # check we have some rate calculated for each transaction - - assert global_rate == used_sample_rate - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_boost_low_volume_transactions_partial( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Test the V2 algorithm is used, only specified projects are balanced and the - rest get a global rate - - Create orgs projects & transactions and then check that the task creates rebalancing data - in Redis - """ - BLENDED_RATE = 0.25 - get_blended_sample_rate.return_value = BLENDED_RATE - - with self.options( - { - "dynamic-sampling.prioritise_transactions.num_explicit_large_transactions": 1, - "dynamic-sampling.prioritise_transactions.rebalance_intensity": 0.7, - } - ): - with self.tasks(): - boost_low_volume_transactions() - - # now redis should contain rebalancing data for our projects - for org in self.orgs_info: - org_id = org["org_id"] - for proj_id in org["project_ids"]: - tran_rate, implicit_rate = get_transactions_resampling_rates( - org_id=org_id, proj_id=proj_id, default_rate=0.1 - ) - # explicit transactions - for transaction_name in ["tl5"]: - assert ( - transaction_name in tran_rate - ) # check we have some rate calculated for each transaction - # implicit transactions - for transaction_name in ["ts1", "ts2", "tm3", "tl4"]: - assert ( - transaction_name not in tran_rate - ) # check we have some rate calculated for each transaction - # we do have some different rate for implicit transactions - assert implicit_rate != BLENDED_RATE - - -@freeze_time(MOCK_DATETIME) -class TestRecalibrateOrgsTasks(TasksTestCase): - @property - def now(self): - return MOCK_DATETIME - - def setUp(self) -> None: - super().setUp() - self.orgs_info = [] - self.orgs = [] - self.num_proj = 2 - self.orgs_sampling = [10, 20, 40] - # create some orgs, projects and transactions - for org_rate in self.orgs_sampling: - org = self.create_old_organization(f"test-org-{org_rate}") - org_info = {"org_id": org.id, "project_ids": [], "projects": []} - self.orgs_info.append(org_info) - self.orgs.append(org) - for proj_idx in range(self.num_proj): - p = self.create_old_project(name=f"test-project-{proj_idx}", organization=org) - org_info["projects"].append(p) - org_info["project_ids"].append(p.id) - self.add_metrics(org, p, org_rate) - - def add_metrics(self, org, project, sample_rate): - base_tags = {"transaction": "trans-x", "is_segment": "true"} - - if sample_rate < 100: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={**base_tags, "decision": "drop"}, - minutes_before_now=2, - value=100 - sample_rate, - project_id=project.id, - org_id=org.id, - ) - if sample_rate > 0: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={**base_tags, "decision": "keep"}, - minutes_before_now=2, - value=sample_rate, - project_id=project.id, - org_id=org.id, - ) - - def add_measure_metrics( - self, - org, - project, - *, - segment_keep: int, - segment_drop: int, - ) -> None: - segment_tags = {"transaction": "trans-x", "is_segment": "true"} - if segment_drop: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={**segment_tags, "decision": "drop"}, - minutes_before_now=2, - value=segment_drop, - project_id=project.id, - org_id=org.id, - ) - if segment_keep: - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={**segment_tags, "decision": "keep"}, - minutes_before_now=2, - value=segment_keep, - project_id=project.id, - org_id=org.id, - ) - - @staticmethod - def set_sliding_window_org_cache_entry(org_id: int, value: str): - redis = get_redis_client_for_ds() - cache_key = generate_sliding_window_org_cache_key(org_id=org_id) - redis.set(cache_key, value) - - def set_sliding_window_org_sample_rate(self, org_id: int, sample_rate: float): - self.set_sliding_window_org_cache_entry(org_id, str(sample_rate)) - - def for_all_orgs(self, block: Callable[[int], None]): - for org in self.orgs_info: - org_id = org["org_id"] - block(org_id) - - def set_sliding_window_org_sample_rate_for_all(self, sample_rate: float): - self.for_all_orgs( - lambda org_id: self.set_sliding_window_org_sample_rate(org_id, sample_rate) - ) - - @patch("sentry.dynamic_sampling.tasks.recalibrate_orgs.compute_adjusted_factor") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_with_no_dynamic_sampling( - self, get_blended_sample_rate, computed_adjusted_factor - ): - """ - Test that the recalibration of orgs doesn't happen if dynamic sampling is not enabled - """ - get_blended_sample_rate.return_value = 0.1 - self.set_sliding_window_org_sample_rate_for_all(0.2) - - with self.tasks(): - recalibrate_orgs() - - computed_adjusted_factor.assert_not_called() - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_with_sliding_window_org( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Test that the org are going to be rebalanced at 20% and that the sample rate used is the one from the sliding - window org. - - The first org is 10%, so we should increase the sampling - The second org is at 20%, so we are spot on - The third is at 40%, so we should decrease the sampling - """ - get_blended_sample_rate.return_value = 0.1 - self.set_sliding_window_org_sample_rate_for_all(0.2) - - redis_client = get_redis_client_for_ds() - - with self.tasks(): - recalibrate_orgs() - - for idx, org in enumerate(self.orgs): - cache_key = generate_recalibrate_orgs_cache_key(org.id) - val = redis_client.get(cache_key) - - if idx == 0: - assert val is not None - # we sampled at 10% half of what we want so we should adjust by 2 - assert float(val) == 2.0 - elif idx == 1: - # we sampled at 20% we should be spot on (no adjustment) - assert val is None - elif idx == 2: - assert val is not None - # we sampled at 40% twice as much as we wanted we should adjust by 0.5 - assert float(val) == 0.5 - - # now if we run it again (with the same data in the database, the algorithm - # should double down... the previous factor didn't do anything so apply it again) - with self.tasks(): - recalibrate_orgs() - - for idx, org in enumerate(self.orgs): - cache_key = generate_recalibrate_orgs_cache_key(org.id) - val = redis_client.get(cache_key) - - if idx == 0: - assert val is not None - # we sampled at 10% when already having a factor of two half of what we want so we - # should double the current factor to 4 - assert float(val) == 4.0 - elif idx == 1: - # we sampled at 20% we should be spot on (no adjustment) - assert val is None - elif idx == 2: - assert val is not None - # we sampled at 40% twice as much as we wanted we already have a factor of 0.5 - # half it again to 0.25 - assert float(val) == 0.25 - - @with_feature("organizations:dynamic-sampling") - @override_options({"dynamic-sampling.legacy.killswitch": True}) - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_killswitch(self, get_blended_sample_rate: MagicMock) -> None: - get_blended_sample_rate.return_value = 0.1 - self.set_sliding_window_org_sample_rate_for_all(0.2) - - with self.tasks(): - recalibrate_orgs() - - redis_client = get_redis_client_for_ds() - for org in self.orgs: - assert redis_client.get(generate_recalibrate_orgs_cache_key(org.id)) is None - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_skips_orgs_served_the_per_org_factor( - self, get_blended_sample_rate: MagicMock - ) -> None: - """An org served the per-org factor keeps the legacy cache untouched. - - Both loops step by previous * target / measured. Writing here as well would step a - factor nothing applies, so it would walk to a rebalance bound. - """ - get_blended_sample_rate.return_value = 0.1 - self.set_sliding_window_org_sample_rate_for_all(0.2) - - served_org = self.orgs[0] - redis_client = get_redis_client_for_ds() - - with ( - override_options({"dynamic-sampling.per_org.serving-org-ids": [served_org.id]}), - self.tasks(), - ): - recalibrate_orgs() - - # The served org sampled at 10% against a 20% target, so the legacy task would have - # written a factor of 2.0 here. - assert redis_client.get(generate_recalibrate_orgs_cache_key(served_org.id)) is None - - # Every other org still has its factor written by the legacy task. - other_factor = redis_client.get(generate_recalibrate_orgs_cache_key(self.orgs[2].id)) - assert other_factor is not None - assert float(other_factor) == 0.5 - - @with_feature("organizations:dynamic-sampling") - def test_recalibrate_orgs_continues_from_the_per_org_factor_after_switching_back( - self, - ) -> None: - """An org switched back to the legacy pipeline steps from the factor it was served. - - The legacy key expired while the per-org pipeline served the org, so without the - carry-over the correction would restart from 1.0 and drop the whole boost. - """ - self.set_sliding_window_org_sample_rate_for_all(0.2) - - # This org stored metrics at a 10% sampling rate, so it measures at 0.1. - switched_back_org = self.orgs[0] - per_org_cache.set_adjusted_factor(switched_back_org.id, 3.0) - - redis_client = get_redis_client_for_ds() - - with self.tasks(): - recalibrate_orgs() - - # 3.0 * (0.2 target / 0.1 measured), instead of the 2.0 a restart from 1.0 gives. - factor = redis_client.get(generate_recalibrate_orgs_cache_key(switched_back_org.id)) - assert factor is not None - assert float(factor) == 6.0 - - @with_feature("organizations:dynamic-sampling") - @with_feature("organizations:dynamic-sampling-custom") - def test_recalibrate_orgs_with_custom_ds(self) -> None: - """ - Test several organizations with mixed sampling mode. - - The first org is 10%, so we should increase the sampling - The second org is at 20%, so we are spot on - The third is at 40%, so we should decrease the sampling - """ - - # First two orgs have a 20% sample rate configured, third one is in project mode - self.orgs[0].update_option("sentry:target_sample_rate", 0.2) - self.orgs[1].update_option("sentry:target_sample_rate", 0.2) - self.orgs[2].update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) - - # First project gets same 20% sample rate, the other one stays at implicit 100% - p1, p2 = self.orgs_info[2]["projects"] - p1.update_option("sentry:target_sample_rate", 0.2) - - with self.tasks(): - recalibrate_orgs() - - redis_client = get_redis_client_for_ds() - - # First org was sampled at 10%, should be recalibrated at 2x to 20%. - assert redis_client.get(generate_recalibrate_orgs_cache_key(self.orgs[0].id)) == "2.0" - # Second org was sampled at 20%, should not be recalibrated. - assert redis_client.get(generate_recalibrate_orgs_cache_key(self.orgs[1].id)) is None - - # Third org should not have org-level recalibration. - assert redis_client.get(generate_recalibrate_orgs_cache_key(self.orgs[2].id)) is None - # First project was sampled at 40%, should be recalibrated at 0.5x to 20%. - assert redis_client.get(generate_recalibrate_projects_cache_key(p1.id)) == "0.5" - # Second project was sampled at 40%, should be recalibrated at 2.5x to 100%. - assert redis_client.get(generate_recalibrate_projects_cache_key(p2.id)) == "2.5" - - assert RecalibrationBias().generate_rules(p1, base_sample_rate=1.0) == [ - { - "samplingValue": {"type": "factor", "value": 0.5}, - "type": "trace", - "condition": {"op": "and", "inner": []}, - "id": 1004, - } - ] - - assert RecalibrationBias().generate_rules(p2, base_sample_rate=1.0) == [ - { - "samplingValue": {"type": "factor", "value": 2.5}, - "type": "trace", - "condition": {"op": "and", "inner": []}, - "id": 1004, - } - ] - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_rules_generation_with_recalibrate_orgs( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Test that we pass rebalancing values all the way to the rules. - """ - get_blended_sample_rate.return_value = 0.20 - - with self.tasks(): - recalibrate_orgs() - - for org_idx, org in enumerate(self.orgs): - for project in org.project_set.all(): - rules = RecalibrationBias().generate_rules(project, base_sample_rate=0.5) - if org_idx == 0: - # we sampled at 10% half of what we want so we should adjust by 2 - assert rules == [ - { - "samplingValue": {"type": "factor", "value": 2.0}, - "type": "trace", - "condition": {"op": "and", "inner": []}, - "id": 1004, - } - ] - elif org_idx == 1: - # we sampled at 20% we should be spot on (no rule) - assert rules == [] - elif org_idx == 2: - # we sampled at 40% twice as much as we wanted we should adjust by 0.5 - assert rules == [ - { - "samplingValue": {"type": "factor", "value": 0.5}, - "type": "trace", - "condition": {"op": "and", "inner": []}, - "id": 1004, - } - ] - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_uses_segments_measure( - self, get_blended_sample_rate: MagicMock - ) -> None: - """ - Test that all orgs use segment metrics (SEGMENTS measure) for recalibration. - """ - get_blended_sample_rate.return_value = 0.1 - self.set_sliding_window_org_sample_rate_for_all(0.2) - - redis_client = get_redis_client_for_ds() - - with self.tasks(): - recalibrate_orgs() - - # First org should be recalibrated (sampled at 10%, target 20% -> factor 2.0) - cache_key = generate_recalibrate_orgs_cache_key(self.orgs[0].id) - val = redis_client.get(cache_key) - assert val is not None - assert float(val) == 2.0 - - # Second org sampled at 20%, target 20% -> no adjustment needed - cache_key = generate_recalibrate_orgs_cache_key(self.orgs[1].id) - val = redis_client.get(cache_key) - assert val is None - - # Third org sampled at 40%, target 20% -> factor 0.5 - cache_key = generate_recalibrate_orgs_cache_key(self.orgs[2].id) - val = redis_client.get(cache_key) - assert val is not None - assert float(val) == 0.5 - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.quotas.backend.get_blended_sample_rate") - def test_recalibrate_orgs_multiple_orgs_with_different_volumes( - self, get_blended_sample_rate: MagicMock - ) -> None: - get_blended_sample_rate.return_value = 0.2 - org1 = self.create_old_organization("org-1") - org2 = self.create_old_organization("org-2") - project1 = self.create_old_project(name="project-1", organization=org1) - project2 = self.create_old_project(name="project-2", organization=org2) - - self.add_measure_metrics( - org1, - project1, - segment_keep=10, - segment_drop=90, - ) - self.add_measure_metrics( - org2, - project2, - segment_keep=50, - segment_drop=50, - ) - - self.set_sliding_window_org_sample_rate(org1.id, 0.2) - self.set_sliding_window_org_sample_rate(org2.id, 0.2) - - with self.tasks(): - recalibrate_orgs() - - redis_client = get_redis_client_for_ds() - org1_value = redis_client.get(generate_recalibrate_orgs_cache_key(org1.id)) - org2_value = redis_client.get(generate_recalibrate_orgs_cache_key(org2.id)) - - assert org1_value is not None - assert org2_value is not None - assert float(org1_value) == pytest.approx(2.0) - assert float(org2_value) == pytest.approx(0.4) - - -@freeze_time(MOCK_DATETIME) -class TestSlidingWindowOrgTask(TasksTestCase): - """Tests for the sliding_window_org task with measure parameter support.""" - - @property - def now(self): - return MOCK_DATETIME - - def setUp(self) -> None: - super().setUp() - self.orgs = [] - # Create orgs with different volumes - for i, volume in enumerate([100, 500, 1000]): - org = self.create_old_organization(f"test-org-{i}") - self.orgs.append(org) - project = self.create_old_project(name=f"test-project-{i}", organization=org) - - self.store_performance_metric( - name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, - tags={"transaction": "foo", "is_segment": "true"}, - minutes_before_now=30, - value=volume, - project_id=project.id, - org_id=org.id, - ) - - @with_feature("organizations:dynamic-sampling") - @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume") - @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume") - def test_sliding_window_org_processes_all_orgs_by_default( - self, - get_transaction_sampling_tier_for_volume: MagicMock, - extrapolate_monthly_volume: MagicMock, - ) -> None: - """ - Test that sliding_window_org processes all orgs using SEGMENTS measure by default. - """ - extrapolate_monthly_volume.side_effect = lambda volume, hours: volume - get_transaction_sampling_tier_for_volume.return_value = (1000, 0.25) - redis_client = get_redis_client_for_ds() - - with self.tasks(): - sliding_window_org() - - # All orgs should have cache entries - for org in self.orgs: - cache_key = generate_sliding_window_org_cache_key(org.id) - val = redis_client.get(cache_key) - assert val is not None, f"Org {org.id} should have sliding window cache entry" - - @with_feature("organizations:dynamic-sampling") - @override_options({"dynamic-sampling.legacy.killswitch": True}) - @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume") - @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume") - def test_sliding_window_org_killswitch( - self, - get_transaction_sampling_tier_for_volume: MagicMock, - extrapolate_monthly_volume: MagicMock, - ) -> None: - extrapolate_monthly_volume.side_effect = lambda volume, hours: volume - get_transaction_sampling_tier_for_volume.return_value = (1000, 0.25) - redis_client = get_redis_client_for_ds() - - with self.tasks(): - sliding_window_org() - - for org in self.orgs: - assert redis_client.get(generate_sliding_window_org_cache_key(org.id)) is None From 54a0f6c17f8b7bceb71aedfbef0435ea78e8eeba Mon Sep 17 00:00:00 2001 From: Simon Hellmayr Date: Thu, 10 Sep 2026 11:04:26 +0200 Subject: [PATCH 2/2] ref(dynamic-sampling): Keep the endpoint's project balancing out of the jobs removal The organization details endpoint still runs the legacy project balancing when the sampling mode or the target rate changes. Moving it onto the per-org pipeline is its own change, so leave the endpoint and the query and task it uses in place. Only the scheduled entry task, its fan-out task and the active-orgs scan go. Co-Authored-By: Claude Fable 5.1 --- src/sentry/conf/server.py | 1 + .../core/endpoints/organization_details.py | 60 +- .../tasks/boost_low_volume_projects.py | 383 +++++++++++ .../dynamic_sampling/tasks/constants.py | 1 + src/sentry/dynamic_sampling/tasks/utils.py | 41 ++ src/sentry/snuba/referrer.py | 3 + .../endpoints/test_organization_details.py | 52 +- .../tasks/test_boost_low_volume_projects.py | 622 ++++++++++++++++++ 8 files changed, 1108 insertions(+), 55 deletions(-) create mode 100644 src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py create mode 100644 src/sentry/dynamic_sampling/tasks/utils.py create mode 100644 tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py diff --git a/src/sentry/conf/server.py b/src/sentry/conf/server.py index fcf74b91a808..42c18e0cb315 100644 --- a/src/sentry/conf/server.py +++ b/src/sentry/conf/server.py @@ -881,6 +881,7 @@ def SOCIAL_AUTH_DEFAULT_USERNAME() -> str: "sentry.demo_mode.tasks", "sentry.dynamic_sampling.per_org.feature_cache", "sentry.dynamic_sampling.per_org.scheduler", + "sentry.dynamic_sampling.tasks.boost_low_volume_projects", "sentry.feedback.tasks.update_user_reports", "sentry.hybridcloud.tasks.deliver_from_outbox", "sentry.hybridcloud.tasks.deliver_webhooks", diff --git a/src/sentry/core/endpoints/organization_details.py b/src/sentry/core/endpoints/organization_details.py index 4eb32300db64..ec45db4b76fc 100644 --- a/src/sentry/core/endpoints/organization_details.py +++ b/src/sentry/core/endpoints/organization_details.py @@ -73,16 +73,16 @@ SEER_AUTOMATED_RUN_STOPPING_POINT_DEFAULT, SEER_DEFAULT_CODING_AGENT_DEFAULT, TARGET_SAMPLE_RATE_DEFAULT, + ObjectStatus, ) from sentry.core.endpoints.project_details import MAX_SENSITIVE_FIELD_CHARS from sentry.deletions.models.scheduleddeletion import CellScheduledDeletion -from sentry.dynamic_sampling.per_org.calculations import run_project_balancing -from sentry.dynamic_sampling.per_org.configuration import ( - CustomDynamicSamplingOrganizationConfiguration, +from sentry.dynamic_sampling.tasks.boost_low_volume_projects import ( + boost_low_volume_projects_of_org_with_query, + calculate_sample_rates_of_projects, + query_project_counts_by_org, ) -from sentry.dynamic_sampling.per_org.queries import get_eap_project_volumes -from sentry.dynamic_sampling.per_org.scheduler import run_calculations_per_org_task_entry -from sentry.dynamic_sampling.types import DynamicSamplingMode +from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure from sentry.dynamic_sampling.utils import ( has_custom_dynamic_sampling, is_organization_mode_sampling, @@ -101,6 +101,7 @@ from sentry.models.options.project_option import ProjectOption from sentry.models.organization import Organization, OrganizationStatus from sentry.models.organizationmember import OrganizationMember +from sentry.models.project import Project from sentry.organizations.services.organization import organization_service from sentry.organizations.services.organization.model import ( RpcOrganization, @@ -1217,7 +1218,7 @@ def put( if "samplingMode" in changed_data: with transaction.atomic(router.db_for_write(ProjectOption)): if is_project_mode_sampling(organization): - self._compute_project_target_sample_rates(organization) + self._compute_project_target_sample_rates(request, organization) organization.delete_option("sentry:target_sample_rate") changed_data["samplingMode"] = "to Advanced Mode" @@ -1245,7 +1246,9 @@ def put( if is_org_mode and ( "samplingMode" in changed_data or "targetSampleRate" in changed_data ): - run_calculations_per_org_task_entry.delay(organization.id) + boost_low_volume_projects_of_org_with_query.delay( + organization.id, + ) if is_org_mode and "defaultAutofixAutomationTuning" in changed_data: organization.update_option( @@ -1311,23 +1314,36 @@ def put( return self.respond(context) return self.respond(as_validation_errors(serializer), status=status.HTTP_400_BAD_REQUEST) - def _compute_project_target_sample_rates(self, organization: Organization) -> None: - """Seed every active project with the rate it had under organization mode. - - Balances the last 30 days of volume at the organization's target rate, which is - still set when this runs, so that the switch to project mode keeps the stored - volume per project unchanged. - """ + def _compute_project_target_sample_rates(self, request: Request, organization: Organization): # TODO: this will take a long time for organizations with a lot of projects # so we need to refactor this into an async task we can run and observe - config = CustomDynamicSamplingOrganizationConfiguration(organization) - project_volumes = get_eap_project_volumes(config, time_interval=timedelta(days=30)) - for rebalanced_item in run_project_balancing(config, project_volumes): - ProjectOption.objects.update_or_create( - project_id=rebalanced_item.id, - key="sentry:target_sample_rate", - defaults={"value": round(rebalanced_item.new_sample_rate, 4)}, + org_id = organization.id + measure = SamplingMeasure.SEGMENTS + projects_with_tx_count_and_rates = [] + for chunk in query_project_counts_by_org( + [org_id], measure, query_interval=timedelta(days=30) + ): + for row in chunk: + projects_with_tx_count_and_rates.append(row[1:]) + + rebalanced_projects = calculate_sample_rates_of_projects( + org_id, projects_with_tx_count_and_rates + ) + + project_ids = set( + Project.objects.filter(organization_id=org_id, status=ObjectStatus.ACTIVE).values_list( + "id", flat=True ) + ) + + if rebalanced_projects is not None: + for rebalanced_item in rebalanced_projects: + if int(rebalanced_item.id) in project_ids: + ProjectOption.objects.update_or_create( + project_id=rebalanced_item.id, + key="sentry:target_sample_rate", + defaults={"value": round(rebalanced_item.new_sample_rate, 4)}, + ) def handle_delete(self, request: Request, organization: Organization): """ diff --git a/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py b/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py new file mode 100644 index 000000000000..e013341dd74a --- /dev/null +++ b/src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py @@ -0,0 +1,383 @@ +from __future__ import annotations + +import logging +from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence +from datetime import timedelta + +import sentry_sdk +from snuba_sdk import ( + Column, + Condition, + Direction, + Entity, + Function, + Granularity, + Limit, + LimitBy, + Op, + OrderBy, + Query, + Request, +) +from taskbroker_client.retry import Retry + +from sentry import quotas +from sentry.constants import ObjectStatus +from sentry.dynamic_sampling.models.common import RebalancedItem, guarded_run +from sentry.dynamic_sampling.models.projects_rebalancing import ( + ProjectsRebalancingInput, + ProjectsRebalancingModel, +) +from sentry.dynamic_sampling.rules.utils import ( + DecisionDropCount, + DecisionKeepCount, + OrganizationId, + ProjectId, + get_redis_client_for_ds, +) +from sentry.dynamic_sampling.tasks.common import ( + MEASURE_CONFIGS, + are_equal_with_epsilon, + sample_rate_to_float, +) +from sentry.dynamic_sampling.tasks.constants import ( + CHUNK_SIZE, + DEFAULT_REDIS_CACHE_KEY_TTL, + MAX_TRANSACTIONS_PER_PROJECT, +) +from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( + generate_boost_low_volume_projects_cache_key, +) +from sentry.dynamic_sampling.tasks.helpers.sample_rate import get_org_sample_rate +from sentry.dynamic_sampling.tasks.utils import dynamic_sampling_task +from sentry.dynamic_sampling.types import SamplingMeasure +from sentry.dynamic_sampling.utils import has_dynamic_sampling, is_project_mode_sampling +from sentry.models.organization import Organization +from sentry.models.project import Project +from sentry.sentry_metrics import indexer +from sentry.silo.base import SiloMode +from sentry.snuba.dataset import Dataset, EntityKey +from sentry.snuba.referrer import Referrer +from sentry.tasks.base import instrumented_task +from sentry.tasks.relay import schedule_invalidate_project_config +from sentry.taskworker.namespaces import telemetry_experience_tasks +from sentry.utils import metrics +from sentry.utils.dates import deprecated_utcnow +from sentry.utils.snuba import raw_snql_query + +# This set contains all the projects for which we want to start extracting the sample rate over time. This is done +# as a temporary solution to dogfood our own product without exploding the cardinality of the project_id tag. +PROJECTS_WITH_METRICS = {1, 11276} # sentry # javascript +logger = logging.getLogger(__name__) + +# a tuple type alias of project_id, root_count, keep_count, drop_count, to be used in extraction of metrics for a specific project +ProjectVolumes = tuple[ProjectId, int, DecisionKeepCount, DecisionDropCount] + +# the same as ProjectVolumes, but with the organization ID added +OrgProjectVolumes = tuple[OrganizationId, ProjectId, int, DecisionKeepCount, DecisionDropCount] + + +@instrumented_task( + name="sentry.dynamic_sampling.boost_low_volume_projects_of_org_with_query", + namespace=telemetry_experience_tasks, + processing_deadline_duration=3 * 60 + 5, + retry=Retry(times=5, delay=5), + silo_mode=SiloMode.CELL, +) +@dynamic_sampling_task +def boost_low_volume_projects_of_org_with_query(org_id: OrganizationId) -> None: + """ + Task to adjust the sample rates of the projects of a single organization specified by an + organization ID. Transaction counts and rates are fetched within this task. + """ + logger.info( + "boost_low_volume_projects_of_org_with_query", + extra={"traceparent": sentry_sdk.get_traceparent(), "baggage": sentry_sdk.get_baggage()}, + ) + + org = Organization.objects.get_from_cache(id=org_id) + if is_project_mode_sampling(org): + return + + projects_with_tx_count_and_rates = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org_id], + measure=SamplingMeasure.SEGMENTS, + )[org_id] + rebalanced_projects = calculate_sample_rates_of_projects( + org_id, projects_with_tx_count_and_rates + ) + if rebalanced_projects is not None: + store_rebalanced_projects(org_id, rebalanced_projects) + + +@metrics.wraps("dynamic_sampling.fetch_projects_with_total_root_transaction_count_and_rates") +def fetch_projects_with_total_root_transaction_count_and_rates( + org_ids: list[int], + measure: SamplingMeasure, + query_interval: timedelta | None = None, +) -> Mapping[OrganizationId, Sequence[ProjectVolumes]]: + """ + Fetches for each org and each project the total root transaction count and how many transactions were kept and + dropped. + """ + aggregated_projects = defaultdict(list) + project_count_query_iter = query_project_counts_by_org(org_ids, measure, query_interval) + for chunk in project_count_query_iter: + for org_id, project_id, root_count_value, keep_count, drop_count in chunk: + aggregated_projects[org_id].append( + ( + project_id, + root_count_value, + keep_count, + drop_count, + ) + ) + + return aggregated_projects + + +@dynamic_sampling_task +def query_project_counts_by_org( + org_ids: list[int], measure: SamplingMeasure, query_interval: timedelta | None = None +) -> Iterator[Sequence[OrgProjectVolumes]]: + """Queries the total root transaction count and how many transactions were kept and dropped + for each project in a given interval (defaults to the last hour). + + Yields chunks of result rows, to allow timeouts to be handled in the caller. + """ + if not org_ids: + return + + if query_interval is None: + query_interval = timedelta(hours=1) + + if query_interval > timedelta(days=1): + granularity = Granularity(24 * 3600) + else: + granularity = Granularity(60) + + metrics.incr( + "dynamic_sampling.query_project_counts_by_org.count", + amount=len(org_ids), + tags={"measure": str(measure.value)}, + ) + + org_ids = list(org_ids) + project_ids = list( + Project.objects.filter(organization_id__in=org_ids, status=ObjectStatus.ACTIVE).values_list( + "id", flat=True + ) + ) + decision_string_id = indexer.resolve_shared_org("decision") + decision_tag = f"tags_raw[{decision_string_id}]" + + config = MEASURE_CONFIGS.get(measure) + if config is None: + raise ValueError(f"Unsupported measure: {measure}") + + metric_id = indexer.resolve_shared_org(str(config["mri"])) + use_case_id = config["use_case_id"] + + where_conditions = [ + Condition(Column("timestamp"), Op.GTE, deprecated_utcnow() - query_interval), + Condition(Column("timestamp"), Op.LT, deprecated_utcnow()), + Condition(Column("metric_id"), Op.EQ, metric_id), + Condition(Column("org_id"), Op.IN, org_ids), + Condition(Column("project_id"), Op.IN, project_ids), + ] + + # Add tag filters from config + for tag_name, tag_value in config["tags"].items(): + tag_string_id = indexer.resolve_shared_org(tag_name) + tag_column = f"tags_raw[{tag_string_id}]" + where_conditions.append(Condition(Column(tag_column), Op.EQ, tag_value)) + + query = Query( + match=Entity(EntityKey.GenericOrgMetricsCounters.value), + select=[ + Function("sum", [Column("value")], "root_count_value"), + Column("org_id"), + Column("project_id"), + Function( + "sumIf", + [ + Column("value"), + Function("equals", [Column(decision_tag), "keep"]), + ], + alias="keep_count", + ), + Function( + "sumIf", + [ + Column("value"), + Function("equals", [Column(decision_tag), "drop"]), + ], + alias="drop_count", + ), + ], + groupby=[Column("org_id"), Column("project_id")], + where=where_conditions, + granularity=granularity, + orderby=[ + OrderBy(Column("org_id"), Direction.ASC), + OrderBy(Column("project_id"), Direction.ASC), + ], + limitby=LimitBy( + columns=[Column("org_id"), Column("project_id")], + count=MAX_TRANSACTIONS_PER_PROJECT, + ), + # we are fetching one more than the chunk size to determine if there are more results + limit=Limit(CHUNK_SIZE + 1), + ) + + offset = 0 + more_results: bool = True + while more_results: + with metrics.timer( + "dynamic_sampling.query_project_counts_by_org.query_time", + tags={"measure": str(measure.value)}, + ): + request = Request( + dataset=Dataset.PerformanceMetrics.value, + app_id="dynamic_sampling", + query=query.set_offset(offset), + tenant_ids={"use_case_id": use_case_id.value, "cross_org_query": 1}, + ) + data = raw_snql_query( + request, + referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value, + )["data"] + + more_results = len(data) > CHUNK_SIZE + offset += CHUNK_SIZE + + # re-adjust, for the extra row we fetched + if more_results: + data = data[:-1] + + yield [ + ( + row["org_id"], + row["project_id"], + row["root_count_value"], + row["keep_count"], + row["drop_count"], + ) + for row in data + ] + + +@dynamic_sampling_task +def calculate_sample_rates_of_projects( + org_id: int, + projects_with_tx_count: Sequence[ProjectVolumes], +) -> list[RebalancedItem] | None: + """ + Calculates the sample rates of projects belonging to a specific org. + """ + try: + # We need the organization object for the feature flag. + organization = Organization.objects.get_from_cache(id=org_id) + except Organization.DoesNotExist: + # In case an org is not found, it might be that it has been deleted in the time between + # the query triggering this job and the actual execution of the job. + organization = None + + # If the org doesn't have dynamic sampling, we want to early return to avoid unnecessary work. + if not has_dynamic_sampling(organization): + return 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. + + default_sample_rate = quotas.backend.get_blended_sample_rate(organization_id=org_id) + sample_rate, success = get_org_sample_rate( + org_id=org_id, + default_sample_rate=default_sample_rate, + ) + + # If we didn't find any sample rate, it doesn't make sense to run the adjustment model. + if sample_rate is None: + sentry_sdk.capture_message( + "Sample rate of org not found when trying to adjust the sample rates of its projects" + ) + return None + + projects_with_counts = { + project_id: count_per_root for project_id, count_per_root, _, _ in projects_with_tx_count + } + + # The rebalancing will not work (or would make sense) when we have only projects with zero-counts. + if not any(projects_with_counts.values()): + return None + + # Since we don't mind about strong consistency, we query a replica of the main database with the possibility of + # having out of date information. This is a trade-off we accept, since we work under the assumption that eventually + # the projects of an org will be replicated consistently across replicas, because no org should continue to create + # new projects. + all_projects_ids = ( + Project.objects.using_replica() + .filter(organization=organization) + .values_list("id", flat=True) + ) + for project_id in all_projects_ids: + # In case a specific project has not been considered in the count query, it means that no metrics were extracted + # for it, thus we consider it as having 0 transactions for the query's time window. + if project_id not in projects_with_counts: + projects_with_counts[project_id] = 0 + + projects = [] + for project_id, count_per_root in projects_with_counts.items(): + projects.append( + RebalancedItem( + id=project_id, + count=count_per_root, + ) + ) + + model = ProjectsRebalancingModel() + rebalanced_projects: list[RebalancedItem] | None = guarded_run( + model, ProjectsRebalancingInput(classes=projects, sample_rate=sample_rate) + ) + + return rebalanced_projects + + +@dynamic_sampling_task +def store_rebalanced_projects(org_id: int, rebalanced_projects: list[RebalancedItem]) -> None: + """Stores the rebalanced projects in the cache and invalidates the project configs.""" + redis_client = get_redis_client_for_ds() + with redis_client.pipeline(transaction=False) as pipeline: + for rebalanced_project in rebalanced_projects: + cache_key = generate_boost_low_volume_projects_cache_key(org_id=org_id) + # We want to get the old sample rate, which will be None in case it was not set. + old_sample_rate = sample_rate_to_float( + redis_client.hget(cache_key, str(rebalanced_project.id)) + ) + + if rebalanced_project.id in PROJECTS_WITH_METRICS: + metrics.gauge( + "dynamic_sampling.project_sample_rate", + rebalanced_project.new_sample_rate * 100, + tags={"project_id": rebalanced_project.id}, + unit="percent", + ) + + # We want to store the new sample rate as a string. + pipeline.hset( + cache_key, + str(rebalanced_project.id), + rebalanced_project.new_sample_rate, # redis stores is as string + ) + pipeline.pexpire(cache_key, DEFAULT_REDIS_CACHE_KEY_TTL) + + # We invalidate the caches only if there was a change in the sample rate. This is to avoid flooding the + # system with project config invalidations, especially for projects with no volume. + if not are_equal_with_epsilon(old_sample_rate, rebalanced_project.new_sample_rate): + schedule_invalidate_project_config( + project_id=rebalanced_project.id, + trigger="dynamic_sampling_boost_low_volume_projects", + ) + + pipeline.execute() diff --git a/src/sentry/dynamic_sampling/tasks/constants.py b/src/sentry/dynamic_sampling/tasks/constants.py index 933ab1e799b5..c771786623b4 100644 --- a/src/sentry/dynamic_sampling/tasks/constants.py +++ b/src/sentry/dynamic_sampling/tasks/constants.py @@ -15,6 +15,7 @@ def adjusted_factor_ttl_ms() -> int: # Parameters to bound the queries run in Snuba. MAX_ORGS_PER_QUERY = 80 +MAX_TRANSACTIONS_PER_PROJECT = 20 # MIN and MAX rebalance factor in order to make sure we don't go crazy when rebalancing orgs. MIN_REBALANCE_FACTOR = 0.1 diff --git a/src/sentry/dynamic_sampling/tasks/utils.py b/src/sentry/dynamic_sampling/tasks/utils.py new file mode 100644 index 000000000000..93797c749ad9 --- /dev/null +++ b/src/sentry/dynamic_sampling/tasks/utils.py @@ -0,0 +1,41 @@ +from collections.abc import Callable +from functools import wraps +from random import random +from typing import Any + +import sentry_sdk + +from sentry.utils import metrics + + +def sample_function(function: Callable[..., Any], _sample_rate: float = 1.0, **kwargs: Any) -> None: + """ + Calls the supplied function with a uniform probability of `_sample_rate`. + """ + if _sample_rate >= 1.0 or 0.0 <= random() <= _sample_rate: + function(**kwargs) + + +def _compute_task_name(function_name: str) -> str: + return f"sentry.tasks.dynamic_sampling.{function_name}" + + +def dynamic_sampling_task(func: Callable[..., Any]) -> Callable[..., Any]: + """ + Decorator to wrap dynamic sampling related tasks to record metrics for the execution of + the task, durations associated with it as metrics, and capture all exceptions in sentry. + """ + + @wraps(func) + def _wrapper(*args: Any, **kwargs: Any) -> Any: + function_name = func.__name__ + task_name = _compute_task_name(function_name) + metrics.incr(f"{task_name}.start", sample_rate=1.0) + with metrics.timer(task_name, sample_rate=1.0): + try: + return func(*args, **kwargs) + except Exception as e: + sentry_sdk.capture_exception(e) + raise + + return _wrapper diff --git a/src/sentry/snuba/referrer.py b/src/sentry/snuba/referrer.py index 16901b96d089..5da166495282 100644 --- a/src/sentry/snuba/referrer.py +++ b/src/sentry/snuba/referrer.py @@ -648,6 +648,9 @@ class Referrer(StrEnum): DYNAMIC_SAMPLING_COUNTERS_GET_ORG_TRANSACTION_VOLUMES = ( "dynamic_sampling.counters.get_org_transaction_volumes" ) + DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT = ( + "dynamic_sampling.distribution.fetch_projects_with_count_per_root_total_volumes" + ) DYNAMIC_SAMPLING_PER_ORG_GET_EAP_ORG_VOLUME = "dynamic_sampling.per_org.get_eap_org_volume" DYNAMIC_SAMPLING_PER_ORG_GET_EAP_PROJECT_VOLUMES = ( "dynamic_sampling.per_org.get_eap_project_volumes" diff --git a/tests/sentry/core/endpoints/test_organization_details.py b/tests/sentry/core/endpoints/test_organization_details.py index 7fc357a48f7e..19a0b6fb5613 100644 --- a/tests/sentry/core/endpoints/test_organization_details.py +++ b/tests/sentry/core/endpoints/test_organization_details.py @@ -43,8 +43,8 @@ from sentry.replays.models import OrganizationMemberReplayAccess from sentry.signals import project_created from sentry.silo.safety import unguarded_write -from sentry.testutils.cases import APITestCase, SnubaTestCase, SpanTestCase, TwoFactorAPITestCase -from sentry.testutils.helpers.datetime import before_now +from sentry.snuba.metrics import SpanMRI +from sentry.testutils.cases import APITestCase, BaseMetricsLayerTestCase, TwoFactorAPITestCase from sentry.testutils.helpers.features import with_feature from sentry.testutils.outbox import outbox_runner from sentry.testutils.pytest.fixtures import django_db_all @@ -95,7 +95,7 @@ def has_scope(self, scope): @cell_silo_test(cells=cells, include_monolith_run=True) -class OrganizationDetailsTest(OrganizationDetailsTestBase, SnubaTestCase, SpanTestCase): +class OrganizationDetailsTest(OrganizationDetailsTestBase, BaseMetricsLayerTestCase): @property def now(self): return datetime.now().replace(microsecond=0) @@ -568,29 +568,11 @@ def test_sampling_mode_change_requires_write_scope(self) -> None: assert response.status_code == 403 - @django_db_all - def test_change_org_target_sample_rate_schedules_per_org_calculation(self) -> None: - self.organization.update_option( - "sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION.value - ) - - with ( - self.feature("organizations:dynamic-sampling-custom"), - patch( - "sentry.core.endpoints.organization_details.run_calculations_per_org_task_entry" - ) as task, - ): - response = self.get_response(self.organization.slug, method="put", targetSampleRate=0.1) - - assert response.status_code == 200 - task.delay.assert_called_once_with(self.organization.id) - @django_db_all @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) - def test_sampling_mode_change_with_deleted_projects_that_had_spans(self) -> None: + def test_sampling_mode_change_with_deleted_projects_that_had_metrics(self) -> None: project_1 = self.create_project(organization=self.organization) project_2 = self.create_project(organization=self.organization) - self.organization.update_option("sentry:target_sample_rate", 0.5) # Create a team member for project_1 only team_1 = self.create_team(organization=self.organization) @@ -601,17 +583,21 @@ def test_sampling_mode_change_with_deleted_projects_that_had_spans(self) -> None ) self.login_as(user=member_user) - timestamp = before_now(days=12) - self.store_spans( - [ - self.create_span( - {"is_segment": True, "sentry_tags": {"dsc.project_id": str(project.id)}}, - organization=self.organization, - project=project, - start_ts=timestamp, - ) - for project in (project_1, project_2) - ] + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"is_segment": "true", "decision": "keep"}, + minutes_before_now=60 * 24 * 12, + value=1, + project_id=project_1.id, + org_id=self.organization.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"is_segment": "true", "decision": "keep"}, + minutes_before_now=60 * 24 * 12, + value=1, + project_id=project_2.id, + org_id=self.organization.id, ) project_2.delete() diff --git a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py new file mode 100644 index 000000000000..21e421df21fb --- /dev/null +++ b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py @@ -0,0 +1,622 @@ +from datetime import datetime, timedelta +from typing import cast +from unittest.mock import patch + +from django.utils import timezone + +from sentry.dynamic_sampling.rules.base import get_guarded_project_sample_rate +from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds +from sentry.dynamic_sampling.tasks.boost_low_volume_projects import ( + boost_low_volume_projects_of_org_with_query, + fetch_projects_with_total_root_transaction_count_and_rates, + query_project_counts_by_org, +) +from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( + get_boost_low_volume_projects_sample_rate, +) +from sentry.dynamic_sampling.tasks.helpers.sliding_window import ( + generate_sliding_window_org_cache_key, +) +from sentry.dynamic_sampling.types import DynamicSamplingMode, SamplingMeasure +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 SpanMRI +from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase +from sentry.testutils.helpers.datetime import freeze_time +from sentry.testutils.helpers.features import with_feature +from sentry.testutils.helpers.options import override_options + +MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace( + hour=0, minute=0, second=0, microsecond=0 +) + + +@freeze_time(MOCK_DATETIME) +class PrioritiseProjectsSnubaQueryTest(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): + @property + def now(self) -> datetime: + return MOCK_DATETIME + + def test_simple_one_org_one_project(self) -> None: + org1 = self.create_organization("test-org") + p1 = self.create_project(organization=org1) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=p1.id, + org_id=org1.id, + ) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=p1.id, + org_id=org1.id, + ) + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org1.id], measure=SamplingMeasure.SEGMENTS + ) + assert results[org1.id] == [(p1.id, 4.0, 1, 3)] + + def test_deleted_projects_are_not_queried(self) -> None: + org1 = self.create_organization("test-org") + p1 = self.create_project(organization=org1) + p2 = self.create_project(organization=org1) + + for p in [p1, p2]: + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=p.id, + org_id=org1.id, + ) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=p.id, + org_id=org1.id, + ) + p2.delete() + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org1.id], measure=SamplingMeasure.SEGMENTS + ) + 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_sliding_window_sample_rate(self) -> None: + org1 = self.create_organization("test-org") + p1 = self.create_project(organization=org1) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=p1.id, + org_id=org1.id, + ) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=p1.id, + org_id=org1.id, + ) + + # simulate having a sliding window sample rate for the org + redis_client = get_redis_client_for_ds() + cache_key = generate_sliding_window_org_cache_key(org1.id) + redis_client.set(cache_key, 1.0) + + 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 == 1.0 + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_simple_one_org_one_project_task_target_sample_rate(self) -> None: + 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=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=p1.id, + org_id=org1.id, + ) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + 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 (sample_rate, got_value) == (0.5, True) + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_per_project_sample_rate_override(self) -> None: + # A per-project override configured via options hard-replaces the rate the + # custom dynamic sampling path would otherwise resolve for that project -- + # winning even over the recently-added 100% boost -- and leaves other projects + # untouched. + org1 = self.create_organization("am3-override-org") + org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) + org1.update_option("sentry:target_sample_rate", 0.5) + overridden = self.create_project(organization=org1) + normal = self.create_project(organization=org1) + + # Baseline: freshly-created projects are boosted to 1.0 by the recently-added + # rule, so neither resolves to the org target yet. + assert get_guarded_project_sample_rate(org1, overridden) == 1.0 + + with override_options( + {"dynamic-sampling.sample-rate-override-per-project": {str(overridden.id): 0.9}} + ): + assert get_guarded_project_sample_rate(org1, overridden) == 0.9 + # Not in the override map -> unaffected by the override. + assert get_guarded_project_sample_rate(org1, normal) == 1.0 + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_per_project_sample_rate_override_ignores_out_of_range(self) -> None: + org1 = self.create_organization("am3-override-org-bad") + org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) + org1.update_option("sentry:target_sample_rate", 0.5) + project = self.create_project(organization=org1) + + baseline = get_guarded_project_sample_rate(org1, project) + with override_options( + {"dynamic-sampling.sample-rate-override-per-project": {str(project.id): 2.0}} + ): + # Out-of-range override is ignored; the resolved rate is unchanged. + assert get_guarded_project_sample_rate(org1, project) == baseline + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_project_mode_sampling_with_query(self) -> None: + org1 = self.create_organization("test-org") + p1 = self.create_project(organization=org1) + + org1.update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) + p1.update_option("sentry:target_sample_rate", 0.2) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=p1.id, + org_id=org1.id, + ) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + 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 (sample_rate, got_value) == (None, False) + + assert get_guarded_project_sample_rate(org1, p1) == 0.2 + + def test_complex(self) -> None: + org1 = self.create_organization("test-org1") + p1_1 = self.create_project(organization=org1, name="p1_1") + p1_2 = self.create_project(organization=org1, name="p1_2") + org2 = self.create_organization("test-org2") + p2_1 = self.create_project(organization=org2, name="p2_1") + p2_2 = self.create_project(organization=org2, name="p2_2") + + proj_orgs = [ + {"org": org1, "projects": [p1_1, p1_2]}, + {"org": org2, "projects": [p2_1, p2_2]}, + ] + + proj_counts = {"p1_1": (1, 2), "p1_2": (3, 4), "p2_1": (5, 6), "p2_2": (7, 8)} # keep,drop + + for org_info in proj_orgs: + org = cast(Organization, org_info.get("org")) + projects = cast(list[Project], org_info.get("projects")) + for project in projects: + keep, drop = proj_counts[project.name] + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={ + "transaction": "foo_transaction", + "decision": "keep", + "is_segment": "true", + }, + minutes_before_now=29, + value=keep, + project_id=project.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={ + "transaction": "foo_transaction", + "decision": "drop", + "is_segment": "true", + }, + minutes_before_now=29, + value=drop, + project_id=project.id, + org_id=org.id, + ) + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org1.id, org2.id], measure=SamplingMeasure.SEGMENTS + ) + + assert len(results) == 2 # two orgs + + org_1_results = results[org1.id] + + assert len(org_1_results) == 2 + + # (p.id, total, keep, drop) == result + assert (p1_1.id, 3, 1, 2) in org_1_results + assert (p1_2.id, 7, 3, 4) in org_1_results + + org_2_results = results[org2.id] + assert len(org_2_results) == 2 + assert (p2_1.id, 11, 5, 6) in org_2_results + assert (p2_2.id, 15, 7, 8) in org_2_results + + +@freeze_time(MOCK_DATETIME) +class TestQueryProjectCountsByOrgEmptyOrgIds(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): + """ + Test that query_project_counts_by_org correctly skips Snuba queries + when org_ids is empty, avoiding unnecessary queries. + """ + + @property + def now(self) -> datetime: + return MOCK_DATETIME + + def test_query_skips_for_empty_org_ids_when_option_enabled(self) -> None: + """ + Confirms that query_project_counts_by_org does NOT make a Snuba query + when called with an empty org_ids list. + """ + with patch( + "sentry.dynamic_sampling.tasks.boost_low_volume_projects.raw_snql_query" + ) as mock_query: + mock_query.return_value = {"data": []} + + list(query_project_counts_by_org([], SamplingMeasure.SEGMENTS)) + + assert mock_query.call_count == 0 + + def test_fetch_projects_only_queries_measures_with_orgs(self) -> None: + """ + Confirms that fetch_projects_with_total_root_transaction_count_and_rates + does NOT make a Snuba query when called with an empty org_ids list. + """ + org = self.create_organization("test-org") + self.create_project(organization=org) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=1, + project_id=org.project_set.first().id, + org_id=org.id, + ) + + with patch( + "sentry.dynamic_sampling.tasks.boost_low_volume_projects.raw_snql_query" + ) as mock_query: + mock_query.return_value = {"data": []} + + # Query with org should make one call + fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org.id], measure=SamplingMeasure.SEGMENTS + ) + assert mock_query.call_count == 1 + + # Query with empty list should not make any additional calls + fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[], measure=SamplingMeasure.SEGMENTS + ) + assert mock_query.call_count == 1 + + +@freeze_time(MOCK_DATETIME) +class TestSpanMetricQuery(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): + """ + Tests that verify the span metric query works correctly with is_segment filter. + """ + + @property + def now(self) -> datetime: + return MOCK_DATETIME + + def test_span_metric_with_is_segment_filter(self) -> None: + """ + Test that span metric queries only count spans with is_segment=true. + """ + org = self.create_organization("test-org") + project = self.create_project(organization=org) + + # Store span metrics with is_segment=true (should be counted) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=5, + project_id=project.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo_transaction", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=10, + project_id=project.id, + org_id=org.id, + ) + + # Store span metrics without is_segment tag (should NOT be counted) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "bar_transaction", "decision": "keep"}, + minutes_before_now=30, + value=100, + project_id=project.id, + org_id=org.id, + ) + + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org.id], measure=SamplingMeasure.SEGMENTS + ) + + # Should only count the is_segment=true metrics (5 + 10 = 15) + assert results[org.id] == [(project.id, 15.0, 5, 10)] + + def test_span_metric_multiple_projects(self) -> None: + """ + Test span metric query with multiple projects. + """ + org = self.create_organization("test-org") + p1 = self.create_project(organization=org) + p2 = self.create_project(organization=org) + + # Project 1: 3 keep, 7 drop + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=p1.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=7, + project_id=p1.id, + org_id=org.id, + ) + + # Project 2: 2 keep, 8 drop + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "bar", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=2, + project_id=p2.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "bar", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=8, + project_id=p2.id, + org_id=org.id, + ) + + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org.id], measure=SamplingMeasure.SPANS + ) + + assert len(results[org.id]) == 2 + assert (p1.id, 10.0, 3, 7) in results[org.id] + assert (p2.id, 10.0, 2, 8) in results[org.id] + + +@freeze_time(MOCK_DATETIME) +class TestEndToEndMeasureDispatching(BaseMetricsLayerTestCase, TestCase, SnubaTestCase): + """ + End-to-end tests verifying that the per-org query task and the project count + query use the right measure. + """ + + @property + def now(self) -> datetime: + return MOCK_DATETIME + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_org_uses_segments_measure_in_with_query_task(self) -> None: + """ + boost_low_volume_projects_of_org_with_query should use SEGMENTS measure. + """ + org = self.create_organization("test-org") + p1 = self.create_project(organization=org) + + # Store span metrics with is_segment=true (used by SEGMENTS measure) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=5, + project_id=p1.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=10, + project_id=p1.id, + org_id=org.id, + ) + + redis_client = get_redis_client_for_ds() + cache_key = generate_sliding_window_org_cache_key(org.id) + redis_client.set(cache_key, 0.5) + + with self.tasks(): + boost_low_volume_projects_of_org_with_query.delay(org.id) + + sample_rate, got_value = get_boost_low_volume_projects_sample_rate( + org.id, p1.id, error_sample_rate_fallback=None + ) + assert got_value + assert sample_rate is not None + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_segments_query_uses_span_mri_with_is_segment_tag(self) -> None: + """ + When processing an org with SEGMENTS measure, the Snuba query should use + SpanMRI and filter by is_segment=true, not TransactionMRI. + """ + org = self.create_organization("test-org") + project = self.create_project(organization=org) + + # Store ONLY span metrics with is_segment=true + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=project.id, + org_id=org.id, + ) + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "drop", "is_segment": "true"}, + minutes_before_now=30, + value=7, + project_id=project.id, + org_id=org.id, + ) + + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org.id], measure=SamplingMeasure.SEGMENTS + ) + + # Should only see the span/segment metrics (3 + 7 = 10), not the transaction metrics (100) + assert results[org.id] == [(project.id, 10.0, 3, 7)] + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_spans_query_uses_span_mri_without_is_segment(self) -> None: + """ + When processing an org with SPANS measure, the Snuba query should use + SpanMRI but NOT filter by is_segment (counts all spans). + """ + org = self.create_organization("test-org") + project = self.create_project(organization=org) + + # Store span metrics WITH is_segment=true + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=3, + project_id=project.id, + org_id=org.id, + ) + + # Store span metrics WITHOUT is_segment tag + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "bar", "decision": "keep"}, + minutes_before_now=30, + value=7, + project_id=project.id, + org_id=org.id, + ) + + results = fetch_projects_with_total_root_transaction_count_and_rates( + org_ids=[org.id], measure=SamplingMeasure.SPANS + ) + + # SPANS measure should count ALL spans (both with and without is_segment) + # Total = 3 + 7 = 10, all keeps + assert results[org.id] == [(project.id, 10.0, 10, 0)] + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_with_query_task_skips_project_mode_orgs(self) -> None: + """ + boost_low_volume_projects_of_org_with_query should early-return for + project-mode orgs without storing any rebalanced rates. + """ + org = self.create_organization("test-org") + p1 = self.create_project(organization=org) + org.update_option("sentry:sampling_mode", DynamicSamplingMode.PROJECT) + + self.store_performance_metric( + name=SpanMRI.COUNT_PER_ROOT_PROJECT.value, + tags={"transaction": "foo", "decision": "keep", "is_segment": "true"}, + minutes_before_now=30, + value=5, + project_id=p1.id, + org_id=org.id, + ) + + redis_client = get_redis_client_for_ds() + cache_key = generate_sliding_window_org_cache_key(org.id) + redis_client.set(cache_key, 0.5) + + with self.tasks(): + boost_low_volume_projects_of_org_with_query.delay(org.id) + + sample_rate, got_value = get_boost_low_volume_projects_sample_rate( + org.id, p1.id, error_sample_rate_fallback=None + ) + assert not got_value + assert sample_rate is None