feat(threading): Add new relay-threading module (#4500)

Author: iambriccardo

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+**Features**:
+
+- Add new `relay-threading` crate with asynchronous thread pool. ([#4500](https://github.com/getsentry/relay/pull/4500))
+
 ## 25.2.0
 
 - Allow log ingestion behind a flag, only for internal use currently. ([#4471](https://github.com/getsentry/relay/pull/4471))

Cargo.lock

@@ -56,6 +56,7 @@ relay-server = { path = "relay-server" }
 relay-spans = { path = "relay-spans" }
 relay-statsd = { path = "relay-statsd" }
 relay-system = { path = "relay-system" }
+relay-threading = { path = "relay-threading" }
 relay-ua = { path = "relay-ua" }
 relay-test = { path = "relay-test" }
 relay-protocol-derive = { path = "relay-protocol-derive" }
@@ -98,6 +99,7 @@ enumset = "1.0.13"
 flate2 = "1.0.35"
 fnv = "1.0.7"
 futures = { version = "0.3", default-features = false, features = ["std"] }
+flume = { version = "0.11.1", default-features = false }
 globset = "0.4.15"
 hash32 = "0.3.1"
 hashbrown = "0.14.5"

Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "relay-threading"
+authors = ["Sentry <[email protected]>"]
+description = "Threading code that is used by Relay"
+homepage = "https://getsentry.github.io/relay/"
+repository = "https://github.com/getsentry/relay"
+version = "25.2.0"
+edition = "2021"
+license-file = "../LICENSE.md"
+publish = false
+
+[dependencies]
+flume = { workspace = true }
+futures = { workspace = true }
+tokio = { workspace = true }
+pin-project-lite = { workspace = true }
+
+[dev-dependencies]
+criterion = { workspace = true, features = ["async_tokio"] }
+futures = { workspace = true, features = ["executor"] }
+tokio = { workspace = true, features = ["rt", "rt-multi-thread", "time", "sync", "macros"] }
+
+[[bench]]
+name = "pool"
+harness = false

relay-threading/Cargo.toml

@@ -0,0 +1,159 @@
+use std::future::Future;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use futures::future::{BoxFuture, FutureExt};
+use relay_threading::AsyncPoolBuilder;
+use tokio::runtime::Runtime;
+use tokio::sync::Semaphore;
+
+struct BenchBarrier {
+    semaphore: Arc<Semaphore>,
+    count: usize,
+}
+
+impl BenchBarrier {
+    fn new(count: usize) -> Self {
+        Self {
+            semaphore: Arc::new(Semaphore::new(count)),
+            count,
+        }
+    }
+
+    async fn spawn<F, Fut>(&self, pool: &relay_threading::AsyncPool<BoxFuture<'static, ()>>, f: F)
+    where
+        F: FnOnce() -> Fut + Send + 'static,
+        Fut: Future<Output = ()> + Send + 'static,
+    {
+        let semaphore = self.semaphore.clone();
+        let permit = semaphore.acquire_owned().await.unwrap();
+        pool.spawn_async(
+            async move {
+                f().await;
+                drop(permit);
+            }
+            .boxed(),
+        )
+        .await;
+    }
+
+    async fn wait(&self) {
+        let _ = self
+            .semaphore
+            .acquire_many(self.count as u32)
+            .await
+            .unwrap();
+    }
+}
+
+fn create_runtime() -> Runtime {
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(4)
+        .enable_all()
+        .build()
+        .unwrap()
+}
+
+async fn run_benchmark(pool: &relay_threading::AsyncPool<BoxFuture<'static, ()>>, count: usize) {
+    let counter = Arc::new(AtomicUsize::new(0));
+    let barrier = BenchBarrier::new(count);
+
+    // Spawn tasks
+    for _ in 0..count {
+        let counter = counter.clone();
+        barrier
+            .spawn(pool, move || async move {
+                // Simulate some work
+                tokio::time::sleep(Duration::from_micros(50)).await;
+                counter.fetch_add(1, Ordering::SeqCst);
+            })
+            .await;
+    }
+
+    // Wait for all tasks to complete
+    barrier.wait().await;
+    assert_eq!(counter.load(Ordering::SeqCst), count);
+}
+
+fn bench_pool_scaling(c: &mut Criterion) {
+    let runtime = create_runtime();
+    let mut group = c.benchmark_group("pool_scaling");
+    group.sampling_mode(criterion::SamplingMode::Flat);
+    group.measurement_time(Duration::from_secs(10));
+
+    // Test with different numbers of threads
+    for threads in [1, 2, 4, 8].iter() {
+        let pool = AsyncPoolBuilder::new(runtime.handle().clone())
+            .num_threads(*threads)
+            .max_concurrency(100)
+            .build()
+            .unwrap();
+
+        // Test with different task counts
+        for tasks in [100, 1000, 10000].iter() {
+            group.bench_with_input(
+                BenchmarkId::new(format!("threads_{}", threads), tasks),
+                tasks,
+                |b, &tasks| {
+                    b.to_async(&runtime).iter(|| run_benchmark(&pool, tasks));
+                },
+            );
+        }
+    }
+
+    group.finish();
+}
+
+fn bench_multi_threaded_spawn(c: &mut Criterion) {
+    let runtime = create_runtime();
+    let mut group = c.benchmark_group("multi_threaded_spawn");
+    group.sampling_mode(criterion::SamplingMode::Flat);
+    group.measurement_time(Duration::from_secs(10));
+
+    // Test with different numbers of spawning threads
+    for spawn_threads in [2, 4, 8].iter() {
+        // Test with different task counts
+        for tasks in [1000, 10000].iter() {
+            group.bench_with_input(
+                BenchmarkId::new(format!("spawn_threads_{}", spawn_threads), tasks),
+                tasks,
+                |b, &tasks| {
+                    b.to_async(&runtime).iter(|| async {
+                        let pool = Arc::new(
+                            AsyncPoolBuilder::new(runtime.handle().clone())
+                                .num_threads(4) // Fixed number of worker threads
+                                .max_concurrency(100)
+                                .build()
+                                .unwrap(),
+                        );
+
+                        let tasks_per_thread = tasks / spawn_threads;
+                        let mut handles = Vec::new();
+
+                        // Spawn tasks from multiple threads
+                        for _ in 0..*spawn_threads {
+                            let runtime = runtime.handle().clone();
+                            let pool = pool.clone();
+                            let handle = std::thread::spawn(move || {
+                                runtime.block_on(run_benchmark(&pool, tasks_per_thread));
+                            });
+                            handles.push(handle);
+                        }
+
+                        // Wait for all spawning threads to complete
+                        for handle in handles {
+                            handle.join().unwrap();
+                        }
+                    });
+                },
+            );
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_pool_scaling, bench_multi_threaded_spawn);
+criterion_main!(benches);

relay-threading/benches/pool.rs

@@ -0,0 +1,129 @@
+use std::any::Any;
+use std::future::Future;
+use std::io;
+use std::sync::Arc;
+
+use crate::pool::{AsyncPool, Thread};
+use crate::pool::{CustomSpawn, DefaultSpawn, ThreadSpawn};
+
+/// Type alias for a thread safe closure that is used for panic handling across the code.
+pub(crate) type PanicHandler = dyn Fn(Box<dyn Any + Send>) + Send + Sync;
+
+/// [`AsyncPoolBuilder`] provides a flexible way to configure and build an [`AsyncPool`] for executing
+/// asynchronous tasks concurrently on dedicated threads.
+///
+/// This builder enables you to customize the number of threads, concurrency limits, thread naming,
+/// and panic handling strategies.
+pub struct AsyncPoolBuilder<S = DefaultSpawn> {
+    pub(crate) runtime: tokio::runtime::Handle,
+    pub(crate) thread_name: Option<Box<dyn FnMut(usize) -> String>>,
+    pub(crate) thread_panic_handler: Option<Arc<PanicHandler>>,
+    pub(crate) task_panic_handler: Option<Arc<PanicHandler>>,
+    pub(crate) spawn_handler: S,
+    pub(crate) num_threads: usize,
+    pub(crate) max_concurrency: usize,
+}
+
+impl AsyncPoolBuilder<DefaultSpawn> {
+    /// Initializes a new [`AsyncPoolBuilder`] with default settings.
+    ///
+    /// The builder is tied to the provided [`tokio::runtime::Handle`] and prepares to configure an [`AsyncPool`].
+    pub fn new(runtime: tokio::runtime::Handle) -> AsyncPoolBuilder<DefaultSpawn> {
+        AsyncPoolBuilder {
+            runtime,
+            thread_name: None,
+            thread_panic_handler: None,
+            task_panic_handler: None,
+            spawn_handler: DefaultSpawn,
+            num_threads: 1,
+            max_concurrency: 1,
+        }
+    }
+}
+
+impl<S> AsyncPoolBuilder<S>
+where
+    S: ThreadSpawn,
+{
+    /// Specifies a custom naming convention for threads in the [`AsyncPool`].
+    ///
+    /// The provided closure receives the thread's index and returns a name,
+    /// which can be useful for debugging and logging.
+    pub fn thread_name<F>(mut self, thread_name: F) -> Self
+    where
+        F: FnMut(usize) -> String + 'static,
+    {
+        self.thread_name = Some(Box::new(thread_name));
+        self
+    }
+
+    /// Sets a custom panic handler for threads in the [`AsyncPool`].
+    ///
+    /// If a thread panics, the provided handler will be invoked so that you can perform
+    /// custom error handling or cleanup.
+    pub fn thread_panic_handler<F>(mut self, panic_handler: F) -> Self
+    where
+        F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
+    {
+        self.thread_panic_handler = Some(Arc::new(panic_handler));
+        self
+    }
+
+    /// Sets a custom panic handler for tasks executed by the [`AsyncPool`].
+    ///
+    /// This handler is used to manage panics that occur during task execution, allowing for graceful
+    /// error handling.
+    pub fn task_panic_handler<F>(mut self, panic_handler: F) -> Self
+    where
+        F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
+    {
+        self.task_panic_handler = Some(Arc::new(panic_handler));
+        self
+    }
+
+    /// Configures a custom thread spawning procedure for the [`AsyncPool`].
+    ///
+    /// This method allows you to adjust thread settings (e.g. naming, stack size) before thread creation,
+    /// making it possible to apply application-specific configurations.
+    pub fn spawn_handler<F>(self, spawn_handler: F) -> AsyncPoolBuilder<CustomSpawn<F>>
+    where
+        F: FnMut(Thread) -> io::Result<()>,
+    {
+        AsyncPoolBuilder {
+            runtime: self.runtime,
+            thread_name: self.thread_name,
+            thread_panic_handler: self.thread_panic_handler,
+            task_panic_handler: self.task_panic_handler,
+            spawn_handler: CustomSpawn::new(spawn_handler),
+            num_threads: self.num_threads,
+            max_concurrency: self.max_concurrency,
+        }
+    }
+
+    /// Sets the number of worker threads for the [`AsyncPool`].
+    ///
+    /// This determines how many dedicated threads will be available for running tasks concurrently.
+    pub fn num_threads(mut self, num_threads: usize) -> Self {
+        self.num_threads = num_threads;
+        self
+    }
+
+    /// Sets the maximum number of concurrent tasks per thread in the [`AsyncPool`].
+    ///
+    /// This controls how many futures can be polled simultaneously on each worker thread.
+    pub fn max_concurrency(mut self, max_concurrency: usize) -> Self {
+        self.max_concurrency = max_concurrency;
+        self
+    }
+
+    /// Constructs an [`AsyncPool`] based on the configured settings.
+    ///
+    /// Finalizing the builder sets up dedicated worker threads and configures the executor
+    /// to enforce the specified concurrency limits.
+    pub fn build<F>(self) -> Result<AsyncPool<F>, io::Error>
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        AsyncPool::new(self)
+    }
+}