diff --git a/synapse/util/logcontext.py b/synapse/util/logcontext.py index 192e3f49f0..3dce8d2bf3 100644 --- a/synapse/util/logcontext.py +++ b/synapse/util/logcontext.py @@ -134,13 +134,25 @@ class PreserveLoggingContext(object): def preserve_context_over_fn(fn, *args, **kwargs): - with PreserveLoggingContext(): - deferred = fn(*args, **kwargs) + """Takes a function and invokes it with the given arguments, but removes + and restores the current logging context while doing so. - return preserve_context_over_deferred(deferred) + If the result is a deferred, call preserve_context_over_deferred before + returning it. + """ + with PreserveLoggingContext(): + res = fn(*args, **kwargs) + + if isinstance(res, defer.Deferred): + return preserve_context_over_deferred(res) + else: + return res def preserve_context_over_deferred(deferred): + """Given a deferred wrap it such that any callbacks added later to it will + be invoked with the current context. + """ d = defer.Deferred() current_context = LoggingContext.current_context()