Cleaner interface.

Author: JSOD11

src/execution_plans/distributed.rs

@@ -1,7 +1,7 @@
 use crate::common::require_one_child;
 use crate::config_extension_ext::get_config_extension_propagation_headers;
 use crate::distributed_planner::NetworkBoundaryExt;
-use crate::networking::{get_distributed_task_router, get_distributed_worker_resolver};
+use crate::networking::{RouterInfo, get_distributed_task_router, get_distributed_worker_resolver};
 use crate::passthrough_headers::get_passthrough_headers;
 use crate::protobuf::{DistributedCodec, tonic_status_to_datafusion_error};
 use crate::stage::{ExecutionTask, Stage};
@@ -131,7 +131,11 @@ impl DistributedExec {
                 .iter()
                 .enumerate()
                 .map(|(i, _)| {
-                    let url = task_router.route_task(start_idx, i, &urls);
+                    let url = task_router.route(RouterInfo {
+                        task_number: i,
+                        urls: urls.clone(),
+                        stage_seed: start_idx,
+                    });
                     let execution_task = ExecutionTask {
                         url: Some(url.clone()),
                     };

src/networking/mod.rs

@@ -7,7 +7,7 @@ pub use channel_resolver::{
     get_distributed_channel_resolver,
 };
 pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver};
-pub use task_router::{TaskRouter, get_distributed_task_router};
+pub use task_router::{RouterInfo, TaskRouter, get_distributed_task_router};
 pub(crate) use task_router::{TaskRouterExtension, set_distributed_task_router};
 
 pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver};

src/networking/task_router.rs

@@ -3,17 +3,19 @@ use datafusion::prelude::SessionConfig;
 use std::sync::Arc;
 use url::Url;
 
-/// Allows users to choose which worker should execute a given distributed task.
-pub trait TaskRouter {
-    /// Returns the index of the worker that should execute the task. The default implementation
-    /// assigns tasks to workers round-robin starting from start_index.
-    fn route_fn(&self, start_index: usize, task_number: usize, num_urls: usize) -> usize {
-        (start_index + task_number) % num_urls
-    }
+pub struct RouterInfo {
+    pub task_number: usize,
+    pub urls: Vec<Url>,
+    pub stage_seed: usize,
+}
 
-    /// Returns the url of the worker that should execute the task.
-    fn route_task(&self, start_index: usize, task_number: usize, urls: &[Url]) -> Url {
-        urls[self.route_fn(start_index, task_number, urls.len())].clone()
+/// Allows users to route tasks to worker nodes.
+pub trait TaskRouter {
+    /// Returns the url of the worker that should execute the task. The default implementation
+    /// assigns tasks to workers round-robin starting from stage_seed.
+    fn route(&self, router_info: RouterInfo) -> Url {
+        router_info.urls[router_info.stage_seed + router_info.task_number % router_info.urls.len()]
+            .clone()
     }
 }