Reconcile automations when Postgres notifications are unavailable (#21089)

Author: zzstoatzz

src/prefect/server/events/triggers.py

@@ -75,6 +75,9 @@
 
 AutomationID: TypeAlias = UUID
 TriggerID: TypeAlias = UUID
+AutomationStateSnapshot: TypeAlias = Tuple[
+    Optional[prefect.types._datetime.DateTime], int
+]
 
 
 AUTOMATION_BUCKET_BATCH_SIZE = 500
@@ -679,6 +682,8 @@ async def periodic_evaluation(now: prefect.types._datetime.DateTime) -> None:
 
     logger.debug("Running periodic evaluation as of %s (offset %ss)", as_of, offset)
 
+    await reconcile_automations()
+
     # Any followers that have been sitting around longer than our lookback are never
     # going to see their leader event (maybe it was lost or took too long to arrive).
     # These events can just be evaluated now in the order they occurred.
@@ -713,6 +718,7 @@ async def evaluate_periodically(periodic_granularity: timedelta) -> None:
 automations_by_id: Dict[UUID, Automation] = {}
 triggers: Dict[TriggerID, EventTrigger] = {}
 next_proactive_runs: Dict[TriggerID, prefect.types._datetime.DateTime] = {}
+automation_state_snapshot: Optional[AutomationStateSnapshot] = None
 
 # This lock governs any changes to the set of loaded automations; any routine that will
 # add/remove automations must be holding this lock when it does so.  It's best to use
@@ -732,6 +738,12 @@ def find_interested_triggers(event: ReceivedEvent) -> Collection[EventTrigger]:
     return [trigger for trigger in candidates if trigger.covers(event)]
 
 
+def clear_loaded_automations() -> None:
+    automations_by_id.clear()
+    triggers.clear()
+    next_proactive_runs.clear()
+
+
 def load_automation(automation: Optional[Automation]) -> None:
     """Loads the given automation into memory so that it is available for evaluations"""
     if not automation:
@@ -762,14 +774,17 @@ async def automation_changed(
     automation_id: UUID,
     event: Literal["automation__created", "automation__updated", "automation__deleted"],
 ) -> None:
+    global automation_state_snapshot
+
     async with _automations_lock():
         if event in ("automation__deleted", "automation__updated"):
             forget_automation(automation_id)
 
-        if event in ("automation__created", "automation__updated"):
-            async with automations_session() as session:
+        async with automations_session() as session:
+            if event in ("automation__created", "automation__updated"):
                 automation = await read_automation(session, automation_id)
                 load_automation(automation)
+            automation_state_snapshot = await read_automation_state_snapshot(session)
 
 
 @db_injector
@@ -788,6 +803,53 @@ async def load_automations(db: PrefectDBInterface, session: AsyncSession):
     )
 
 
+@db_injector
+async def read_automation_state_snapshot(
+    db: PrefectDBInterface, session: AsyncSession
+) -> AutomationStateSnapshot:
+    query = sa.select(
+        sa.func.max(db.Automation.updated),
+        sa.func.count(db.Automation.id),
+    ).select_from(db.Automation)
+
+    latest_updated, count = (await session.execute(query)).one()
+
+    return (
+        prefect.types._datetime.create_datetime_instance(latest_updated)
+        if latest_updated
+        else None,
+        count or 0,
+    )
+
+
+async def reconcile_automations(force: bool = False) -> bool:
+    global automation_state_snapshot
+
+    async with _automations_lock():
+        async with automations_session() as session:
+            current_snapshot = await read_automation_state_snapshot(session)
+            if not force and current_snapshot == automation_state_snapshot:
+                return False
+
+            previous_automations = automations_by_id.copy()
+            previous_triggers = triggers.copy()
+            previous_next_proactive_runs = next_proactive_runs.copy()
+
+            clear_loaded_automations()
+
+            try:
+                await load_automations(session)
+            except Exception:
+                clear_loaded_automations()
+                automations_by_id.update(previous_automations)
+                triggers.update(previous_triggers)
+                next_proactive_runs.update(previous_next_proactive_runs)
+                raise
+
+            automation_state_snapshot = current_snapshot
+            return True
+
+
 @db_injector
 async def remove_buckets_exceeding_threshold(
     db: PrefectDBInterface, session: AsyncSession, trigger: EventTrigger
@@ -1068,10 +1130,11 @@ async def sweep_closed_buckets(
 
 async def reset() -> None:
     """Resets the in-memory state of the service"""
+    global automation_state_snapshot
+
     await reset_events_clock()
-    automations_by_id.clear()
-    triggers.clear()
-    next_proactive_runs.clear()
+    clear_loaded_automations()
+    automation_state_snapshot = None
 
 
 async def listen_for_automation_changes() -> None:
@@ -1096,6 +1159,8 @@ async def listen_for_automation_changes() -> None:
                 f"Listening for automation changes on {AUTOMATION_CHANGES_CHANNEL}"
             )
 
+            await reconcile_automations()
+
             async for payload in pg_listen(
                 conn,
                 AUTOMATION_CHANGES_CHANNEL,
@@ -1152,8 +1217,7 @@ async def consumer(
     # Start the automation change listener task
     sync_task = asyncio.create_task(listen_for_automation_changes())
 
-    async with automations_session() as session:
-        await load_automations(session)
+    await reconcile_automations(force=True)
 
     proactive_task = asyncio.create_task(evaluate_periodically(periodic_granularity))
 

tests/events/server/triggers/test_service.py

@@ -103,6 +103,17 @@ def load_automations(monkeypatch: pytest.MonkeyPatch) -> mock.AsyncMock:
     return m
 
 
+@pytest.fixture(autouse=True)
+def read_automation_state_snapshot(
+    monkeypatch: pytest.MonkeyPatch,
+) -> mock.AsyncMock:
+    m = mock.AsyncMock(return_value=(None, 0))
+    monkeypatch.setattr(
+        "prefect.server.events.triggers.read_automation_state_snapshot", m
+    )
+    return m
+
+
 @pytest.fixture(autouse=True)
 def periodic_evaluation(monkeypatch: pytest.MonkeyPatch) -> mock.AsyncMock:
     m = mock.AsyncMock(spec=triggers.periodic_evaluation)
@@ -243,6 +254,45 @@ async def test_loads_automations_at_startup(
         load_automations.assert_awaited_once_with(open_automations_session)
 
 
+async def test_reconcile_automations_skips_reload_when_snapshot_matches(
+    open_automations_session: mock.Mock,
+    load_automations: mock.AsyncMock,
+    read_automation_state_snapshot: mock.AsyncMock,
+):
+    snapshot = (now("UTC"), 2)
+    triggers.automation_state_snapshot = snapshot
+    read_automation_state_snapshot.return_value = snapshot
+
+    changed = await triggers.reconcile_automations()
+
+    assert changed is False
+    load_automations.assert_not_awaited()
+
+
+async def test_reconcile_automations_reloads_when_snapshot_changes(
+    arachnophobia: Automation,
+    open_automations_session: mock.Mock,
+    load_automations: mock.AsyncMock,
+    read_automation_state_snapshot: mock.AsyncMock,
+):
+    previous_snapshot = (now("UTC"), 1)
+    current_snapshot = (now("UTC") + timedelta(seconds=1), 2)
+    triggers.automation_state_snapshot = previous_snapshot
+    triggers.load_automation(arachnophobia)
+    (trigger,) = arachnophobia.triggers_of_type(EventTrigger)
+    triggers.next_proactive_runs[trigger.id] = now("UTC")
+    read_automation_state_snapshot.return_value = current_snapshot
+
+    changed = await triggers.reconcile_automations()
+
+    assert changed is True
+    load_automations.assert_awaited_once_with(open_automations_session)
+    assert triggers.automation_state_snapshot == current_snapshot
+    assert triggers.automations_by_id == {}
+    assert triggers.triggers == {}
+    assert triggers.next_proactive_runs == {}
+
+
 async def test_only_considers_messages_with_attributes(
     effective_automations,
     reactive_evaluation: mock.AsyncMock,