Merge branch 'rav/faster_joins/07_partial_state_events_tracker' into …

…rav/faster_joins/08_no_wait_for_filtered_state

synapse/storage/util/partial_state_events_tracker.py

@@ -31,7 +31,9 @@ class PartialStateEventsTracker:
 
     def __init__(self, store: EventsWorkerStore):
         self._store = store
-        self._observers: Dict[str, Set[Deferred]] = defaultdict(set)
+        # a map from event id to a set of Deferreds which are waiting for that event to be
+        # un-partial-stated.
+        self._observers: Dict[str, Set[Deferred[None]]] = defaultdict(set)
 
     def notify_un_partial_stated(self, event_id: str) -> None:
         """Notify that we now have full state for a given event
@@ -92,7 +94,9 @@ async def await_full_state(self, event_ids: Collection[str]) -> None:
             for event_id, partial in (
                 await self._store.get_partial_state_events(observers.keys())
             ).items():
-                if not partial:
+                # there may have been a call to notify_un_partial_stated during the
+                # db query, so the observers may already have been called.
+                if not partial and not observers[event_id].called:
                     observers[event_id].callback(None)
 
             await make_deferred_yieldable(

tests/storage/util/test_partial_state_events_tracker.py

@@ -75,6 +75,29 @@ async def get_partial_state_events(events):
             ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
         )
 
+    def test_un_partial_state_during_get_partial_state_events(self):
+        # we should correctly handle a call to notify_un_partial_stated during the
+        # second call to get_partial_state_events.
+
+        self._events_dict = {"event1": True, "event2": False}
+
+        async def get_partial_state_events1(events):
+            self.mock_store.get_partial_state_events.side_effect = (
+                get_partial_state_events2
+            )
+            return {e: self._events_dict[e] for e in events}
+
+        async def get_partial_state_events2(events):
+            self.tracker.notify_un_partial_stated("event1")
+            self._events_dict["event1"] = False
+            return {e: self._events_dict[e] for e in events}
+
+        self.mock_store.get_partial_state_events.side_effect = get_partial_state_events1
+
+        self.successResultOf(
+            ensureDeferred(self.tracker.await_full_state(["event1", "event2"]))
+        )
+
     def test_cancellation(self):
         self._events_dict = {"event1": True, "event2": False}