tracing: Add --trace-json for visualising traces and add more spans

Cargo.toml

@@ -188,6 +188,7 @@ tokio-util = { version = "0.7", features = ["codec"] }
 toml = "0.8"
 tower = "0.4"
 tracing = { version = "0.1", features = ["attributes", "log"] }
+tracing-chrome = "0.7.2"
 tracing-subscriber = { version = "0.3", features = ["time", "env-filter"] }
 try-traits = "0.1"
 tungstenite = "0.26.1"

crates/common/mqtt_channel/Cargo.toml

@@ -18,6 +18,7 @@ rumqttc = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt", "time", "rt-multi-thread"] }
+tracing = { workspace = true }
 zeroize = { workspace = true }
 
 [dev-dependencies]

crates/common/mqtt_channel/src/connection.rs

@@ -8,8 +8,6 @@ use futures::channel::mpsc;
 use futures::channel::oneshot;
 use futures::SinkExt;
 use futures::StreamExt;
-use log::error;
-use log::info;
 use rumqttc::AsyncClient;
 use rumqttc::Event;
 use rumqttc::EventLoop;
@@ -21,6 +19,10 @@ use std::time::Duration;
 use tokio::sync::OwnedSemaphorePermit;
 use tokio::sync::Semaphore;
 use tokio::time::sleep;
+use tracing::error;
+use tracing::info;
+use tracing::instrument;
+use tracing::trace;
 
 /// A connection to some MQTT server
 pub struct Connection {
@@ -123,6 +125,7 @@ impl Connection {
         let _ = self.pub_done.await;
     }
 
+    #[instrument(level = "trace", skip(message_sender, error_sender))]
     async fn open(
         config: &Config,
         mut message_sender: mpsc::UnboundedSender<MqttMessage>,
@@ -201,6 +204,7 @@ impl Connection {
         Ok((mqtt_client, event_loop))
     }
 
+    #[instrument(skip_all, level = "trace")]
     async fn receiver_loop(
         mqtt_client: AsyncClient,
         config: Config,
@@ -306,17 +310,21 @@ impl Connection {
         let _ = message_sender.close().await;
         let _ = error_sender.close().await;
         let _ = done.send(());
+        trace!("terminating receiver loop");
         Ok(())
     }
 
+    #[instrument(skip_all, level = "trace")]
     async fn sender_loop(
         mqtt_client: AsyncClient,
         mut messages_receiver: mpsc::UnboundedReceiver<MqttMessage>,
         mut error_sender: mpsc::UnboundedSender<MqttError>,
         last_will: Option<MqttMessage>,
         _disconnect_permit: OwnedSemaphorePermit,
     ) {
+        trace!("waiting for message");
         while let Some(message) = messages_receiver.next().await {
+            trace!(msg = ?message, "received message");
             let payload = Vec::from(message.payload_bytes());
             if let Err(err) = mqtt_client
                 .publish(message.topic, message.qos, message.retain, payload)

crates/common/tedge_config/Cargo.toml

@@ -36,6 +36,7 @@ tedge_utils = { workspace = true, features = ["timestamp"] }
 thiserror = { workspace = true }
 toml = { workspace = true }
 tracing = { workspace = true }
+tracing-chrome = { workspace = true }
 tracing-subscriber = { workspace = true }
 url = { workspace = true }
 which = { workspace = true }

crates/common/tedge_config/src/cli.rs

@@ -41,6 +41,10 @@ pub struct LogConfigArgs {
     #[clap(long, global = true)]
     #[clap(add(ArgValueCandidates::new(log_level_completions)))]
     pub log_level: Option<tracing::Level>,
+
+    /// Emit a JSON trace file that can be opened by ui.perfetto.dev
+    #[clap(long, global = true)]
+    pub trace_json: bool,
 }
 
 fn log_level_completions() -> Vec<CompletionCandidate> {

crates/common/tedge_config/src/system_services/log_config.rs

@@ -1,9 +1,14 @@
 use camino::Utf8Path;
+use tracing_chrome::ChromeLayerBuilder;
+use tracing_subscriber::fmt::format::FmtSpan;
+use tracing_subscriber::prelude::*;
+use tracing_subscriber::EnvFilter;
 
 use crate::cli::LogConfigArgs;
 use crate::system_services::SystemConfig;
 use crate::system_services::SystemServiceError;
 use std::io::IsTerminal;
+use std::path::PathBuf;
 use std::str::FromStr;
 
 /// Configures and enables logging taking into account flags, env variables and file config.
@@ -18,34 +23,75 @@ pub fn log_init(
     sname: &str,
     flags: &LogConfigArgs,
     config_dir: &Utf8Path,
-) -> Result<(), SystemServiceError> {
-    let subscriber = tracing_subscriber::fmt()
+) -> Result<Option<tracing_chrome::FlushGuard>, SystemServiceError> {
+    let print_file_and_line = std::env::var("RUST_LOG").is_ok();
+    let file_level = get_log_level(sname, config_dir)?;
+
+    let filter_layer = filter_layer(flags, file_level);
+
+    // print code location if RUST_LOG is used
+    let fmt_layer = tracing_subscriber::fmt::layer()
         .with_writer(std::io::stderr)
         .with_ansi(std::io::stderr().is_terminal())
-        .with_timer(tracing_subscriber::fmt::time::UtcTime::rfc_3339());
+        .with_timer(tracing_subscriber::fmt::time::UtcTime::rfc_3339())
+        .with_span_events(FmtSpan::NONE)
+        .with_file(print_file_and_line)
+        .with_line_number(print_file_and_line)
+        .with_filter(filter_layer);
+
+    // chrome layer if `--trace-json` OR `TEDGE_TRACE_DIR` envvar present
+
+    // `TEDGE_TRACE_DIR` is helpful if we want to enable tracing for all the services; turns out using a config is
+    // tricky because we set up logging before loading the config to be able to log config loading and parsing logic! So
+    // we can either 1) try to change and reinstall the subscriber after we've loaded desired logging config, which i'm
+    // sure is tricky, or 2) the simpler and selected approach of just bypassing the config and using an envvar.
+    let trace_json_enabled = std::env::var("TEDGE_TRACE_DIR").is_ok() || flags.trace_json;
+    let trace_filename = format!(
+        "{sname}-trace-{}.json",
+        std::time::SystemTime::UNIX_EPOCH
+            .elapsed()
+            .unwrap()
+            .as_micros()
+    );
+    let trace_path_file: PathBuf = match std::env::var("TEDGE_TRACE_DIR") {
+        Ok(dir) => [&dir, &trace_filename].iter().collect(),
+        Err(_) => ["./", &trace_filename].iter().collect(),
+    };
+
+    let (chrome_layer, guard) = if trace_json_enabled {
+        let (chrome_layer, guard) = ChromeLayerBuilder::new()
+            .file(trace_path_file)
+            .include_args(true)
+            .build();
+        (Some(chrome_layer), Some(guard))
+    } else {
+        (None, None)
+    };
+
+    tracing_subscriber::registry()
+        .with(chrome_layer)
+        .with(fmt_layer)
+        .init();
+
+    Ok(guard)
+}
 
+fn filter_layer(flags: &LogConfigArgs, file_level: tracing::Level) -> EnvFilter {
+    // 1. use level from flags if they're present
     let log_level = flags
         .log_level
         .or(flags.debug.then_some(tracing::Level::DEBUG));
-
     if let Some(log_level) = log_level {
-        subscriber.with_max_level(log_level).init();
-        return Ok(());
+        return EnvFilter::new(log_level.to_string());
     }
 
+    // 2. if not, use RUST_LOG
     if std::env::var("RUST_LOG").is_ok() {
-        subscriber
-            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-            .with_file(true)
-            .with_line_number(true)
-            .init();
-        return Ok(());
+        return EnvFilter::from_default_env();
     }
 
-    let log_level = get_log_level(sname, config_dir)?;
-    subscriber.with_max_level(log_level).init();
-
-    Ok(())
+    // 3. if not, use file content (info if no logging preferences in file)
+    EnvFilter::new(file_level.to_string())
 }
 
 pub fn get_log_level(

crates/core/tedge_actors/Cargo.toml

@@ -25,6 +25,7 @@ tokio = { workspace = true, default_features = false, features = [
     "macros",
     "time",
 ] }
+tracing = { workspace = true }
 
 [dev-dependencies]
 env_logger = { workspace = true } # TODO: remove me

crates/core/tedge_actors/src/channels.rs

@@ -5,6 +5,8 @@ use async_trait::async_trait;
 use futures::channel::mpsc;
 use futures::channel::oneshot;
 use futures::SinkExt;
+use tracing::debug;
+use tracing::instrument;
 
 /// A sender of messages of type `M`
 ///
@@ -75,13 +77,24 @@ impl<M: Message, N: Message + Into<M>> CloneSender<N> for DynSender<M> {
 #[async_trait]
 impl<M: Message, N: Message + Into<M>> Sender<N> for mpsc::Sender<M> {
     async fn send(&mut self, message: N) -> Result<(), ChannelError> {
-        Ok(SinkExt::send(&mut self, message.into()).await?)
+        let message = message.into();
+        if let Err(err) = self.try_send(message) {
+            if err.is_full() {
+                debug!("Sender is full");
+            }
+            let message = err.into_inner();
+            SinkExt::send(&mut self, message).await?;
+            debug!("Blocked send completed");
+        }
+
+        Ok(())
     }
 }
 
 /// An `mpsc::UnboundedSender<M>` is a `DynSender<N>` provided `N` implements `Into<M>`
 #[async_trait]
 impl<M: Message, N: Message + Into<M>> Sender<N> for mpsc::UnboundedSender<M> {
+    #[instrument(name = "mpsc::UnboundedSender::send", skip_all)]
     async fn send(&mut self, message: N) -> Result<(), ChannelError> {
         Ok(SinkExt::send(&mut self, message.into()).await?)
     }

crates/core/tedge_actors/src/message_boxes.rs

@@ -95,6 +95,7 @@ use futures::channel::mpsc;
 use futures::StreamExt;
 use log::debug;
 use std::fmt::Debug;
+use tracing::instrument;
 
 #[async_trait]
 pub trait MessageReceiver<Input> {
@@ -163,13 +164,15 @@ impl<Input: Debug> LoggingReceiver<Input> {
 impl<Input: Send + Debug> MessageReceiver<Input> for LoggingReceiver<Input> {
     async fn try_recv(&mut self) -> Result<Option<Input>, RuntimeRequest> {
         let message = self.receiver.try_recv().await;
-        debug!(target: &self.name, "recv {:?}", message);
+        debug!("recv {:?}", message);
         message
     }
 
+    #[instrument(name = "LoggingReceiver::recv", skip(self), fields(name = self.name))]
     async fn recv(&mut self) -> Option<Input> {
+        debug!("attempting recv");
         let message = self.receiver.recv().await;
-        debug!(target: &self.name, "recv {:?}", message);
+        debug!("recv");
         message
     }
 
@@ -202,6 +205,7 @@ impl<Output> LoggingSender<Output> {
 
 #[async_trait]
 impl<Output: Message> Sender<Output> for LoggingSender<Output> {
+    #[instrument(name = "LoggingSender::send", skip(self, message), fields(name = self.name))]
     async fn send(&mut self, message: Output) -> Result<(), ChannelError> {
         log_message_sent(&self.name, &message);
         self.sender.send(message).await

crates/core/tedge_actors/src/run_actor.rs

@@ -1,3 +1,6 @@
+use tracing::debug_span;
+use tracing::Instrument;
+
 use crate::Actor;
 use crate::Builder;
 use crate::DynSender;
@@ -39,7 +42,11 @@ impl RunActor {
     }
 
     pub async fn run(self) -> Result<(), RuntimeError> {
-        self.actor.run_boxed().await
+        let name = self.actor.name().to_string();
+        self.actor
+            .run_boxed()
+            .instrument(debug_span!("Actor::run", actor_name = name))
+            .await
     }
 }