Trace Rust state resolution, event loads and HTTP requests

- `state.resolve` around a state resolution on the database, tagged
  `room_id`, `state_sets` and, once done, the resolution's database and
  auth chain time. It covers the handler's resolutions and the
  persister's, under `persist.state_res`. A room without a chain cover
  index is tagged `fallback`, not an error: the caller resolves its
  state in Python.
- `events.load` around an event load, tagged with the number of events
  asked for, `reload`, and on success how many this load read from the
  database (`fetched`) and how many it did not (`cached`, which counts
  events another load was already reading). Its `db.*` spans nest
  under it.
- `outgoing-client-request` around a request of the Rust HTTP client
  (MAS introspection), with the tags Python's `SimpleHttpClient` span
  has: `span.kind`, `http.method`, `http.url` (without credentials or
  query) and `http.status_code`. As in Python, an error status is the
  caller's to judge, so only a failure to get a response is marked as
  an error, before the conversion to a Python exception.

A trial test loads an uncached event with tracing on and checks that
`events.load` is under the caller's Python span, with the event fetched
and its reads under it. Rust tests check `state.resolve` on a failed
resolution and the HTTP client span against a local server.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W4zxygfJC93SsugxskhoJG

rust/src/http_client.rs

@@ -18,9 +18,11 @@ use anyhow::Context;
 use http_body_util::BodyExt;
 use pyo3::prelude::*;
 use reqwest::RequestBuilder;
+use tracing::Instrument;
 
 use crate::deferred::create_deferred;
 use crate::errors::HttpResponseException;
+use crate::logging::opentracing::record_error;
 use crate::runtime::RustRuntime;
 
 /// Called when registering modules with python.
@@ -104,35 +106,151 @@ impl HttpClient {
         builder: RequestBuilder,
         response_limit: usize,
     ) -> PyResult<Bound<'a, PyAny>> {
+        let client = self.client.clone();
         create_deferred(py, &self.runtime, async move {
-            let response = builder.send().await.context("sending request")?;
-
-            let status = response.status();
-
-            // A light-weight way to read the response up until the `response_limit`. We
-            // want to avoid allocating a giant response object on the server above our
-            // expected `response_limit` to avoid out-of-memory DOS problems.
-            let body = reqwest::Body::from(response);
-            let limited_body = http_body_util::Limited::new(body, response_limit);
-            let collected = limited_body
-                .collect()
-                .await
-                .map_err(anyhow::Error::from_boxed)
-                .with_context(|| {
-                    format!(
-                        "Response body exceeded response limit ({} bytes)",
-                        response_limit
-                    )
-                })?;
-            let bytes: bytes::Bytes = collected.to_bytes();
-
-            if !status.is_success() {
-                return Err(HttpResponseException::new(status, bytes));
-            }
-
-            // Because of the `pyo3` `bytes` feature, we can pass this back to Python
-            // land efficiently
-            Ok(bytes)
+            let request = builder.build().context("building request")?;
+            traced_request(client, request, response_limit).await
         })
     }
 }
+
+/// Send `request` in Python's `SimpleHttpClient` span, created in the
+/// task so that the caller's Python span is its parent.
+async fn traced_request(
+    client: reqwest::Client,
+    request: reqwest::Request,
+    response_limit: usize,
+) -> PyResult<bytes::Bytes> {
+    // Without credentials or the query, which may carry secrets.
+    let mut url = request.url().clone();
+    let _ = url.set_password(None);
+    let _ = url.set_username("");
+    url.set_query(None);
+    let span = tracing::info_span!(
+        "outgoing-client-request",
+        span.kind = "client",
+        http.method = request.method().as_str(),
+        http.url = url.as_str(),
+        http.status_code = tracing::field::Empty,
+        error = tracing::field::Empty,
+    );
+    let result = read_response(client, request, response_limit)
+        .instrument(span.clone())
+        .await;
+    match result {
+        Ok(bytes) => Ok(bytes),
+        // As Python's span, an error status is not a failure of the
+        // request: the caller decides.
+        Err(ReadError::Status(status, bytes)) => Err(HttpResponseException::new(status, bytes)),
+        Err(ReadError::Other(err)) => {
+            // Before the conversion to a Python error, which needs the GIL.
+            record_error(&span, &format!("{err:#}"));
+            Err(err.into())
+        }
+    }
+}
+
+/// Why [`read_response`] failed.
+enum ReadError {
+    /// The server answered with an error status.
+    Status(reqwest::StatusCode, bytes::Bytes),
+    Other(anyhow::Error),
+}
+
+impl From<anyhow::Error> for ReadError {
+    fn from(err: anyhow::Error) -> Self {
+        ReadError::Other(err)
+    }
+}
+
+/// Send `request` and read its response, up to `response_limit` bytes.
+async fn read_response(
+    client: reqwest::Client,
+    request: reqwest::Request,
+    response_limit: usize,
+) -> Result<bytes::Bytes, ReadError> {
+    let response = client.execute(request).await.context("sending request")?;
+
+    let status = response.status();
+    tracing::Span::current().record("http.status_code", status.as_u16());
+
+    // A light-weight way to read the response up until the `response_limit`. We
+    // want to avoid allocating a giant response object on the server above our
+    // expected `response_limit` to avoid out-of-memory DOS problems.
+    let body = reqwest::Body::from(response);
+    let limited_body = http_body_util::Limited::new(body, response_limit);
+    let collected = limited_body
+        .collect()
+        .await
+        .map_err(anyhow::Error::from_boxed)
+        .with_context(|| {
+            format!(
+                "Response body exceeded response limit ({} bytes)",
+                response_limit
+            )
+        })?;
+    let bytes: bytes::Bytes = collected.to_bytes();
+
+    if !status.is_success() {
+        return Err(ReadError::Status(status, bytes));
+    }
+
+    // Because of the `pyo3` `bytes` feature, we can pass this back to Python
+    // land efficiently
+    Ok(bytes)
+}
+
+#[cfg(test)]
+mod tests {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    use super::*;
+    use crate::logging::opentracing::layer::testing::{tag, traced_until_dropped};
+    use crate::logging::opentracing::thrift::TagValue;
+    use crate::logging::opentracing::SamplerConfig;
+
+    /// The request is Python's `outgoing-client-request` span, with its
+    /// status, and a URL without credentials or query.
+    #[tokio::test]
+    async fn traced_request_span() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let port = listener.local_addr().unwrap().port();
+        let server = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut buf = [0; 4096];
+            let _ = socket.read(&mut buf).await.unwrap();
+            socket
+                .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\nconnection: close\r\n\r\nok")
+                .await
+                .unwrap();
+        });
+
+        let (guard, sink) = traced_until_dropped(SamplerConfig::Const(true));
+        let client = reqwest::Client::builder().no_proxy().build().unwrap();
+        let request = client
+            .get(format!(
+                "http://user:secret@127.0.0.1:{port}/introspect?token=abc"
+            ))
+            .build()
+            .unwrap();
+        let body = traced_request(client, request, 1024).await.unwrap();
+        drop(guard);
+        server.await.unwrap();
+
+        assert_eq!(&body[..], b"ok");
+        let spans = sink.take();
+        let [span] = &spans[..] else {
+            panic!("expected one span: {spans:#?}")
+        };
+        assert_eq!(span.operation, "outgoing-client-request");
+        assert_eq!(tag(span, "http.method"), Some(&TagValue::Str("GET".into())));
+        assert_eq!(
+            tag(span, "http.url"),
+            Some(&TagValue::Str(
+                format!("http://127.0.0.1:{port}/introspect").into()
+            ))
+        );
+        assert_eq!(tag(span, "http.status_code"), Some(&TagValue::Long(200)));
+        assert_eq!(tag(span, "error"), None);
+    }
+}

rust/src/state/mod.rs

@@ -45,12 +45,14 @@ use pyo3::{
     prelude::*,
     types::{PyDict, PyModule, PyTuple},
 };
+use tracing::Instrument;
 
 use crate::caches::{python::PyCacheKey, Cache};
 use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool};
 use crate::deferred::create_deferred;
 use crate::event_auth::AuthFailure;
 use crate::events::Event;
+use crate::logging::opentracing::record_error;
 use crate::room_versions::{RoomVersion, StateResolutionVersions};
 use crate::runtime::RustRuntime;
 use crate::storage::{caches::StoreCaches, store::Store};
@@ -700,6 +702,56 @@ pub(crate) async fn resolve_on_pool(
     separated: (StateMap, v2::Conflicted),
     events: EventMap,
     stats: StateResStats,
+) -> anyhow::Result<ResolveOutput> {
+    let span = tracing::info_span!(
+        "state.resolve",
+        room_id = room_id,
+        state_sets = state_sets.len() as u64,
+        db_ms = tracing::field::Empty,
+        auth_chain_ms = tracing::field::Empty,
+        fallback = tracing::field::Empty,
+        error = tracing::field::Empty,
+    );
+    let output = resolve_on_pool_in_span(
+        pool,
+        handle,
+        caches,
+        cache,
+        room_id,
+        room_version,
+        state_sets,
+        separated,
+        events,
+        stats,
+    )
+    .instrument(span.clone())
+    .await;
+    match &output {
+        Ok(output) => {
+            span.record("db_ms", output.stats.db_time * 1000.0);
+            span.record("auth_chain_ms", output.stats.auth_chain_time * 1000.0);
+        }
+        // Expected: the caller resolves the state in Python instead.
+        Err(err) if err.downcast_ref::<store::NoChainCoverIndex>().is_some() => {
+            span.record("fallback", true);
+        }
+        Err(err) => record_error(&span, err),
+    }
+    output
+}
+
+#[allow(clippy::too_many_arguments)]
+async fn resolve_on_pool_in_span(
+    pool: &ConnectionPool,
+    handle: tokio::runtime::Handle,
+    caches: Arc<StoreCaches>,
+    cache: &ConflictCache,
+    room_id: &str,
+    room_version: &'static RoomVersion,
+    state_sets: &[Arc<StateMap>],
+    separated: (StateMap, v2::Conflicted),
+    events: EventMap,
+    stats: StateResStats,
 ) -> anyhow::Result<ResolveOutput> {
     let start = Instant::now();
     let store = PostgresStateResStore::begin(pool, handle, Some(caches)).await?;

rust/src/state/store.rs

@@ -576,4 +576,74 @@ mod tracing_tests {
             ]
         );
     }
+
+    /// A resolution that fails reports `state.resolve` marked failed, with
+    /// its reads under it.
+    #[test]
+    fn traced_resolve_failure() {
+        use crate::caches::registry::CacheRegistry;
+        use crate::room_versions::RoomVersion;
+        use crate::state::{conflict_cache_for_test, resolve_on_pool, v2, StateKey, StateMap};
+        use crate::storage::caches::StoreCaches;
+
+        let Some((rt, pool)) = test_pool() else {
+            return;
+        };
+        let registry = CacheRegistry::new();
+        let caches = Arc::new(StoreCaches::new(&registry, 10));
+        let cache = conflict_cache_for_test(10);
+        let key = StateKey::new("m.room.name", "");
+        let state_sets: Vec<Arc<StateMap>> = ["$a", "$b"]
+            .into_iter()
+            .map(|event_id| Arc::new(StateMap::from([(key.clone(), Arc::from(event_id))])))
+            .collect();
+        let separated = v2::separate(&state_sets);
+
+        let (result, spans) = traced(SamplerConfig::Const(true), || {
+            let result = rt.block_on(resolve_on_pool(
+                &pool,
+                rt.handle().clone(),
+                caches,
+                &cache,
+                "!unknown:test",
+                &RoomVersion::V11,
+                &state_sets,
+                separated,
+                Default::default(),
+                Default::default(),
+            ));
+            // The failed store rolls back on a task of its own, which holds
+            // the spans open: wait for it to give the connection back.
+            rt.block_on(async {
+                let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
+                loop {
+                    let status = pool.status();
+                    if status.available == status.size {
+                        break;
+                    }
+                    assert!(
+                        std::time::Instant::now() < deadline,
+                        "rollback never finished"
+                    );
+                    tokio::time::sleep(std::time::Duration::from_millis(5)).await;
+                }
+            });
+            result
+        });
+        assert!(result.is_err());
+
+        let resolve = spans
+            .iter()
+            .find(|span| span.operation == "state.resolve")
+            .expect("a state.resolve span");
+        assert_eq!(
+            tag(resolve, "room_id"),
+            Some(&TagValue::Str("!unknown:test".into()))
+        );
+        assert_eq!(tag(resolve, "state_sets"), Some(&TagValue::Long(2)));
+        assert_eq!(tag(resolve, "error"), Some(&TagValue::Bool(true)));
+        assert!(spans
+            .iter()
+            .any(|span| span.operation == "db.connection" && span.parent_id == resolve.span_id));
+    }
 }

rust/src/storage/python.rs

@@ -56,7 +56,9 @@ use pyo3::{
     },
     Bound, IntoPyObject, Py, PyAny, PyClassInitializer, PyErr, PyResult, Python,
 };
+use tracing::Instrument;
 
+use crate::logging::opentracing::record_error;
 use crate::{
     caches::{
         python::{PyCacheKey, RustCache},
@@ -213,6 +215,14 @@ impl RustStore {
         let ids: Vec<Arc<str>> = event_ids.into_iter().map(Arc::from).collect();
         Ok(async move {
             let mut stats = EventLoadStats::default();
+            let span = tracing::info_span!(
+                "events.load",
+                events = ids.len() as u64,
+                reload,
+                cached = tracing::field::Empty,
+                fetched = tracing::field::Empty,
+                error = tracing::field::Empty,
+            );
             let events = async {
                 let querier = PoolQuerier::checkout(&pool).await?;
                 let loader = EventLoader::new(&querier, Some(&store.caches.events));
@@ -222,7 +232,17 @@ impl RustStore {
                     loader.load(&ids, &mut stats).await
                 }
             }
+            .instrument(span.clone())
             .await;
+            match &events {
+                Ok(_) => {
+                    // Events another caller's load read count as cached:
+                    // this load did not read them.
+                    span.record("cached", stats.cached);
+                    span.record("fetched", stats.fetched);
+                }
+                Err(err) => record_error(&span, err),
+            }
             events
                 .map(|events| LoadedEvents {
                     events,

tests/logging/test_rust_tracing.py

@@ -676,3 +676,65 @@ class SyncTracingTestCase(_AgentTestCase):
         # Encoded on the reactor, under the request's Python span.
         self.assertEqual(encode.traceIdLow, entries.traceIdLow)
         self.assertIn("SyncRestServlet", self._python_ancestors(encode, python_by_id))
+
+
+class EventLoadTracingTestCase(_AgentTestCase):
+    """An event load in Rust reports `events.load` under the caller's Python
+    span, with the event fetched from the database.
+
+    Needs the native Rust Postgres driver:
+
+        SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.logging.test_rust_tracing
+    """
+
+    servlets = [
+        admin.register_servlets,
+        room.register_servlets,
+        login.register_servlets,
+    ]
+
+    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
+        hs = self.setup_test_homeserver()
+        if hs.get_datastores().main._rust_event_loader is None:
+            raise SkipTest("needs the native Rust Postgres driver")
+        return hs
+
+    def prepare_before_tracing(self) -> None:
+        self.user = self.register_user("user", "pass")
+        self.token = self.login("user", "pass")
+        room_id = self.helper.create_room_as(self.user, tok=self.token)
+        self.event_id = self.helper.send(room_id, "hello", tok=self.token)["event_id"]
+
+    def test_event_load_span(self) -> None:
+        store = self.hs.get_datastores().main
+        # Out of the cache, so that the load reads the database.
+        store._invalidate_local_get_event_cache(self.event_id)
+
+        tracer = opentracing.global_tracer()
+        with LoggingContext(name="ctx", server_name="test_server"):
+            with start_active_span("load", tracer=tracer) as scope:
+                self.get_success(store.get_event(self.event_id))
+
+        spans = self._received_spans()
+        trace_id = _signed(_context(scope.span).trace_id)
+        loads = [
+            s
+            for s in spans
+            if s.operationName == "events.load" and s.traceIdLow == trace_id
+        ]
+        self.assertEqual(len(loads), 1, [s.operationName for s in spans])
+        (load,) = loads
+        # Under the caller's span, through the store's own traced calls.
+        ancestors = self._python_ancestors(load, self._python_spans_by_id())
+        self.assertEqual(ancestors[-1], "load", ancestors)
+        tags = {t.key: t.vLong for t in load.tags if t.vType == ttypes.TagType.LONG}
+        self.assertEqual(tags["fetched"], 1)
+        self.assertEqual(tags["events"], 1)
+
+        # Its reads nest under it.
+        self.assertTrue(
+            any(
+                s.operationName == "db.connection" and s.parentSpanId == load.spanId
+                for s in spans
+            )
+        )