Wrap the Rust HTTP client with make_deferred_yieldable

changelog.d/18903.misc

@@ -0,0 +1 @@
+Wrap the Rust HTTP client with `make_deferred_yieldable` so it follows Synapse logcontext rules.

synapse/api/auth/mas.py

@@ -33,15 +33,14 @@
     UnrecognizedRequestError,
 )
 from synapse.http.site import SynapseRequest
-from synapse.logging.context import PreserveLoggingContext
 from synapse.logging.opentracing import (
     active_span,
     force_tracing,
     inject_request_headers,
     start_active_span,
 )
 from synapse.metrics import SERVER_NAME_LABEL
-from synapse.synapse_rust.http_client import HttpClient
+from synapse.synapse_rust_wrapper.http_client import HttpClient
 from synapse.types import JsonDict, Requester, UserID, create_requester
 from synapse.util import json_decoder
 from synapse.util.caches.cached_call import RetryOnExceptionCachedCall
@@ -229,13 +228,12 @@ async def _introspect_token(
         try:
             with start_active_span("mas-introspect-token"):
                 inject_request_headers(raw_headers)
-                with PreserveLoggingContext():
-                    resp_body = await self._rust_http_client.post(
-                        url=self._introspection_endpoint,
-                        response_limit=1 * 1024 * 1024,
-                        headers=raw_headers,
-                        request_body=body,
-                    )
+                resp_body = await self._rust_http_client.post(
+                    url=self._introspection_endpoint,
+                    response_limit=1 * 1024 * 1024,
+                    headers=raw_headers,
+                    request_body=body,
+                )
         except HttpResponseException as e:
             end_time = self._clock.time()
             introspection_response_timer.labels(

synapse/api/auth/msc3861_delegated.py

@@ -38,15 +38,14 @@
     UnrecognizedRequestError,
 )
 from synapse.http.site import SynapseRequest
-from synapse.logging.context import PreserveLoggingContext
 from synapse.logging.opentracing import (
     active_span,
     force_tracing,
     inject_request_headers,
     start_active_span,
 )
 from synapse.metrics import SERVER_NAME_LABEL
-from synapse.synapse_rust.http_client import HttpClient
+from synapse.synapse_rust_wrapper.http_client import HttpClient
 from synapse.types import Requester, UserID, create_requester
 from synapse.util import json_decoder
 from synapse.util.caches.cached_call import RetryOnExceptionCachedCall
@@ -327,13 +326,12 @@ async def _introspect_token(
         try:
             with start_active_span("mas-introspect-token"):
                 inject_request_headers(raw_headers)
-                with PreserveLoggingContext():
-                    resp_body = await self._rust_http_client.post(
-                        url=uri,
-                        response_limit=1 * 1024 * 1024,
-                        headers=raw_headers,
-                        request_body=body,
-                    )
+                resp_body = await self._rust_http_client.post(
+                    url=uri,
+                    response_limit=1 * 1024 * 1024,
+                    headers=raw_headers,
+                    request_body=body,
+                )
         except HttpResponseException as e:
             end_time = self._clock.time()
             introspection_response_timer.labels(

synapse/synapse_rust/http_client.pyi

@@ -17,6 +17,13 @@ from twisted.internet.defer import Deferred
 from synapse.types import ISynapseReactor
 
 class HttpClient:
+    """
+    Since the returned deferreds don't follow Synapse logcontext rules,
+    this is not meant to be used by Synapse code directly.
+
+    Use `synapse.synapse_rust_wrapper.http_client.HttpClient` instead.
+    """
+
     def __init__(self, reactor: ISynapseReactor, user_agent: str) -> None: ...
     def get(self, url: str, response_limit: int) -> Deferred[bytes]: ...
     def post(

synapse/synapse_rust_wrapper/__init__.py

@@ -0,0 +1,11 @@
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2025 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.

synapse/synapse_rust_wrapper/http_client.py

@@ -0,0 +1,44 @@
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2025 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+
+from typing import Mapping
+
+from twisted.internet.defer import Deferred
+
+from synapse.logging.context import make_deferred_yieldable
+from synapse.synapse_rust.http_client import HttpClient as RustHttpClient
+from synapse.types import ISynapseReactor
+
+
+class HttpClient:
+    """
+    Wrap `synapse.synapse_rust.http_client.HttpClient` to ensure the returned
+    deferreds follow Synapse logcontext rules.
+    """
+
+    def __init__(self, reactor: ISynapseReactor, user_agent: str) -> None:
+        self._http_client = RustHttpClient(reactor, user_agent)
+
+    def get(self, url: str, response_limit: int) -> Deferred[bytes]:
+        deferred = self._http_client.get(url, response_limit)
+        return make_deferred_yieldable(deferred)
+
+    def post(
+        self,
+        url: str,
+        response_limit: int,
+        headers: Mapping[str, str],
+        request_body: str,
+    ) -> Deferred[bytes]:
+        deferred = self._http_client.post(url, response_limit, headers, request_body)
+        return make_deferred_yieldable(deferred)

tests/synapse_rust/__init__.py

@@ -0,0 +1,11 @@
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2025 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.

tests/synapse_rust/test_http_client.py

@@ -0,0 +1,87 @@
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright (C) 2025 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+import time
+from typing import Any, Coroutine, Generator, TypeVar, Union
+
+from twisted.internet.defer import Deferred, ensureDeferred
+from twisted.internet.testing import MemoryReactor
+
+from synapse.logging.context import LoggingContext
+from synapse.server import HomeServer
+from synapse.synapse_rust.http_client import HttpClient
+from synapse.util import Clock
+
+from tests.unittest import HomeserverTestCase
+
+T = TypeVar("T")
+
+
+class HttpClientTestCase(HomeserverTestCase):
+    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
+        hs = self.setup_test_homeserver()
+        # This triggers the server startup hooks, which starts the Tokio thread pool
+        reactor.run()
+        return hs
+
+    def tearDown(self) -> None:
+        # MemoryReactor doesn't trigger the shutdown phases, and we want the
+        # Tokio thread pool to be stopped
+        # XXX: This logic should probably get moved somewhere else
+        shutdown_triggers = self.reactor.triggers.get("shutdown", {})
+        for phase in ["before", "during", "after"]:
+            triggers = shutdown_triggers.get(phase, [])
+            for callbable, args, kwargs in triggers:
+                callbable(*args, **kwargs)
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        self._http_client = hs.get_proxied_http_client()
+        self._rust_http_client = HttpClient(
+            reactor=hs.get_reactor(),
+            user_agent=self._http_client.user_agent.decode("utf8"),
+        )
+
+    def till_deferred_has_result(
+        self,
+        awaitable: Union[
+            "Coroutine[Deferred[Any], Any, T]",
+            "Generator[Deferred[Any], Any, T]",
+            "Deferred[T]",
+        ],
+    ) -> "Deferred[T]":
+        """Wait until a deferred has a result.
+
+        This is useful because the Rust HTTP client will resolve the deferred
+        using reactor.callFromThread, which are only run when we call
+        reactor.advance.
+        """
+        deferred = ensureDeferred(awaitable)
+        tries = 0
+        while not deferred.called:
+            time.sleep(0.1)
+            self.reactor.advance(0)
+            tries += 1
+            if tries > 100:
+                raise Exception("Timed out waiting for deferred to resolve")
+
+        return deferred
+
+    def test_logging_context(self) -> None:
+        async def asdf() -> None:
+            with LoggingContext("test"):
+                # TODO: Test logging context before/after this call
+                await self._rust_http_client.get(
+                    url="http://localhost",
+                    response_limit=1 * 1024 * 1024,
+                )
+
+        self.get_success(self.till_deferred_has_result(asdf()))