Add a cache to state res to cache more cases (#20185)

This is an attempt to better cache the cases where there are a large
number of extremities to resolve over, which keep slightly changing.
This spawns from seeing issues on matrix.org.

We already have a cache over the exact state groups being resolved.
However, we can do better by caching the inputs into state res (i.e. the
conflicted sets), which are more likely to be constant across repeated
state res in a room. We key this cache based on a sha256 hash, on the
assumption that this will never conflict.

Also includes a commit that removes needless copying of the state.

changelog.d/20185.misc

@@ -0,0 +1 @@
+Add a cache to state resolution keyed off the conflicted events.

synapse/state/__init__.py

@@ -32,7 +32,6 @@ from typing import (
 )
 
 import attr
-from immutabledict import immutabledict
 from prometheus_client import Counter, Histogram
 
 from synapse.api.constants import EventTypes
@@ -53,6 +52,7 @@ from synapse.storage.databases.main.event_federation import StateDifference
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.types import StateMap, StrCollection
 from synapse.types.state import StateFilter
+from synapse.util import MutableOverlayMapping
 from synapse.util.async_helpers import Linearizer
 from synapse.util.caches.expiringcache import ExpiringCache
 from synapse.util.duration import Duration
@@ -113,18 +113,13 @@ class _StateCacheEntry:
         #
         # This can be None if we have a `state_group` (as then we can fetch the
         # state from the DB.)
-        self._state: StateMap[str] | None = (
-            immutabledict(state) if state is not None else None
-        )
-
+        self._state = state
         # the ID of a state group if one and only one is involved.
         # otherwise, None otherwise?
         self.state_group = state_group
 
         self.prev_group = prev_group
-        self.delta_ids: StateMap[str] | None = (
-            immutabledict(delta_ids) if delta_ids is not None else None
-        )
+        self.delta_ids = delta_ids
 
     async def get_state(
         self,
@@ -174,14 +169,29 @@ class _StateCacheEntry:
         length = 0
 
         if self._state:
-            length += len(self._state)
+            length += _state_map_size(self._state)
 
         if self.delta_ids:
-            length += len(self.delta_ids)
+            length += _state_map_size(self.delta_ids)
 
         return length or 1  # Make sure its not 0.
 
 
+def _state_map_size(state_map: Mapping[Any, Any]) -> int:
+    """Estimate a proxy for the memory a state map holds, for sizing caches.
+
+    Since state maps are often combinations of `ChainMap` and
+    `MutableOverlayMapping`, we look at the total number of entries across all
+    layers rather than just the number of distinct keys. This is both faster and
+    a more accurate proxy for memory usage.
+    """
+    if isinstance(state_map, ChainMap):
+        return sum(_state_map_size(layer) for layer in state_map.maps)
+    if isinstance(state_map, MutableOverlayMapping):
+        return state_map.total_entries()
+    return len(state_map)
+
+
 class StateHandler:
     """Fetches bits of state from the stores, and does state resolution
     where necessary
@@ -636,6 +646,31 @@ class StateResolutionHandler:
             )
         )
 
+        # The result of resolving a conflicted set of state, keyed on a digest
+        # of the inputs to `_resolve_conflicted_set`. See
+        # `v2._conflict_cache_key`.
+        #
+        # This is different to `_state_cache` above, which caches the resolved
+        # state based on the state groups. This cache aims to address the case
+        # where resolving across different state groups often produces the same
+        # conflicted set, which we can then cache.
+        #
+        # We bound the size of the cache based on the size calculated by
+        # `_state_map_size`, which calculates a proxy for a rough estimate of
+        # the memory footprint of a state map.
+        self._conflict_resolution_cache: ExpiringCache[bytes, StateMap[str]] = (
+            ExpiringCache(
+                cache_name="state_conflict_resolution_cache",
+                server_name=self.server_name,
+                hs=hs,
+                clock=self.clock,
+                max_len=100000,
+                expiry_ms=EVICTION_TIMEOUT_SECONDS * 1000,
+                size_callback=_state_map_size,
+                reset_expiry_on_get=True,
+            )
+        )
+
         #
         # stuff for tracking time spent on state-res by room
         #
@@ -798,6 +833,7 @@ class StateResolutionHandler:
                         state_sets,
                         event_map,
                         state_res_store,
+                        conflict_cache=self._conflict_resolution_cache,
                     )
         finally:
             self._record_state_res_metrics(room_id, m.get_resource_usage())

synapse/state/v2.py

@@ -18,18 +18,24 @@
 #
 #
 
+import hashlib
 import heapq
 import itertools
+import json
 import logging
+from collections import ChainMap
 from typing import (
+    AbstractSet,
     Any,
     Awaitable,
     Callable,
     Generator,
     Iterable,
     Literal,
+    MutableMapping,
     Protocol,
     Sequence,
+    cast,
     overload,
 )
 
@@ -39,7 +45,8 @@ from synapse.api.errors import AuthError
 from synapse.api.room_versions import RoomVersion, StateResolutionVersions
 from synapse.events import EventBase, is_creator
 from synapse.storage.databases.main.event_federation import StateDifference
-from synapse.types import MutableStateMap, StateMap, StrCollection
+from synapse.types import MutableStateMap, StateKey, StateMap, StrCollection
+from synapse.util import MutableOverlayMapping
 from synapse.util.duration import Duration
 
 logger = logging.getLogger(__name__)
@@ -86,6 +93,7 @@ async def resolve_events_with_store(
     state_sets: Sequence[StateMap[str]],
     event_map: dict[str, EventBase] | None,
     state_res_store: StateResolutionStore,
+    conflict_cache: "ConflictCache | None" = None,
 ) -> StateMap[str]:
     """Resolves the state using the v2 state resolution algorithm
 
@@ -104,6 +112,10 @@ async def resolve_events_with_store(
             If None, all events will be fetched via state_res_store.
 
         state_res_store:
+        conflict_cache:
+            if given, the resolution of the conflicted set is looked up here
+            before it is computed, and stored here afterwards. See
+            `_conflict_cache_key` for what the key covers.
 
     Returns:
         A map from (type, state_key) to event_id.
@@ -163,6 +175,145 @@ async def resolve_events_with_store(
 
     logger.debug("%d full_conflicted_set entries", len(full_conflicted_set))
 
+    # Calculate the base state.
+    #
+    # v2 uses the unconflicted state as the base state, but v2.1 uses the empty
+    # set.
+    base_state: StateMap[str] = {}
+    if room_version.state_res != StateResolutionVersions.V2_1:
+        # Resolving conflicted sets requires the following types from the base
+        # state:
+        #   - the `auth_types_for_event(..)` of the conflicted events for
+        #     `_iterative_auth_checks`
+        #   - the power levels for `_mainline_sort`
+        #
+        # We can therefore safely restrict the base state to those keys, which
+        # keeps keys that cannot affect the outcome out of the cache key. The
+        # rest of the unconflicted state is layered back on below.
+        base_state_keys = {(EventTypes.PowerLevels, "")}
+        for event_id in full_conflicted_set:
+            base_state_keys.update(
+                event_auth.auth_types_for_event(room_version, event_map[event_id])
+            )
+
+        base_state = {
+            key: unconflicted_state[key]
+            for key in base_state_keys
+            if key in unconflicted_state
+        }
+
+    resolved_state: StateMap[str] | None = None
+    if conflict_cache is not None:
+        cache_key = _conflict_cache_key(
+            room_id, room_version, full_conflicted_set, base_state
+        )
+        resolved_state = conflict_cache.get(cache_key)
+
+    if resolved_state is None:
+        resolved_state = await _resolve_conflicted_set(
+            clock,
+            room_id,
+            room_version,
+            full_conflicted_set,
+            base_state,
+            event_map,
+            state_res_store,
+        )
+        if conflict_cache is not None:
+            conflict_cache[cache_key] = resolved_state
+    else:
+        logger.debug(
+            "Reusing the resolution of %d conflicted events",
+            len(full_conflicted_set),
+        )
+
+    logger.debug("done")
+
+    # Finally, we copy the unconflicted state over the resolved state.
+    #
+    # We use a `ChainMap` here to avoid a copy.
+    #
+    # `ChainMap` expects mutable mappings as it is a mutable mapping. However,
+    # the return type of this function is an immutable mapping so it is safe to
+    # cast the underlying mappings to mutable mappings to satify `ChainMap`'s
+    # signature.
+    return ChainMap(
+        cast(MutableMapping[StateKey, str], unconflicted_state),
+        cast(MutableMapping[StateKey, str], resolved_state),
+    )
+
+
+class ConflictCache(Protocol):
+    """What `resolve_events_with_store` needs from a `conflict_cache`.
+
+    Satisfied by a plain `dict` and by `ExpiringCache`.
+
+    Mainly used to allow unit tests to pass a `dict` rather than building a full
+    cache.
+    """
+
+    def get(self, key: bytes) -> StateMap[str] | None: ...
+
+    def __setitem__(self, key: bytes, value: StateMap[str]) -> None: ...
+
+
+def _conflict_cache_key(
+    room_id: str,
+    room_version: RoomVersion,
+    full_conflicted_set: AbstractSet[str],
+    base_state: StateMap[str],
+) -> bytes:
+    """Create a key for the conflict cache. Incorporates everything
+    that would affect the result of `_resolve_conflicted_set`.
+
+    We use a digest as the key as the inputs can be very large.
+    """
+
+    # We use JSON as the serialization format for ease, we could use a
+    # hand-rolled format but one has to be careful to ensure that you can't get
+    # collisions (given in some room versions the room ID is arbitrary bytes
+    # rather than a hash).
+    key_material = json.dumps(
+        {
+            "room_id": room_id,
+            "room_version": room_version.identifier,
+            "conflicted_set": sorted(full_conflicted_set),
+            "base_state": sorted(base_state.values()),
+        }
+    )
+    return hashlib.sha256(key_material.encode("utf-8")).digest()
+
+
+async def _resolve_conflicted_set(
+    clock: Clock,
+    room_id: str,
+    room_version: RoomVersion,
+    full_conflicted_set: set[str],
+    base_state: StateMap[str],
+    event_map: dict[str, EventBase],
+    state_res_store: StateResolutionStore,
+) -> MutableStateMap[str]:
+    """Apply the conflicted events to the base state.
+
+    This is the expensive part of `resolve_events_with_store`: the reverse
+    topological power sort, the two iterative auth check passes and the
+    mainline sort.
+
+    Args:
+        clock
+        room_id: the room we are working in
+        room_version: the room version
+        full_conflicted_set: the conflicted events plus the auth chain
+            difference. All must be present in `event_map`.
+        base_state: the state to apply the conflicted events to
+        event_map: updated in place with the events fetched along the way
+        state_res_store
+
+    Returns:
+        The base state with the conflicted events that passed auth applied
+        over it. The unconflicted state is not merged in.
+    """
+
     # Get and sort all the power events (kicks/bans/etc)
     power_events = (
         eid for eid in full_conflicted_set if _is_power_event(event_map[eid])
@@ -174,17 +325,6 @@ async def resolve_events_with_store(
 
     logger.debug("sorted %d power events", len(sorted_power_events))
 
-    # v2.1 starts iterative auth checks from the empty set and not the unconflicted state.
-    # It relies on IAC behaviour which populates the base state with the events from auth_events
-    # if the state tuple is missing from the base state. This ensures the base state is only
-    # populated from auth_events rather than whatever the unconflicted state is (which could be
-    # completely bogus).
-    base_state = (
-        {}
-        if room_version.state_res == StateResolutionVersions.V2_1
-        else unconflicted_state
-    )
-
     # Now sequentially auth each one
     resolved_state = await _iterative_auth_checks(
         clock,
@@ -227,11 +367,6 @@ async def resolve_events_with_store(
 
     logger.debug("resolved")
 
-    # We make sure that unconflicted state always still applies.
-    resolved_state.update(unconflicted_state)
-
-    logger.debug("done")
-
     return resolved_state
 
 
@@ -691,7 +826,7 @@ async def _iterative_auth_checks(
     Returns:
         Returns the final updated state
     """
-    resolved_state = dict(base_state)
+    resolved_state = MutableOverlayMapping(base_state)
 
     for idx, event_id in enumerate(event_ids, start=1):
         event = event_map[event_id]

synapse/util/__init__.py

@@ -151,19 +151,35 @@ class MutableOverlayMapping(collections.abc.MutableMapping[K, V]):
                 yield key
 
     def __len__(self) -> int:
-        count = len(self._underlying_map)
-        for key in self._deletions:
-            if key in self._underlying_map:
-                count -= 1
+        # The distinct keys can be calculated via `(underlying ∪ mutable) −
+        # deletions`. A key is never in both `_mutable_map` and `_deletions`, so
+        # the only deletions to subtract are those of keys in the underlying
+        # map:
+        #
+        #   |underlying| + |mutable| − |mutable ∩ underlying| − |underlying ∩ deletions|
+        #
+        # The intersections run at C speed and iterates over the smaller
+        # operand, so this is much cheaper than a Python loop over every key.
+        underlying_keys = self._underlying_map.keys()
+        return (
+            len(self._underlying_map)
+            + len(self._mutable_map)
+            - len(self._mutable_map.keys() & underlying_keys)
+            - len(underlying_keys & self._deletions)
+        )
 
-        for key in self._mutable_map:
-            # `key` should not be in both _mutable_map and _deletions
-            assert key not in self._deletions
+    def total_entries(self) -> int:
+        """The number of entries held across the underlying map, the
+        overrides and the deletions, following nested overlays down.
 
-            if key not in self._underlying_map:
-                count += 1
-
-        return count
+        Useful for estimating the memory usage of the overlay mapping.
+        """
+        underlying = self._underlying_map
+        if isinstance(underlying, MutableOverlayMapping):
+            underlying_size = underlying.total_entries()
+        else:
+            underlying_size = len(underlying)
+        return underlying_size + len(self._mutable_map) + len(self._deletions)
 
     def clear(self) -> None:
         self._underlying_map = {}

synapse/util/caches/expiringcache.py

@@ -24,8 +24,8 @@ from collections import OrderedDict
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Generic,
-    Iterable,
     Literal,
     TypeVar,
     overload,
@@ -66,6 +66,7 @@ class ExpiringCache(Generic[KT, VT]):
         expiry_ms: int = 0,
         reset_expiry_on_get: bool = False,
         iterable: bool = False,
+        size_callback: Callable[[VT], int] | None = None,
     ):
         """
         Args:
@@ -83,7 +84,18 @@ class ExpiringCache(Generic[KT, VT]):
                 an item on access. Defaults to False.
             iterable: If true, the size is calculated by summing the
                 sizes of all entries, rather than the number of entries.
+                Shorthand for `size_callback=len`.
+            size_callback: If given, the size of the cache is the sum of this
+                function over all values, rather than the number of entries.
         """
+        if iterable and size_callback is not None:
+            raise ValueError("`iterable` and `size_callback` are exclusive")
+        if iterable:
+            # type-ignore: if `iterable` is true, then the value type VT should
+            # be Sized (i.e. have a `__len__` method). We don't enforce this via
+            # the type system at present.
+            size_callback = len  # type: ignore[assignment]
+
         self._cache_name = cache_name
 
         self._original_max_size = max_len
@@ -97,7 +109,7 @@ class ExpiringCache(Generic[KT, VT]):
 
         self._cache: OrderedDict[KT, _CacheEntry[VT]] = OrderedDict()
 
-        self.iterable = iterable
+        self._size_callback = size_callback
 
         self.metrics = register_cache(
             cache_type="expiring",
@@ -124,13 +136,7 @@ class ExpiringCache(Generic[KT, VT]):
         # Evict if there are now too many items
         while self._max_size and len(self) > self._max_size:
             _key, value = self._cache.popitem(last=False)
-            if self.iterable:
-                # type-ignore, here and below: if self.iterable is true, then the value
-                # type VT should be Sized (i.e. have a __len__ method). We don't enforce
-                # this via the type system at present.
-                self.metrics.inc_evictions(EvictionReason.size, len(value.value))  # type: ignore[arg-type]
-            else:
-                self.metrics.inc_evictions(EvictionReason.size)
+            self.metrics.inc_evictions(EvictionReason.size, self._size_of(value.value))
 
     def __getitem__(self, key: KT) -> VT:
         try:
@@ -161,10 +167,9 @@ class ExpiringCache(Generic[KT, VT]):
                 raise KeyError(key)
             return default
 
-        if self.iterable:
-            self.metrics.inc_evictions(EvictionReason.invalidation, len(value.value))
-        else:
-            self.metrics.inc_evictions(EvictionReason.invalidation)
+        self.metrics.inc_evictions(
+            EvictionReason.invalidation, self._size_of(value.value)
+        )
 
         return value.value
 
@@ -207,10 +212,7 @@ class ExpiringCache(Generic[KT, VT]):
 
         for k in keys_to_delete:
             value = self._cache.pop(k)
-            if self.iterable:
-                self.metrics.inc_evictions(EvictionReason.time, len(value.value))  # type: ignore[arg-type]
-            else:
-                self.metrics.inc_evictions(EvictionReason.time)
+            self.metrics.inc_evictions(EvictionReason.time, self._size_of(value.value))
 
         logger.debug(
             "[%s] _prune_cache before: %d, after len: %d",
@@ -219,12 +221,16 @@ class ExpiringCache(Generic[KT, VT]):
             len(self),
         )
 
+    def _size_of(self, value: VT) -> int:
+        """How much a value counts towards `max_len`."""
+        if self._size_callback is None:
+            return 1
+        return self._size_callback(value)
+
     def __len__(self) -> int:
-        if self.iterable:
-            g: Iterable[int] = (len(entry.value) for entry in self._cache.values())  # type: ignore[arg-type]
-            return sum(g)
-        else:
+        if self._size_callback is None:
             return len(self._cache)
+        return sum(self._size_callback(entry.value) for entry in self._cache.values())
 
     def set_cache_factor(self, factor: float) -> bool:
         """

tests/state/test_v2.py

@@ -23,18 +23,21 @@ from typing import (
     Collection,
     Iterable,
     Mapping,
+    Sequence,
     TypeVar,
 )
 
 import attr
+from parameterized import parameterized
 
 from twisted.internet import defer
 
 from synapse.api.constants import EventTypes, JoinRules, Membership
-from synapse.api.room_versions import RoomVersions
+from synapse.api.room_versions import RoomVersion, RoomVersions
 from synapse.event_auth import auth_types_for_event
 from synapse.events import EventBase
 from synapse.state.v2 import (
+    ConflictCache,
     _get_auth_chain_difference,
     _get_power_level_for_sender,
     lexicographical_topological_sort,
@@ -181,6 +184,59 @@ INITIAL_EVENTS = [
 
 INITIAL_EDGES = ["START", "IMZ", "IMC", "IMB", "IJR", "IPOWER", "IMA", "CREATE"]
 
+ZARA_KEY = (EventTypes.Member, ZARA)
+TOPIC_KEY = (EventTypes.Topic, "")
+
+
+def _member(node_id: str, sender: str, state_key: str, content: dict) -> FakeEvent:
+    return FakeEvent(
+        id=node_id,
+        sender=sender,
+        type=EventTypes.Member,
+        state_key=state_key,
+        content=content,
+    )
+
+
+# Events for the conflict cache tests. All branch off START.
+#
+# PA is a power levels event that gives Bob PL 50. T1 and T2 are topic changes
+# by Bob. Bob has no power under IPOWER, so they only pass auth if PA is in the
+# state.
+#
+# ZJ1 and ZJ2 are Zara re-joining on two branches, and INV1 and INV2 are
+# invites she sends to Evelyn on each branch. The invites pull the joins into
+# the auth chain difference (see `test_conflict_cache_key_repartitioned`).
+CACHE_TEST_CASE_EVENTS = [
+    FakeEvent(
+        id="PA",
+        sender=ALICE,
+        type=EventTypes.PowerLevels,
+        state_key="",
+        content={"users": {ALICE: 100, BOB: 50}},
+    ),
+    FakeEvent(id="T1", sender=BOB, type=EventTypes.Topic, state_key="", content={}),
+    FakeEvent(id="T2", sender=BOB, type=EventTypes.Topic, state_key="", content={}),
+    _member("ZJ1", ZARA, ZARA, MEMBERSHIP_CONTENT_JOIN),
+    _member("ZJ2", ZARA, ZARA, MEMBERSHIP_CONTENT_JOIN),
+    _member("INV1", ZARA, EVELYN, {"membership": Membership.INVITE}),
+    _member("INV2", ZARA, EVELYN, {"membership": Membership.INVITE}),
+]
+
+CACHE_TEST_CASE_EDGES = [
+    ["PA", "START"],
+    ["T1", "START"],
+    ["T2", "START"],
+    ["INV1", "ZJ1", "START"],
+    ["INV2", "ZJ2", "START"],
+]
+
+# Room versions that use v2 and v2.1 state resolution respectively. Both have
+# the same auth rules. The difference is that v2.1 starts the iterative auth
+# checks from empty state rather than from the unconflicted state.
+V2_ROOM = RoomVersions.V11
+V21_ROOM = RoomVersions.HydraV11
+
 
 class StateTestCase(unittest.TestCase):
     def test_ban_vs_pl(self) -> None:
@@ -453,21 +509,137 @@ class StateTestCase(unittest.TestCase):
 
         self.do_check(events, edges, expected_state_ids)
 
-    def do_check(
+    # Helpers for the conflict cache tests. These use a plain dict as the
+    # cache, so `len(conflict_cache)` after a call tells us whether it was a
+    # hit or a miss.
+
+    def _build_cache_scenario(self) -> None:
+        self.event_map, self.state_at_event = self.build_event_graph(
+            CACHE_TEST_CASE_EVENTS, CACHE_TEST_CASE_EDGES
+        )
+
+    def _state(self, *node_ids: str) -> StateMap[str]:
+        """The state at START with the given events applied on top."""
+        state = dict(self.state_at_event["START"])
+        for node_id in node_ids:
+            event = self.event_map[EventID(node_id, "example.com").to_string()]
+            state[(event.type, event.state_key)] = event.event_id
+        return state
+
+    def _resolve_with_cache(
+        self,
+        room_version: RoomVersion,
+        state_sets: Sequence[StateMap[str]],
+        conflict_cache: ConflictCache,
+    ) -> StateMap[str]:
+        return self.successResultOf(
+            defer.ensureDeferred(
+                resolve_events_with_store(
+                    FakeClock(),
+                    ROOM_ID,
+                    room_version,
+                    state_sets,
+                    event_map=None,
+                    state_res_store=TestStateResolutionStore(self.event_map),
+                    conflict_cache=conflict_cache,
+                )
+            )
+        )
+
+    @parameterized.expand((V2_ROOM, V21_ROOM))
+    def test_conflict_cache_shared_across_unconflicted_state(
+        self, room_version: RoomVersion
+    ) -> None:
+        """Test that two resolutions with the same conflicted set but different
+        unconflicted state share a cache entry, and that each result still
+        includes its own unconflicted state.
+
+        The unconflicted state differs on Zara's membership. The topic events
+        aren't authed against that, so under v2 it isn't part of the cache key.
+        """
+        self._build_cache_scenario()
+        conflict_cache: dict[bytes, StateMap[str]] = {}
+
+        with_zara = [self._state("PA", "T1"), self._state("PA", "T2")]
+        without_zara = [
+            {key: value for key, value in state.items() if key != ZARA_KEY}
+            for state in with_zara
+        ]
+
+        first = self._resolve_with_cache(room_version, with_zara, conflict_cache)
+        second = self._resolve_with_cache(room_version, without_zara, conflict_cache)
+        self.assertEqual(len(conflict_cache), 1, "expected a cache hit")
+
+        self.assertIn(ZARA_KEY, first)
+        self.assertNotIn(ZARA_KEY, second)
+
+        # Everything else agrees.
+        self.assertEqual({k: v for k, v in first.items() if k != ZARA_KEY}, second)
+
+    def test_conflict_cache_keys_on_base_state(self) -> None:
+        """Test that under v2 the cache key includes the unconflicted state the
+        auth checks depend on. Changing the power levels is a cache miss and
+        gives a different result."""
+        self._build_cache_scenario()
+
+        powerless = [self._state("T1"), self._state("T2")]
+        powered = [self._state("PA", "T1"), self._state("PA", "T2")]
+
+        conflict_cache: dict[bytes, StateMap[str]] = {}
+        under_ipower = self._resolve_with_cache(V2_ROOM, powerless, conflict_cache)
+        under_pa = self._resolve_with_cache(V2_ROOM, powered, conflict_cache)
+        self.assertEqual(len(conflict_cache), 2, "expected a cache miss")
+
+        # Bob's topics fail auth under IPOWER and pass under PA.
+        self.assertNotIn(TOPIC_KEY, under_ipower)
+        self.assertIn(TOPIC_KEY, under_pa)
+
+    def test_conflict_cache_key_repartitioned(self) -> None:
+        """Test that a cached result is correct when the same conflicted set is
+        split differently between conflicted and unconflicted keys.
+
+        In the first call Zara's membership is unconflicted, but ZJ1 and ZJ2 are
+        in the auth chain difference (via the invites) and so are in the
+        conflicted set. In the second call Zara's membership is itself
+        conflicted. Both calls have the same cache key, so the cached result
+        must include the resolved Zara membership, even though the first call
+        overrides it with its unconflicted state.
+        """
+        self._build_cache_scenario()
+
+        agreed = [self._state("INV1"), self._state("INV2")]
+        conflicting = [self._state("ZJ1", "INV1"), self._state("ZJ2", "INV2")]
+
+        conflict_cache: dict[bytes, StateMap[str]] = {}
+        agreed_result = self._resolve_with_cache(V21_ROOM, agreed, conflict_cache)
+        warm = self._resolve_with_cache(V21_ROOM, conflicting, conflict_cache)
+        self.assertEqual(len(conflict_cache), 1, "expected a cache hit")
+
+        # The first call's unconflicted state takes precedence over the cached
+        # resolution.
+        self.assertEqual(agreed_result[ZARA_KEY], self._state()[ZARA_KEY])
+
+        # The second call gets the winner from the cached resolution, the same
+        # one it computes from cold.
+        conflict_cache.clear()
+        cold = self._resolve_with_cache(V21_ROOM, conflicting, conflict_cache)
+        self.assertNotEqual(warm[ZARA_KEY], agreed_result[ZARA_KEY])
+        self.assertEqual(warm, cold)
+
+    def build_event_graph(
         self,
         events: list[FakeEvent],
         edges: list[list[str]],
-        expected_state_ids: list[str],
-    ) -> None:
-        """Take a list of events and edges and calculate the state of the
-        graph at END, and asserts it matches `expected_state_ids`
+    ) -> tuple[dict[str, EventBase], dict[str, StateMap[str]]]:
+        """Build the graph of `INITIAL_EVENTS` plus `events`.
 
         Args:
             events
             edges: A list of chains of event edges, e.g.
                 `[[A, B, C]]` are edges A->B and B->C.
-            expected_state_ids: The expected state at END, (excluding
-                the keys that haven't changed since START).
+
+        Returns:
+            The events by event ID, and the state after each node ID.
         """
         # We want to sort the events into topological order for processing.
         graph: dict[str, set[str]] = {}
@@ -539,6 +711,26 @@ class StateTestCase(unittest.TestCase):
             state_at_event[node_id] = state_after
             event_map[event_id] = event
 
+        return event_map, state_at_event
+
+    def do_check(
+        self,
+        events: list[FakeEvent],
+        edges: list[list[str]],
+        expected_state_ids: list[str],
+    ) -> None:
+        """Take a list of events and edges and calculate the state of the
+        graph at END, and asserts it matches `expected_state_ids`
+
+        Args:
+            events
+            edges: A list of chains of event edges, e.g.
+                `[[A, B, C]]` are edges A->B and B->C.
+            expected_state_ids: The expected state at END, (excluding
+                the keys that haven't changed since START).
+        """
+        event_map, state_at_event = self.build_event_graph(events, edges)
+
         expected_state = {}
         for node_id in expected_state_ids:
             # expected_state_ids are node IDs rather than event IDs,

tests/test_state.py

@@ -18,6 +18,7 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
+from collections import ChainMap
 from typing import (
     Any,
     Collection,
@@ -34,9 +35,15 @@ from synapse.api.constants import EventTypes, Membership
 from synapse.api.room_versions import RoomVersions
 from synapse.events import EventBase
 from synapse.events.snapshot import EventContext
-from synapse.state import StateHandler, StateResolutionHandler, _make_state_cache_entry
+from synapse.state import (
+    StateHandler,
+    StateResolutionHandler,
+    _make_state_cache_entry,
+    _state_map_size,
+)
 from synapse.types import JsonDict, MutableStateMap, StateMap
 from synapse.types.state import StateFilter
+from synapse.util import MutableOverlayMapping
 from synapse.util.macaroons import MacaroonGenerator
 
 from tests import unittest
@@ -864,6 +871,31 @@ class StateTestCase(unittest.TestCase):
         result = yield defer.ensureDeferred(self.state.compute_event_context(event))
         return result
 
+    def test_state_map_size(self) -> None:
+        "Cache sizing counts every held entry, not the distinct keys"
+
+        base: StateMap[str] = {("a", ""): "A", ("b", ""): "B", ("c", ""): "C"}
+        self.assertEqual(_state_map_size(base), 3)
+
+        # Overriding and deleting keys leaves `len()` alone but holds entries.
+        overlay = MutableOverlayMapping(base)
+        overlay[("a", "")] = "A2"
+        overlay[("d", "")] = "D"
+        del overlay[("b", "")]
+        self.assertEqual(len(overlay), 3)
+        self.assertEqual(_state_map_size(overlay), 3 + 2 + 1)
+
+        # Nested overlays are followed down.
+        outer = MutableOverlayMapping(overlay)
+        outer[("e", "")] = "E"
+        self.assertEqual(len(outer), 4)
+        self.assertEqual(_state_map_size(outer), 6 + 1)
+
+        # A `ChainMap` holds every layer in full, however much they overlap.
+        chain = ChainMap({("a", ""): "A3", ("f", ""): "F"}, outer)
+        self.assertEqual(len(chain), 5)
+        self.assertEqual(_state_map_size(chain), 2 + 7)
+
     def test_make_state_cache_entry(self) -> None:
         "Test that calculating a prev_group and delta is correct"
 

tests/util/test_expiring_cache.py

@@ -87,6 +87,49 @@ class ExpiringCacheTestCase(unittest.HomeserverTestCase):
         self.assertEqual(cache.get("key3"), [4, 5])
         self.assertEqual(cache.get("key4"), [6, 7])
 
+    def test_size_callback_eviction(self) -> None:
+        reactor, clock = get_clock()
+        cache: ExpiringCache[str, list[int]] = ExpiringCache(
+            cache_name="test",
+            server_name="testserver",
+            hs=self.hs,
+            clock=clock,
+            max_len=5,
+            # Size each value by its sum, so that `len()` and the reported
+            # size disagree.
+            size_callback=sum,
+        )
+
+        cache["key"] = [1]
+        cache["key2"] = [2]
+        cache["key3"] = [2]
+
+        self.assertEqual(cache.get("key"), [1])
+        self.assertEqual(cache.get("key2"), [2])
+        self.assertEqual(cache.get("key3"), [2])
+        self.assertEqual(len(cache), 5)
+
+        # Three entries of `len()` 1 each, but a size of 6, so the oldest two
+        # go.
+        cache["key4"] = [3]
+        self.assertEqual(cache.get("key"), None)
+        self.assertEqual(cache.get("key2"), None)
+        self.assertEqual(cache.get("key3"), [2])
+        self.assertEqual(cache.get("key4"), [3])
+        self.assertEqual(len(cache), 5)
+
+    def test_iterable_and_size_callback_are_exclusive(self) -> None:
+        reactor, clock = get_clock()
+        with self.assertRaises(ValueError):
+            ExpiringCache(
+                cache_name="test",
+                server_name="testserver",
+                hs=self.hs,
+                clock=clock,
+                iterable=True,
+                size_callback=len,
+            )
+
     def test_time_eviction(self) -> None:
         reactor, clock = get_clock()
         cache: ExpiringCache[str, int] = ExpiringCache(

tests/util/test_mutable_overlay_mapping.py

@@ -187,3 +187,76 @@ class TestMutableOverlayMapping(unittest.TestCase):
         self.assertNotIn("c", mapping)
         self.assertIn("d", mapping)
         self.assertNotIn("e", mapping)
+
+    def test_len_after_reset_and_redelete(self) -> None:
+        """len() must follow keys that move between the underlying map, the
+        overrides and the deletions."""
+        underlying = {"a": 1, "b": 2, "c": 3}
+        mapping = MutableOverlayMapping(underlying)
+
+        # Delete an underlying key, then set it again.
+        del mapping["a"]
+        self.assertEqual(len(mapping), 2)
+        mapping["a"] = 10
+        self.assertEqual(len(mapping), 3)
+
+        # Add a key only the overlay knows about, delete it, and add it back.
+        mapping["d"] = 4
+        self.assertEqual(len(mapping), 4)
+        del mapping["d"]
+        self.assertEqual(len(mapping), 3)
+        mapping["d"] = 40
+        self.assertEqual(len(mapping), 4)
+
+        # Override then delete an underlying key.
+        mapping["b"] = 20
+        del mapping["b"]
+        self.assertEqual(len(mapping), 3)
+
+        self.assertEqual(len(mapping), len(dict(mapping)))
+
+    def test_len_nested(self) -> None:
+        """An overlay over an overlay reports the right length."""
+        inner = MutableOverlayMapping({"a": 1, "b": 2})
+        inner["c"] = 3
+        del inner["a"]
+
+        outer = MutableOverlayMapping(inner)
+        self.assertEqual(len(outer), 2)
+
+        outer["a"] = 10  # deleted in inner, so new to outer
+        outer["b"] = 20  # present in inner's underlying map
+        outer["c"] = 30  # present in inner's overrides
+        outer["d"] = 40  # new
+        self.assertEqual(len(outer), 4)
+
+        del outer["b"]
+        del outer["d"]
+        self.assertEqual(len(outer), 2)
+        self.assertEqual(len(outer), len(dict(outer)))
+
+    def test_total_entries(self) -> None:
+        """total_entries() counts the base map plus every override and
+        deletion, unlike len()."""
+        mapping = MutableOverlayMapping({"a": 1, "b": 2, "c": 3})
+        self.assertEqual(mapping.total_entries(), 3)
+
+        mapping["a"] = 10  # override: +1 entry, same length
+        mapping["d"] = 4  # new key: +1 entry, +1 length
+        del mapping["b"]  # deletion: +1 entry, -1 length
+        self.assertEqual(len(mapping), 3)
+        self.assertEqual(mapping.total_entries(), 6)
+
+        # Deleting an override drops it from the overrides and records the
+        # deletion, so the entry count is unchanged.
+        del mapping["d"]
+        self.assertEqual(len(mapping), 2)
+        self.assertEqual(mapping.total_entries(), 6)
+
+        # Nested overlays are counted all the way down.
+        outer = MutableOverlayMapping(mapping)
+        outer["e"] = 5
+        self.assertEqual(outer.total_entries(), 7)
+
+        mapping.clear()
+        self.assertEqual(mapping.total_entries(), 0)