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.

rust/src/http_client.rs

@@ -19,6 +19,7 @@ use futures::TryStreamExt;
 use once_cell::sync::OnceCell;
 use pyo3::{create_exception, exceptions::PyException, prelude::*};
 use reqwest::RequestBuilder;
+use std::sync::OnceLock;
 use tokio::runtime::Runtime;
 
 use crate::errors::HttpResponseException;
@@ -218,29 +219,34 @@ impl HttpClient {
         builder: RequestBuilder,
         response_limit: usize,
     ) -> PyResult<Bound<'a, PyAny>> {
-        create_deferred(py, self.reactor.bind(py), async move {
-            let response = builder.send().await.context("sending request")?;
-
-            let status = response.status();
-
-            let mut stream = response.bytes_stream();
-            let mut buffer = Vec::new();
-            while let Some(chunk) = stream.try_next().await.context("reading body")? {
-                if buffer.len() + chunk.len() > response_limit {
-                    Err(anyhow::anyhow!("Response size too large"))?;
+        // We use `make_deferred_yieldable` to make the returned deferred follow Synapse
+        // logcontext rules.
+        Ok(make_deferred_yieldable(
+            py,
+            &create_deferred(py, self.reactor.bind(py), async move {
+                let response = builder.send().await.context("sending request")?;
+
+                let status = response.status();
+
+                let mut stream = response.bytes_stream();
+                let mut buffer = Vec::new();
+                while let Some(chunk) = stream.try_next().await.context("reading body")? {
+                    if buffer.len() + chunk.len() > response_limit {
+                        Err(anyhow::anyhow!("Response size too large"))?;
+                    }
+
+                    buffer.extend_from_slice(&chunk);
                 }
 
-                buffer.extend_from_slice(&chunk);
-            }
-
-            if !status.is_success() {
-                return Err(HttpResponseException::new(status, buffer));
-            }
+                if !status.is_success() {
+                    return Err(HttpResponseException::new(status, buffer));
+                }
 
-            let r = Python::with_gil(|py| buffer.into_pyobject(py).map(|o| o.unbind()))?;
+                let r = Python::with_gil(|py| buffer.into_pyobject(py).map(|o| o.unbind()))?;
 
-            Ok(r)
-        })
+                Ok(r)
+            })?,
+        ))
     }
 }
 
@@ -301,3 +307,23 @@ where
 
     Ok(deferred)
 }
+
+static MAKE_DEFERRED_YIELDABLE: OnceLock<pyo3::Py<pyo3::PyAny>> = OnceLock::new();
+
+/// Given a deferred, make it follow the Synapse logcontext rules
+fn make_deferred_yieldable<'py>(
+    py: Python<'py>,
+    deferred: &Bound<'py, PyAny>,
+) -> Bound<'py, PyAny> {
+    let make_deferred_yieldable = MAKE_DEFERRED_YIELDABLE.get_or_init(|| {
+        let sys = PyModule::import(py, "synapse.logging.context").unwrap();
+        let func = sys.getattr("make_deferred_yieldable").unwrap().unbind();
+        func
+    });
+
+    make_deferred_yieldable
+        .call1(py, (deferred,))
+        .unwrap()
+        .extract(py)
+        .unwrap()
+}

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.caches.cached_call import RetryOnExceptionCachedCall
 from synapse.util.caches.response_cache import ResponseCache, ResponseCacheContext
@@ -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.caches.cached_call import RetryOnExceptionCachedCall
 from synapse.util.caches.response_cache import ResponseCache, ResponseCacheContext
@@ -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,10 @@ from twisted.internet.defer import Deferred
 from synapse.types import ISynapseReactor
 
 class HttpClient:
+    """
+    The returned deferreds follow Synapse logcontext rules.
+    """
+
     def __init__(self, reactor: ISynapseReactor, user_agent: str) -> None: ...
     def get(self, url: str, response_limit: int) -> Deferred[bytes]: ...
     def post(

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>.