From 132279a46fd76de8f767bc6977192900c450fec9 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 27 Sep 2019 15:11:14 +0100
Subject: [PATCH 01/13] Patch inlinecallbacks for log contexts

---
 synapse/__init__.py             |  6 +++
 synapse/handlers/room_member.py |  4 +-
 synapse/storage/_base.py        |  7 +--
 synapse/storage/push_rule.py    |  2 +-
 tests/patch_inline_callbacks.py | 86 +++++++++++++++++++++++++++++++--
 5 files changed, 95 insertions(+), 10 deletions(-)

diff --git a/synapse/__init__.py b/synapse/__init__.py
index ddfe9ec542..4401fd52f0 100644
--- a/synapse/__init__.py
+++ b/synapse/__init__.py
@@ -17,8 +17,11 @@
 """ This is a reference implementation of a Matrix home server.
 """
 
+import os
 import sys
 
+from tests.patch_inline_callbacks import do_patch
+
 # Check that we're not running on an unsupported Python version.
 if sys.version_info < (3, 5):
     print("Synapse requires Python 3.5 or above.")
@@ -36,3 +39,6 @@ except ImportError:
     pass
 
 __version__ = "1.4.0rc1"
+
+if bool(os.environ.get("SYNAPSE_TEST_PATCH_LOG_CONTEXTS", False)):
+    do_patch()
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 8abdb1b6e6..19e44b5460 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -213,11 +213,11 @@ class RoomMemberHandler(object):
 
             if predecessor:
                 # It is an upgraded room. Copy over old tags
-                self.copy_room_tags_and_direct_to_room(
+                yield self.copy_room_tags_and_direct_to_room(
                     predecessor["room_id"], room_id, user_id
                 )
                 # Move over old push rules
-                self.store.move_push_rules_from_room_to_room_for_user(
+                yield self.store.move_push_rules_from_room_to_room_for_user(
                     predecessor["room_id"], room_id, user_id
                 )
         elif event.membership == Membership.LEAVE:
diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py
index abe16334ec..06cc14fcd1 100644
--- a/synapse/storage/_base.py
+++ b/synapse/storage/_base.py
@@ -30,7 +30,7 @@ from prometheus_client import Histogram
 from twisted.internet import defer
 
 from synapse.api.errors import StoreError
-from synapse.logging.context import LoggingContext, PreserveLoggingContext
+from synapse.logging.context import LoggingContext, make_deferred_yieldable
 from synapse.metrics.background_process_metrics import run_as_background_process
 from synapse.storage.engines import PostgresEngine, Sqlite3Engine
 from synapse.types import get_domain_from_id
@@ -550,8 +550,9 @@ class SQLBaseStore(object):
 
                 return func(conn, *args, **kwargs)
 
-        with PreserveLoggingContext():
-            result = yield self._db_pool.runWithConnection(inner_func, *args, **kwargs)
+        result = yield make_deferred_yieldable(
+            self._db_pool.runWithConnection(inner_func, *args, **kwargs)
+        )
 
         return result
 
diff --git a/synapse/storage/push_rule.py b/synapse/storage/push_rule.py
index a6517c4cf3..1f878e6575 100644
--- a/synapse/storage/push_rule.py
+++ b/synapse/storage/push_rule.py
@@ -235,7 +235,7 @@ class PushRulesWorkerStore(
                 (c.get("key") == "room_id" and c.get("pattern") == old_room_id)
                 for c in conditions
             ):
-                self.move_push_rule_from_room_to_room(new_room_id, user_id, rule)
+                yield self.move_push_rule_from_room_to_room(new_room_id, user_id, rule)
 
     @defer.inlineCallbacks
     def bulk_get_push_rules_for_room(self, event, context):
diff --git a/tests/patch_inline_callbacks.py b/tests/patch_inline_callbacks.py
index 220884311c..5ef0aff0c3 100644
--- a/tests/patch_inline_callbacks.py
+++ b/tests/patch_inline_callbacks.py
@@ -15,6 +15,7 @@
 
 from __future__ import print_function
 
+import inspect
 import functools
 import sys
 
@@ -33,17 +34,19 @@ def do_patch():
     orig_inline_callbacks = defer.inlineCallbacks
 
     def new_inline_callbacks(f):
-
-        orig = orig_inline_callbacks(f)
-
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
             start_context = LoggingContext.current_context()
+            changes = []
+            orig = orig_inline_callbacks(_check_yield_points(f, changes, start_context))
 
             try:
                 res = orig(*args, **kwargs)
             except Exception:
                 if LoggingContext.current_context() != start_context:
+                    for err in changes:
+                        print(err, file=sys.stderr)
+
                     err = "%s changed context from %s to %s on exception" % (
                         f,
                         start_context,
@@ -55,7 +58,10 @@ def do_patch():
 
             if not isinstance(res, Deferred) or res.called:
                 if LoggingContext.current_context() != start_context:
-                    err = "%s changed context from %s to %s" % (
+                    for err in changes:
+                        print(err, file=sys.stderr)
+
+                    err = "Completed %s changed context from %s to %s" % (
                         f,
                         start_context,
                         LoggingContext.current_context(),
@@ -76,6 +82,8 @@ def do_patch():
 
             def check_ctx(r):
                 if LoggingContext.current_context() != start_context:
+                    for err in changes:
+                        print(err, file=sys.stderr)
                     err = "%s completion of %s changed context from %s to %s" % (
                         "Failure" if isinstance(r, Failure) else "Success",
                         f,
@@ -92,3 +100,73 @@ def do_patch():
         return wrapped
 
     defer.inlineCallbacks = new_inline_callbacks
+
+
+def _check_yield_points(f, changes, start_context):
+    from synapse.logging.context import LoggingContext
+
+    @functools.wraps(f)
+    def check_yield_points_inner(*args, **kwargs):
+        gen = f(*args, **kwargs)
+
+        last_yield_line_no = 1
+        result = None
+        while True:
+            try:
+                isFailure = isinstance(result, Failure)
+                if isFailure:
+                    d = result.throwExceptionIntoGenerator(gen)
+                else:
+                    d = gen.send(result)
+            except (StopIteration, defer._DefGen_Return) as e:
+                if LoggingContext.current_context() != start_context:
+                    # This happens when the context is lost sometime *after* the
+                    # final yield and returning. E.g. we forgot to yield on a
+                    # function that returns a deferred.
+                    err = (
+                        "%s returned and changed context from %s to %s, in %s between %d and end of func"
+                        % (
+                            f.__qualname__,
+                            start_context,
+                            LoggingContext.current_context(),
+                            f.__code__.co_filename,
+                            last_yield_line_no,
+                        )
+                    )
+                    changes.append(err)
+                    # print(err, file=sys.stderr)
+                    # raise Exception(err)
+                return getattr(e, "value", None)
+
+            try:
+                result = yield d
+            except Exception as e:
+                result = Failure(e)
+
+            frame = gen.gi_frame
+            if frame.f_code.co_name == "check_yield_points_inner":
+                frame = inspect.getgeneratorlocals(gen)["gen"].gi_frame
+
+            if LoggingContext.current_context() != start_context:
+                # This happens because the context is lost sometime *after* the
+                # previous yield and *after* the current yield. E.g. the
+                # deferred we waited on didn't follow the rules, or we forgot to
+                # yield on a function between the two yield points.
+                err = (
+                    "%s changed context from %s to %s, happened between lines %d and %d in %s"
+                    % (
+                        frame.f_code.co_name,
+                        start_context,
+                        LoggingContext.current_context(),
+                        last_yield_line_no,
+                        frame.f_lineno,
+                        frame.f_code.co_filename,
+                    )
+                )
+                changes.append(err)
+                # print(err, file=sys.stderr)
+                # raise Exception(err)
+
+            last_yield_line_no = frame.f_lineno
+
+    return check_yield_points_inner

From 6374ca40c2ff3d2eaa41af5ebf4f8324522ecb84 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 27 Sep 2019 15:58:14 +0100
Subject: [PATCH 02/13] Update

---
 tests/patch_inline_callbacks.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/tests/patch_inline_callbacks.py b/tests/patch_inline_callbacks.py
index 5ef0aff0c3..a35a1d3305 100644
--- a/tests/patch_inline_callbacks.py
+++ b/tests/patch_inline_callbacks.py
@@ -15,7 +15,6 @@
 
 from __future__ import print_function
 
-import inspect
 import functools
 import sys
 
@@ -32,6 +31,8 @@ def do_patch():
     from synapse.logging.context import LoggingContext
 
     orig_inline_callbacks = defer.inlineCallbacks
+    if hasattr(orig_inline_callbacks, "patched_by_synapse"):
+        return
 
     def new_inline_callbacks(f):
         @functools.wraps(f)
@@ -100,13 +101,20 @@ def do_patch():
         return wrapped
 
     defer.inlineCallbacks = new_inline_callbacks
+    new_inline_callbacks.patched_by_synapse = True
 
 
 def _check_yield_points(f, changes, start_context):
+    """Wraps a generator that is about to passed to defer.inlineCallbacks
+    checking that after every yield the log contexts are correct.
+    """
+
     from synapse.logging.context import LoggingContext
 
     @functools.wraps(f)
     def check_yield_points_inner(*args, **kwargs):
+        expected_context = start_context
+
         gen = f(*args, **kwargs)
 
         last_yield_line_no = 1
@@ -119,12 +127,13 @@ def _check_yield_points(f, changes, start_context):
                 else:
                     d = gen.send(result)
             except (StopIteration, defer._DefGen_Return) as e:
-                if LoggingContext.current_context() != start_context:
+                if LoggingContext.current_context() != expected_context:
                     # This happens when the context is lost sometime *after* the
                     # final yield and returning. E.g. we forgot to yield on a
                     # function that returns a deferred.
                     err = (
-                        "%s returned and changed context from %s to %s, in %s between %d and end of func"
+                        "Function %r returned and changed context from %s to %s,"
+                        " in %s between %d and end of func"
                         % (
                             f.__qualname__,
                             start_context,
@@ -134,7 +143,6 @@ def _check_yield_points(f, changes, start_context):
                         )
                     )
                     changes.append(err)
-                    # print(err, file=sys.stderr)
                     # raise Exception(err)
                 return getattr(e, "value", None)
 
@@ -144,10 +152,8 @@ def _check_yield_points(f, changes, start_context):
                 result = Failure(e)
 
             frame = gen.gi_frame
-            if frame.f_code.co_name == "check_yield_points_inner":
-                frame = inspect.getgeneratorlocals(gen)["gen"].gi_frame
 
-            if LoggingContext.current_context() != start_context:
+            if LoggingContext.current_context() != expected_context:
                 # This happens because the context is lost sometime *after* the
                 # previous yield and *after* the current yield. E.g. the
                 # deferred we waited on didn't follow the rules, or we forgot to
@@ -164,9 +170,10 @@ def _check_yield_points(f, changes, start_context):
                     )
                 )
                 changes.append(err)
-                # print(err, file=sys.stderr)
                 # raise Exception(err)
 
+                expected_context = LoggingContext.current_context()
+
             last_yield_line_no = frame.f_lineno
 
     return check_yield_points_inner

From 1d6dd1c2944c22147258dda8ccf2777c68b38fba Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 10:53:06 +0100
Subject: [PATCH 03/13] Move patch_inline_callbacks into synapse/

---
 synapse/__init__.py                               | 2 +-
 {tests => synapse/util}/patch_inline_callbacks.py | 2 +-
 tests/__init__.py                                 | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)
 rename {tests => synapse/util}/patch_inline_callbacks.py (98%)

diff --git a/synapse/__init__.py b/synapse/__init__.py
index 1055f54e00..bf102244a9 100644
--- a/synapse/__init__.py
+++ b/synapse/__init__.py
@@ -20,7 +20,7 @@
 import os
 import sys
 
-from tests.patch_inline_callbacks import do_patch
+from synapse.util.patch_inline_callbacks import do_patch
 
 # Check that we're not running on an unsupported Python version.
 if sys.version_info < (3, 5):
diff --git a/tests/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
similarity index 98%
rename from tests/patch_inline_callbacks.py
rename to synapse/util/patch_inline_callbacks.py
index a35a1d3305..4fb49b0b2b 100644
--- a/tests/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -105,7 +105,7 @@ def do_patch():
 
 
 def _check_yield_points(f, changes, start_context):
-    """Wraps a generator that is about to passed to defer.inlineCallbacks
+    """Wraps a generator that is about to be passed to defer.inlineCallbacks
     checking that after every yield the log contexts are correct.
     """
 
diff --git a/tests/__init__.py b/tests/__init__.py
index f7fc502f01..ed805db1c2 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -16,9 +16,9 @@
 
 from twisted.trial import util
 
-import tests.patch_inline_callbacks
+from synapse.util.patch_inline_callbacks import do_patch
 
 # attempt to do the patch before we load any synapse code
-tests.patch_inline_callbacks.do_patch()
+do_patch()
 
 util.DEFAULT_TIMEOUT_DURATION = 20

From 3e4272961a4cb659513bccd981cbd42f4e506362 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 10:58:32 +0100
Subject: [PATCH 04/13] Test for sentinel commit

---
 synapse/util/patch_inline_callbacks.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 4fb49b0b2b..4a45824f52 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -146,13 +146,30 @@ def _check_yield_points(f, changes, start_context):
                     # raise Exception(err)
                 return getattr(e, "value", None)
 
+            frame = gen.gi_frame
+
+            if isinstance(d, defer.Deferred):
+                # This happens if we yield on a deferred that doesn't follow
+                # the log context rules without wrappin in a `make_deferred_yieldable`
+                if LoggingContext.current_context() != LoggingContext.Sentinel:
+                    err = (
+                        "%s yielded with context %s rather than Sentinel,"
+                        " yielded on line %d in %s"
+                        % (
+                            frame.f_code.co_name,
+                            start_context,
+                            LoggingContext.current_context(),
+                            frame.f_lineno,
+                            frame.f_code.co_filename,
+                        )
+                    )
+                    changes.append(err)
+
             try:
                 result = yield d
             except Exception as e:
                 result = Failure(e)
 
-            frame = gen.gi_frame
-
             if LoggingContext.current_context() != expected_context:
                 # This happens because the context is lost sometime *after* the
                 # previous yield and *after* the current yield. E.g. the

From ec0596f2ab4502c9a6183813a7e5dc2a5bfedd48 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 10:59:07 +0100
Subject: [PATCH 05/13] Log correct context

---
 synapse/util/patch_inline_callbacks.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 4a45824f52..5ef7190b14 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -117,7 +117,7 @@ def _check_yield_points(f, changes, start_context):
 
         gen = f(*args, **kwargs)
 
-        last_yield_line_no = 1
+        last_yield_line_no = gen.gi_frame.f_lineno
         result = None
         while True:
             try:
@@ -136,7 +136,7 @@ def _check_yield_points(f, changes, start_context):
                         " in %s between %d and end of func"
                         % (
                             f.__qualname__,
-                            start_context,
+                            expected_context,
                             LoggingContext.current_context(),
                             f.__code__.co_filename,
                             last_yield_line_no,
@@ -148,22 +148,22 @@ def _check_yield_points(f, changes, start_context):
 
             frame = gen.gi_frame
 
-            if isinstance(d, defer.Deferred):
+            if isinstance(d, defer.Deferred) and not d.called:
                 # This happens if we yield on a deferred that doesn't follow
                 # the log context rules without wrappin in a `make_deferred_yieldable`
-                if LoggingContext.current_context() != LoggingContext.Sentinel:
+                if LoggingContext.current_context() is not LoggingContext.sentinel:
                     err = (
-                        "%s yielded with context %s rather than Sentinel,"
+                        "%s yielded with context %s rather than sentinel,"
                         " yielded on line %d in %s"
                         % (
                             frame.f_code.co_name,
-                            start_context,
                             LoggingContext.current_context(),
                             frame.f_lineno,
                             frame.f_code.co_filename,
                         )
                     )
                     changes.append(err)
+                    # raise Exception(err)
 
             try:
                 result = yield d

From 128d5948c4e0066f1263b347198f4754e72010c8 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 11:16:26 +0100
Subject: [PATCH 06/13] Fix packaging

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

diff --git a/synapse/__init__.py b/synapse/__init__.py
index bf102244a9..56df3f5ac6 100644
--- a/synapse/__init__.py
+++ b/synapse/__init__.py
@@ -20,8 +20,6 @@
 import os
 import sys
 
-from synapse.util.patch_inline_callbacks import do_patch
-
 # Check that we're not running on an unsupported Python version.
 if sys.version_info < (3, 5):
     print("Synapse requires Python 3.5 or above.")
@@ -41,4 +39,8 @@ except ImportError:
 __version__ = "1.4.0"
 
 if bool(os.environ.get("SYNAPSE_TEST_PATCH_LOG_CONTEXTS", False)):
+    # We import here so that we don't have to install a bunch of deps when
+    # running the packaging tox test.
+    from synapse.util.patch_inline_callbacks import do_patch
+
     do_patch()

From f4571a0497791285ff044d293a6287ed90575b71 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 11:17:00 +0100
Subject: [PATCH 07/13] Newsfile

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

diff --git a/changelog.d/6127.misc b/changelog.d/6127.misc
new file mode 100644
index 0000000000..7bfbcfc252
--- /dev/null
+++ b/changelog.d/6127.misc
@@ -0,0 +1 @@
+Add env var to turn on tracking of log context changes.

From 791a8c559bf4ea984637c047fad7d6097e34ce99 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 11:53:57 +0100
Subject: [PATCH 08/13] Add coments

---
 synapse/util/patch_inline_callbacks.py | 30 +++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 5ef7190b14..b518dae256 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -107,6 +107,19 @@ def do_patch():
 def _check_yield_points(f, changes, start_context):
     """Wraps a generator that is about to be passed to defer.inlineCallbacks
     checking that after every yield the log contexts are correct.
+
+    Its perfectly valid for log contexts to change within a function, e.g. due
+    to new Measure blocks, so such changes are added to the given `changes`
+    list instead of triggering an exception.
+
+    Args:
+        f: generator function to wrap
+        changes (list[str]): A list of strings detailing how the contexts
+            changed within a function.
+        start_context (LoggingContext): The initial context we're expecting
+
+    Returns:
+        function
     """
 
     from synapse.logging.context import LoggingContext
@@ -131,6 +144,10 @@ def _check_yield_points(f, changes, start_context):
                     # This happens when the context is lost sometime *after* the
                     # final yield and returning. E.g. we forgot to yield on a
                     # function that returns a deferred.
+                    #
+                    # We don't raise here as its perfectly valid for contexts to
+                    # change in a function, as long as it sets the correct context
+                    # on resolving (which is checked separately).
                     err = (
                         "Function %r returned and changed context from %s to %s,"
                         " in %s between %d and end of func"
@@ -143,14 +160,14 @@ def _check_yield_points(f, changes, start_context):
                         )
                     )
                     changes.append(err)
-                    # raise Exception(err)
                 return getattr(e, "value", None)
 
             frame = gen.gi_frame
 
             if isinstance(d, defer.Deferred) and not d.called:
                 # This happens if we yield on a deferred that doesn't follow
-                # the log context rules without wrappin in a `make_deferred_yieldable`
+                # the log context rules without wrappin in a `make_deferred_yieldable`.
+                # We raise here as this should never happen.
                 if LoggingContext.current_context() is not LoggingContext.sentinel:
                     err = (
                         "%s yielded with context %s rather than sentinel,"
@@ -162,8 +179,7 @@ def _check_yield_points(f, changes, start_context):
                             frame.f_code.co_filename,
                         )
                     )
-                    changes.append(err)
-                    # raise Exception(err)
+                    raise Exception(err)
 
             try:
                 result = yield d
@@ -171,10 +187,15 @@ def _check_yield_points(f, changes, start_context):
                 result = Failure(e)
 
             if LoggingContext.current_context() != expected_context:
+
                 # This happens because the context is lost sometime *after* the
                 # previous yield and *after* the current yield. E.g. the
                 # deferred we waited on didn't follow the rules, or we forgot to
                 # yield on a function between the two yield points.
+                #
+                # We don't raise here as its perfectly valid for contexts to
+                # change in a function, as long as it sets the correct context
+                # on resolving (which is checked separately).
                 err = (
                     "%s changed context from %s to %s, happened between lines %d and %d in %s"
                     % (
@@ -187,7 +208,6 @@ def _check_yield_points(f, changes, start_context):
                     )
                 )
                 changes.append(err)
-                # raise Exception(err)
 
                 expected_context = LoggingContext.current_context()
 

From 941edad58355a829b49b0a43d382bbb0bf9ba021 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 12:15:17 +0100
Subject: [PATCH 09/13] Appease mypy

---
 synapse/util/patch_inline_callbacks.py | 31 +++++++++++++++-----------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index b518dae256..64a2c891c3 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -18,11 +18,17 @@ from __future__ import print_function
 import functools
 import sys
 
+from typing import List, Callable, Any
+
 from twisted.internet import defer
 from twisted.internet.defer import Deferred
 from twisted.python.failure import Failure
 
 
+# Tracks if we've already patched inlineCallbacks
+_already_patched = False
+
+
 def do_patch():
     """
     Patch defer.inlineCallbacks so that it checks the state of the logcontext on exit
@@ -30,16 +36,18 @@ def do_patch():
 
     from synapse.logging.context import LoggingContext
 
+    global _already_patched
+
     orig_inline_callbacks = defer.inlineCallbacks
-    if hasattr(orig_inline_callbacks, "patched_by_synapse"):
+    if _already_patched:
         return
 
     def new_inline_callbacks(f):
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
             start_context = LoggingContext.current_context()
-            changes = []
-            orig = orig_inline_callbacks(_check_yield_points(f, changes, start_context))
+            changes: List[str] = []
+            orig = orig_inline_callbacks(_check_yield_points(f, changes))
 
             try:
                 res = orig(*args, **kwargs)
@@ -101,10 +109,10 @@ def do_patch():
         return wrapped
 
     defer.inlineCallbacks = new_inline_callbacks
-    new_inline_callbacks.patched_by_synapse = True
+    _already_patched = True
 
 
-def _check_yield_points(f, changes, start_context):
+def _check_yield_points(f: Callable, changes: List[str]):
     """Wraps a generator that is about to be passed to defer.inlineCallbacks
     checking that after every yield the log contexts are correct.
 
@@ -114,9 +122,8 @@ def _check_yield_points(f, changes, start_context):
 
     Args:
         f: generator function to wrap
-        changes (list[str]): A list of strings detailing how the contexts
+        changes: A list of strings detailing how the contexts
             changed within a function.
-        start_context (LoggingContext): The initial context we're expecting
 
     Returns:
         function
@@ -126,13 +133,13 @@ def _check_yield_points(f, changes, start_context):
 
     @functools.wraps(f)
     def check_yield_points_inner(*args, **kwargs):
-        expected_context = start_context
-
         gen = f(*args, **kwargs)
 
         last_yield_line_no = gen.gi_frame.f_lineno
-        result = None
+        result: Any = None
         while True:
+            expected_context = LoggingContext.current_context()
+
             try:
                 isFailure = isinstance(result, Failure)
                 if isFailure:
@@ -200,7 +207,7 @@ def _check_yield_points(f, changes, start_context):
                     "%s changed context from %s to %s, happened between lines %d and %d in %s"
                     % (
                         frame.f_code.co_name,
-                        start_context,
+                        expected_context,
                         LoggingContext.current_context(),
                         last_yield_line_no,
                         frame.f_lineno,
@@ -209,8 +216,6 @@ def _check_yield_points(f, changes, start_context):
                 )
                 changes.append(err)
 
-                expected_context = LoggingContext.current_context()
-
             last_yield_line_no = frame.f_lineno
 
     return check_yield_points_inner

From f735aeec65a5117c71cf0f1e5f61cb900683533a Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 12:20:29 +0100
Subject: [PATCH 10/13] sort

---
 synapse/util/patch_inline_callbacks.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 64a2c891c3..3b78451dc0 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -17,14 +17,12 @@ from __future__ import print_function
 
 import functools
 import sys
-
-from typing import List, Callable, Any
+from typing import Any, Callable, List
 
 from twisted.internet import defer
 from twisted.internet.defer import Deferred
 from twisted.python.failure import Failure
 
-
 # Tracks if we've already patched inlineCallbacks
 _already_patched = False
 

From c349e3ebafbe044022b93ca5c04d8b2fcb640c0a Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 12:29:38 +0100
Subject: [PATCH 11/13] Fix py3.5

---
 synapse/util/patch_inline_callbacks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 3b78451dc0..812dc88835 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -44,7 +44,7 @@ def do_patch():
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
             start_context = LoggingContext.current_context()
-            changes: List[str] = []
+            changes = []  # type: List[str]
             orig = orig_inline_callbacks(_check_yield_points(f, changes))
 
             try:

From 59e0ed83065874775be350e25bb9f87da67b87c2 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 12:47:07 +0100
Subject: [PATCH 12/13] Fix py3.5

---
 synapse/util/patch_inline_callbacks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 812dc88835..66c3d47519 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -134,7 +134,7 @@ def _check_yield_points(f: Callable, changes: List[str]):
         gen = f(*args, **kwargs)
 
         last_yield_line_no = gen.gi_frame.f_lineno
-        result: Any = None
+        result = None  # type: Any
         while True:
             expected_context = LoggingContext.current_context()
 

From fe1c1e6c28e09f88b30e0587161f9b1dbd6e8acf Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 10 Oct 2019 13:17:19 +0100
Subject: [PATCH 13/13] Fixup comments

Co-Authored-By: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 synapse/util/patch_inline_callbacks.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index 66c3d47519..3925927f9f 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -114,7 +114,7 @@ def _check_yield_points(f: Callable, changes: List[str]):
     """Wraps a generator that is about to be passed to defer.inlineCallbacks
     checking that after every yield the log contexts are correct.
 
-    Its perfectly valid for log contexts to change within a function, e.g. due
+    It's perfectly valid for log contexts to change within a function, e.g. due
     to new Measure blocks, so such changes are added to the given `changes`
     list instead of triggering an exception.
 
@@ -150,7 +150,7 @@ def _check_yield_points(f: Callable, changes: List[str]):
                     # final yield and returning. E.g. we forgot to yield on a
                     # function that returns a deferred.
                     #
-                    # We don't raise here as its perfectly valid for contexts to
+                    # We don't raise here as it's perfectly valid for contexts to
                     # change in a function, as long as it sets the correct context
                     # on resolving (which is checked separately).
                     err = (
@@ -171,7 +171,7 @@ def _check_yield_points(f: Callable, changes: List[str]):
 
             if isinstance(d, defer.Deferred) and not d.called:
                 # This happens if we yield on a deferred that doesn't follow
-                # the log context rules without wrappin in a `make_deferred_yieldable`.
+                # the log context rules without wrapping in a `make_deferred_yieldable`.
                 # We raise here as this should never happen.
                 if LoggingContext.current_context() is not LoggingContext.sentinel:
                     err = (