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

api/app/settings/common.py

@@ -1483,6 +1483,20 @@
 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_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

@@ -175,12 +175,43 @@ def refresh_all_segment_counts() -> None:
         return
 
     project_ids = Segment.live_objects.values_list("project_id", flat=True)
-    for project in (
+    num_orgs = (
+        Project.objects.filter(id__in=project_ids)
+        .values("organisation_id")
+        .distinct()
+        .count()
+    )
+    if not num_orgs:
+        return
+
+    org_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
+    )
+    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")
-        .iterator()
-    ):
-        enqueue_membership_refresh(project)
+        .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)
 
 
 @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,6 +376,22 @@ def test_refresh_all_segment_counts__no_clickhouse_creds__skips(
     enqueue.assert_not_called()
 
 
+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__delegates_to_enqueue(
     mocker: MockerFixture,
     settings: SettingsWrapper,
@@ -394,8 +411,85 @@ def test_refresh_all_segment_counts__live_segment_projects__delegates_to_enqueue
     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}
+    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
+    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
+    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]
 
 
 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:283`
 
 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:256`
 
 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:227`
+ - `api/segment_membership/tasks.py:236`
 
 Attributes:
  - `project.id`