fix(log): adding span metadata to tokio spawned futures (#4118)

ShankarSinghC · web-flow · commit 070622125f49 · 2024-03-27T06:16:22.000Z
diff --git a/crates/common_utils/src/macros.rs b/crates/common_utils/src/macros.rs
@@ -47,13 +47,6 @@ macro_rules! newtype {
     };
 }
 
-#[macro_export]
-macro_rules! async_spawn {
-    ($t:block) => {
-        tokio::spawn(async move { $t });
-    };
-}
-
 /// Use this to ensure that the corresponding
 /// openapi route has been implemented in the openapi crate
 #[macro_export]
diff --git a/crates/diesel_models/src/lib.rs b/crates/diesel_models/src/lib.rs
@@ -24,7 +24,6 @@ pub mod gsm;
 #[cfg(feature = "kv_store")]
 pub mod kv;
 pub mod locker_mock_up;
-pub mod macros;
 pub mod mandate;
 pub mod merchant_account;
 pub mod merchant_connector_account;
diff --git a/crates/diesel_models/src/macros.rs b/crates/diesel_models/src/macros.rs
diff --git a/crates/drainer/src/handler.rs b/crates/drainer/src/handler.rs
@@ -1,5 +1,6 @@
 use std::sync::{atomic, Arc};
 
+use router_env::tracing::Instrument;
 use tokio::{
     sync::{mpsc, oneshot},
     time::{self, Duration},
@@ -68,13 +69,16 @@ impl Handler {
         while self.running.load(atomic::Ordering::SeqCst) {
             metrics::DRAINER_HEALTH.add(&metrics::CONTEXT, 1, &[]);
             if self.store.is_stream_available(stream_index).await {
-                tokio::spawn(drainer_handler(
-                    self.store.clone(),
-                    stream_index,
-                    self.conf.max_read_count,
-                    self.active_tasks.clone(),
-                    jobs_picked.clone(),
-                ));
+                let _task_handle = tokio::spawn(
+                    drainer_handler(
+                        self.store.clone(),
+                        stream_index,
+                        self.conf.max_read_count,
+                        self.active_tasks.clone(),
+                        jobs_picked.clone(),
+                    )
+                    .in_current_span(),
+                );
             }
             stream_index = utils::increment_stream_index(
                 (stream_index, jobs_picked.clone()),
@@ -116,10 +120,12 @@ impl Handler {
         let redis_conn_clone = self.store.redis_conn.clone();
 
         // Spawn a task to monitor if redis is down or not
-        tokio::spawn(async move { redis_conn_clone.on_error(redis_error_tx).await });
+        let _task_handle = tokio::spawn(
+            async move { redis_conn_clone.on_error(redis_error_tx).await }.in_current_span(),
+        );
 
         //Spawns a task to send shutdown signal if redis goes down
-        tokio::spawn(redis_error_receiver(redis_error_rx, tx));
+        let _task_handle = tokio::spawn(redis_error_receiver(redis_error_rx, tx).in_current_span());
 
         Ok(())
     }
diff --git a/crates/drainer/src/lib.rs b/crates/drainer/src/lib.rs
@@ -18,7 +18,10 @@ use common_utils::signals::get_allowed_signals;
 use diesel_models::kv;
 use error_stack::{IntoReport, ResultExt};
 use hyperswitch_interfaces::secrets_interface::secret_state::RawSecret;
-use router_env::{instrument, tracing};
+use router_env::{
+    instrument,
+    tracing::{self, Instrument},
+};
 use tokio::sync::mpsc;
 
 pub(crate) type Settings = crate::settings::Settings<RawSecret>;
@@ -39,7 +42,8 @@ pub async fn start_drainer(store: Arc<Store>, conf: DrainerSettings) -> errors::
                 "Failed while getting allowed signals".to_string(),
             ))?;
     let handle = signal.handle();
-    let task_handle = tokio::spawn(common_utils::signals::signal_handler(signal, tx.clone()));
+    let task_handle =
+        tokio::spawn(common_utils::signals::signal_handler(signal, tx.clone()).in_current_span());
 
     let handler_clone = drainer_handler.clone();
 
diff --git a/crates/router/src/bin/scheduler.rs b/crates/router/src/bin/scheduler.rs
@@ -16,7 +16,10 @@ use router::{
     services::{self, api},
     workflows,
 };
-use router_env::{instrument, tracing};
+use router_env::{
+    instrument,
+    tracing::{self, Instrument},
+};
 use scheduler::{
     consumer::workflows::ProcessTrackerWorkflow, errors::ProcessTrackerError,
     workflows::ProcessTrackerWorkflows, SchedulerAppState,
@@ -49,10 +52,9 @@ async fn main() -> CustomResult<(), ProcessTrackerError> {
     .await;
     // channel to shutdown scheduler gracefully
     let (tx, rx) = mpsc::channel(1);
-    tokio::spawn(router::receiver_for_error(
-        redis_shutdown_signal_rx,
-        tx.clone(),
-    ));
+    let _task_handle = tokio::spawn(
+        router::receiver_for_error(redis_shutdown_signal_rx, tx.clone()).in_current_span(),
+    );
 
     #[allow(clippy::expect_used)]
     let scheduler_flow_str =
@@ -81,10 +83,13 @@ async fn main() -> CustomResult<(), ProcessTrackerError> {
     .await
     .expect("Failed to create the server");
 
-    tokio::spawn(async move {
-        let _ = web_server.await;
-        logger::error!("The health check probe stopped working!");
-    });
+    let _task_handle = tokio::spawn(
+        async move {
+            let _ = web_server.await;
+            logger::error!("The health check probe stopped working!");
+        }
+        .in_current_span(),
+    );
 
     logger::debug!(startup_config=?state.conf);
 
diff --git a/crates/router/src/core/payments/flows/authorize_flow.rs b/crates/router/src/core/payments/flows/authorize_flow.rs
@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use error_stack;
+use router_env::tracing::Instrument;
 
 use super::{ConstructFlowSpecificData, Feature};
 use crate::{
@@ -125,26 +126,32 @@ impl Feature<api::Authorize, types::PaymentsAuthorizeData> for types::PaymentsAu
                 let state = state.clone();
 
                 logger::info!("Call to save_payment_method in locker");
-                tokio::spawn(async move {
-                    logger::info!("Starting async call to save_payment_method in locker");
-
-                    let result = Box::pin(tokenization::save_payment_method(
-                        &state,
-                        &connector,
-                        response,
-                        &maybe_customer,
-                        &merchant_account,
-                        self.request.payment_method_type,
-                        &key_store,
-                        Some(resp.request.amount),
-                        Some(resp.request.currency),
-                    ))
-                    .await;
-
-                    if let Err(err) = result {
-                        logger::error!("Asynchronously saving card in locker failed : {:?}", err);
+                let _task_handle = tokio::spawn(
+                    async move {
+                        logger::info!("Starting async call to save_payment_method in locker");
+
+                        let result = Box::pin(tokenization::save_payment_method(
+                            &state,
+                            &connector,
+                            response,
+                            &maybe_customer,
+                            &merchant_account,
+                            self.request.payment_method_type,
+                            &key_store,
+                            Some(resp.request.amount),
+                            Some(resp.request.currency),
+                        ))
+                        .await;
+
+                        if let Err(err) = result {
+                            logger::error!(
+                                "Asynchronously saving card in locker failed : {:?}",
+                                err
+                            );
+                        }
                     }
-                });
+                    .in_current_span(),
+                );
 
                 Ok(resp)
             }
diff --git a/crates/router/src/core/payments/operations/payment_confirm.rs b/crates/router/src/core/payments/operations/payment_confirm.rs
@@ -160,18 +160,21 @@ impl<F: Send + Clone, Ctx: PaymentMethodRetrieve>
 
         let store = state.store.clone();
 
-        let business_profile_fut = tokio::spawn(async move {
-            store
-                .find_business_profile_by_profile_id(&profile_id)
-                .map(|business_profile_result| {
-                    business_profile_result.to_not_found_response(
-                        errors::ApiErrorResponse::BusinessProfileNotFound {
-                            id: profile_id.to_string(),
-                        },
-                    )
-                })
-                .await
-        });
+        let business_profile_fut = tokio::spawn(
+            async move {
+                store
+                    .find_business_profile_by_profile_id(&profile_id)
+                    .map(|business_profile_result| {
+                        business_profile_result.to_not_found_response(
+                            errors::ApiErrorResponse::BusinessProfileNotFound {
+                                id: profile_id.to_string(),
+                            },
+                        )
+                    })
+                    .await
+            }
+            .in_current_span(),
+        );
 
         let store = state.store.clone();
 
@@ -498,13 +501,17 @@ impl<F: Send + Clone, Ctx: PaymentMethodRetrieve>
 
         let store = state.clone().store;
 
-        let additional_pm_data_fut = tokio::spawn(async move {
-            Ok(n_request_payment_method_data
-                .async_map(|payment_method_data| async move {
-                    helpers::get_additional_payment_data(&payment_method_data, store.as_ref()).await
-                })
-                .await)
-        });
+        let additional_pm_data_fut = tokio::spawn(
+            async move {
+                Ok(n_request_payment_method_data
+                    .async_map(|payment_method_data| async move {
+                        helpers::get_additional_payment_data(&payment_method_data, store.as_ref())
+                            .await
+                    })
+                    .await)
+            }
+            .in_current_span(),
+        );
 
         let store = state.clone().store;
 
diff --git a/crates/router/src/lib.rs b/crates/router/src/lib.rs
@@ -31,6 +31,7 @@ use actix_web::{
 };
 use http::StatusCode;
 use hyperswitch_interfaces::secrets_interface::secret_state::SecuredSecret;
+use router_env::tracing::Instrument;
 use routes::AppState;
 use storage_impl::errors::ApplicationResult;
 use tokio::sync::{mpsc, oneshot};
@@ -192,7 +193,7 @@ pub async fn start_server(conf: settings::Settings<SecuredSecret>) -> Applicatio
         .workers(server.workers)
         .shutdown_timeout(server.shutdown_timeout)
         .run();
-    tokio::spawn(receiver_for_error(rx, server.handle()));
+    let _task_handle = tokio::spawn(receiver_for_error(rx, server.handle()).in_current_span());
     Ok(server)
 }
 
diff --git a/crates/router/tests/utils.rs b/crates/router/tests/utils.rs
@@ -12,6 +12,7 @@ use actix_web::{
 };
 use derive_deref::Deref;
 use router::{configs::settings::Settings, routes::AppState, services};
+use router_env::tracing::Instrument;
 use serde::{de::DeserializeOwned, Deserialize};
 use serde_json::{json, Value};
 use tokio::sync::{oneshot, OnceCell};
@@ -24,7 +25,7 @@ async fn spawn_server() -> bool {
         .await
         .expect("failed to create server");
 
-    let _server = tokio::spawn(server);
+    let _server = tokio::spawn(server.in_current_span());
     true
 }
 
diff --git a/crates/scheduler/src/consumer.rs b/crates/scheduler/src/consumer.rs
@@ -10,7 +10,10 @@ pub use diesel_models::{self, process_tracker as storage};
 use error_stack::{IntoReport, ResultExt};
 use futures::future;
 use redis_interface::{RedisConnectionPool, RedisEntryId};
-use router_env::{instrument, tracing};
+use router_env::{
+    instrument,
+    tracing::{self, Instrument},
+};
 use time::PrimitiveDateTime;
 use tokio::sync::mpsc;
 use uuid::Uuid;
@@ -64,7 +67,8 @@ pub async fn start_consumer<T: SchedulerAppState + 'static>(
         .into_report()
         .attach_printable("Failed while creating a signals handler")?;
     let handle = signal.handle();
-    let task_handle = tokio::spawn(common_utils::signals::signal_handler(signal, tx));
+    let task_handle =
+        tokio::spawn(common_utils::signals::signal_handler(signal, tx).in_current_span());
 
     'consumer: loop {
         match rx.try_recv() {
diff --git a/crates/scheduler/src/producer.rs b/crates/scheduler/src/producer.rs
@@ -3,7 +3,10 @@ use std::sync::Arc;
 use common_utils::errors::CustomResult;
 use diesel_models::enums::ProcessTrackerStatus;
 use error_stack::{report, IntoReport, ResultExt};
-use router_env::{instrument, tracing};
+use router_env::{
+    instrument,
+    tracing::{self, Instrument},
+};
 use time::Duration;
 use tokio::sync::mpsc;
 
@@ -57,7 +60,8 @@ where
         .into_report()
         .attach_printable("Failed while creating a signals handler")?;
     let handle = signal.handle();
-    let task_handle = tokio::spawn(common_utils::signals::signal_handler(signal, tx));
+    let task_handle =
+        tokio::spawn(common_utils::signals::signal_handler(signal, tx).in_current_span());
 
     loop {
         match rx.try_recv() {
diff --git a/crates/storage_impl/src/redis.rs b/crates/storage_impl/src/redis.rs
@@ -6,7 +6,7 @@ use std::sync::{atomic, Arc};
 
 use error_stack::{IntoReport, ResultExt};
 use redis_interface::PubsubInterface;
-use router_env::logger;
+use router_env::{logger, tracing::Instrument};
 
 use self::{kv_store::RedisConnInterface, pub_sub::PubSubInterface};
 
@@ -35,9 +35,12 @@ impl RedisStore {
 
     pub fn set_error_callback(&self, callback: tokio::sync::oneshot::Sender<()>) {
         let redis_clone = self.redis_conn.clone();
-        tokio::spawn(async move {
-            redis_clone.on_error(callback).await;
-        });
+        let _task_handle = tokio::spawn(
+            async move {
+                redis_clone.on_error(callback).await;
+            }
+            .in_current_span(),
+        );
     }
 
     pub async fn subscribe_to_channel(
@@ -54,11 +57,14 @@ impl RedisStore {
             .change_context(redis_interface::errors::RedisError::SubscribeError)?;
 
         let redis_clone = self.redis_conn.clone();
-        tokio::spawn(async move {
-            if let Err(e) = redis_clone.on_message().await {
-                logger::error!(pubsub_err=?e);
+        let _task_handle = tokio::spawn(
+            async move {
+                if let Err(e) = redis_clone.on_message().await {
+                    logger::error!(pubsub_err=?e);
+                }
             }
-        });
+            .in_current_span(),
+        );
         Ok(())
     }
 }
