From 88aa56c9ed0820ac1e0487d17a031893d559c611 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 19 Jul 2022 13:57:12 +0100
Subject: [PATCH] Make `LruCacheget_multi` return something sane.

---
 synapse/util/caches/dictionary_cache.py |  8 ++++----
 synapse/util/caches/lrucache.py         | 27 +++++++++++++++++++------
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/synapse/util/caches/dictionary_cache.py b/synapse/util/caches/dictionary_cache.py
index 51d8d88eaa42..62bb6689505d 100644
--- a/synapse/util/caches/dictionary_cache.py
+++ b/synapse/util/caches/dictionary_cache.py
@@ -20,7 +20,7 @@
 from typing_extensions import Literal
 
 from synapse.util.caches.lrucache import LruCache
-from synapse.util.caches.treecache import TreeCache, iterate_tree_cache_items
+from synapse.util.caches.treecache import TreeCache
 
 logger = logging.getLogger(__name__)
 
@@ -228,9 +228,9 @@ def _get_full_dict(
         # and `_PerKeyValue` into the `DictionaryEntry`.
         values = {}
         known_absent = set()
-        for key_tuple, node in iterate_tree_cache_items((), all_entries):
-            dict_key = key_tuple[0]
-            dict_value = node.value
+        for cache_key, dict_value in all_entries:
+            # The key used for the `TreeCache` is `(key, dict_key)`
+            dict_key = cache_key[1]
 
             # We have explicitly looked for a full cache key, so we
             # shouldn't see one.
diff --git a/synapse/util/caches/lrucache.py b/synapse/util/caches/lrucache.py
index c26dce94747c..a3fc0940e968 100644
--- a/synapse/util/caches/lrucache.py
+++ b/synapse/util/caches/lrucache.py
@@ -24,9 +24,11 @@
     Callable,
     Collection,
     Dict,
+    Generator,
     Generic,
     List,
     Optional,
+    Tuple,
     Type,
     TypeVar,
     Union,
@@ -46,8 +48,8 @@
 from synapse.util.caches import CacheMetric, EvictionReason, register_cache
 from synapse.util.caches.treecache import (
     TreeCache,
-    TreeCacheNode,
     iterate_tree_cache_entry,
+    iterate_tree_cache_items,
 )
 from synapse.util.linked_list import ListNode
 
@@ -596,7 +598,7 @@ def cache_get_multi(
             key: tuple,
             default: Literal[None] = None,
             update_metrics: bool = True,
-        ) -> Union[None, TreeCacheNode]:
+        ) -> Union[None, Generator[Tuple[KT, VT], None, None]]:
             ...
 
         @overload
@@ -604,7 +606,7 @@ def cache_get_multi(
             key: tuple,
             default: T,
             update_metrics: bool = True,
-        ) -> Union[T, TreeCacheNode]:
+        ) -> Union[T, Generator[Tuple[KT, VT], None, None]]:
             ...
 
         @synchronized
@@ -612,8 +614,15 @@ def cache_get_multi(
             key: tuple,
             default: Optional[T] = None,
             update_metrics: bool = True,
-        ) -> Union[None, T, TreeCacheNode]:
-            """Used only for `TreeCache` to fetch a subtree."""
+        ) -> Union[None, T, Generator[Tuple[KT, VT], None, None]]:
+            """Returns a generator yielding all entries under the given key.
+
+            Can only be used if backed by a tree cache.
+
+            Returns:
+                Either default if the key doesn't exist, or a generator of the
+                key/value pairs.
+            """
 
             assert isinstance(cache, TreeCache)
 
@@ -621,7 +630,13 @@ def cache_get_multi(
             if node is not None:
                 if update_metrics and metrics:
                     metrics.inc_hits()
-                return node
+
+                # Iterating over the node will return values of type `_Node`,
+                # which we need to unwrap.
+                return (
+                    (full_key, lru_node.value)
+                    for full_key, lru_node in iterate_tree_cache_items(key, node)
+                )
             else:
                 if update_metrics and metrics:
                     metrics.inc_misses()