refactor(Segment Membership): Distribute count refreshes evenly across the window

Per review, drop the per-organisation grouping (and its cumulative-delay overflow
handling and STAGGER_SECONDS setting) in favour of spreading all live-segment
projects evenly across the stagger window, ordered by organisation so an org's
projects still refresh contiguously. Simpler, and no org-slot overflow to guard.

beep boop
Claude-Session: https://claude.ai/code/session_01EgZ5iHpDASZzCapiHRxHLB

Author: khvn26

api/app/settings/common.py

@@ -1494,9 +1494,6 @@
         "SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS must not exceed "
         "SEGMENT_MEMBERSHIP_REFRESH_INTERVAL_HOURS."
     )
-SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_SECONDS = env.int(
-    "SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_SECONDS", default=10
-)
 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,44 +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)
-    num_orgs = (
-        Project.objects.filter(id__in=project_ids)
-        .values("organisation_id")
-        .distinct()
-        .count()
+    projects = (
+        Project.objects.filter(
+            Exists(Segment.live_objects.filter(project=OuterRef("pk")))
+        )
+        .select_related("organisation")
+        .order_by("organisation_id", "id")
     )
-    if not num_orgs:
+    total = projects.count()
+    if not total:
         return
 
-    org_spacing = timedelta(
+    spacing = timedelta(
         hours=settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS
-    ) / (num_orgs + 1)
-    project_spacing = timedelta(
-        seconds=settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_SECONDS
-    )
+    ) / (total + 1)
     now = timezone.now()
-
-    # Accumulate the delay so it never moves backwards: each organisation starts
-    # at its window slot or after the previous organisation's last project,
-    # whichever is later. This keeps an organisation with many projects from
-    # spilling its stagger into the next organisation's slot.
-    cumulative_delay = timedelta(0)
-    org_index = -1
-    current_org_id = None
-    projects = (
-        Project.objects.filter(id__in=project_ids)
-        .select_related("organisation")
-        .order_by("organisation_id", "id")
-    )
-    for project in projects.iterator():
-        if project.organisation_id != current_org_id:
-            current_org_id = project.organisation_id
-            org_index += 1
-            cumulative_delay = max(cumulative_delay, org_spacing * org_index)
-        else:
-            cumulative_delay += project_spacing
-        enqueue_membership_refresh(project, delay_until=now + cumulative_delay)
+    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

@@ -392,104 +392,39 @@ def test_refresh_all_segment_counts__no_live_segments__does_nothing(
     enqueue.assert_not_called()
 
 
-def test_refresh_all_segment_counts__live_segment_projects__delegates_to_enqueue(
+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
-    )
-    Segment.objects.create(name="seg-b", project=project_b)
-    settings.CLICKHOUSE_ENABLED = True
-    enqueue = mocker.patch.object(tasks, "enqueue_membership_refresh")
-
-    # When
-    refresh_all_segment_counts()
-
-    # Then
-    calls = enqueue.call_args_list
-    assert {call.args[0].id for call in calls} == {project.id, project_b.id}
-    # enqueues are staggered across the refresh window, not fired all at once
-    delays = sorted(call.kwargs["delay_until"] for call in calls)
-    window = timedelta(
-        hours=settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS
-    )
-    assert delays[0] < delays[1]
-    assert delays[-1] - delays[0] < window
-
-
-def test_refresh_all_segment_counts__multiple_orgs__staggers_by_org_then_project(
-    mocker: MockerFixture,
-    settings: SettingsWrapper,
-    project: Project,
-    segment: Segment,
-) -> None:
-    # Given org A with two live-segment projects and org B with one
+    # 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_INTERVAL_HOURS = 6
-    settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS = 3
-    settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_SECONDS = 10
+    settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_WINDOW_HOURS = 4
     enqueue = mocker.patch.object(tasks, "enqueue_membership_refresh")
 
     # When
     refresh_all_segment_counts()
 
-    # Then
-    delay_by_id = {
-        call.args[0].id: call.kwargs["delay_until"] for call in enqueue.call_args_list
-    }
-    assert set(delay_by_id) == {project.id, project_a2.id, project_b.id}
-    # org A's projects cluster a project-stagger apart
-    org_a = sorted([delay_by_id[project.id], delay_by_id[project_a2.id]])
-    assert org_a[1] - org_a[0] == timedelta(seconds=10)
-    # org B sits a full window-slot (window / (orgs + 1)) away, not a project-stagger
-    org_spacing = timedelta(hours=3) / 3
-    assert abs(delay_by_id[project_b.id] - org_a[0]) == org_spacing
-
-
-def test_refresh_all_segment_counts__org_overflows_slot__no_cross_org_collision(
-    mocker: MockerFixture,
-    settings: SettingsWrapper,
-    project: Project,
-    segment: Segment,
-) -> None:
-    # Given org A with three live-segment projects and org B with one, and a
-    # project stagger wide enough that A's projects overflow its window slot
-    # (org_spacing = 1h / 3 = 20min, but the projects are 1h apart).
-    project_a2 = Project.objects.create(name="a2", organisation=project.organisation)
-    Segment.objects.create(name="seg-a2", project=project_a2)
-    project_a3 = Project.objects.create(name="a3", organisation=project.organisation)
-    Segment.objects.create(name="seg-a3", project=project_a3)
-    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 = 1
-    settings.SEGMENT_MEMBERSHIP_REFRESH_PROJECT_STAGGER_SECONDS = 3600
-    enqueue = mocker.patch.object(tasks, "enqueue_membership_refresh")
-
-    # When
-    refresh_all_segment_counts()
-
-    # Then
-    delay_by_id = {
-        call.args[0].id: call.kwargs["delay_until"] for call in enqueue.call_args_list
-    }
-    org_a = sorted(delay_by_id[p] for p in (project.id, project_a2.id, project_a3.id))
-    # org A's projects stay a project-stagger apart, overflowing the 20min slot
-    assert org_a[1] - org_a[0] == timedelta(seconds=3600)
-    assert org_a[2] - org_a[1] == timedelta(seconds=3600)
-    # org B never schedules before org A's last project, so the two don't collide
-    assert delay_by_id[project_b.id] >= org_a[-1]
+    # 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:283`
+ - `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:256`
+ - `api/segment_membership/tasks.py:235`
 
 Attributes:
  - `project.id`
 
 ### `segment_membership.refresh.project.skipped`
 
 Logged at `info` from:
- - `api/segment_membership/tasks.py:227`
- - `api/segment_membership/tasks.py:236`
+ - `api/segment_membership/tasks.py:206`
+ - `api/segment_membership/tasks.py:215`
 
 Attributes:
  - `project.id`