|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import logging |
| 4 | +from collections.abc import Sequence |
| 5 | +from datetime import timedelta |
| 6 | + |
| 7 | +import sentry_sdk |
| 8 | + |
| 9 | +from sentry import features, options |
| 10 | +from sentry.models.organization import Organization, OrganizationStatus |
| 11 | +from sentry.tasks.base import instrumented_task |
| 12 | +from sentry.taskworker.namespaces import seer_tasks |
| 13 | +from sentry.utils.iterators import chunked |
| 14 | +from sentry.utils.query import RangeQuerySetWrapper |
| 15 | + |
| 16 | +logger = logging.getLogger("sentry.tasks.seer.night_shift") |
| 17 | + |
| 18 | +NIGHT_SHIFT_DISPATCH_STEP_SECONDS = 37 |
| 19 | +NIGHT_SHIFT_SPREAD_DURATION = timedelta(hours=4) |
| 20 | + |
| 21 | +FEATURE_NAMES = [ |
| 22 | + "organizations:seer-night-shift", |
| 23 | + "organizations:gen-ai-features", |
| 24 | +] |
| 25 | + |
| 26 | + |
| 27 | +@instrumented_task( |
| 28 | + name="sentry.tasks.seer.night_shift.schedule_night_shift", |
| 29 | + namespace=seer_tasks, |
| 30 | + processing_deadline_duration=15 * 60, |
| 31 | +) |
| 32 | +def schedule_night_shift() -> None: |
| 33 | + """ |
| 34 | + Nightly scheduler: iterates active orgs in batches, checks feature flags |
| 35 | + in bulk, and dispatches per-org worker tasks with jitter. |
| 36 | + """ |
| 37 | + if not options.get("seer.night_shift.enable"): |
| 38 | + return |
| 39 | + |
| 40 | + spread_seconds = int(NIGHT_SHIFT_SPREAD_DURATION.total_seconds()) |
| 41 | + batch_index = 0 |
| 42 | + |
| 43 | + for org_batch in chunked( |
| 44 | + RangeQuerySetWrapper[Organization]( |
| 45 | + Organization.objects.filter(status=OrganizationStatus.ACTIVE), |
| 46 | + step=1000, |
| 47 | + ), |
| 48 | + 100, |
| 49 | + ): |
| 50 | + for org in _get_eligible_orgs_from_batch(org_batch): |
| 51 | + if bool(org.get_option("sentry:hide_ai_features")): |
| 52 | + continue |
| 53 | + |
| 54 | + delay = (batch_index * NIGHT_SHIFT_DISPATCH_STEP_SECONDS) % spread_seconds |
| 55 | + |
| 56 | + run_night_shift_for_org.apply_async( |
| 57 | + args=[org.id], |
| 58 | + countdown=delay, |
| 59 | + ) |
| 60 | + batch_index += 1 |
| 61 | + |
| 62 | + logger.info( |
| 63 | + "night_shift.schedule_complete", |
| 64 | + extra={"orgs_dispatched": batch_index}, |
| 65 | + ) |
| 66 | + |
| 67 | + |
| 68 | +@instrumented_task( |
| 69 | + name="sentry.tasks.seer.night_shift.run_night_shift_for_org", |
| 70 | + namespace=seer_tasks, |
| 71 | + processing_deadline_duration=5 * 60, |
| 72 | +) |
| 73 | +def run_night_shift_for_org(organization_id: int) -> None: |
| 74 | + try: |
| 75 | + organization = Organization.objects.get( |
| 76 | + id=organization_id, status=OrganizationStatus.ACTIVE |
| 77 | + ) |
| 78 | + except Organization.DoesNotExist: |
| 79 | + return |
| 80 | + |
| 81 | + sentry_sdk.set_tags( |
| 82 | + { |
| 83 | + "organization_id": organization.id, |
| 84 | + "organization_slug": organization.slug, |
| 85 | + } |
| 86 | + ) |
| 87 | + |
| 88 | + logger.info( |
| 89 | + "night_shift.org_dispatched", |
| 90 | + extra={ |
| 91 | + "organization_id": organization_id, |
| 92 | + "organization_slug": organization.slug, |
| 93 | + }, |
| 94 | + ) |
| 95 | + |
| 96 | + |
| 97 | +def _get_eligible_orgs_from_batch( |
| 98 | + orgs: Sequence[Organization], |
| 99 | +) -> list[Organization]: |
| 100 | + """ |
| 101 | + Check feature flags for a batch of orgs using batch_has_for_organizations. |
| 102 | + Returns orgs that have all required feature flags enabled. |
| 103 | + """ |
| 104 | + eligible = list(orgs) |
| 105 | + |
| 106 | + for feature_name in FEATURE_NAMES: |
| 107 | + batch_result = features.batch_has_for_organizations(feature_name, eligible) |
| 108 | + if batch_result is None: |
| 109 | + raise RuntimeError(f"batch_has_for_organizations returned None for {feature_name}") |
| 110 | + |
| 111 | + eligible = [org for org in eligible if batch_result.get(f"organization:{org.id}", False)] |
| 112 | + |
| 113 | + if not eligible: |
| 114 | + return [] |
| 115 | + |
| 116 | + return eligible |
0 commit comments