Move Threadpool class to separate module (#151)

Move the ThreadPool class to a separate module for organization
purposes. 

This PR is the first step in a larger effort to update the multithreaded runtime 
with the following goals in mind:

- Allow us to isolate the async part of our codebase to only the
portions that really need it.
- Abstract away the specific async runtime (now tokio) to allow us to
swap it out for lighter weight runtime down the road if desired.
- Expose compute focused non-async primitives (e.g. par_for, JoinSet,
JoinHandle) that we use frequently in the async world, but that could
easily be done in the multithreaded sync world instead.
- Expose reliable ways to call into async routines from the non-async
runtime.
- Lay the groundwork for better CTRL-C cancellations.

Cargo.toml

@@ -15,7 +15,7 @@ members = [
     "utils",
     "cas_object",
     "cas_types",
-    "chunk_cache",
+    "chunk_cache", "xet_threadpool",
 ]
 
 exclude = ["hf_xet", "chunk_cache_bench"]

cas_client/Cargo.toml

@@ -18,6 +18,7 @@ utils = { path = "../utils" }
 merkledb = { path = "../merkledb" }
 mdb_shard = { path = "../mdb_shard" }
 merklehash = { path = "../merklehash" }
+xet_threadpool = { path = "../xet_threadpool" }
 thiserror = "2.0"
 tokio = { version = "1.41", features = ["full"] }
 async-trait = "0.1.9"

cas_client/src/remote_client.rs

@@ -21,7 +21,7 @@ use tracing::{debug, error, trace};
 use utils::auth::AuthConfig;
 use utils::progress::ProgressUpdater;
 use utils::singleflight::Group;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::error::Result;
 use crate::interface::*;

data/Cargo.toml

@@ -21,6 +21,7 @@ utils = { path = "../utils" }
 parutils = { path = "../parutils" }
 file_utils = { path = "../file_utils" }
 error_printer = { path = "../error_printer" }
+xet_threadpool = { path = "../xet_threadpool" }
 thiserror = "2.0"
 tokio = { version = "1.36", features = ["full"] }
 anyhow = "1"

data/src/bin/example.rs

@@ -10,7 +10,7 @@ use cas_client::CacheConfig;
 use clap::{Args, Parser, Subcommand};
 use data::configurations::*;
 use data::{PointerFile, PointerFileTranslator};
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 #[derive(Parser)]
 struct XCommand {

data/src/cas_interface.rs

@@ -7,7 +7,7 @@ use cas_client::{CacheConfig, RemoteClient};
 use mdb_shard::ShardFileManager;
 use tracing::info;
 use utils::auth::AuthConfig;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::configurations::*;
 use crate::errors::Result;

data/src/chunking.rs

@@ -7,7 +7,7 @@ use merklehash::compute_data_hash;
 use tokio::sync::mpsc::{Receiver, Sender};
 use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use super::clean::BufferItem;
 use crate::errors::{DataProcessingError, Result};

data/src/clean.rs

@@ -23,7 +23,7 @@ use tokio::sync::Mutex;
 use tokio::task::{JoinHandle, JoinSet};
 use tracing::{debug, info, warn};
 use utils::progress::ProgressUpdater;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::chunking::{chunk_target_default, ChunkYieldType};
 use crate::constants::MIN_SPACING_BETWEEN_GLOBAL_DEDUP_QUERIES;

data/src/data_client.rs

@@ -14,7 +14,7 @@ use parutils::{tokio_par_for_each, ParallelError};
 use tempfile::{tempdir_in, TempDir};
 use utils::auth::{AuthConfig, TokenRefresher};
 use utils::progress::ProgressUpdater;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::configurations::*;
 use crate::errors::DataProcessingError;

data/src/data_processing.rs

@@ -13,7 +13,7 @@ use mdb_shard::ShardFileManager;
 use merklehash::MerkleHash;
 use tokio::sync::{Mutex, Semaphore};
 use utils::progress::ProgressUpdater;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::cas_interface::create_cas_client;
 use crate::clean::Cleaner;

data/src/parallel_xorb_uploader.rs

@@ -9,7 +9,7 @@ use merklehash::MerkleHash;
 use tokio::sync::{Mutex, Semaphore};
 use tokio::task::JoinSet;
 use utils::progress::ProgressUpdater;
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::data_processing::CASDataAggregator;
 use crate::errors::DataProcessingError::*;

data/src/remote_shard_interface.rs

@@ -15,7 +15,7 @@ use merklehash::MerkleHash;
 use parutils::tokio_par_for_each;
 use tokio::task::JoinHandle;
 use tracing::{debug, info};
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use super::configurations::{FileQueryPolicy, StorageConfig};
 use super::errors::{DataProcessingError, Result};

hf_xet/Cargo.toml

@@ -19,6 +19,7 @@ pyo3 = { version = "0.23.3", features = [
 error_printer = { path = "../error_printer" }
 data = { path = "../data" }
 utils = { path = "../utils" }
+xet_threadpool = { path = "../xet_threadpool" }
 tokio = { version = "1.36", features = ["full"] }
 parutils = { path = "../parutils" }
 tracing = "0.1.*"

hf_xet/src/log.rs

@@ -6,7 +6,7 @@ use tracing_subscriber::filter::FilterFn;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::{EnvFilter, Layer};
-use utils::ThreadPool;
+use xet_threadpool::ThreadPool;
 
 use crate::log_buffer::{get_telemetry_task, LogBufferLayer, TelemetryTaskInfo, TELEMETRY_PRE_ALLOC_BYTES};
 

hf_xet/src/runtime.rs

@@ -5,7 +5,8 @@ use std::time::Duration;
 use lazy_static::lazy_static;
 use pyo3::exceptions::{PyKeyboardInterrupt, PyRuntimeError};
 use pyo3::prelude::*;
-use utils::threadpool::{MultithreadedRuntimeError, ThreadPool};
+use xet_threadpool::errors::MultithreadedRuntimeError;
+use xet_threadpool::ThreadPool;
 
 use crate::log;
 

utils/Cargo.toml

@@ -9,6 +9,7 @@ path = "src/lib.rs"
 
 [dependencies]
 merklehash = { path = "../merklehash" }
+xet_threadpool = { path = "../xet_threadpool" }
 thiserror = "2.0"
 futures = "0.3.28"
 

utils/src/lib.rs

@@ -4,8 +4,6 @@ pub mod auth;
 pub mod errors;
 pub mod serialization_utils;
 pub mod singleflight;
-pub mod threadpool;
-pub use threadpool::ThreadPool;
 
 mod async_read;
 mod output_bytes;

utils/src/singleflight.rs

@@ -10,6 +10,7 @@
 //!
 //! use futures::future::join_all;
 //! use utils::singleflight::Group;
+//! use xet_threadpool;
 //!
 //! const RES: usize = 7;
 //!
@@ -20,7 +21,7 @@
 //!
 //! #[tokio::main]
 //! async fn main() {
-//!     let threadpool = Arc::new(utils::ThreadPool::new().unwrap());
+//!     let threadpool = Arc::new(xet_threadpool::ThreadPool::new().unwrap());
 //!     let g = Arc::new(Group::<_, ()>::new(threadpool.clone()));
 //!     let mut handlers = Vec::new();
 //!     for _ in 0..10 {
@@ -50,9 +51,9 @@ use parking_lot::RwLock;
 use pin_project::{pin_project, pinned_drop};
 use tokio::sync::{Mutex, Notify};
 use tracing::debug;
+use xet_threadpool::ThreadPool;
 
 pub use crate::errors::SingleflightError;
-use crate::ThreadPool;
 
 type SingleflightResult<T, E> = Result<T, SingleflightError<E>>;
 type CallMap<T, E> = HashMap<String, Arc<Call<T, E>>>;
@@ -366,11 +367,11 @@ mod tests {
     use tokio::sync::{Mutex, Notify};
     use tokio::task::JoinHandle;
     use tokio::time::timeout;
+    use xet_threadpool::ThreadPool;
 
     use super::Group;
     use crate::errors::SingleflightError;
     use crate::singleflight::{Call, OwnerTask};
-    use crate::ThreadPool;
 
     /// A period of time for waiters to wait for a notification from the owner
     /// task. This is expected to be sufficient time for the test futures to

xet_threadpool/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "xet_threadpool"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+tokio = { version = "1.41", features = ["full"] }
+thiserror = "2.0"
+tracing = "0.1.31"
+lazy_static = "1"

xet_threadpool/src/errors.rs

@@ -0,0 +1,18 @@
+use thiserror::Error;
+
+/// Define an error time for spawning external threads.
+#[derive(Debug, Error)]
+#[non_exhaustive]
+pub enum MultithreadedRuntimeError {
+    #[error("Error Initializing Multithreaded Runtime: {0:?}")]
+    RuntimeInitializationError(std::io::Error),
+
+    #[error("Task Panic: {0:?}.")]
+    TaskPanic(tokio::task::JoinError),
+
+    #[error("Task cancelled; possible runtime shutdown in progress ({0}).")]
+    TaskCanceled(String),
+
+    #[error("Unknown task runtime error: {0}")]
+    Other(String),
+}

xet_threadpool/src/lib.rs

@@ -0,0 +1,4 @@
+pub mod errors;
+pub mod threadpool;
+
+pub use threadpool::ThreadPool;

utils/src/threadpool.rs → xet_threadpool/src/threadpool.rs

@@ -3,7 +3,6 @@ use std::future::Future;
 use std::sync::atomic::Ordering::SeqCst;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 
-use thiserror::Error;
 /// This module provides a simple wrapper around Tokio's runtime to create a thread pool
 /// with some default settings. It is intended to be used as a singleton thread pool for
 /// the entire application.
@@ -14,7 +13,7 @@ use thiserror::Error;
 /// # Example
 ///
 /// ```rust
-/// use utils::ThreadPool;
+/// use xet_threadpool::ThreadPool;
 ///
 /// let pool = ThreadPool::new().expect("Error initializing runtime.");
 ///
@@ -51,31 +50,16 @@ use thiserror::Error;
 ///
 /// - `new_threadpool`: Creates a new Tokio runtime with the specified settings.
 use tokio;
-use tokio::task::{JoinError, JoinHandle};
+use tokio::task::JoinHandle;
 use tracing::{debug, error};
 
+use crate::errors::MultithreadedRuntimeError;
+
 const THREADPOOL_NUM_WORKER_THREADS: usize = 4; // 4 active threads
 const THREADPOOL_THREAD_ID_PREFIX: &str = "hf-xet"; // thread names will be hf-xet-0, hf-xet-1, etc.
 const THREADPOOL_STACK_SIZE: usize = 8_000_000; // 8MB stack size
 const THREADPOOL_MAX_BLOCKING_THREADS: usize = 100; // max 100 threads can block IO
 
-/// Define an error time for spawning external threads.
-#[derive(Debug, Error)]
-#[non_exhaustive]
-pub enum MultithreadedRuntimeError {
-    #[error("Error Initializing Multithreaded Runtime: {0:?}")]
-    RuntimeInitializationError(std::io::Error),
-
-    #[error("Task Panic: {0:?}.")]
-    TaskPanic(JoinError),
-
-    #[error("Task cancelled; possible runtime shutdown in progress ({0}).")]
-    TaskCanceled(String),
-
-    #[error("Unknown task runtime error: {0}")]
-    Other(String),
-}
-
 #[derive(Debug)]
 pub struct ThreadPool {
     // This has to allow for exclusive access to enable shutdown when