fix(Segment Membership): Stagger the recurring count refresh fan-out (#7943)

Author: khvn26

api/app/settings/common.py

@@ -1483,6 +1483,17 @@
 SEGMENT_MEMBERSHIP_REFRESH_INTERVAL_HOURS = env.int(
     "SEGMENT_MEMBERSHIP_REFRESH_INTERVAL_HOURS", default=6
 )
+SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS = env.int(
+    "SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS", default=1
+)
+if (
+    SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS
+    > SEGMENT_MEMBERSHIP_REFRESH_INTERVAL_HOURS
+):
+    raise ImproperlyConfigured(
+        "SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS must not exceed "
+        "SEGMENT_MEMBERSHIP_REFRESH_INTERVAL_HOURS."
+    )
 SEGMENT_MEMBERSHIP_DELETE_REFRESH_DELAY_SECONDS = env.int(
     "SEGMENT_MEMBERSHIP_DELETE_REFRESH_DELAY_SECONDS",
     default=120,  # We can expect the identity deletion to propagate by T+120 seconds based on Edge CDC SLO.

api/segment_membership/tasks.py

@@ -174,13 +174,23 @@ def refresh_all_segment_counts() -> None:
     if not settings.CLICKHOUSE_ENABLED:
         return
 
-    project_ids = Segment.live_objects.values_list("project_id", flat=True)
-    for project in (
-        Project.objects.filter(id__in=project_ids)
+    projects = (
+        Project.objects.filter(
+            Exists(Segment.live_objects.filter(project=OuterRef("pk")))
+        )
         .select_related("organisation")
-        .iterator()
-    ):
-        enqueue_membership_refresh(project)
+        .order_by("organisation_id", "id")
+    )
+    total = projects.count()
+    if not total:
+        return
+
+    spacing = timedelta(
+        hours=settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS
+    ) / (total + 1)
+    now = timezone.now()
+    for index, project in enumerate(projects.iterator()):
+        enqueue_membership_refresh(project, delay_until=now + spacing * index)
 
 
 @register_task_handler(

api/tests/unit/segment_membership/test_unit_segment_membership_tasks.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from datetime import timezone as dt_timezone
 from unittest.mock import MagicMock
 
@@ -13,6 +13,7 @@
 from task_processor.task_run_method import TaskRunMethod
 
 from environments.models import Environment
+from organisations.models import Organisation
 from projects.models import Project
 from segment_membership import tasks
 from segment_membership.models import SegmentMembershipCount, SegmentMembershipSeed
@@ -375,27 +376,55 @@ def test_refresh_all_segment_counts__no_clickhouse_creds__skips(
     enqueue.assert_not_called()
 
 
-def test_refresh_all_segment_counts__live_segment_projects__delegates_to_enqueue(
+def test_refresh_all_segment_counts__no_live_segments__does_nothing(
+    mocker: MockerFixture,
+    settings: SettingsWrapper,
+    project: Project,
+) -> None:
+    # Given a project but no live segments, so there are no organisations to seed
+    settings.CLICKHOUSE_ENABLED = True
+    enqueue = mocker.patch.object(tasks, "enqueue_membership_refresh")
+
+    # When
+    refresh_all_segment_counts()
+
+    # Then
+    enqueue.assert_not_called()
+
+
+def test_refresh_all_segment_counts__live_segment_projects__staggers_evenly_by_org(
     mocker: MockerFixture,
     settings: SettingsWrapper,
     project: Project,
     segment: Segment,
 ) -> None:
-    # Given two projects with live segments (flag + debounce gating is the
-    # helper's job, so refresh_all just delegates one call per project)
-    project_b = Project.objects.create(
-        name="project-b", organisation=project.organisation
-    )
+    # Given three live-segment projects across two organisations
+    project_a2 = Project.objects.create(name="a2", organisation=project.organisation)
+    Segment.objects.create(name="seg-a2", project=project_a2)
+    org_b = Organisation.objects.create(name="org-b")
+    project_b = Project.objects.create(name="b1", organisation=org_b)
     Segment.objects.create(name="seg-b", project=project_b)
     settings.CLICKHOUSE_ENABLED = True
+    settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS = 4
     enqueue = mocker.patch.object(tasks, "enqueue_membership_refresh")
 
     # When
     refresh_all_segment_counts()
 
-    # Then
-    dispatched_ids = {call.args[0].id for call in enqueue.call_args_list}
-    assert dispatched_ids == {project.id, project_b.id}
+    # Then one call per project, ordered by (organisation_id, id) so an
+    # organisation's projects refresh contiguously
+    called = [
+        (call.args[0].id, call.kwargs["delay_until"]) for call in enqueue.call_args_list
+    ]
+    expected_order = sorted(
+        (project, project_a2, project_b), key=lambda p: (p.organisation_id, p.id)
+    )
+    assert [pid for pid, _ in called] == [p.id for p in expected_order]
+    # spread evenly across the window: spacing = window / (total + 1)
+    spacing = timedelta(hours=4) / 4
+    delays = [delay for _, delay in called]
+    assert delays[1] - delays[0] == spacing
+    assert delays[2] - delays[1] == spacing
 
 
 def test_refresh_project_segment_counts__no_clickhouse_creds__skips(

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -388,7 +388,7 @@ Attributes:
 ### `segment_membership.refresh.project.completed`
 
 Logged at `info` from:
- - `api/segment_membership/tasks.py:252`
+ - `api/segment_membership/tasks.py:262`
 
 Attributes:
  - `membership_counts.count`
@@ -398,16 +398,16 @@ Attributes:
 ### `segment_membership.refresh.project.failed`
 
 Logged at `exception` from:
- - `api/segment_membership/tasks.py:225`
+ - `api/segment_membership/tasks.py:235`
 
 Attributes:
  - `project.id`
 
 ### `segment_membership.refresh.project.skipped`
 
 Logged at `info` from:
- - `api/segment_membership/tasks.py:196`
- - `api/segment_membership/tasks.py:205`
+ - `api/segment_membership/tasks.py:206`
+ - `api/segment_membership/tasks.py:215`
 
 Attributes:
  - `project.id`