Fix push badge count ignoring a room's notifications when another room's summary is up to date (#20237)

The count of unread notifications is computed in two phases:

1. From the summaries (`event_push_summary`) when those are up to date
with the user's last read receipt
2. Otherwise from counting the push actions (`event_push_actions`) after
that receipt

The problem in this process is that the threads whose summary is up to
date are identified with only their `thread_id`. But all main timelines
share the same `"main"` thread_id, so as soon as one room has an
up-to-date summary, phase 2 skips the main timeline of every room.

The fix identifies the up-to-date summaries with their `(room_id,
thread_id)` pair.

Authored by @sandhose. Found while investigating the
`TestThreadedReceipts` Complement flake (#15517, #18537), but it is a
bug on its own.

### Pull Request Checklist

<!-- Please read
https://element-hq.github.io/synapse/latest/development/contributing_guide.html
before submitting your pull request -->

* [x] Pull request is based on the develop branch
* [x] Pull request includes a [changelog
file](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#changelog).
The entry should:
- Be a short description of your change which makes sense to users.
"Fixed a bug that prevented receiving messages from other servers."
instead of "Moved X method from `EventStore` to `EventWorkerStore`.".
  - Use markdown where necessary, mostly for `code blocks`.
  - End with either a period (.) or an exclamation mark (!).
  - Start with a capital letter.
- Feel free to credit yourself, by adding a sentence "Contributed by
@github_username." or "Contributed by [Your Name]." to the end of the
entry.
* [x] [Code
style](https://element-hq.github.io/synapse/latest/code_style.html) is
correct (run the
[linters](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#run-the-linters))

---------

Co-authored-by: Quentin Gliech <quentingliech@gmail.com>
Co-authored-by: Devon Hudson <devonhudson@librem.one>

changelog.d/20237.bugfix

@@ -0,0 +1 @@
+Fix a bug where push badge counts could ignore notifications in one room based on another room's summarised counts.

synapse/storage/databases/main/event_push_actions.py

@@ -435,12 +435,13 @@ class EventPushActionsWorkerStore(ReceiptsWorkerStore, StreamWorkerStore, SQLBas
         """
         txn.execute(sql, args)
 
-        seen_thread_ids = set()
+        # The (room ID, thread ID) pairs we found an up-to-date summary for.
+        seen_room_thread_ids = set()
         room_to_count: dict[str, int] = defaultdict(int)
 
         for room_id, thread_id, notif_count in txn:
             room_to_count[room_id] += notif_count
-            seen_thread_ids.add(thread_id)
+            seen_room_thread_ids.add((room_id, thread_id))
 
         # Now get any event push actions that haven't been rotated using the same OR
         # join and filter by receipt and event push summary rotated up to stream ordering.
@@ -460,35 +461,31 @@ class EventPushActionsWorkerStore(ReceiptsWorkerStore, StreamWorkerStore, SQLBas
 
         for room_id, thread_id, notif_count in txn:
             # Note: only count push actions we have valid summaries for with up to date receipt.
-            if thread_id not in seen_thread_ids:
+            if (room_id, thread_id) not in seen_room_thread_ids:
                 continue
             room_to_count[room_id] += notif_count
 
-        thread_id_clause, thread_ids_args = make_in_list_sql_clause(
-            self.database_engine, "epa.thread_id", seen_thread_ids
-        )
-
-        # Finally re-check event_push_actions for any rooms not in the summary, ignoring
-        # the rotated up-to position. This handles the case where a read receipt has arrived
-        # but not been rotated meaning the summary table is out of date, so we go back to
-        # the push actions table.
+        # Finally re-check event_push_actions for any room/threads not in the summary,
+        # ignoring the rotated up-to position. This handles the case where a read receipt
+        # has arrived but not been rotated meaning the summary table is out of date, so we
+        # go back to the push actions table.
         sql = f"""
             {receipts_cte}
-            SELECT epa.room_id, COUNT(CASE WHEN epa.notif = 1 THEN 1 END) AS notif_count
+            SELECT epa.room_id, epa.thread_id, COUNT(CASE WHEN epa.notif = 1 THEN 1 END) AS notif_count
             FROM event_push_actions AS epa
             {receipts_joins}
             WHERE user_id = ?
-            AND NOT {thread_id_clause}
             AND epa.notif = 1
             AND (threaded_receipt_stream_ordering IS NULL OR stream_ordering > threaded_receipt_stream_ordering)
             AND (unthreaded_receipt_stream_ordering IS NULL OR stream_ordering > unthreaded_receipt_stream_ordering)
-            GROUP BY epa.room_id
+            GROUP BY epa.room_id, epa.thread_id
         """
 
-        args.extend(thread_ids_args)
         txn.execute(sql, args)
 
-        for room_id, notif_count in txn:
+        for room_id, thread_id, notif_count in txn:
+            if (room_id, thread_id) in seen_room_thread_ids:
+                continue
             room_to_count[room_id] += notif_count
 
         return room_to_count

tests/storage/test_event_push_actions.py

@@ -472,6 +472,82 @@ class EventPushActionsStoreTestCase(HomeserverTestCase):
         self.get_success(self.store._rotate_notifs())
         _assert_badge(1)
 
+    def test_count_aggregation_badge_recount_is_scoped_per_room(self) -> None:
+        """
+        Regression test: a room whose summary row is out of date must be recounted
+        from `event_push_actions`, even when another room has an up-to-date summary
+        for the same thread ID.
+
+        The set of threads a valid summary was found for used to be keyed on the
+        thread ID alone, so a single room with an up-to-date `main` summary excluded
+        `main` from the recount in *every* room, dropping those rooms' counts.
+        """
+        user_id, token, other_id, other_token, room_id = self._create_users_and_room()
+
+        stale_room_id = self.helper.create_room_as(user_id, tok=token)
+        self.helper.join(stale_room_id, other_id, tok=other_token)
+
+        def _send(room: str) -> str:
+            return self.helper.send_event(
+                room,
+                type="m.room.message",
+                content={"msgtype": "m.text", "body": "msg"},
+                tok=other_token,
+            )["event_id"]
+
+        def _read(room: str, event_id: str) -> None:
+            self.get_success(
+                self.store.insert_receipt(
+                    room,
+                    "m.read",
+                    user_id=user_id,
+                    event_ids=[event_id],
+                    thread_id=None,
+                    data={},
+                )
+            )
+
+        def _badge(room: str) -> int:
+            counts = self.get_success(
+                self.store.db_pool.runInteraction(
+                    "get-aggregate-unread-counts",
+                    self.store._get_unread_counts_by_room_for_user_txn,
+                    user_id,
+                )
+            )
+            return counts.get(room, 0)
+
+        # `room_id` keeps an up-to-date summary throughout, so its `main` thread is
+        # always one we found a valid summary for.
+        first = _send(room_id)
+        _send(room_id)
+
+        stale_first = _send(stale_room_id)
+        stale_second = _send(stale_room_id)
+        _send(stale_room_id)
+
+        # Read one event in each room and rotate, so that both summary rows record
+        # the receipt they were calculated against.
+        _read(room_id, first)
+        _read(stale_room_id, stale_first)
+        self.get_success(self.store._rotate_notifs())
+
+        self.assertEqual(_badge(room_id), 1)
+        self.assertEqual(_badge(stale_room_id), 2)
+
+        # A second receipt, which rotation has not processed yet: `stale_room_id`'s
+        # summary row no longer matches it, so its count has to be recovered from
+        # `event_push_actions`.
+        _read(stale_room_id, stale_second)
+
+        self.assertEqual(_badge(room_id), 1)
+        self.assertEqual(_badge(stale_room_id), 1)
+        # A new event, not yet rotated, while the summary row is still stale.
+        _send(stale_room_id)
+
+        self.assertEqual(_badge(room_id), 1)
+        self.assertEqual(_badge(stale_room_id), 2)
+
     def test_count_aggregation_threads(self) -> None:
         """
         This is essentially the same test as test_count_aggregation, but adds