diff --git a/libsignal-service-actix/Cargo.toml b/libsignal-service-actix/Cargo.toml index 782816358..315b36536 100644 --- a/libsignal-service-actix/Cargo.toml +++ b/libsignal-service-actix/Cargo.toml @@ -19,6 +19,7 @@ mpart-async = "0.6" serde_json = "1.0" futures = "0.3" tracing = "0.1" +tracing-futures = "0.2" bytes = "1" rustls = "0.21" rustls-pemfile = "0.3" diff --git a/libsignal-service-actix/src/push_service.rs b/libsignal-service-actix/src/push_service.rs index 0ef21ae4f..e6b30e31e 100644 --- a/libsignal-service-actix/src/push_service.rs +++ b/libsignal-service-actix/src/push_service.rs @@ -13,6 +13,7 @@ use libsignal_service::{ websocket::SignalWebSocket, }; use serde::{Deserialize, Serialize}; +use tracing_futures::Instrument; use crate::websocket::AwcWebSocket; @@ -85,6 +86,7 @@ impl AwcPushService { }) } + #[instrument(name = "extracting error", skip(response))] async fn from_response( response: &mut ClientResponse, ) -> Result<(), ServiceError> @@ -201,8 +203,7 @@ impl PushService for AwcPushService { }, })?; - let _span = - tracing::debug_span!("processing response", ?response).entered(); + let span = tracing::debug_span!("processing response", ?response); Self::from_response(&mut response).await?; @@ -646,6 +647,7 @@ impl PushService for AwcPushService { additional_headers: &[(&str, &str)], credentials: Option, ) -> Result { + let span = tracing::debug_span!("websocket"); let (ws, stream) = AwcWebSocket::with_client( &mut self.client, self.cfg.base_url(Endpoint::Service), @@ -653,13 +655,14 @@ impl PushService for AwcPushService { additional_headers, credentials.as_ref(), ) + .instrument(span.clone()) .await?; let (ws, task) = SignalWebSocket::from_socket( ws, stream, keep_alive_path.to_owned(), ); - actix_rt::spawn(task); + actix_rt::spawn(task.instrument(span)); Ok(ws) } } diff --git a/libsignal-service-hyper/Cargo.toml b/libsignal-service-hyper/Cargo.toml index 547d4359e..d874aac8d 100644 --- a/libsignal-service-hyper/Cargo.toml +++ b/libsignal-service-hyper/Cargo.toml @@ -13,6 +13,7 @@ async-trait = "0.1" bytes = "1.0" futures = "0.3" tracing = "0.1" +tracing-futures = "0.2" mpart-async = "0.6" serde = "1.0" serde_json = "1.0" diff --git a/libsignal-service/Cargo.toml b/libsignal-service/Cargo.toml index 03bc31a1e..45285513e 100644 --- a/libsignal-service/Cargo.toml +++ b/libsignal-service/Cargo.toml @@ -35,6 +35,7 @@ url = { version = "2.1", features = ["serde"] } uuid = { version = "1", features = ["serde"] } tracing = "0.1" +tracing-futures = "0.2" [build-dependencies] prost-build = "0.10" diff --git a/libsignal-service/src/account_manager.rs b/libsignal-service/src/account_manager.rs index de18bde59..cc2e2d1a2 100644 --- a/libsignal-service/src/account_manager.rs +++ b/libsignal-service/src/account_manager.rs @@ -13,6 +13,7 @@ use libsignal_protocol::{ use prost::Message; use serde::{Deserialize, Serialize}; use sha2::Sha256; +use tracing_futures::Instrument; use zkgroup::profiles::ProfileKey; use crate::pre_keys::KyberPreKeyEntity; @@ -102,29 +103,26 @@ impl AccountManager { pq_pre_keys_offset_id: u32, use_last_resort_key: bool, ) -> Result<(u32, u32, u32), ServiceError> { - let prekey_status = { - let _span = tracing::span!( + let prekey_status = match self + .service + .get_pre_key_status(ServiceIdType::AccountIdentity) + .instrument(tracing::span!( tracing::Level::DEBUG, "Fetching pre key status" - ) - .entered(); - match self - .service - .get_pre_key_status(ServiceIdType::AccountIdentity) - .await - { - Ok(status) => status, - Err(ServiceError::Unauthorized) => { - tracing::info!("Got Unauthorized when fetching pre-key status. Assuming first installment."); - // Additionally, the second PUT request will fail if this really comes down to an - // authorization failure. - crate::push_service::PreKeyStatus { - count: 0, - pq_count: 0, - } - }, - Err(e) => return Err(e), - } + )) + .await + { + Ok(status) => status, + Err(ServiceError::Unauthorized) => { + tracing::info!("Got Unauthorized when fetching pre-key status. Assuming first installment."); + // Additionally, the second PUT request will fail if this really comes down to an + // authorization failure. + crate::push_service::PreKeyStatus { + count: 0, + pq_count: 0, + } + }, + Err(e) => return Err(e), }; tracing::trace!("Remaining pre-keys on server: {:?}", prekey_status); @@ -140,11 +138,11 @@ impl AccountManager { } let pre_key_state = { - let _span = - tracing::span!(tracing::Level::DEBUG, "Generating pre keys") - .entered(); + let span = + tracing::span!(tracing::Level::DEBUG, "Generating pre keys"); + let identity_key_pair = - protocol_store.get_identity_key_pair().await?; + protocol_store.get_identity_key_pair().instrument(tracing::trace_span!(parent: &span, "get identity key pair")).await?; let mut pre_key_entities = vec![]; let mut pq_pre_key_entities = vec![]; @@ -159,7 +157,7 @@ impl AccountManager { let pre_key_record = PreKeyRecord::new(pre_key_id, &key_pair); protocol_store .save_pre_key(pre_key_id, &pre_key_record) - .await?; + .instrument(tracing::trace_span!(parent: &span, "save pre key", ?pre_key_id)).await?; // TODO: Shouldn't this also remove the previous pre-keys from storage? // I think we might want to update the storage, and then sync the storage to the // server. @@ -180,7 +178,7 @@ impl AccountManager { )?; protocol_store .save_kyber_pre_key(pre_key_id, &pre_key_record) - .await?; + .instrument(tracing::trace_span!(parent: &span, "save kyber pre key", ?pre_key_id)).await?; // TODO: Shouldn't this also remove the previous pre-keys from storage? // I think we might want to update the storage, and then sync the storage to the // server. @@ -214,7 +212,7 @@ impl AccountManager { next_signed_pre_key_id.into(), &signed_prekey_record, ) - .await?; + .instrument(tracing::trace_span!(parent: &span, "save signed pre key", signed_pre_key_id = ?next_signed_pre_key_id)).await?; PreKeyState { pre_keys: pre_key_entities, @@ -240,17 +238,13 @@ impl AccountManager { } }; - { - let _span = - tracing::span!(tracing::Level::DEBUG, "Uploading pre keys") - .entered(); - self.service - .register_pre_keys( - ServiceIdType::AccountIdentity, - pre_key_state, - ) - .await?; - } + self.service + .register_pre_keys(ServiceIdType::AccountIdentity, pre_key_state) + .instrument(tracing::span!( + tracing::Level::DEBUG, + "Uploading pre keys" + )) + .await?; Ok(( pre_keys_offset_id + PRE_KEY_BATCH_SIZE, diff --git a/libsignal-service/src/sender.rs b/libsignal-service/src/sender.rs index 658c30c2b..09c56eb02 100644 --- a/libsignal-service/src/sender.rs +++ b/libsignal-service/src/sender.rs @@ -7,6 +7,7 @@ use libsignal_protocol::{ }; use rand::{CryptoRng, Rng}; use tracing::{info, trace}; +use tracing_futures::Instrument; use uuid::Uuid; use crate::{ @@ -195,19 +196,16 @@ where }); // Request upload attributes - let attrs = { - let _span = - tracing::trace_span!("requesting upload attributes").entered(); - self.identified_ws - .get_attachment_v2_upload_attributes() - .await? - }; - let (id, digest) = { - let _span = tracing::trace_span!("Uploading attachment").entered(); - self.service - .upload_attachment(&attrs, &mut std::io::Cursor::new(&contents)) - .await? - }; + let attrs = self + .identified_ws + .get_attachment_v2_upload_attributes() + .instrument(tracing::trace_span!("requesting upload attributes")) + .await?; + let (id, digest) = self + .service + .upload_attachment(&attrs, &mut std::io::Cursor::new(&contents)) + .instrument(tracing::trace_span!("Uploading attachment")) + .await?; Ok(AttachmentPointer { content_type: Some(spec.content_type), @@ -758,10 +756,10 @@ where } } - let _span = tracing::trace_span!("encrypting message").entered(); let message = self .cipher .encrypt(&recipient_protocol_address, unidentified_access, content) + .instrument(tracing::trace_span!("encrypting message")) .await?; Ok(message)