Add --threads and --workers options in tpch benchmarks

Author: gabotechs

benchmarks/src/bin/dfbench.rs

@@ -30,12 +30,14 @@ enum Options {
 }
 
 // Main benchmark runner entrypoint
-#[tokio::main]
-pub async fn main() -> Result<()> {
+pub fn main() -> Result<()> {
     env_logger::init();
 
     match Options::from_args() {
-        Options::Tpch(opt) => Box::pin(opt.run()).await,
-        Options::TpchConvert(opt) => opt.run().await,
+        Options::Tpch(opt) => opt.run(),
+        Options::TpchConvert(opt) => {
+            let rt = tokio::runtime::Runtime::new()?;
+            rt.block_on(async { opt.run().await })
+        }
     }
 }

benchmarks/src/tpch/run.rs

@@ -20,9 +20,6 @@ use super::{
     TPCH_QUERY_START_ID, TPCH_TABLES,
 };
 use async_trait::async_trait;
-use std::path::PathBuf;
-use std::sync::Arc;
-
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::arrow::util::pretty::{self, pretty_format_batches};
 use datafusion::common::instant::Instant;
@@ -40,18 +37,25 @@ use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
+use datafusion_distributed::MappedDistributedSessionBuilderExt;
+use std::path::PathBuf;
+use std::sync::Arc;
 
 use crate::util::{
     BenchmarkRun, CommonOpt, InMemoryCacheExecCodec, InMemoryDataSourceRule, QueryResult,
     WarmingUpMarker,
 };
-use datafusion_distributed::test_utils::localhost::start_localhost_context;
+use datafusion_distributed::test_utils::localhost::{
+    get_free_ports, spawn_flight_service, start_localhost_context, LocalHostChannelResolver,
+};
 use datafusion_distributed::{
     DistributedExt, DistributedPhysicalOptimizerRule, DistributedSessionBuilder,
     DistributedSessionBuilderContext,
 };
 use log::info;
 use structopt::StructOpt;
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
 
 // hack to avoid `default_value is meaningless for bool` errors
 type BoolDefaultTrue = bool;
@@ -113,6 +117,14 @@ pub struct RunOpt {
     /// Number of partitions per task.
     #[structopt(long = "ppt")]
     partitions_per_task: Option<usize>,
+
+    /// Number of physical threads per worker (default 1)
+    #[structopt(long, default_value = "1")]
+    workers: usize,
+
+    /// Number of physical threads per worker
+    #[structopt(long)]
+    threads: Option<usize>,
 }
 
 #[async_trait]
@@ -156,7 +168,50 @@ impl DistributedSessionBuilder for RunOpt {
 }
 
 impl RunOpt {
-    pub async fn run(mut self) -> Result<()> {
+    pub fn spawn_workers(self) -> Vec<(tokio::runtime::Runtime, JoinHandle<()>)> {
+        let ports = get_free_ports(self.workers);
+        let channel_resolver = LocalHostChannelResolver::new(ports.clone());
+        let threads_per_worker = self.threads;
+        let session_builder = self.map(move |builder: SessionStateBuilder| {
+            let channel_resolver = channel_resolver.clone();
+            Ok(builder
+                .with_distributed_channel_resolver(channel_resolver)
+                .build())
+        });
+        let mut handles = vec![];
+        for port in ports {
+            let session_builder = session_builder.clone();
+            let rt = tokio::runtime::Builder::new_multi_thread()
+                .worker_threads(threads_per_worker.unwrap_or(get_available_parallelism()))
+                .enable_all()
+                .build()
+                .unwrap();
+            let handle = rt.spawn(async move {
+                let listener = TcpListener::bind(format!("127.0.0.1:{port}"))
+                    .await
+                    .unwrap();
+                spawn_flight_service(session_builder, listener)
+                    .await
+                    .unwrap();
+            });
+
+            handles.push((rt, handle));
+        }
+        handles
+    }
+
+    pub fn run(self) -> Result<()> {
+        let _handle = self.clone().spawn_workers();
+
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(self.threads.unwrap_or(get_available_parallelism()))
+            .enable_all()
+            .build()?;
+
+        rt.block_on(async move { self._run().await })
+    }
+
+    pub async fn _run(mut self) -> Result<()> {
         let (ctx, _guard) = start_localhost_context(1, self.clone()).await;
         println!("Running benchmarks with the following options: {self:?}");
         let query_range = match self.query {

src/test_utils/localhost.rs

@@ -17,6 +17,17 @@ use tokio::net::TcpListener;
 use tonic::transport::{Channel, Server};
 use url::Url;
 
+pub fn get_free_ports(n: usize) -> Vec<u16> {
+    let listeners = (0..n)
+        .map(|_| std::net::TcpListener::bind("127.0.0.1:0"))
+        .collect::<Result<Vec<_>, _>>()
+        .expect("Failed to bind to address");
+    listeners
+        .iter()
+        .map(|listener| listener.local_addr().unwrap().port())
+        .collect()
+}
+
 pub async fn start_localhost_context<B>(
     num_workers: usize,
     session_builder: B,
@@ -25,23 +36,7 @@ where
     B: DistributedSessionBuilder + Send + Sync + 'static,
     B: Clone,
 {
-    let listeners = futures::future::try_join_all(
-        (0..num_workers)
-            .map(|_| TcpListener::bind("127.0.0.1:0"))
-            .collect::<Vec<_>>(),
-    )
-    .await
-    .expect("Failed to bind to address");
-
-    let ports: Vec<u16> = listeners
-        .iter()
-        .map(|listener| {
-            listener
-                .local_addr()
-                .expect("Failed to get local address")
-                .port()
-        })
-        .collect();
+    let ports = get_free_ports(num_workers);
 
     let channel_resolver = LocalHostChannelResolver::new(ports.clone());
     let session_builder = session_builder.map(move |builder: SessionStateBuilder| {
@@ -51,8 +46,11 @@ where
             .build())
     });
     let mut join_set = JoinSet::new();
-    for listener in listeners {
+    for port in ports {
         let session_builder = session_builder.clone();
+        let listener = TcpListener::bind(format!("127.0.0.1:{port}"))
+            .await
+            .unwrap();
         join_set.spawn(async move {
             spawn_flight_service(session_builder, listener)
                 .await