wip2

2 files changed, 96 insertions, 65 deletions

diff --git a/synapse/util/caches/dual_lookup_cache.py b/synapse/util/caches/dual_lookup_cache.py
index 7529a5b8f9..6ee72de705 100644
--- a/synapse/util/caches/dual_lookup_cache.py
+++ b/synapse/util/caches/dual_lookup_cache.py
@@ -17,12 +17,14 @@ from typing import (
     Dict,
     Generic,
     ItemsView,
+    List,
     Optional,
-    Set,
     TypeVar,
+    Union,
     ValuesView,
 )
 
+# Used to discern between a value not existing in a map, or the value being 'None'.
 SENTINEL = object()
 
 # The type of the primary dict's keys.
@@ -35,6 +37,13 @@ SKT = TypeVar("SKT")
 logger = logging.getLogger(__name__)
 
 
+class SecondarySet(set):
+    """
+    Used to differentiate between an entry in the secondary_dict, and a set stored
+    in the primary_dict. This is necessary as pop() can return either.
+    """
+
+
 class DualLookupCache(Generic[PKT, PVT, SKT]):
     """
     A backing store for LruCache that supports multiple entry points.
@@ -79,7 +88,7 @@ class DualLookupCache(Generic[PKT, PVT, SKT]):
 
     def __init__(self, secondary_key_function: Callable[[PVT], SKT]) -> None:
         self._primary_dict: Dict[PKT, PVT] = {}
-        self._secondary_dict: Dict[SKT, Set[PKT]] = {}
+        self._secondary_dict: Dict[SKT, SecondarySet] = {}
         self._secondary_key_function = secondary_key_function
 
     def __setitem__(self, key: PKT, value: PVT) -> None:
@@ -108,7 +117,9 @@ class DualLookupCache(Generic[PKT, PVT, SKT]):
 
         # And create a mapping in the secondary_dict to a set containing the
         # primary_key, creating the set if necessary.
-        secondary_key_set = self._secondary_dict.setdefault(secondary_key, set())
+        secondary_key_set = self._secondary_dict.setdefault(
+            secondary_key, SecondarySet()
+        )
         secondary_key_set.add(key)
 
         logger.info("*** Insert into primary_dict: %s: %s", key, value)
@@ -138,69 +149,84 @@ class DualLookupCache(Generic[PKT, PVT, SKT]):
         self._primary_dict.clear()
         self._secondary_dict.clear()
 
-    def pop(self, key: PKT, default: Optional[PVT] = None) -> Optional[PVT]:
-        """Remove the given key, from the cache if it exists, and return the associated
-        value.
+    def pop(
+        self, key: Union[PKT, SKT], default: Optional[Union[Dict[PKT, PVT], PVT]] = None
+    ) -> Optional[Union[Dict[PKT, PVT], PVT]]:
+        """Remove an entry from either the primary_dict or secondary_dict.
 
-        Evicts an entry from both the primary_dict and secondary_dict.
+        The primary_dict is checked first for the key. If an entry is found, it is
+        removed from the primary_dict and returned.
+
+        If no entry in the primary_dict exists, then the secondary_dict is checked.
+        If an entry exists, all associated entries in the primary_dict will be
+        deleted, and all primary_dict keys returned from this function in a SecondarySet.
 
         Args:
-            key: The key to remove from the cache.
-            default: The value to return if the given key is not found.
+            key: A key to drop from either the primary_dict or secondary_dict.
+            default: The default value if the key does not exist in either dict.
 
         Returns:
-            The value associated with the given key if it is found. Otherwise, the value
-            of `default`.
+            Either a matched value from the primary_dict or the secondary_dict. If no
+            value is found for the key, then None.
         """
-        # Exit immediately if the key is not found
-        if key not in self._primary_dict:
-            return default
-
-        # Pop the entry from the primary_dict to retrieve the desired value
-        primary_value = self._primary_dict.pop(key)
-
-        logger.info("*** Popping from primary_dict: %s: %s", key, primary_value)
-
-        # Derive the secondary_key from the primary_value
-        secondary_key = self._secondary_key_function(primary_value)
-
-        # Pop the entry from the secondary_dict
-        secondary_key_set = self._secondary_dict[secondary_key]
-        if len(secondary_key_set) > 1:
-            # Delete just the set entry for the given key.
-            secondary_key_set.remove(key)
-            logger.info("*** Popping from secondary_dict: %s: %s", secondary_key, key)
-
-        else:
-            # Delete the entire soon-to-be-empty set referenced by the secondary_key.
-            del self._secondary_dict[secondary_key]
-            logger.info("*** Popping from secondary_dict: %s", secondary_key)
+        # Attempt to remove from the primary_dict first.
+        primary_value = self._primary_dict.pop(key, SENTINEL)
+        if primary_value is not SENTINEL:
+            # We found a value in the primary_dict. Remove it from the corresponding
+            # entry in the secondary_dict, and then return it.
+            logger.info(
+                "*** Popped entry from primary_dict: %s: %s", key, primary_value
+            )
 
-        return primary_value
+            # Derive the secondary_key from the primary_value
+            secondary_key = self._secondary_key_function(primary_value)
+
+            # Pop the entry from the secondary_dict
+            secondary_key_set = self._secondary_dict[secondary_key]
+            if len(secondary_key_set) > 1:
+                # Delete just the set entry for the given key.
+                secondary_key_set.remove(key)
+                logger.info(
+                    "*** Popping from secondary_dict: %s: %s", secondary_key, key
+                )
+            else:
+                # Delete the entire set referenced by the secondary_key, as it only
+                # has one entry.
+                del self._secondary_dict[secondary_key]
+                logger.info("*** Popping from secondary_dict: %s", secondary_key)
+
+            return primary_value
+
+        # There was no matching value in the primary_dict. Attempt the secondary_dict.
+        primary_key_set = self._secondary_dict.pop(key, SENTINEL)
+        if primary_key_set is not SENTINEL:
+            # We found a set in the secondary_dict.
+            logger.info(
+                "*** Found '%s' in secondary_dict: %s: ",
+                key,
+                primary_key_set,
+            )
 
-    def del_multi(self, secondary_key: SKT) -> None:
-        """Remove an entry from the secondary_dict, removing all associated entries
-        in the primary_dict as well.
+            popped_primary_dict_values: List[PVT] = []
 
-        Args:
-            secondary_key: A secondary_key to drop. May be associated with zero or more
-                primary keys. If any associated primary keys are found, they will be
-                dropped as well.
-        """
-        primary_key_set = self._secondary_dict.pop(secondary_key, None)
-        if not primary_key_set:
+            # We found an entry in the secondary_dict. Delete all related entries in the
+            # primary_dict.
             logger.info(
-                "*** Did not find '%s' in secondary_dict: %s",
-                secondary_key,
-                self._secondary_dict,
+                "*** Found key in secondary_dict to pop: %s. "
+                "Popping primary_dict entries",
+                key,
             )
-            return
+            for primary_key in primary_key_set:
+                logger.info("*** Popping entry from primary_dict: %s", primary_key)
+                logger.info("*** primary_dict: %s", self._primary_dict)
+                popped_primary_dict_values = self._primary_dict[primary_key]
+                del self._primary_dict[primary_key]
+
+            # Now return the unmodified copy of the set.
+            return popped_primary_dict_values
 
-        logger.info("*** Popping whole key from secondary_dict: %s", secondary_key)
-        for primary_key in primary_key_set:
-            logger.info("*** Popping entry from primary_dict: %s", primary_key)
-            logger.info("*** primary_dict: %s", self._primary_dict)
-            del self._primary_dict[primary_key]
+        # No match in either dict.
+        return default
 
     def values(self) -> ValuesView:
         return self._primary_dict.values()
diff --git a/synapse/util/caches/lrucache.py b/synapse/util/caches/lrucache.py
index 30765e630d..6345735680 100644
--- a/synapse/util/caches/lrucache.py
+++ b/synapse/util/caches/lrucache.py
@@ -46,9 +46,10 @@ from synapse.metrics.background_process_metrics import wrap_as_background_proces
 from synapse.metrics.jemalloc import get_jemalloc_stats
 from synapse.util import Clock, caches
 from synapse.util.caches import CacheMetric, EvictionReason, register_cache
-from synapse.util.caches.dual_lookup_cache import DualLookupCache
+from synapse.util.caches.dual_lookup_cache import DualLookupCache, SecondarySet
 from synapse.util.caches.treecache import (
     TreeCache,
+    TreeCacheNode,
     iterate_tree_cache_entry,
     iterate_tree_cache_items,
 )
@@ -751,21 +752,25 @@ class LruCache(Generic[KT, VT]):
             may be of lower cardinality than the TreeCache - in which case the whole
             subtree is deleted.
             """
-            if isinstance(cache, DualLookupCache):
-                # Make use of DualLookupCache's del_multi feature
-                cache.del_multi(key)
-                return
-
             # Remove an entry from the cache.
             # In the case of a 'dict' cache type, we're just removing an entry from the
             # dict. For a TreeCache, we're removing a subtree which has children.
-            popped_entry = cache.pop(key, None)
-            if popped_entry is not None and cache_type is TreeCache:
-                # We've popped a subtree - now we need to clean up each child node.
-                # For each deleted node, we remove it from the linked list and run
-                # its callbacks.
+            popped_entry: _Node[KT, VT] = cache.pop(key, None)
+            if popped_entry is None:
+                return
+
+            if isinstance(popped_entry, TreeCacheNode):
+                # We've popped a subtree from a TreeCache - now we need to clean up
+                # each child node.
                 for leaf in iterate_tree_cache_entry(popped_entry):
+                    # For each deleted child node, we remove it from the linked list and
+                    # run its callbacks.
+                    delete_node(leaf)
+            elif isinstance(popped_entry, SecondarySet):
+                for leaf in popped_entry:
                     delete_node(leaf)
+            else:
+                delete_node(popped_entry)
 
         @synchronized
         def cache_clear() -> None: