From 9a2223d4c8c2a46f7ab072597bdba2614bc3e6ea Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 10 Dec 2019 11:22:12 +0000
Subject: [PATCH 1/3] Fix make_deferred_yieldable to work with coroutines

---
 synapse/logging/context.py    |  9 ++++++++-
 tests/util/test_logcontext.py | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/synapse/logging/context.py b/synapse/logging/context.py
index 2c1fb9dda..9f484ce59 100644
--- a/synapse/logging/context.py
+++ b/synapse/logging/context.py
@@ -23,6 +23,7 @@ them.
 See doc/log_contexts.rst for details on how this works.
 """
 
+import inspect
 import logging
 import threading
 import types
@@ -612,7 +613,8 @@ def run_in_background(f, *args, **kwargs):
 
 
 def make_deferred_yieldable(deferred):
-    """Given a deferred, make it follow the Synapse logcontext rules:
+    """Given a deferred (or coroutine), make it follow the Synapse logcontext
+    rules:
 
     If the deferred has completed (or is not actually a Deferred), essentially
     does nothing (just returns another completed deferred with the
@@ -624,6 +626,11 @@ def make_deferred_yieldable(deferred):
 
     (This is more-or-less the opposite operation to run_in_background.)
     """
+    if inspect.isawaitable(deferred):
+        # If we're given a coroutine we need to convert it to a deferred so that
+        # we can attach callbacks (and not immediately return).
+        deferred = defer.ensureDeferred(deferred)
+
     if not isinstance(deferred, defer.Deferred):
         return deferred
 
diff --git a/tests/util/test_logcontext.py b/tests/util/test_logcontext.py
index 8b8455c8b..281b32c4b 100644
--- a/tests/util/test_logcontext.py
+++ b/tests/util/test_logcontext.py
@@ -179,6 +179,30 @@ class LoggingContextTestCase(unittest.TestCase):
             nested_context = nested_logging_context(suffix="bar")
             self.assertEqual(nested_context.request, "foo-bar")
 
+    @defer.inlineCallbacks
+    def test_make_deferred_yieldable_with_await(self):
+        # an async function which retuns an incomplete coroutine, but doesn't
+        # follow the synapse rules.
+
+        async def blocking_function():
+            d = defer.Deferred()
+            reactor.callLater(0, d.callback, None)
+            await d
+
+        sentinel_context = LoggingContext.current_context()
+
+        with LoggingContext() as context_one:
+            context_one.request = "one"
+
+            d1 = make_deferred_yieldable(blocking_function())
+            # make sure that the context was reset by make_deferred_yieldable
+            self.assertIs(LoggingContext.current_context(), sentinel_context)
+
+            yield d1
+
+            # now it should be restored
+            self._check_test_key("one")
+
 
 # a function which returns a deferred which has been "called", but
 # which had a function which returned another incomplete deferred on

From f5bb1531b7307bc2d826789746e7c82fa4dbf36c Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 10 Dec 2019 11:23:52 +0000
Subject: [PATCH 2/3] Newsfile

---
 changelog.d/6505.misc | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/6505.misc

diff --git a/changelog.d/6505.misc b/changelog.d/6505.misc
new file mode 100644
index 000000000..3a75b2d9d
--- /dev/null
+++ b/changelog.d/6505.misc
@@ -0,0 +1 @@
+Make `make_deferred_yieldable` to work with async/await.

From ffeafade4879a1135ab6795d4ba0c11131f82f01 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 10 Dec 2019 13:17:39 +0000
Subject: [PATCH 3/3] Update comment

---
 synapse/logging/context.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/synapse/logging/context.py b/synapse/logging/context.py
index 9f484ce59..6747f29e6 100644
--- a/synapse/logging/context.py
+++ b/synapse/logging/context.py
@@ -627,8 +627,10 @@ def make_deferred_yieldable(deferred):
     (This is more-or-less the opposite operation to run_in_background.)
     """
     if inspect.isawaitable(deferred):
-        # If we're given a coroutine we need to convert it to a deferred so that
-        # we can attach callbacks (and not immediately return).
+        # If we're given a coroutine we convert it to a deferred so that we
+        # run it and find out if it immediately finishes, it it does then we
+        # don't need to fiddle with log contexts at all and can return
+        # immediately.
         deferred = defer.ensureDeferred(deferred)
 
     if not isinstance(deferred, defer.Deferred):