Add an approximate difference method to StateFilters

synapse/storage/databases/state/store.py

@@ -24,6 +24,7 @@
 from synapse.storage.types import Cursor
 from synapse.storage.util.sequence import build_sequence_generator
 from synapse.types import MutableStateMap, StateMap
+from synapse.util.async_helpers import ObservableDeferred
 from synapse.util.caches.descriptors import cached
 from synapse.util.caches.dictionary_cache import DictionaryCache
 
@@ -91,6 +92,12 @@ def __init__(self, database: DatabasePool, db_conn, hs):
             500000,
         )
 
+        # Current ongoing get_state_for_groups in-flight requests
+        # {group ID -> {StateFilter -> ObservableDeferred}}
+        self._state_group_inflight_requests: Dict[
+            int, Dict[StateFilter, ObservableDeferred[StateMap[str]]]
+        ] = {}
+
         def get_max_state_group_txn(txn: Cursor):
             txn.execute("SELECT COALESCE(max(id), 0) FROM state_groups")
             return txn.fetchone()[0]
@@ -249,6 +256,44 @@ async def _get_state_for_groups(
         if not incomplete_groups:
             return state
 
+        # try and rely on in-flight requests to complete this request without
+        # needing to spawn additional queries.
+
+        # (group ID, ObservableDeferred of request result)
+        reusable_inflight_requests: List[
+            Tuple[int, ObservableDeferred[StateMap[str]]]
+        ] = []
+        # group ID -> left over StateFilter to request
+        requests_to_spawn: Dict[int, StateFilter] = {}
+
+        for group in incomplete_groups:
+            requests_in_flight_for_group = self._state_group_inflight_requests.get(
+                group
+            )
+            if requests_in_flight_for_group is None:
+                continue
+
+            state_filter_left_over = state_filter
+            for (
+                request_state_filter,
+                request_deferred,
+            ) in requests_in_flight_for_group.items():
+                new_state_filter_left_over = state_filter_left_over.approx_difference(
+                    request_state_filter
+                )
+                if new_state_filter_left_over != state_filter_left_over:
+                    # reusing this request narrows our StateFilter down a bit.
+                    reusable_inflight_requests.append((group, request_deferred))
+                state_filter_left_over = new_state_filter_left_over
+                if state_filter_left_over == StateFilter.none():
+                    # we have managed to collect enough of the in-flight requests
+                    # to cover our StateFilter and give us the state we need.
+                    break
+            else:
+                # we have some of the state filter left over, so need to spawn
+                # a request
+                requests_to_spawn[group] = state_filter_left_over
+
         cache_sequence_nm = self._state_group_cache.sequence
         cache_sequence_m = self._state_group_members_cache.sequence
 

synapse/storage/state.py

@@ -361,18 +361,20 @@ def approx_difference(self, subtrahend: "StateFilter") -> "StateFilter":
         Returns a state filter which represents `self - subtrahend`.
 
         The resultant state filter MUST admit all state events that are admitted
-        by only this filter (`self`) and not `subtrahend`. (1)
+        by only this filter (`self`) and not `subtrahend`.
         The resultant filter MAY be an over-approximation: the resultant state
-        filter MAY additionally admit other state events.
+        filter MAY additionally admit some state events from `subtrahend`.
 
 
         Formally, if the set of state events admitted by a state filter F are
         written as E(F), then the resultant state filter bears this property:
 
-            E(difference(self, subtrahend)) ⊇ E(self) ∖ E(subtrahend)
+            E(self) ∖ E(subtrahend)
+                ⊆ E(approx_difference(self, subtrahend))
+                ⊆ E(self)
 
 
-        This function attempts to return the narrowest such state filter.
+        This implementation attempts to return the narrowest such state filter.
         In the case that `self` contains wildcards for state types where
         `subtrahend` contains specific state keys, an approximation must be made:
         the resultant state filter keeps the wildcard, as state filters are not