change strategy to ResponseCache.wrap_conditional()

synapse/handlers/sync.py

@@ -40,7 +40,7 @@
 from synapse.util.async_helpers import concurrently_execute
 from synapse.util.caches.expiringcache import ExpiringCache
 from synapse.util.caches.lrucache import LruCache
-from synapse.util.caches.response_cache import NoTimedCache, ResponseCache
+from synapse.util.caches.response_cache import ResponseCache
 from synapse.util.metrics import Measure, measure_func
 from synapse.visibility import filter_events_for_client
 
@@ -277,8 +277,9 @@ async def wait_for_sync_for_user(
         user_id = sync_config.user.to_string()
         await self.auth.check_auth_blocking(requester=requester)
 
-        res = await self.response_cache.wrap(
+        res = await self.response_cache.wrap_conditional(
             sync_config.request_key,
+            lambda result: since_token != result.next_batch,
             self._wait_for_sync_for_user,
             sync_config,
             since_token,
@@ -331,11 +332,6 @@ def current_sync_callback(before_token, after_token):
                 lazy_loaded = "false"
             non_empty_sync_counter.labels(sync_type, lazy_loaded).inc()
 
-        # fixme hack: sanity check to invalidate cache, prevent a self-referral loop
-        #  replace with contextvars approach once it gets working on twisted (twisted/twisted#1262)
-        if since_token == result.next_batch:
-            result = NoTimedCache(result)  # type: ignore # because it is unwrapped in ResponseCache.set.<remove>()
-
         return result
 
     async def current_sync_for_user(

synapse/util/caches/response_cache.py

@@ -13,9 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, TypeVar
-
-import attr
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Set, TypeVar
 
 from twisted.internet import defer
 
@@ -31,11 +29,6 @@
 T = TypeVar("T")
 
 
-@attr.s(slots=True, frozen=True)
-class NoTimedCache:
-    obj = attr.ib()
-
-
 class ResponseCache(Generic[T]):
     """
     This caches a deferred response. Until the deferred completes it will be
@@ -47,6 +40,7 @@ class ResponseCache(Generic[T]):
     def __init__(self, hs: "HomeServer", name: str, timeout_ms: float = 0):
         # Requests that haven't finished yet.
         self.pending_result_cache = {}  # type: Dict[T, ObservableDeferred]
+        self.pending_conditionals = {}  # type: Dict[T, Set[Callable[[Any], bool]]]
 
         self.clock = hs.get_clock()
         self.timeout_sec = timeout_ms / 1000.0
@@ -108,17 +102,45 @@ def set(self, key: T, deferred: defer.Deferred) -> defer.Deferred:
         self.pending_result_cache[key] = result
 
         def remove(r):
-            if self.timeout_sec and not isinstance(r, NoTimedCache):
+            should_cache = all(
+                func(r) for func in self.pending_conditionals.pop(key, [])
+            )
+
+            if self.timeout_sec and should_cache:
                 self.clock.call_later(
                     self.timeout_sec, self.pending_result_cache.pop, key, None
                 )
             else:
                 self.pending_result_cache.pop(key, None)
-            return r.obj if isinstance(r, NoTimedCache) else r
+            return r
 
         result.addBoth(remove)
         return result.observe()
 
+    def add_conditional(self, key: T, conditional: Callable[[Any], bool]):
+        self.pending_conditionals.setdefault(key, set()).add(conditional)
+
+    def wrap_conditional(
+        self,
+        key: T,
+        should_cache: Callable[[Any], bool],
+        callback: "Callable[..., Any]",
+        *args: Any,
+        **kwargs: Any
+    ) -> defer.Deferred:
+        """The same as wrap(), but adds a conditional to the final execution.
+
+        When the final execution completes, *all* conditionals need to return True for it to properly cache,
+        else it'll not be cached in a timed fashion."""
+
+        # See if there's already a result on this key that hasn't yet completed, due to the single-threaded nature of
+        # python, adding a key immediately in the same execution thread will not cause a race condition.
+        result = self.get(key)
+        if not result or isinstance(result, defer.Deferred) and not result.called:
+            self.add_conditional(key, should_cache)
+
+        return self.wrap(key, callback, *args, **kwargs)
+
     def wrap(
         self, key: T, callback: "Callable[..., Any]", *args: Any, **kwargs: Any
     ) -> defer.Deferred: