From ded0896df5bf4095c8efc86f315be7bc50a02c59 Mon Sep 17 00:00:00 2001
From: Aditya Pandey <prblmslvr.aditya@gmail.com>
Date: Fri, 5 May 2023 23:18:45 +0530
Subject: [PATCH] Task Executor Metrics (#2580)

Co-authored-by: Aditya Pandey <aditya.p@Aditya-P.local>
---
 Cargo.lock                  |  2 ++
 crates/tasks/Cargo.toml     |  4 ++++
 crates/tasks/src/lib.rs     | 11 ++++++++++-
 crates/tasks/src/metrics.rs | 24 ++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 crates/tasks/src/metrics.rs

diff --git a/Cargo.lock b/Cargo.lock
index e1fa534d6..acaf989fe 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5455,6 +5455,8 @@ version = "0.1.0"
 dependencies = [
  "dyn-clone",
  "futures-util",
+ "metrics",
+ "reth-metrics-derive",
  "thiserror",
  "tokio",
  "tracing",
diff --git a/crates/tasks/Cargo.toml b/crates/tasks/Cargo.toml
index 6d03dfcca..64609ef8b 100644
--- a/crates/tasks/Cargo.toml
+++ b/crates/tasks/Cargo.toml
@@ -19,5 +19,9 @@ tracing = { version = "0.1", default-features = false }
 thiserror = "1.0"
 dyn-clone = "1.0"
 
+## rpc/metrics
+metrics = "0.20.1"
+reth-metrics-derive = { path = "../metrics/metrics-derive" }
+
 [dev-dependencies]
 tokio = { version = "1", features = ["sync", "rt", "rt-multi-thread", "time", "macros"] }
diff --git a/crates/tasks/src/lib.rs b/crates/tasks/src/lib.rs
index 02954b7fc..4f7e61b6b 100644
--- a/crates/tasks/src/lib.rs
+++ b/crates/tasks/src/lib.rs
@@ -7,7 +7,10 @@
 
 //! reth task management
 
-use crate::shutdown::{signal, Shutdown, Signal};
+use crate::{
+    metrics::TaskExecutorMetrics,
+    shutdown::{signal, Shutdown, Signal},
+};
 use dyn_clone::DynClone;
 use futures_util::{
     future::{select, BoxFuture},
@@ -25,6 +28,7 @@ use tokio::{
 use tracing::error;
 use tracing_futures::Instrument;
 
+pub mod metrics;
 pub mod shutdown;
 
 /// A type that can spawn tasks.
@@ -160,6 +164,7 @@ impl TaskManager {
             handle: self.handle.clone(),
             on_shutdown: self.on_shutdown.clone(),
             panicked_tasks_tx: self.panicked_tasks_tx.clone(),
+            metrics: Default::default(),
         }
     }
 }
@@ -192,6 +197,8 @@ pub struct TaskExecutor {
     on_shutdown: Shutdown,
     /// Sender half for sending panic signals to this type
     panicked_tasks_tx: UnboundedSender<&'static str>,
+    // Task Executor Metrics
+    metrics: TaskExecutorMetrics,
 }
 
 // === impl TaskExecutor ===
@@ -358,10 +365,12 @@ impl TaskExecutor {
 
 impl TaskSpawner for TaskExecutor {
     fn spawn(&self, fut: BoxFuture<'static, ()>) -> JoinHandle<()> {
+        self.metrics.inc_regular_task();
         self.spawn(fut)
     }
 
     fn spawn_critical(&self, name: &'static str, fut: BoxFuture<'static, ()>) -> JoinHandle<()> {
+        self.metrics.inc_critical_tasks();
         TaskExecutor::spawn_critical(self, name, fut)
     }
 
diff --git a/crates/tasks/src/metrics.rs b/crates/tasks/src/metrics.rs
new file mode 100644
index 000000000..a56ccfbbe
--- /dev/null
+++ b/crates/tasks/src/metrics.rs
@@ -0,0 +1,24 @@
+//! Task Executor Metrics
+use metrics::Counter;
+use reth_metrics_derive::Metrics;
+
+/// Task Executor Metrics
+#[derive(Metrics, Clone)]
+#[metrics(scope = "executor.spawn")]
+pub struct TaskExecutorMetrics {
+    /// Number of spawned critical tasks
+    pub(crate) critical_tasks: Counter,
+
+    /// Number of spawned regular tasks
+    pub(crate) regular_tasks: Counter,
+}
+
+impl TaskExecutorMetrics {
+    pub(crate) fn inc_critical_tasks(&self) {
+        self.critical_tasks.increment(1);
+    }
+
+    pub(crate) fn inc_regular_task(&self) {
+        self.regular_tasks.increment(1);
+    }
+}