add tracing instruments

Author: br0kenpixel

src/logging.rs

@@ -1,6 +1,10 @@
 use crate::error::Error;
 use std::time::SystemTime;
-use tracing_subscriber::fmt::{format::Writer, time::FormatTime};
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::fmt::{
+    format::{FmtSpan, Writer},
+    time::FormatTime,
+};
 
 struct DateTimeFormatter;
 
@@ -12,11 +16,19 @@ impl FormatTime for DateTimeFormatter {
     }
 }
 
-pub fn setup() -> Result<(), Error> {
+pub fn setup(debug: bool) -> Result<(), Error> {
+    let level = if cfg!(debug_assertions) | debug {
+        LevelFilter::DEBUG
+    } else {
+        LevelFilter::INFO
+    };
+
     let subscriber = tracing_subscriber::fmt()
         .compact()
         .with_timer(DateTimeFormatter)
         .with_target(false)
+        .with_max_level(level)
+        .with_span_events(FmtSpan::CLOSE)
         .finish();
     tracing::subscriber::set_global_default(subscriber)?;
 

src/main.rs

@@ -19,7 +19,7 @@ pub static MIGRATOR: Migrator = sqlx::migrate!();
 async fn main() -> Result<(), error::Error> {
     let args = cli::Cli::parse();
 
-    logging::setup()?;
+    logging::setup(args.debug)?;
 
     info!("PixelWeather Server v{}", env!("CARGO_PKG_VERSION"));
     debug!("Arguments: {args:?}");

src/server/client.rs

@@ -64,6 +64,7 @@ impl<S> Client<S> {
         message.take_request().ok_or(Error::NotRequest)
     }
 
+    #[tracing::instrument(name = "Client::send_message()", skip(self), level = "debug")]
     async fn send_message(&mut self, msg: Message) -> Result<()> {
         // Serialize the message.
         let raw = msg.serialize();
@@ -87,6 +88,13 @@ impl<S> Client<S> {
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "Client::receive_message()",
+        skip(self),
+        level = "debug",
+        err,
+        ret
+    )]
     async fn receive_message(&mut self) -> Result<Message> {
         // First read the message size.
         self.stream
@@ -183,6 +191,13 @@ impl Client<Unathenticated> {
         }
     }
 
+    #[tracing::instrument(
+        name = "Client<Unathenticated>::receive_handshake()",
+        skip(self),
+        level = "debug",
+        err,
+        ret
+    )]
     async fn receive_handshake(&mut self) -> Result<Mac> {
         let req = self.receive_request().await?;
         let Request::Handshake { mac } = req else {
@@ -230,6 +245,11 @@ impl Client<Authenticated> {
         self.state.update_state = UpdateState::UpToDate;
     }
 
+    #[tracing::instrument(
+        name = "Client<Authenticated>::store_update_check_result()",
+        skip(self, blob),
+        level = "debug"
+    )]
     pub fn store_update_check_result(
         &mut self,
         current_version: Version,

src/server/client_handle.rs

@@ -5,10 +5,10 @@ use super::{
     rate_limit::RateLimiter,
 };
 use crate::error::Error;
-use tracing::{debug, error, warn};
 use pwmp_client::pwmp_msg::{request::Request, response::Response};
 use std::{io::Read, net::SocketAddr, sync::Arc, time::Duration};
 use tokio::{net::TcpStream, time::timeout};
+use tracing::{debug, error, warn};
 
 #[allow(clippy::needless_pass_by_value)]
 pub async fn handle_client(
@@ -78,6 +78,7 @@ pub async fn handle_client(
     Ok(())
 }
 
+#[tracing::instrument(name = "handle_request()", skip_all, level = "debug", err, ret)]
 async fn handle_request(
     req: Request,
     client: &mut Client<Authenticated>,

src/server/db.rs

@@ -10,6 +10,7 @@ use sqlx::{
     Pool, Postgres,
     postgres::{PgConnectOptions, PgPoolOptions, PgSslMode},
 };
+use std::fmt::Debug;
 
 pub type NodeId = i32;
 pub type MeasurementId = i32;
@@ -25,6 +26,7 @@ macro_rules! query {
 }
 
 impl DatabaseClient {
+    #[tracing::instrument(name = "DatabaseClient::init()", level = "debug", err, skip_all)]
     pub async fn new(config: &Config) -> Result<Self, Error> {
         let mut opts = PgConnectOptions::new()
             .host(&config.database.host)
@@ -45,6 +47,12 @@ impl DatabaseClient {
         Ok(Self(pool))
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::setup_timezone()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn setup_timezone(&self, tz: &str) -> Result<(), Error> {
         if !self.validate_timezone(tz).await? {
             return Err(Error::InvalidTimeZone(tz.into()));
@@ -57,6 +65,13 @@ impl DatabaseClient {
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::authorize_device()",
+        level = "debug",
+        skip(self),
+        err,
+        ret
+    )]
     pub async fn authorize_device(&self, mac: &Mac) -> Result<Option<NodeId>, Error> {
         let mac = mac.to_string();
 
@@ -68,6 +83,12 @@ impl DatabaseClient {
         }
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::create_notification()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn create_notification(&self, node_id: NodeId, content: &str) -> Result<(), Error> {
         query!(
             self.0,
@@ -80,6 +101,13 @@ impl DatabaseClient {
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::get_settings()",
+        level = "debug",
+        skip(self),
+        err,
+        ret
+    )]
     pub async fn get_settings(&self, node_id: NodeId) -> Result<Option<NodeSettings>, Error> {
         let result = query!(
             self.0,
@@ -101,6 +129,13 @@ impl DatabaseClient {
         }
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::post_results()",
+        level = "debug",
+        skip(self),
+        err,
+        ret
+    )]
     pub async fn post_results(
         &self,
         node: NodeId,
@@ -126,6 +161,12 @@ impl DatabaseClient {
         Ok(result.id)
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::post_stats()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn post_stats(
         &self,
         measurement: MeasurementId,
@@ -146,12 +187,24 @@ impl DatabaseClient {
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::run_migrations()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn run_migrations(&self) -> Result<(), Error> {
         crate::MIGRATOR.run(&self.0).await?;
 
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::check_os_update()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn check_os_update(
         &self,
         node: NodeId,
@@ -183,6 +236,13 @@ impl DatabaseClient {
         }
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::send_os_update_stat()",
+        level = "debug",
+        skip(self),
+        err,
+        ret
+    )]
     pub async fn send_os_update_stat(
         &self,
         node_id: NodeId,
@@ -208,6 +268,12 @@ impl DatabaseClient {
         Ok(result.id)
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::mark_os_update_stat()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     pub async fn mark_os_update_stat(&self, node_id: NodeId, success: bool) -> Result<(), Error> {
         let last_update_id = self.get_last_os_update_stat(node_id).await?;
         query!(
@@ -221,6 +287,7 @@ impl DatabaseClient {
         Ok(())
     }
 
+    #[tracing::instrument(name = "DatabaseClient::erase()", level = "debug", skip(self), err)]
     pub async fn erase(&self, content_only: bool, keep_devices: bool) -> Result<(), Error> {
         if content_only {
             if keep_devices {
@@ -247,6 +314,13 @@ impl DatabaseClient {
         Ok(())
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::get_last_os_update_stat()",
+        level = "debug",
+        skip(self),
+        err,
+        ret
+    )]
     async fn get_last_os_update_stat(&self, node_id: NodeId) -> Result<UpdateStatId, Error> {
         Ok(query!(
             self.0,
@@ -257,6 +331,12 @@ impl DatabaseClient {
         .id)
     }
 
+    #[tracing::instrument(
+        name = "DatabaseClient::get_supported_time_zones()",
+        level = "debug",
+        skip(self),
+        err
+    )]
     async fn get_supported_time_zones(&self) -> Result<Vec<String>, Error> {
         let results = query!(self.0, "queries/get_tz_names.sql", fetch_all,)?;
         let names: Vec<String> = results
@@ -267,7 +347,14 @@ impl DatabaseClient {
         Ok(names)
     }
 
-    async fn validate_timezone<S: PartialEq<String>>(&self, tz: S) -> Result<bool, Error> {
+    #[tracing::instrument(
+        name = "DatabaseClient::validate_timezone()",
+        level = "debug",
+        skip(self, tz),
+        err,
+        ret /* this will print `tz` too */
+    )]
+    async fn validate_timezone<S: PartialEq<String> + Debug>(&self, tz: S) -> Result<bool, Error> {
         let supported = self.get_supported_time_zones().await?;
 
         Ok(supported.iter().any(|candidate| tz.eq(candidate)))