ref(processor): Switch processor to a rayon thread pool

CHANGELOG.md

@@ -2,10 +2,14 @@
 
 ## Unreleased
 
-**Features**
+**Features**:
 
 - "Cardinality limit" outcomes now report which limit was exceeded. ([#3825](https://github.com/getsentry/relay/pull/3825))
 
+**Internal**:
+
+- Use a dedicated thread pool for CPU intensive workloads. ([#3833](https://github.com/getsentry/relay/pull/3833))
+
 ## 24.7.0
 
 **Bug Fixes**:

Cargo.toml

@@ -128,6 +128,7 @@ quote = "1.0.2"
 r2d2 = "0.8.10"
 rand = "0.8.5"
 rand_pcg = "0.3.1"
+rayon = "1.10"
 rdkafka = "0.29.0"
 rdkafka-sys = "4.3.0"
 # Git revision until https://github.com/redis-rs/redis-rs/pull/1097 (merged) and https://github.com/redis-rs/redis-rs/pull/1253 are released.

relay-server/Cargo.toml

@@ -64,6 +64,7 @@ multer = { workspace = true }
 once_cell = { workspace = true }
 pin-project-lite = { workspace = true }
 rand = { workspace = true }
+rayon = { workspace = true }
 regex = { workspace = true }
 relay-auth = { workspace = true }
 relay-base-schema = { workspace = true }

relay-server/src/service.rs

@@ -7,6 +7,7 @@ use crate::services::stats::RelayStats;
 use anyhow::{Context, Result};
 use axum::extract::FromRequestParts;
 use axum::http::request::Parts;
+use rayon::ThreadPool;
 use relay_cogs::Cogs;
 use relay_config::{Config, RedisConnection};
 use relay_redis::RedisPool;
@@ -81,6 +82,24 @@ pub fn create_runtime(name: &str, threads: usize) -> Runtime {
         .unwrap()
 }
 
+fn create_processor_pool(config: &Config) -> Result<ThreadPool> {
+    // Adjust thread count for small cpu counts to not have too many idle cores
+    // and distribute workload better.
+    let thread_count = match config.cpu_concurrency() {
+        conc @ 0..=2 => conc.max(1),
+        conc @ 3..=4 => conc - 1,
+        conc => conc - 2,
+    };
+    relay_log::info!("starting {thread_count} envelope processing workers");
+
+    let pool = crate::utils::ThreadPoolBuilder::new("processor")
+        .num_threads(thread_count)
+        .runtime(tokio::runtime::Handle::current())
+        .build()?;
+
+    Ok(pool)
+}
+
 #[derive(Debug)]
 struct StateInner {
     config: Arc<Config>,
@@ -171,6 +190,7 @@ impl ServiceState {
         let cogs = Cogs::new(CogsServiceRecorder::new(&config, cogs.start()));
 
         EnvelopeProcessorService::new(
+            create_processor_pool(&config)?,
             config.clone(),
             global_config_handle,
             cogs,

relay-server/src/services/processor.rs

@@ -84,7 +84,7 @@ use crate::services::upstream::{
 };
 use crate::statsd::{RelayCounters, RelayHistograms, RelayTimers};
 use crate::utils::{
-    self, InvalidProcessingGroupType, ManagedEnvelope, SamplingResult, TypedEnvelope,
+    self, InvalidProcessingGroupType, ManagedEnvelope, SamplingResult, ThreadPool, TypedEnvelope,
 };
 use crate::{http, metrics};
 
@@ -1113,6 +1113,7 @@ impl Default for Addrs {
 }
 
 struct InnerProcessor {
+    pool: ThreadPool,
     config: Arc<Config>,
     global_config: GlobalConfigHandle,
     cogs: Cogs,
@@ -1132,6 +1133,7 @@ struct InnerProcessor {
 impl EnvelopeProcessorService {
     /// Creates a multi-threaded envelope processor.
     pub fn new(
+        pool: ThreadPool,
         config: Arc<Config>,
         global_config: GlobalConfigHandle,
         cogs: Cogs,
@@ -1150,6 +1152,7 @@ impl EnvelopeProcessorService {
         });
 
         let inner = InnerProcessor {
+            pool,
             global_config,
             cogs,
             #[cfg(feature = "processing")]
@@ -2811,18 +2814,9 @@ impl Service for EnvelopeProcessorService {
     type Interface = EnvelopeProcessor;
 
     fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) {
-        // Adjust thread count for small cpu counts to not have too many idle cores
-        // and distribute workload better.
-        let thread_count = match self.inner.config.cpu_concurrency() {
-            conc @ 0..=2 => conc.max(1),
-            conc @ 3..=4 => conc - 1,
-            conc => conc - 2,
-        };
-        relay_log::info!("starting {thread_count} envelope processing workers");
+        let semaphore = Arc::new(Semaphore::new(self.inner.pool.current_num_threads()));
 
         tokio::spawn(async move {
-            let semaphore = Arc::new(Semaphore::new(thread_count));
-
             loop {
                 let next_msg = async {
                     let permit_result = semaphore.clone().acquire_owned().await;
@@ -2837,7 +2831,7 @@ impl Service for EnvelopeProcessorService {
 
                     (Some(message), Ok(permit)) = next_msg => {
                         let service = self.clone();
-                        tokio::task::spawn_blocking(move || {
+                        self.inner.pool.spawn(move || {
                             service.handle_message(message);
                             drop(permit);
                         });

relay-server/src/testutils.rs

@@ -23,6 +23,7 @@ use crate::services::outcome::TrackOutcome;
 use crate::services::processor::{self, EnvelopeProcessorService};
 use crate::services::project::ProjectState;
 use crate::services::test_store::TestStore;
+use crate::utils::{ThreadPool, ThreadPoolBuilder};
 
 pub fn state_with_rule_and_condition(
     sample_rate: Option<f64>,
@@ -138,6 +139,7 @@ pub fn create_test_processor(config: Config) -> EnvelopeProcessorService {
 
     let config = Arc::new(config);
     EnvelopeProcessorService::new(
+        create_processor_pool(),
         Arc::clone(&config),
         GlobalConfigHandle::fixed(Default::default()),
         Cogs::noop(),
@@ -177,6 +179,7 @@ pub fn create_test_processor_with_addrs(
 
     let config = Arc::new(config);
     EnvelopeProcessorService::new(
+        create_processor_pool(),
         Arc::clone(&config),
         GlobalConfigHandle::fixed(Default::default()),
         Cogs::noop(),
@@ -192,3 +195,11 @@ pub fn processor_services() -> (Addr<TrackOutcome>, Addr<TestStore>) {
     let (test_store, _) = mock_service("test_store", (), |&mut (), _| {});
     (outcome_aggregator, test_store)
 }
+
+fn create_processor_pool() -> ThreadPool {
+    ThreadPoolBuilder::new("processor")
+        .num_threads(1)
+        .runtime(tokio::runtime::Handle::current())
+        .build()
+        .unwrap()
+}

relay-server/src/utils/mod.rs

@@ -13,6 +13,7 @@ mod sizes;
 mod sleep_handle;
 mod split_off;
 mod statsd;
+mod thread_pool;
 
 #[cfg(feature = "processing")]
 mod native;
@@ -38,5 +39,6 @@ pub use self::sizes::*;
 pub use self::sleep_handle::*;
 pub use self::split_off::*;
 pub use self::statsd::*;
+pub use self::thread_pool::*;
 #[cfg(feature = "processing")]
 pub use self::unreal::*;

relay-server/src/utils/thread_pool.rs

@@ -0,0 +1,97 @@
+use std::thread;
+use tokio::runtime::Handle;
+
+pub use rayon::{ThreadPool, ThreadPoolBuildError};
+
+/// Used to create a new [`ThreadPool`] thread pool.
+pub struct ThreadPoolBuilder {
+    name: &'static str,
+    runtime: Option<Handle>,
+    num_threads: usize,
+}
+
+impl ThreadPoolBuilder {
+    /// Creates a new named thread pool builder.
+    pub fn new(name: &'static str) -> Self {
+        Self {
+            name,
+            runtime: None,
+            num_threads: 0,
+        }
+    }
+
+    /// Sets the number of threads to be used in the rayon threadpool.
+    ///
+    /// See also [`rayon::ThreadPoolBuilder::num_threads`].
+    pub fn num_threads(mut self, num_threads: usize) -> Self {
+        self.num_threads = num_threads;
+        self
+    }
+
+    /// Sets the tokio runtime which will be made available in the workers.
+    pub fn runtime(mut self, runtime: Handle) -> Self {
+        self.runtime = Some(runtime);
+        self
+    }
+
+    /// Creates and returns the thread pool.
+    pub fn build(self) -> Result<ThreadPool, ThreadPoolBuildError> {
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(self.num_threads)
+            .thread_name(move |id| format!("pool-{name}-{id}", name = self.name))
+            // In case of panic, log that there was a panic but keep the thread alive and don't
+            // exist.
+            .panic_handler(move |_panic| {
+                relay_log::error!("thread in pool {name} paniced!", name = self.name)
+            })
+            .spawn_handler(|thread| {
+                let mut b = thread::Builder::new();
+                if let Some(name) = thread.name() {
+                    b = b.name(name.to_owned());
+                }
+                if let Some(stack_size) = thread.stack_size() {
+                    b = b.stack_size(stack_size);
+                }
+                let runtime = self.runtime.clone();
+                b.spawn(move || {
+                    let _guard = runtime.as_ref().map(|runtime| runtime.enter());
+                    thread.run()
+                })?;
+                Ok(())
+            })
+            .build()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_num_threads() {
+        let pool = ThreadPoolBuilder::new("s").num_threads(3).build().unwrap();
+        assert_eq!(pool.current_num_threads(), 3);
+    }
+
+    #[test]
+    fn test_runtime() {
+        let rt = tokio::runtime::Runtime::new().unwrap();
+
+        let pool = ThreadPoolBuilder::new("s")
+            .num_threads(1)
+            .runtime(rt.handle().clone())
+            .build()
+            .unwrap();
+
+        let has_runtime = pool.install(|| tokio::runtime::Handle::try_current().is_ok());
+        assert!(has_runtime);
+    }
+
+    #[test]
+    fn test_no_runtime() {
+        let pool = ThreadPoolBuilder::new("s").num_threads(1).build().unwrap();
+
+        let has_runtime = pool.install(|| tokio::runtime::Handle::try_current().is_ok());
+        assert!(!has_runtime);
+    }
+}