Revert back and include old keys

synapse/handlers/sliding_sync/__init__.py

@@ -1381,39 +1381,15 @@ def _required_state_changes(
     only want to re-send that entry down sync if it has changed.
 
     Returns:
-        A 2-tuple of updated required state config and the state filter to use
+        A 2-tuple of updated required state config (or None if there is no update)
-        to fetch extra current state that we need to return.
+        and the state filter to use to fetch extra current state that we need to
+        return.
     """
 
     if prev_required_state_map == request_required_state_map:
         # There has been no change. Return immediately.
         return None, StateFilter.none()
 
-    # Convert the list of state_deltas to map from type to state_keys that have
-    # changed.
-    changed_types_to_state_keys: Dict[str, Set[str]] = {}
-    for event_type, state_key in state_deltas:
-        changed_types_to_state_keys.setdefault(event_type, set()).add(state_key)
-
-    new_required_state_map: Dict[str, Set[str]] = {}
-    for event_type in (
-        prev_required_state_map.keys() | request_required_state_map.keys()
-    ):
-        old_state_keys = prev_required_state_map.get(event_type, set())
-        changed_state_keys = changed_types_to_state_keys.get(event_type, set())
-
-        # We keep entries from the previous required state where the state hasn't
-        # changed. This way we can still keep track that we've already sent down
-        # the state to the client.
-        new_required_state_map[event_type] = request_required_state_map.get(
-            event_type, set()
-        ) | {
-            old_state_key
-            for old_state_key in old_state_keys
-            if old_state_key not in changed_state_keys
-            and old_state_key not in {StateValues.WILDCARD, StateValues.LAZY}
-        }
-
     prev_wildcard = prev_required_state_map.get(StateValues.WILDCARD, set())
     request_wildcard = request_required_state_map.get(StateValues.WILDCARD, set())
 
@@ -1434,16 +1410,33 @@ def _required_state_changes(
         # Keys were only removed, so we don't have to fetch everything.
         return request_required_state_map, StateFilter.none()
 
+    # Contains updates to the required state map compared with the previous room
+    # config. This has the same format as `RoomSyncConfig.required_state`
+    changes: Dict[str, AbstractSet[str]] = {}
+
     # The set of types/state keys that we need to fetch and return to the
     # client. Passed to `StateFilter.from_types(...)`
     added: List[Tuple[str, Optional[str]]] = []
 
+    # Convert the list of state deltas to map from type to state_keys that have
+    # changed.
+    changed_types_to_state_keys: Dict[str, Set[str]] = {}
+    for event_type, state_key in state_deltas:
+        changed_types_to_state_keys.setdefault(event_type, set()).add(state_key)
+
     # First we calculate what, if anything, has been *added*.
     for event_type in (
         prev_required_state_map.keys() | request_required_state_map.keys()
     ):
         old_state_keys = prev_required_state_map.get(event_type, set())
         request_state_keys = request_required_state_map.get(event_type, set())
+        changed_state_keys = changed_types_to_state_keys.get(event_type, set())
+
+        invalidated_state_keys = (
+            old_state_keys
+            - request_state_keys
+            - {StateValues.WILDCARD, StateValues.LAZY}
+        ) & changed_state_keys
 
         if old_state_keys == request_state_keys:
             # No change to this type
@@ -1453,6 +1446,11 @@ def _required_state_changes(
             # Nothing *added*, so we skip. Removals happen below.
             continue
 
+        # Always update changes to include the newly added keys
+        changes[event_type] = request_state_keys | (
+            old_state_keys - invalidated_state_keys
+        )
+
         if StateValues.WILDCARD in old_state_keys:
             # We were previously fetching everything for this type, so we don't need to
             # fetch anything new.
@@ -1467,8 +1465,11 @@ def _required_state_changes(
                 if state_key == StateValues.ME:
                     added.append((event_type, user_id))
                 elif state_key == StateValues.LAZY:
-                    # We handle lazy loading separately (outside this function), so
+                    # We handle lazy loading separately (outside this function),
-                    # don't need to explicitly add anything here.
+                    # so don't need to explicitly add anything here.
+                    #
+                    # LAZY values should also be ignore for event types that are
+                    # not membership.
                     pass
                 else:
                     added.append((event_type, state_key))
@@ -1478,22 +1479,78 @@ def _required_state_changes(
     # Figure out what changes we need to apply to the effective required state
     # config.
     for event_type, changed_state_keys in changed_types_to_state_keys.items():
-        if event_type in new_required_state_map:
+        old_state_keys = prev_required_state_map.get(event_type, set())
-            old_state_keys = prev_required_state_map.get(event_type, set())
+        request_state_keys = request_required_state_map.get(event_type, set())
-            request_state_keys = request_required_state_map.get(event_type, set())
 
-            # We only remove state keys from the effective state if they've been
+        invalidated_state_keys = (
-            # removed from the request *and* the state has changed. This ensures
+            old_state_keys
-            # that if a client removes and then re-adds a state key, we only send
+            - request_state_keys
-            # down the associated current state event if its changed (rather than
+            - {StateValues.WILDCARD, StateValues.LAZY}
-            # sending down the same event twice).
+        ) & changed_state_keys
-            invalidated = (
+
-                old_state_keys
+        if old_state_keys == request_state_keys:
-                - request_state_keys
+            # No change.
-                - {StateValues.WILDCARD, StateValues.LAZY}
+            continue
-            ) & changed_state_keys
+
-            new_required_state_map[event_type] = (
+        if request_state_keys - old_state_keys:
-                new_required_state_map[event_type] - invalidated
+            # We've expanded the set of state keys, so we just clobber the
+            # current set with the new set.
+            #
+            # We could also ensure that we keep entries where the state hasn't
+            # changed, but are no longer in the requested required state, but
+            # that's a sufficient edge case that we can ignore (as its only a
+            # performance optimization).
+            changes[event_type] = request_state_keys | (
+                old_state_keys - invalidated_state_keys
             )
+            continue
 
-    return new_required_state_map, added_state_filter
+        old_state_key_wildcard = StateValues.WILDCARD in old_state_keys
+        request_state_key_wildcard = StateValues.WILDCARD in request_state_keys
+
+        if old_state_key_wildcard != request_state_key_wildcard:
+            # If a state_key wildcard has been added or removed, we always update the
+            # effective room required state config to match the request.
+            changes[event_type] = request_state_keys
+            continue
+
+        if event_type == EventTypes.Member:
+            old_state_key_lazy = StateValues.LAZY in old_state_keys
+            request_state_key_lazy = StateValues.LAZY in request_state_keys
+
+            if old_state_key_lazy != request_state_key_lazy:
+                # If a "$LAZY" has been added or removed we always update the effective room
+                # required state config to match the request.
+                changes[event_type] = request_state_keys
+                continue
+
+        # Handle "$ME" values by adding "$ME" if the state key matches the user
+        # ID.
+        if user_id in changed_state_keys:
+            changed_state_keys.add(StateValues.ME)
+
+        # At this point there are no wildcards and no additions to the set of
+        # state keys requested, only deletions.
+        #
+        # We only remove state keys from the effective state if they've been
+        # removed from the request *and* the state has changed. This ensures
+        # that if a client removes and then re-adds a state key, we only send
+        # down the associated current state event if its changed (rather than
+        # sending down the same event twice).
+        invalidated = (old_state_keys - request_state_keys) & changed_state_keys
+        if invalidated:
+            changes[event_type] = old_state_keys - invalidated
+
+    if changes:
+        # Update the required state config based on the changes.
+        new_required_state_map = dict(prev_required_state_map)
+        for event_type, state_keys in changes.items():
+            if state_keys:
+                new_required_state_map[event_type] = state_keys
+            else:
+                # Remove entries with empty state keys.
+                new_required_state_map.pop(event_type, None)
+
+        return new_required_state_map, added_state_filter
+    else:
+        return None, added_state_filter

tests/handlers/test_sliding_sync.py

@@ -3687,27 +3687,32 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                         }
                     },
                     request_required_state_map={EventTypes.Member: {StateValues.LAZY}},
-                    state_deltas={(EventTypes.Member, "@user:test"): "$event_id"},
+                    state_deltas={(EventTypes.Member, "@user2:test"): "$event_id"},
                     expected_with_state_deltas=(
-                        # If a "$LAZY" has been added or removed we always update the
+                        # Remove "@user2:test" since that state has changed and is no
-                        # required state to what was requested for simplicity.
+                        # longer being requested anymore. Since something was removed,
+                        # we should persist the changed to required state. That way next
+                        # time, they request "@user2:test", we see that we haven't sent
+                        # it before and send the new state. (we should still keep track
+                        # that we've sent specific `EventTypes.Member` before)
                         {
                             EventTypes.Member: {
                                 StateValues.LAZY,
-                                "@user2:test",
                                 "@user3:test",
                             }
                         },
+                        # We don't need to request anything more if they are requesting
+                        # less state now
                         StateFilter.none(),
                     ),
                     expected_without_state_deltas=(
-                        {
+                        # We're not requesting any specific `EventTypes.Member` now but
-                            EventTypes.Member: {
+                        # since that state hasn't changed, nothing should change (we
-                                StateValues.LAZY,
+                        # should still keep track that we've sent specific
-                                "@user2:test",
+                        # `EventTypes.Member` before).
-                                "@user3:test",
+                        None,
-                            }
+                        # We don't need to request anything more if they are requesting
-                        },
+                        # less state now
                         StateFilter.none(),
                     ),
                 ),
@@ -3730,21 +3735,30 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                     request_required_state_map={
                         EventTypes.Member: {StateValues.LAZY, "@user4:test"}
                     },
-                    state_deltas={(EventTypes.Member, "@user:test"): "$event_id"},
+                    state_deltas={(EventTypes.Member, "@user2:test"): "$event_id"},
                     expected_with_state_deltas=(
-                        # If a "$LAZY" has been added or removed we always update the
+                        # Since "@user4:test" was added, we should persist the changed
-                        # required state to what was requested for simplicity.
+                        # required state config.
+                        #
+                        # Also remove "@user2:test" since that state has changed and is no
+                        # longer being requested anymore. Since something was removed,
+                        # we also should persist the changed to required state. That way next
+                        # time, they request "@user2:test", we see that we haven't sent
+                        # it before and send the new state. (we should still keep track
+                        # that we've sent specific `EventTypes.Member` before)
                         {
                             EventTypes.Member: {
                                 StateValues.LAZY,
-                                "@user2:test",
                                 "@user3:test",
                                 "@user4:test",
                             }
                         },
-                        StateFilter.none(),
+                        # We should see the new state_keys added
+                        StateFilter.from_types([(EventTypes.Member, "@user4:test")]),
                     ),
                     expected_without_state_deltas=(
+                        # Since "@user4:test" was added, we should persist the changed
+                        # required state config.
                         {
                             EventTypes.Member: {
                                 StateValues.LAZY,
@@ -3753,34 +3767,45 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                                 "@user4:test",
                             }
                         },
-                        StateFilter.none(),
+                        # We should see the new state_keys added
+                        StateFilter.from_types([(EventTypes.Member, "@user4:test")]),
                     ),
                 ),
             ),
             (
                 "state_key_expand_lazy_keep_previous_memberships",
                 """
-                Test TODO
+                Test expanding the `required_state` to lazy-loading room members.
                 """,
                 RequiredStateChangesTestParameters(
                     previous_required_state_map={
                         EventTypes.Member: {"@user2:test", "@user3:test"}
                     },
                     request_required_state_map={EventTypes.Member: {StateValues.LAZY}},
-                    state_deltas={(EventTypes.Member, "@user:test"): "$event_id"},
+                    state_deltas={(EventTypes.Member, "@user2:test"): "$event_id"},
                     expected_with_state_deltas=(
-                        # If a "$LAZY" has been added or removed we always update the
+                        # Since `StateValues.LAZY` was added, we should persist the
-                        # required state to what was requested for simplicity.
+                        # changed required state config.
+                        #
+                        # Also remove "@user2:test" since that state has changed and is no
+                        # longer being requested anymore. Since something was removed,
+                        # we also should persist the changed to required state. That way next
+                        # time, they request "@user2:test", we see that we haven't sent
+                        # it before and send the new state. (we should still keep track
+                        # that we've sent specific `EventTypes.Member` before)
                         {
                             EventTypes.Member: {
                                 StateValues.LAZY,
-                                "@user2:test",
                                 "@user3:test",
                             }
                         },
+                        # We don't need to request anything more if they are requesting
+                        # less state now
                         StateFilter.none(),
                     ),
                     expected_without_state_deltas=(
+                        # Since `StateValues.LAZY` was added, we should persist the
+                        # changed required state config.
                         {
                             EventTypes.Member: {
                                 StateValues.LAZY,
@@ -3788,10 +3813,92 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                                 "@user3:test",
                             }
                         },
+                        # We don't need to request anything more if they are requesting
+                        # less state now
                         StateFilter.none(),
                     ),
                 ),
             ),
+            (
+                "state_key_retract_lazy_keep_previous_memberships_no_new_memberships",
+                """
+                Test retracting the `required_state` to no longer lazy-loading room members.
+                """,
+                RequiredStateChangesTestParameters(
+                    previous_required_state_map={
+                        EventTypes.Member: {
+                            StateValues.LAZY,
+                            "@user2:test",
+                            "@user3:test",
+                        }
+                    },
+                    request_required_state_map={},
+                    state_deltas={(EventTypes.Member, "@user2:test"): "$event_id"},
+                    expected_with_state_deltas=(
+                        # TODO
+                        {
+                            EventTypes.Member: {
+                                "@user3:test",
+                            }
+                        },
+                        # We don't need to request anything more if they are requesting
+                        # less state now
+                        StateFilter.none(),
+                    ),
+                    expected_without_state_deltas=(
+                        # TODO
+                        {
+                            EventTypes.Member: {
+                                "@user2:test",
+                                "@user3:test",
+                            }
+                        },
+                        # We don't need to request anything more if they are requesting
+                        # less state now
+                        StateFilter.none(),
+                    ),
+                ),
+            ),
+            (
+                "state_key_retract_lazy_keep_previous_memberships_with_new_memberships",
+                """
+                Test retracting the `required_state` to no longer lazy-loading room members.
+                """,
+                RequiredStateChangesTestParameters(
+                    previous_required_state_map={
+                        EventTypes.Member: {
+                            StateValues.LAZY,
+                            "@user2:test",
+                            "@user3:test",
+                        }
+                    },
+                    request_required_state_map={EventTypes.Member: {"@user4:test"}},
+                    state_deltas={(EventTypes.Member, "@user2:test"): "$event_id"},
+                    expected_with_state_deltas=(
+                        # TODO
+                        {
+                            EventTypes.Member: {
+                                "@user3:test",
+                                "@user4:test",
+                            }
+                        },
+                        # We should see the new state_keys added
+                        StateFilter.from_types([(EventTypes.Member, "@user4:test")]),
+                    ),
+                    expected_without_state_deltas=(
+                        # TODO
+                        {
+                            EventTypes.Member: {
+                                "@user2:test",
+                                "@user3:test",
+                                "@user4:test",
+                            }
+                        },
+                        # We should see the new state_keys added
+                        StateFilter.from_types([(EventTypes.Member, "@user4:test")]),
+                    ),
+                ),
+            ),
             (
                 "type_wildcard_with_state_key_wildcard_to_explicit_state_keys",
                 """
@@ -3892,7 +3999,7 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                 ),
             ),
             (
-                "state_key_wildcard_to_explicit_state_keys",
+                "explicit_state_keys_to_wildcard_state_key",
                 """Test switching from a wildcard to explicit state keys with a concrete type""",
                 RequiredStateChangesTestParameters(
                     previous_required_state_map={
@@ -3905,20 +4012,28 @@ class RequiredStateChangesTestCase(unittest.TestCase):
                     # request. And we need to request all of the state for that type
                     # because we previously, only sent down a few keys.
                     expected_with_state_deltas=(
-                        {"type1": {StateValues.WILDCARD, "state_key2", "state_key3"}},
+                        {"type1": {StateValues.WILDCARD}},
                         StateFilter.from_types([("type1", None)]),
                     ),
+                    # expected_with_state_deltas=(
+                    #     {"type1": {StateValues.WILDCARD, "state_key2", "state_key3"}},
+                    #     StateFilter.from_types([("type1", None)]),
+                    # ),
                     expected_without_state_deltas=(
-                        {
+                        {"type1": {StateValues.WILDCARD}},
-                            "type1": {
-                                StateValues.WILDCARD,
-                                "state_key1",
-                                "state_key2",
-                                "state_key3",
-                            }
-                        },
                         StateFilter.from_types([("type1", None)]),
                     ),
+                    # expected_without_state_deltas=(
+                    #     {
+                    #         "type1": {
+                    #             StateValues.WILDCARD,
+                    #             "state_key1",
+                    #             "state_key2",
+                    #             "state_key3",
+                    #         }
+                    #     },
+                    #     StateFilter.from_types([("type1", None)]),
+                    # ),
                 ),
             ),
         ]
@@ -3942,6 +4057,11 @@ class RequiredStateChangesTestCase(unittest.TestCase):
             test_parameters.expected_without_state_deltas[0],
             "changed_required_state_map does not match (without state_deltas)",
         )
+        logger.info("asdf actual added_state_filter: %s", added_state_filter)
+        logger.info(
+            "asdf expected added_state_filter: %s",
+            test_parameters.expected_with_state_deltas[1],
+        )
         self.assertEqual(
             added_state_filter,
             test_parameters.expected_without_state_deltas[1],