Fix logcontext leak in HttpTransactionCache

ONE DAY I WILL PURGE THE WORLD OF THIS EVIL

synapse/rest/client/transactions.py

@@ -19,6 +19,7 @@ import logging
 
 from synapse.api.auth import get_access_token_from_request
 from synapse.util.async import ObservableDeferred
+from synapse.util.logcontext import make_deferred_yieldable, run_in_background
 
 logger = logging.getLogger(__name__)
 
@@ -80,31 +81,26 @@ class HttpTransactionCache(object):
         Returns:
             Deferred which resolves to a tuple of (response_code, response_dict).
         """
-        try:
-            return self.transactions[txn_key][0].observe()
-        except (KeyError, IndexError):
-            pass  # execute the function instead.
-
-        deferred = fn(*args, **kwargs)
-
-        observable = ObservableDeferred(deferred, consumeErrors=False)
-        self.transactions[txn_key] = (observable, self.clock.time_msec())
-
-        # if the request fails with an exception, remove it from the
-        # transaction map. This is done to ensure that we don't cache
-        # transient errors like rate-limiting errors, etc.
-        #
-        # (make sure we add this errback *after* adding the key above, in case
-        # the deferred has already failed and is running errbacks
-        # synchronously)
-        def remove_from_map(err):
-            self.transactions.pop(txn_key, None)
-            # we deliberately do not propagate the error any further, as we
-            # expect the observers to have reported it.
-
-        deferred.addErrback(remove_from_map)
-
-        return observable.observe()
+        if txn_key in self.transactions:
+            observable = self.transactions[txn_key][0]
+        else:
+            # execute the function instead.
+            deferred = run_in_background(fn, *args, **kwargs)
+
+            observable = ObservableDeferred(deferred)
+            self.transactions[txn_key] = (observable, self.clock.time_msec())
+
+            # if the request fails with an exception, remove it
+            # from the transaction map. This is done to ensure that we don't
+            # cache transient errors like rate-limiting errors, etc.
+            def remove_from_map(err):
+                self.transactions.pop(txn_key, None)
+                # we deliberately do not propagate the error any further, as we
+                # expect the observers to have reported it.
+
+            deferred.addErrback(remove_from_map)
+
+        return make_deferred_yieldable(observable.observe())
 
     def _cleanup(self):
         now = self.clock.time_msec()

tests/rest/client/test_transactions.py

@@ -3,6 +3,7 @@ from synapse.rest.client.transactions import CLEANUP_PERIOD_MS
 from twisted.internet import defer
 from mock import Mock, call
 
+from synapse.util import async
 from synapse.util.logcontext import LoggingContext
 from tests import unittest
 from tests.utils import MockClock
@@ -41,6 +42,26 @@ class HttpTransactionCacheTestCase(unittest.TestCase):
         # expect only a single call to do the work
         cb.assert_called_once_with("some_arg", keyword="arg", changing_args=0)
 
+    @defer.inlineCallbacks
+    def test_logcontexts_with_async_result(self):
+        @defer.inlineCallbacks
+        def cb():
+            yield async.sleep(0)
+            defer.returnValue("yay")
+
+        @defer.inlineCallbacks
+        def test():
+            with LoggingContext("c") as c1:
+                res = yield self.cache.fetch_or_execute(self.mock_key, cb)
+                self.assertIs(LoggingContext.current_context(), c1)
+                self.assertEqual(res, "yay")
+
+        # run the test twice in parallel
+        d = defer.gatherResults([test(), test()])
+        self.assertIs(LoggingContext.current_context(), LoggingContext.sentinel)
+        yield d
+        self.assertIs(LoggingContext.current_context(), LoggingContext.sentinel)
+
     @defer.inlineCallbacks
     def test_does_not_cache_exceptions(self):
         """Checks that, if the callback throws an exception, it is called again