Make desired task-count handlers async

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::<ChildrenIsolatorUnionExec>() {
         // 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

src/events/common.rs

@@ -25,6 +25,10 @@ impl<H: ?Sized> Clone for EventHandlerChain<H> {
 }
 
 impl<H: ?Sized> EventHandlerChain<H> {
+    pub(super) fn iter(&self) -> impl Iterator<Item = &H> {
+        self.custom.iter().chain(&self.builtin).map(AsRef::as_ref)
+    }
+
     pub(super) fn find_map<T>(&self, mut f: impl FnMut(&H) -> Option<T>) -> Option<T> {
         // Give priority to custom handlers registered by users.
         if let Some(res) = self.custom.iter().find_map(|handler| f(handler.as_ref())) {

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(())

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<DesiredTaskCountEventResponse>;
+    async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option<DesiredTaskCountEventResponse>;
 }
 
 impl From<TaskCountAnnotation> for usize {
@@ -120,37 +125,48 @@ impl TaskCountAnnotation {
     }
 }
 
+#[async_trait]
 impl<F> DesiredTaskCountHandler for F
 where
     F: Send + Sync + 'static,
     F: for<'a> Fn(DesiredTaskCountEvent<'a>) -> Option<DesiredTaskCountEventResponse>,
 {
-    fn handle(&self, ev: DesiredTaskCountEvent) -> Option<DesiredTaskCountEventResponse> {
+    async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option<DesiredTaskCountEventResponse> {
         self(ev)
     }
 }
 
+#[async_trait]
 impl DesiredTaskCountHandler for usize {
-    fn handle(&self, ev: DesiredTaskCountEvent) -> Option<DesiredTaskCountEventResponse> {
+    async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option<DesiredTaskCountEventResponse> {
         ev.plan
             .children()
             .is_empty()
             .then(|| DesiredTaskCountEventResponse::desired(*self))
     }
 }
 
+#[async_trait]
 impl DesiredTaskCountHandler for Arc<dyn DesiredTaskCountHandler> {
-    fn handle(&self, ev: DesiredTaskCountEvent) -> Option<DesiredTaskCountEventResponse> {
-        self.as_ref().handle(ev)
+    async fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option<DesiredTaskCountEventResponse> {
+        self.as_ref().handle(ev).await
     }
 }
 
 pub(crate) type DesiredTaskCountHandlers = EventHandlerChain<dyn DesiredTaskCountHandler>;
 
 impl DesiredTaskCountHandlers {
-    pub(crate) fn handle(ev: DesiredTaskCountEvent) -> Option<DesiredTaskCountEventResponse> {
-        ev.session_config
-            .get_extension::<DesiredTaskCountHandlers>()?
-            .find_map(|handler| handler.handle(ev))
+    pub(crate) async fn handle(
+        ev: DesiredTaskCountEvent<'_>,
+    ) -> Option<DesiredTaskCountEventResponse> {
+        let handlers = ev
+            .session_config
+            .get_extension::<DesiredTaskCountHandlers>()?;
+        for handler in handlers.iter() {
+            if let Some(response) = handler.handle(ev).await {
+                return Some(response);
+            }
+        }
+        None
     }
 }