WIP

pyproject.toml

@@ -51,11 +51,7 @@ line-length = 88
 # flake8-bugbear compatible checks. Its error codes are described at
 # https://beta.ruff.rs/docs/rules/#flake8-bugbear-b
 #  B023: Functions defined inside a loop must not use variables redefined in the loop
-ignore = [
+ignore = ["B023", "E501", "E731"]
-    "B023",
-    "E501",
-    "E731",
-]
 select = [
     # pycodestyle
     "E",
@@ -80,7 +76,15 @@ select = [
 
 [tool.isort]
 line_length = 88
-sections = ["FUTURE", "STDLIB", "THIRDPARTY", "TWISTED", "FIRSTPARTY", "TESTS", "LOCALFOLDER"]
+sections = [
+    "FUTURE",
+    "STDLIB",
+    "THIRDPARTY",
+    "TWISTED",
+    "FIRSTPARTY",
+    "TESTS",
+    "LOCALFOLDER",
+]
 default_section = "THIRDPARTY"
 known_first_party = ["synapse"]
 known_tests = ["tests"]
@@ -102,9 +106,7 @@ authors = ["Matrix.org Team and Contributors <packages@matrix.org>"]
 license = "AGPL-3.0-or-later"
 readme = "README.rst"
 repository = "https://github.com/element-hq/synapse"
-packages = [
+packages = [{ include = "synapse" }]
-    { include = "synapse" },
-]
 classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Topic :: Communications :: Chat",
@@ -130,9 +132,7 @@ include = [
     { path = "rust/build.rs", format = "sdist" },
     { path = "rust/src/**", format = "sdist" },
 ]
-exclude = [
+exclude = [{ path = "synapse/*.so", format = "sdist" }]
-    { path = "synapse/*.so", format = "sdist"}
-]
 
 [tool.poetry.build]
 script = "build_rust.py"
@@ -171,9 +171,8 @@ canonicaljson = "^2.0.0"
 signedjson = "^1.1.0"
 # validating SSL certs for IP addresses requires service_identity 18.1.
 service-identity = ">=18.1.0"
-# Twisted 18.9 introduces some logger improvements that the structured
+# Twisted 21.2 introduces contextvar support
-# logger utilises
+Twisted = { extras = ["tls", "contextvars"], version = ">=21.2.0" }
-Twisted = {extras = ["tls"], version = ">=18.9.0"}
 treq = ">=15.1"
 # Twisted has required pyopenssl 16.0 since about Twisted 16.6.
 pyOpenSSL = ">=16.0.0"
@@ -290,7 +289,9 @@ all = [
     # matrix-synapse-ldap3
     "matrix-synapse-ldap3",
     # postgres
-    "psycopg2", "psycopg2cffi", "psycopg2cffi-compat",
+    "psycopg2",
+    "psycopg2cffi",
+    "psycopg2cffi-compat",
     # saml2
     "pysaml2",
     # oidc and jwt
@@ -300,9 +301,11 @@ all = [
     # sentry
     "sentry-sdk",
     # opentracing
-    "jaeger-client", "opentracing",
+    "jaeger-client",
+    "opentracing",
     # redis
-    "txredisapi", "hiredis",
+    "txredisapi",
+    "hiredis",
     # cache-memory
     "pympler",
     # improved user search

synapse/logging/context.py

@@ -33,9 +33,12 @@ import logging
 import threading
 import typing
 import warnings
+from collections.abc import Coroutine, Generator
+from contextvars import ContextVar
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Any,
     Awaitable,
     Callable,
     Optional,
@@ -657,13 +660,14 @@ class PreserveLoggingContext:
                 )
 
 
-_thread_local = threading.local()
+_CURRENT_CONTEXT_VAR: ContextVar[LoggingContextOrSentinel] = ContextVar(
-_thread_local.current_context = SENTINEL_CONTEXT
+    "current_context", default=SENTINEL_CONTEXT
+)
 
 
 def current_context() -> LoggingContextOrSentinel:
     """Get the current logging context from thread local storage"""
-    return getattr(_thread_local, "current_context", SENTINEL_CONTEXT)
+    return _CURRENT_CONTEXT_VAR.get()
 
 
 def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSentinel:
@@ -684,7 +688,7 @@ def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSe
     if current is not context:
         rusage = get_thread_resource_usage()
         current.stop(rusage)
-        _thread_local.current_context = context
+        _CURRENT_CONTEXT_VAR.set(context)
         context.start(rusage)
 
     return current
@@ -971,3 +975,74 @@ def defer_to_threadpool(
             return f(*args, **kwargs)
 
     return make_deferred_yieldable(threads.deferToThreadPool(reactor, threadpool, g))
+
+
+_T = TypeVar("_T")
+
+
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class _ResourceTracker(Generator[defer.Deferred[Any], Any, _T]):
+    gen: Generator[defer.Deferred[Any], Any, _T]
+
+    def send(self, val: Any) -> defer.Deferred[_T]:
+        try:
+            return self.gen.send(val)
+        finally:
+            pass
+
+    @overload
+    def throw(
+        self,
+        a: Type[BaseException],
+        b: object = ...,
+        c: Optional[TracebackType] = ...,
+        /,
+    ) -> defer.Deferred[Any]: ...
+
+    @overload
+    def throw(
+        self, a: BaseException, v: None = ..., c: Optional[TracebackType] = ..., /
+    ) -> defer.Deferred[Any]: ...
+
+    def throw(self, a: Any, b: Any = None, c: Any = None) -> defer.Deferred[Any]:
+        try:
+            return self.throw(a, b, c)
+        finally:
+            pass
+
+
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class _ResourceTracker2(Coroutine[defer.Deferred[Any], Any, _T]):
+    gen: Coroutine[defer.Deferred[Any], Any, _T]
+
+    def send(self, val: Any) -> defer.Deferred[_T]:
+        try:
+            return self.gen.send(val)
+        finally:
+            pass
+
+    @overload
+    def throw(
+        self,
+        a: Type[BaseException],
+        b: object = ...,
+        c: Optional[TracebackType] = ...,
+        /,
+    ) -> defer.Deferred[Any]: ...
+
+    @overload
+    def throw(
+        self, a: BaseException, v: None = ..., c: Optional[TracebackType] = ..., /
+    ) -> defer.Deferred[Any]: ...
+
+    def throw(self, a: Any, b: Any = None, c: Any = None) -> defer.Deferred[Any]:
+        try:
+            return self.throw(a, b, c)
+        finally:
+            pass
+
+    def __await__(self) -> Generator[defer.Deferred[Any], Any, _T]:
+        return _ResourceTracker(self.gen.__await__())
+
+    def close(self) -> None:
+        return self.gen.close()