From 291ebf5e2442c318f8630a5dd5540b921271e1e8 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Tue, 14 Jan 2025 02:22:28 -0500
Subject: [PATCH] Have serai-task warnings print with the name of the task

---
 common/task/src/lib.rs       | 10 +++++++---
 common/task/src/type_name.rs | 31 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 common/task/src/type_name.rs

diff --git a/common/task/src/lib.rs b/common/task/src/lib.rs
index 64cf9416..83eac9bf 100644
--- a/common/task/src/lib.rs
+++ b/common/task/src/lib.rs
@@ -10,6 +10,8 @@ use core::{
 
 use tokio::sync::mpsc;
 
+mod type_name;
+
 /// A handle for a task.
 ///
 /// The task will only stop running once all handles for it are dropped.
@@ -49,8 +51,6 @@ impl Task {
 
 impl TaskHandle {
   /// Tell the task to run now (and not whenever its next iteration on a timer is).
-  ///
-  /// Panics if the task has been dropped.
   pub fn run_now(&self) {
     #[allow(clippy::match_same_arms)]
     match self.run_now.try_send(()) {
@@ -58,6 +58,7 @@ impl TaskHandle {
       // NOP on full, as this task will already be ran as soon as possible
       Err(mpsc::error::TrySendError::Full(())) => {}
       Err(mpsc::error::TrySendError::Closed(())) => {
+        // The task should only be closed if all handles are dropped, and this one hasn't been
         panic!("task was unexpectedly closed when calling run_now")
       }
     }
@@ -131,7 +132,10 @@ pub trait ContinuallyRan: Sized + Send {
             }
           }
           Err(e) => {
-            log::warn!("{e:?}");
+            // Get the type name
+            let type_name = type_name::strip_type_name(core::any::type_name::<Self>());
+            // Print the error as a warning, prefixed by the task's type
+            log::warn!("{type_name}: {e:?}");
             increase_sleep_before_next_task(&mut current_sleep_before_next_task);
           }
         }
diff --git a/common/task/src/type_name.rs b/common/task/src/type_name.rs
new file mode 100644
index 00000000..c6ba1658
--- /dev/null
+++ b/common/task/src/type_name.rs
@@ -0,0 +1,31 @@
+/// Strip the modules from a type name.
+// This may be of the form `a::b::C`, in which case we only want `C`
+pub(crate) fn strip_type_name(full_type_name: &'static str) -> String {
+  // It also may be `a::b::C<d::e::F>`, in which case, we only attempt to strip `a::b`
+  let mut by_generics = full_type_name.split('<');
+
+  // Strip to just `C`
+  let full_outer_object_name = by_generics.next().unwrap();
+  let mut outer_object_name_parts = full_outer_object_name.split("::");
+  let mut last_part_in_outer_object_name = outer_object_name_parts.next().unwrap();
+  for part in outer_object_name_parts {
+    last_part_in_outer_object_name = part;
+  }
+
+  // Push back on the generic terms
+  let mut type_name = last_part_in_outer_object_name.to_string();
+  for generic in by_generics {
+    type_name.push('<');
+    type_name.push_str(generic);
+  }
+  type_name
+}
+
+#[test]
+fn test_strip_type_name() {
+  assert_eq!(strip_type_name("core::option::Option"), "Option");
+  assert_eq!(
+    strip_type_name("core::option::Option<alloc::string::String>"),
+    "Option<alloc::string::String>"
+  );
+}