diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 6e433ec6..75e6813d 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -237,7 +237,7 @@ async fn _inject_network_boundaries( plan: &plan, session_config: nb_ctx.cfg, }; - return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev) { + return if let Some(estimate) = DesiredTaskCountHandlers::handle(ev).await { Ok(nb_ctx.plan_with_task_count(plan, estimate.task_count.limit(nb_ctx.max_tasks()?))) } else { // We could not determine how many tasks this leaf node should run on, so @@ -261,7 +261,9 @@ async fn _inject_network_boundaries( plan: &plan, session_config: nb_ctx.cfg, }; - let mut task_count = DesiredTaskCountHandlers::handle(ev).map_or(Desired(1), |v| v.task_count); + let mut task_count = DesiredTaskCountHandlers::handle(ev) + .await + .map_or(Desired(1), |v| v.task_count); if plan.is::() { // Isolating unions have the chance to decide how many tasks they should run on. If there // is a union with a bunch of children, the user might want to increase parallelism and the diff --git a/src/events/common.rs b/src/events/common.rs index 7f715059..fd5c02e4 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -25,6 +25,10 @@ impl Clone for EventHandlerChain { } impl EventHandlerChain { + pub(super) fn iter(&self) -> impl Iterator { + self.custom.iter().chain(&self.builtin).map(AsRef::as_ref) + } + pub(super) fn find_map(&self, mut f: impl FnMut(&H) -> Option) -> Option { // Give priority to custom handlers registered by users. if let Some(res) = self.custom.iter().find_map(|handler| f(handler.as_ref())) { diff --git a/src/events/defaults/file_scan_config.rs b/src/events/defaults/file_scan_config.rs index 32aef96c..f11451b8 100644 --- a/src/events/defaults/file_scan_config.rs +++ b/src/events/defaults/file_scan_config.rs @@ -115,6 +115,7 @@ mod tests { plan: &plan, session_config: &cfg, }) + .await .expect("a handler should respond"); assert_eq!(response.task_count.as_usize(), 10); Ok(()) @@ -131,6 +132,7 @@ mod tests { plan: &plan, session_config: &cfg, }) + .await .expect("a handler should respond"); assert_eq!(response.task_count.as_usize(), 30); Ok(()) diff --git a/src/events/desired_task_count.rs b/src/events/desired_task_count.rs index 85981b79..cd6cf28e 100644 --- a/src/events/desired_task_count.rs +++ b/src/events/desired_task_count.rs @@ -1,5 +1,6 @@ use super::TaskCountAnnotation::{Desired, Maximum}; use super::common::EventHandlerChain; +use async_trait::async_trait; use datafusion::execution::config::SessionConfig; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; @@ -73,10 +74,14 @@ impl DesiredTaskCountEventResponse { } } +#[async_trait] pub trait DesiredTaskCountHandler: Send + Sync + 'static { /// Function applied to each node that returns a [DesiredTaskCountEventResponse] hinting how /// many tasks should be used in the [Stage] containing that node. /// + /// Handlers are asynchronous and may await metadata or external services. Handler functions + /// return a [`DesiredTaskCountFuture`] so their futures can borrow from the event. + /// /// All the [TaskEstimator] registered in the session will be applied to the node /// until one returns an estimation. /// @@ -86,7 +91,7 @@ pub trait DesiredTaskCountHandler: Send + Sync + 'static { /// that the leaf node cannot be distributed across tasks. /// - If the node is a normal node in the plan, then the maximum task count from its children /// is inherited. - fn handle(&self, ev: DesiredTaskCountEvent) -> Option; + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option; } impl From for usize { @@ -120,18 +125,20 @@ impl TaskCountAnnotation { } } +#[async_trait] impl DesiredTaskCountHandler for F where F: Send + Sync + 'static, F: for<'a> Fn(DesiredTaskCountEvent<'a>) -> Option, { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { self(ev) } } +#[async_trait] impl DesiredTaskCountHandler for usize { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { ev.plan .children() .is_empty() @@ -139,18 +146,27 @@ impl DesiredTaskCountHandler for usize { } } +#[async_trait] impl DesiredTaskCountHandler for Arc { - fn handle(&self, ev: DesiredTaskCountEvent) -> Option { - self.as_ref().handle(ev) + async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { + self.as_ref().handle(ev).await } } pub(crate) type DesiredTaskCountHandlers = EventHandlerChain; impl DesiredTaskCountHandlers { - pub(crate) fn handle(ev: DesiredTaskCountEvent) -> Option { - ev.session_config - .get_extension::()? - .find_map(|handler| handler.handle(ev)) + pub(crate) async fn handle( + ev: DesiredTaskCountEvent<'_>, + ) -> Option { + let handlers = ev + .session_config + .get_extension::()?; + for handler in handlers.iter() { + if let Some(response) = handler.handle(ev).await { + return Some(response); + } + } + None } }