Make block_on_dangerous respect exit, metrics

sigp · Jun 14, 2022 · 13e71b3 · 13e71b3
1 parent b2192db
commit 13e71b3
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 4 deletions.
diff --git a/common/task_executor/Cargo.toml b/common/task_executor/Cargo.toml
@@ -5,7 +5,7 @@ authors = ["Sigma Prime <contact@sigmaprime.io>"]
 edition = "2021"
 
 [dependencies]
-tokio = { version = "1.14.0", features = ["rt-multi-thread"] }
+tokio = { version = "1.14.0", features = ["rt-multi-thread", "macros"] }
 slog = "2.5.2"
 futures = "0.3.7"
 exit-future = "0.2.0"

diff --git a/common/task_executor/src/lib.rs b/common/task_executor/src/lib.rs
@@ -322,12 +322,44 @@ impl TaskExecutor {
     pub fn block_on_dangerous<F: Future>(
         &self,
         future: F,
-        _name: &'static str,
+        name: &'static str,
     ) -> Option<F::Output>
 where {
+        let timer = metrics::start_timer_vec(&metrics::BLOCK_ON_TASKS_HISTOGRAM, &[name]);
+        metrics::inc_gauge_vec(&metrics::BLOCK_ON_TASKS_COUNT, &[name]);
+        let log = self.log.clone();
         let handle = self.handle()?;
-        // TODO(paul): respect the shutdown signal and the name.
-        Some(handle.block_on(future))
+        let exit = self.exit.clone();
+
+        debug!(
+            log,
+            "Starting block_on task";
+            "name" => name
+        );
+
+        handle.block_on(async {
+            let output = tokio::select! {
+                output = future => {
+                    debug!(
+                        log,
+                        "Completed block_on task";
+                        "name" => name
+                    );
+                    Some(output)
+                },
+                _ = exit => {
+                    debug!(
+                        log,
+                        "Cancelled block_on task";
+                        "name" => name,
+                    );
+                    None
+                }
+            };
+            metrics::dec_gauge_vec(&metrics::BLOCK_ON_TASKS_COUNT, &[name]);
+            drop(timer);
+            output
+        })
     }
 
     /// Returns a `Handle` to the current runtime.

diff --git a/common/task_executor/src/metrics.rs b/common/task_executor/src/metrics.rs
@@ -18,6 +18,16 @@ lazy_static! {
         "Time taken by blocking tasks",
         &["blocking_task_hist"]
     );
+    pub static ref BLOCK_ON_TASKS_COUNT: Result<IntGaugeVec> = try_create_int_gauge_vec(
+        "block_on_tasks_count",
+        "Total number of block_on_dangers tasks spawned",
+        &["name"]
+    );
+    pub static ref BLOCK_ON_TASKS_HISTOGRAM: Result<HistogramVec> = try_create_histogram_vec(
+        "block_on_tasks_histogram",
+        "Time taken by block_on_dangerous tasks",
+        &["name"]
+    );
     pub static ref TASKS_HISTOGRAM: Result<HistogramVec> = try_create_histogram_vec(
         "async_tasks_time_histogram",
         "Time taken by async tasks",