feat(sdk): Add a root span to our client

Author: poljar

crates/matrix-sdk/src/client/builder.rs

@@ -29,6 +29,11 @@ use ruma::{
 use thiserror::Error;
 #[cfg(not(target_arch = "wasm32"))]
 use tokio::sync::OnceCell;
+use tracing::{
+    debug,
+    field::{self, debug},
+    instrument, span, Level, Span,
+};
 use url::Url;
 
 use super::{Client, ClientInner};
@@ -88,10 +93,19 @@ pub struct ClientBuilder {
     appservice_mode: bool,
     server_versions: Option<Box<[MatrixVersion]>>,
     handle_refresh_tokens: bool,
+    root_span: Span,
 }
 
 impl ClientBuilder {
     pub(crate) fn new() -> Self {
+        let root_span = span!(
+            Level::INFO,
+            "matrix-sdk",
+            user_id = field::Empty,
+            device_id = field::Empty,
+            ed25519_key = field::Empty
+        );
+
         Self {
             homeserver_cfg: None,
             http_cfg: None,
@@ -101,6 +115,7 @@ impl ClientBuilder {
             appservice_mode: false,
             server_versions: None,
             handle_refresh_tokens: false,
+            root_span,
         }
     }
 
@@ -326,8 +341,12 @@ impl ClientBuilder {
     ///   server discovery request is made which can fail; if you didn't set
     ///   [`server_versions(false)`][Self::server_versions], that amounts to
     ///   another request that can fail
+    #[instrument(skip_all, parent = &self.root_span, target = "matrix_sdk::client", fields(homeserver))]
     pub async fn build(self) -> Result<Client, ClientBuildError> {
+        debug!("Starting to build the Client");
+
         let homeserver_cfg = self.homeserver_cfg.ok_or(ClientBuildError::MissingHomeserver)?;
+        Span::current().record("homeserver", debug(&homeserver_cfg));
 
         let inner_http_client = match self.http_cfg.unwrap_or_default() {
             #[allow(unused_mut)]
@@ -364,6 +383,8 @@ impl ClientBuilder {
         let homeserver = match homeserver_cfg {
             HomeserverConfig::Url(url) => url,
             HomeserverConfig::ServerName(server_name) => {
+                debug!("Trying to discover the homeserver");
+
                 let homeserver = homeserver_from_name(&server_name);
                 let well_known = http_client
                     .send(
@@ -387,6 +408,7 @@ impl ClientBuilder {
                 if let Some(proxy) = well_known.sliding_sync_proxy.map(|p| p.url) {
                     sliding_sync_proxy = Url::parse(&proxy).ok();
                 }
+                debug!(homserver_url = well_known.homeserver.base_url, "Discovered the homeserver");
 
                 well_known.homeserver.base_url
             }
@@ -421,7 +443,15 @@ impl ClientBuilder {
             refresh_token_lock: Mutex::new(Ok(())),
         });
 
-        Ok(Client { inner })
+        debug!("Done building the Client");
+
+        // We drop the root span here so it gets pushed to the subscribers, i.e. it gets
+        // only uploaded to a OpenTelemetry collector if the span gets dropped.
+        // We still want it around so other methods that get called by this
+        // client instance are connected to it, so we clone.
+        drop(self.root_span.clone());
+
+        Ok(Client { inner, root_span: self.root_span })
     }
 }
 

crates/matrix-sdk/src/client/login_builder.rs

@@ -142,6 +142,7 @@ impl LoginBuilder {
     /// Instead of calling this function and `.await`ing its return value, you
     /// can also `.await` the `LoginBuilder` directly.
     #[instrument(
+        parent = &self.client.root_span,
         target = "matrix_sdk::client",
         name = "login",
         skip_all,
@@ -278,7 +279,13 @@ where
     ///
     /// Instead of calling this function and `.await`ing its return value, you
     /// can also `.await` the `SsoLoginBuilder` directly.
-    #[instrument(target = "matrix_sdk::client", name = "login", skip_all, fields(method = "sso"))]
+    #[instrument(
+        parent = &self.client.root_span,
+        target = "matrix_sdk::client",
+        name = "login",
+        skip_all,
+        fields(method = "sso"),
+    )]
     pub async fn send(self) -> Result<login::v3::Response> {
         use std::{
             convert::Infallible,

crates/matrix-sdk/src/client/mod.rs

@@ -69,7 +69,7 @@ use serde::de::DeserializeOwned;
 use tokio::sync::OnceCell;
 #[cfg(feature = "e2e-encryption")]
 use tracing::error;
-use tracing::{debug, info, instrument};
+use tracing::{debug, field::display, info, instrument, trace, Instrument, Span};
 use url::Url;
 
 #[cfg(feature = "e2e-encryption")]
@@ -128,6 +128,7 @@ pub enum LoopCtrl {
 #[derive(Clone)]
 pub struct Client {
     pub(crate) inner: Arc<ClientInner>,
+    pub(crate) root_span: Span,
 }
 
 pub(crate) struct ClientInner {
@@ -1152,6 +1153,15 @@ impl Client {
             }
         }
 
+        self.root_span
+            .record("user_id", display(&response.user_id))
+            .record("device_id", display(&response.device_id));
+
+        #[cfg(feature = "e2e-encryption")]
+        if let Some(key) = self.encryption().ed25519_key().await {
+            self.root_span.record("ed25519_key", key);
+        }
+
         self.inner.base_client.receive_login_response(response).await?;
 
         Ok(())
@@ -1216,10 +1226,27 @@ impl Client {
     /// ```
     ///
     /// [`login`]: #method.login
+    #[instrument(skip_all, parent = &self.root_span)]
     pub async fn restore_session(&self, session: Session) -> Result<()> {
+        debug!("Restoring session");
+
         let (meta, tokens) = session.into_parts();
+
+        self.root_span
+            .record("user_id", display(&meta.user_id))
+            .record("device_id", display(&meta.device_id));
+
         self.base_client().set_session_tokens(tokens);
-        Ok(self.base_client().set_session_meta(meta).await?)
+        self.base_client().set_session_meta(meta).await?;
+
+        #[cfg(feature = "e2e-encryption")]
+        if let Some(key) = self.encryption().ed25519_key().await {
+            self.root_span.record("ed25519_key", key);
+        }
+
+        debug!("Done restoring session");
+
+        Ok(())
     }
 
     /// Refresh the access token.
@@ -1401,7 +1428,7 @@ impl Client {
     /// client.register(request).await;
     /// # })
     /// ```
-    #[instrument(skip_all)]
+    #[instrument(skip_all, parent = &self.root_span)]
     pub async fn register(
         &self,
         request: register::v3::Request,
@@ -1464,7 +1491,7 @@ impl Client {
     ///
     /// let response = client.sync_once(sync_settings).await.unwrap();
     /// # });
-    #[instrument(skip(self, definition))]
+    #[instrument(skip(self, definition), parent = &self.root_span)]
     pub async fn get_or_upload_filter(
         &self,
         filter_name: &str,
@@ -2157,7 +2184,7 @@ impl Client {
     ///     .await;
     /// })
     /// ```
-    #[instrument(skip(self, callback))]
+    #[instrument(skip_all, parent = &self.root_span)]
     pub async fn sync_with_callback<C>(
         &self,
         sync_settings: crate::config::SyncSettings,
@@ -2254,11 +2281,15 @@ impl Client {
         }
 
         loop {
+            trace!("Syncing");
             let result = self.sync_loop_helper(&mut sync_settings).await;
 
+            trace!("Running callback");
             if callback(result).await? == LoopCtrl::Break {
+                trace!("Callback told us to stop");
                 break;
             }
+            trace!("Done running callback");
 
             Client::delay_sync(&mut last_sync_time).await
         }
@@ -2308,7 +2339,7 @@ impl Client {
     ///
     /// # anyhow::Ok(()) });
     /// ```
-    #[instrument(skip(self))]
+    #[instrument(skip(self), parent = &self.root_span)]
     pub async fn sync_stream(
         &self,
         mut sync_settings: crate::config::SyncSettings,
@@ -2319,9 +2350,11 @@ impl Client {
             sync_settings.token = self.sync_token().await;
         }
 
+        let parent_span = Span::current();
+
         async_stream::stream! {
             loop {
-                yield self.sync_loop_helper(&mut sync_settings).await;
+                yield self.sync_loop_helper(&mut sync_settings).instrument(parent_span.clone()).await;
 
                 Client::delay_sync(&mut last_sync_time).await
             }