Adding input param of num_cores to ThreadPool::new()

constellation-rs · Mar 7, 2021 · ce88de4 · ce88de4
1 parent fb9ee3b
commit ce88de4
Show file tree

Hide file tree

Showing 30 changed files with 70 additions and 39 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
 /target
 **/*.rs.bk
 Cargo.lock
+/.idea
diff --git a/README.md b/README.md
@@ -98,7 +98,7 @@ struct LogLine {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
-    let pool = ThreadPool::new(None)?;
+    let pool = ThreadPool::new(None, None)?;
 
     let rows = Parquet::new(ParquetDirectory::new(S3Directory::new_with(
         AwsRegion::UsEast1,
@@ -133,7 +133,7 @@ use std::error::Error;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
-    let pool = ThreadPool::new(None)?;
+    let pool = ThreadPool::new(None, None)?;
 
     let rows = Parquet::new(ParquetDirectory::new(S3Directory::new_with(
         AwsRegion::UsEast1,
@@ -169,7 +169,7 @@ use std::error::Error;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
-    let pool = ThreadPool::new(None)?;
+    let pool = ThreadPool::new(None, None)?;
 
     let rows = Parquet::new(ParquetDirectory::new(S3Directory::new_with(
         AwsRegion::UsEast1,

diff --git a/benches/csv.rs b/benches/csv.rs
@@ -19,7 +19,7 @@ static RT: Lazy<Runtime> = Lazy::new(|| {
 		.build()
 		.unwrap()
 });
-static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None).unwrap());
+static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None, None).unwrap());
 
 #[derive(Data, Clone, Deserialize, PartialEq, PartialOrd, Debug)]
 struct GameDerived {

diff --git a/benches/in_memory.rs b/benches/in_memory.rs
@@ -19,7 +19,7 @@ static RT: Lazy<Runtime> = Lazy::new(|| {
 		.build()
 		.unwrap()
 });
-static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None).unwrap());
+static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None, None).unwrap());
 
 #[bench]
 fn vec(b: &mut Bencher) {

diff --git a/benches/parquet.rs b/benches/parquet.rs
@@ -19,7 +19,7 @@ static RT: Lazy<Runtime> = Lazy::new(|| {
 		.build()
 		.unwrap()
 });
-static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None).unwrap());
+static POOL: Lazy<ThreadPool> = Lazy::new(|| ThreadPool::new(None, None).unwrap());
 
 #[derive(Data, Clone, PartialEq, Debug)]
 struct TenKayVeeTwo {

diff --git a/examples/cloudfront_logs.rs b/examples/cloudfront_logs.rs
@@ -16,7 +16,7 @@ use amadeus::prelude::*;
 #[allow(unreachable_code)]
 #[tokio::main]
 async fn main() {
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let rows = Cloudfront::new_with(
 		AwsRegion::UsEast1,

diff --git a/examples/columnar.rs b/examples/columnar.rs
@@ -2,7 +2,7 @@ use amadeus::prelude::*;
 
 #[tokio::main]
 async fn main() {
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let rows = Cloudfront::new_with(
 		AwsRegion::UsEast1,

diff --git a/examples/commoncrawl.rs b/examples/commoncrawl.rs
@@ -19,7 +19,7 @@ use amadeus::{data::Webpage, dist::prelude::*};
 async fn main() {
 	return; // TODO: runs for a long time
 
-	let pool = ThreadPool::new(None).unwrap();
+	let pool = ThreadPool::new(None, None).unwrap();
 
 	let webpages = CommonCrawl::new("CC-MAIN-2020-24").await.unwrap();
 

diff --git a/src/pool/process.rs b/src/pool/process.rs
@@ -130,7 +130,7 @@ impl ProcessPoolInner {
 							let receiver = Receiver::<Option<Request>>::new(parent);
 							let sender = Sender::<Result<Response, Panicked>>::new(parent);
 
-							let thread_pool = ThreadPool::new(tasks_per_core).unwrap();
+							let thread_pool = ThreadPool::new(tasks_per_core, None).unwrap();
 
 							while let Some(work) = receiver.recv().await.unwrap() {
 								let ret = panic::catch_unwind(panic::AssertUnwindSafe(|| {

diff --git a/src/pool/thread.rs b/src/pool/thread.rs
@@ -26,14 +26,21 @@ struct ThreadPoolInner {
 #[derive(Debug)]
 pub struct ThreadPool(Arc<ThreadPoolInner>);
 impl ThreadPool {
-	pub fn new(tasks_per_core: Option<usize>) -> io::Result<Self> {
-		let logical_cores = if !cfg!(target_arch = "wasm32") {
-			num_cpus::get()
-		} else {
-			1
-		};
-		let tasks_per_core = tasks_per_core.unwrap_or(DEFAULT_TASKS_PER_CORE);
-		#[cfg(not(target_arch = "wasm32"))]
+	pub fn new(tasks_per_core: Option<usize>, num_cores: Option<usize>) -> io::Result<Self> {
+        let logical_cores = if let Some(cores) = num_cores {
+          cores
+        } else if !cfg!(target_arch = "wasm32") {
+            num_cpus::get()
+        } else {
+            1
+        };
+
+        let tasks_per_core = tasks_per_core.unwrap_or(DEFAULT_TASKS_PER_CORE);
+
+        println!("Logical cores: {}", logical_cores);
+        println!("Tasks per core: {}", tasks_per_core);
+
+        #[cfg(not(target_arch = "wasm32"))]
 		let pool = Pool::new(logical_cores);
 		Ok(ThreadPool(Arc::new(ThreadPoolInner {
 			logical_cores,

diff --git a/tests/cloudfront.rs b/tests/cloudfront.rs
@@ -6,7 +6,7 @@ use std::time::SystemTime;
 #[tokio::test(threaded_scheduler)]
 #[cfg_attr(miri, ignore)]
 async fn cloudfront() {
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let start = SystemTime::now();
 

diff --git a/tests/cloudfront_dist.rs b/tests/cloudfront_dist.rs
@@ -17,7 +17,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/commoncrawl.rs b/tests/commoncrawl.rs
@@ -20,7 +20,7 @@ use amadeus::{data::Webpage, prelude::*};
 async fn commoncrawl() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let webpages = CommonCrawl::new("CC-MAIN-2020-24").await.unwrap();
 	let _ = webpages

diff --git a/tests/commoncrawl_dist.rs b/tests/commoncrawl_dist.rs
@@ -31,7 +31,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/csv.rs b/tests/csv.rs
@@ -9,7 +9,7 @@ use amadeus::prelude::*;
 async fn csv() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	#[derive(Data, Clone, PartialEq, PartialOrd, Debug)]
 	struct GameDerived {

diff --git a/tests/csv_dist.rs b/tests/csv_dist.rs
@@ -20,7 +20,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/csv_wasm.rs b/tests/csv_wasm.rs
@@ -44,7 +44,7 @@ async fn csv() {
 	let timer = web_sys::window().unwrap().performance().unwrap();
 	let start = timer.now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	#[derive(Data, Clone, PartialEq, PartialOrd, Debug)]
 	struct GameDerived {

diff --git a/tests/into_par_stream.rs b/tests/into_par_stream.rs
@@ -5,7 +5,7 @@ use amadeus::prelude::*;
 #[tokio::test(threaded_scheduler)]
 #[cfg_attr(miri, ignore)]
 async fn into_par_stream() {
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	<&[usize] as IntoParallelStream>::into_par_stream(&[1, 2, 3])
 		.map(|a: usize| a)

diff --git a/tests/into_par_stream_dist.rs b/tests/into_par_stream_dist.rs
@@ -19,7 +19,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/json.rs b/tests/json.rs
@@ -9,7 +9,7 @@ use amadeus::prelude::*;
 async fn json() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 	let tasks = 100;
 
 	#[derive(Data, Clone, PartialEq, PartialOrd, Debug)]

diff --git a/tests/json_dist.rs b/tests/json_dist.rs
@@ -20,7 +20,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool, 100).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/panic.rs b/tests/panic.rs
@@ -8,7 +8,7 @@ use amadeus::prelude::*;
 async fn panic() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let res = AssertUnwindSafe((0i32..1_000).into_par_stream().for_each(pool, |i| {
 		if i == 500 {

diff --git a/tests/panic_dist.rs b/tests/panic_dist.rs
@@ -21,7 +21,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/parquet.rs b/tests/parquet.rs
@@ -14,7 +14,7 @@ use amadeus::prelude::*;
 async fn parquet() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let rows = Parquet::<_, Value>::new(vec![
 		PathBuf::from("amadeus-testing/parquet/cf-accesslogs/year=2018/month=11/day=02/part-00176-17868f39-cd99-4b60-bb48-8daf9072122e.c000.snappy.parquet"),

diff --git a/tests/parquet_dist.rs b/tests/parquet_dist.rs
@@ -23,7 +23,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/parquet_wasm.rs b/tests/parquet_wasm.rs
@@ -44,7 +44,7 @@ async fn parquet() {
 	let timer = web_sys::window().unwrap().performance().unwrap();
 	let start = timer.now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	let rows = Parquet::<_, Value>::new(vec![
 		PathBuf::from("amadeus-testing/parquet/cf-accesslogs/year=2018/month=11/day=02/part-00176-17868f39-cd99-4b60-bb48-8daf9072122e.c000.snappy.parquet"),

diff --git a/tests/postgres.rs b/tests/postgres.rs
@@ -9,7 +9,7 @@ use amadeus::prelude::*;
 async fn postgres() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 
 	#[derive(Data, Clone, PartialEq, PartialOrd, Debug)]
 	struct Weather {

diff --git a/tests/postgres_dist.rs b/tests/postgres_dist.rs
@@ -18,7 +18,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool).await
 			};
 			#[cfg(feature = "constellation")]

diff --git a/tests/threads.rs b/tests/threads.rs
@@ -1,7 +1,8 @@
 use futures::future::join_all;
 use rand::{rngs::SmallRng, Rng, SeedableRng};
 use std::{
-	convert::TryInto, time::{Duration, SystemTime}
+	convert::TryInto, time::{Duration, SystemTime},
+    collections::HashSet,
 };
 use tokio::time::delay_for as sleep;
 
@@ -12,7 +13,7 @@ use amadeus::dist::prelude::*;
 async fn threads() {
 	let start = SystemTime::now();
 
-	let pool = &ThreadPool::new(None).unwrap();
+	let pool = &ThreadPool::new(None, None).unwrap();
 	let parallel = 1000;
 
 	join_all((0..parallel).map(|i| async move {
@@ -29,3 +30,25 @@ async fn threads() {
 
 	println!("in {:?}", start.elapsed().unwrap());
 }
+
+#[tokio::test(threaded_scheduler)]
+#[cfg_attr(miri, ignore)]
+async fn user_set_core_count() {
+    let num_cores = 4;
+
+    let pool = &ThreadPool::new(Some(1), Some(num_cores)).unwrap();
+    let parallel = 1000;
+
+    let ret = join_all((0..parallel).map(|_| async move {
+        pool
+            .spawn(move || async move {
+                format!("{:?}", std::thread::current().id())
+            })
+            .await.unwrap()
+    })).await;
+
+    let unique_thread_ids: HashSet<String> = ret.into_iter().collect();
+
+    println!("Number of cores used: {}", unique_thread_ids.len());
+    assert_eq!(unique_thread_ids.len(), num_cores);
+}
diff --git a/tests/threads_dist.rs b/tests/threads_dist.rs
@@ -23,7 +23,7 @@ fn main() {
 		.unwrap()
 		.block_on(async {
 			let thread_pool_time = {
-				let thread_pool = ThreadPool::new(None).unwrap();
+				let thread_pool = ThreadPool::new(None, None).unwrap();
 				run(&thread_pool, 1000).await
 			};
 			#[cfg(feature = "constellation")]