Run reconstruction inside a scoped rayon pool (#8075)

Co-Authored-By: Jimmy Chen <jchen.tc@gmail.com>

Co-Authored-By: Eitan Seri- Levi <eserilev@gmail.com>

Co-Authored-By: Eitan Seri-Levi <eserilev@ucsc.edu>

common/task_executor/Cargo.toml

@@ -8,6 +8,8 @@ edition = { workspace = true }
 async-channel = { workspace = true }
 futures = { workspace = true }
 metrics = { workspace = true }
+num_cpus = { workspace = true }
+rayon = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
 tracing = { workspace = true }
 

common/task_executor/src/lib.rs

@@ -1,12 +1,15 @@
 mod metrics;
+mod rayon_pool_provider;
 pub mod test_utils;
 
 use futures::channel::mpsc::Sender;
 use futures::prelude::*;
-use std::sync::Weak;
+use std::sync::{Arc, Weak};
 use tokio::runtime::{Handle, Runtime};
 use tracing::debug;
 
+use crate::rayon_pool_provider::RayonPoolProvider;
+pub use crate::rayon_pool_provider::RayonPoolType;
 pub use tokio::task::JoinHandle;
 
 /// Provides a reason when Lighthouse is shut down.
@@ -84,6 +87,8 @@ pub struct TaskExecutor {
     // FIXME(sproul): delete?
     #[allow(dead_code)]
     service_name: String,
+
+    rayon_pool_provider: Arc<RayonPoolProvider>,
 }
 
 impl TaskExecutor {
@@ -105,6 +110,7 @@ impl TaskExecutor {
             exit,
             signal_tx,
             service_name,
+            rayon_pool_provider: Arc::new(RayonPoolProvider::default()),
         }
     }
 
@@ -115,6 +121,7 @@ impl TaskExecutor {
             exit: self.exit.clone(),
             signal_tx: self.signal_tx.clone(),
             service_name,
+            rayon_pool_provider: self.rayon_pool_provider.clone(),
         }
     }
 
@@ -226,6 +233,47 @@ impl TaskExecutor {
         }
     }
 
+    /// Spawns a blocking task on a dedicated tokio thread pool and installs a rayon context within it.
+    pub fn spawn_blocking_with_rayon<F>(
+        self,
+        task: F,
+        rayon_pool_type: RayonPoolType,
+        name: &'static str,
+    ) where
+        F: FnOnce() + Send + 'static,
+    {
+        let thread_pool = self.rayon_pool_provider.get_thread_pool(rayon_pool_type);
+        self.spawn_blocking(
+            move || {
+                thread_pool.install(|| {
+                    task();
+                });
+            },
+            name,
+        )
+    }
+
+    /// Spawns a blocking computation on a rayon thread pool and awaits the result.
+    pub async fn spawn_blocking_with_rayon_async<F, R>(
+        &self,
+        rayon_pool_type: RayonPoolType,
+        task: F,
+    ) -> Result<R, tokio::sync::oneshot::error::RecvError>
+    where
+        F: FnOnce() -> R + Send + 'static,
+        R: Send + 'static,
+    {
+        let thread_pool = self.rayon_pool_provider.get_thread_pool(rayon_pool_type);
+        let (tx, rx) = tokio::sync::oneshot::channel();
+
+        thread_pool.spawn(move || {
+            let result = task();
+            let _ = tx.send(result);
+        });
+
+        rx.await
+    }
+
     /// Spawn a future on the tokio runtime wrapped in an `async-channel::Receiver` returning an optional
     /// join handle to the future.
     /// The task is cancelled when the corresponding async-channel is dropped.

common/task_executor/src/rayon_pool_provider.rs

@@ -0,0 +1,58 @@
+use rayon::{ThreadPool, ThreadPoolBuilder};
+use std::sync::Arc;
+
+const DEFAULT_LOW_PRIORITY_CPU_PERCENTAGE: usize = 25;
+const DEFAULT_HIGH_PRIORITY_CPU_PERCENTAGE: usize = 80;
+const MINIMUM_THREAD_COUNT: usize = 1;
+
+pub enum RayonPoolType {
+    HighPriority,
+    LowPriority,
+}
+
+pub struct RayonPoolProvider {
+    /// Smaller rayon thread pool for lower-priority, compute-intensive tasks.
+    /// By default ~25% of CPUs or a minimum of 1 thread.
+    low_priority_thread_pool: Arc<ThreadPool>,
+    /// Larger rayon thread pool for high-priority, compute-intensive tasks.
+    /// By default ~80% of CPUs or a minimum of 1 thread. Citical/highest
+    /// priority tasks should use the global pool instead.
+    high_priority_thread_pool: Arc<ThreadPool>,
+}
+
+impl Default for RayonPoolProvider {
+    fn default() -> Self {
+        let low_prio_threads =
+            (num_cpus::get() * DEFAULT_LOW_PRIORITY_CPU_PERCENTAGE / 100).max(MINIMUM_THREAD_COUNT);
+        let low_priority_thread_pool = Arc::new(
+            ThreadPoolBuilder::new()
+                .num_threads(low_prio_threads)
+                .build()
+                .expect("failed to build low-priority rayon pool"),
+        );
+
+        let high_prio_threads = (num_cpus::get() * DEFAULT_HIGH_PRIORITY_CPU_PERCENTAGE / 100)
+            .max(MINIMUM_THREAD_COUNT);
+        let high_priority_thread_pool = Arc::new(
+            ThreadPoolBuilder::new()
+                .num_threads(high_prio_threads)
+                .build()
+                .expect("failed to build high-priority rayon pool"),
+        );
+        Self {
+            low_priority_thread_pool,
+            high_priority_thread_pool,
+        }
+    }
+}
+
+impl RayonPoolProvider {
+    /// Get a scoped thread pool by priority level.
+    /// For critical/highest priority tasks, use the global pool instead.
+    pub fn get_thread_pool(&self, rayon_pool_type: RayonPoolType) -> Arc<ThreadPool> {
+        match rayon_pool_type {
+            RayonPoolType::HighPriority => self.high_priority_thread_pool.clone(),
+            RayonPoolType::LowPriority => self.low_priority_thread_pool.clone(),
+        }
+    }
+}