datafusion_distributed/events/
route_tasks.rs

1use super::common::EventHandlerChain;
2use datafusion::error::Result;
3use datafusion::execution::TaskContext;
4use datafusion::physical_plan::ExecutionPlan;
5use std::sync::Arc;
6use url::Url;
7
8/// Information supplied when the coordinator assigns a stage's tasks to workers.
9///
10/// A routing response contains one worker URL per task, in task-index order. The coordinator
11/// validates that a response has exactly [`Self::handle`] URLs before it starts the stage.
12#[derive(Clone)]
13pub struct RouteTasksEvent<'a> {
14    /// The task context active for the query being coordinated.
15    pub task_ctx: Arc<TaskContext>,
16    /// The head execution plan of the stage whose tasks are being routed.
17    /// WARNING: this is never going to be a custom leaf node, this is the head node of the fragment
18    ///  of the plan that contains the custom leaf node.
19    pub plan: &'a Arc<dyn ExecutionPlan>,
20    /// The number of task slots that need a worker assignment.
21    pub task_count: usize,
22}
23
24/// Worker assignments returned by a [`RouteTasksHandler`].
25pub struct RouteTasksEventResponse {
26    /// One worker URL per task, ordered by task index.
27    pub urls: Vec<Url>,
28}
29
30impl RouteTasksEventResponse {
31    /// Returns a response assigning tasks to `urls` in order.
32    ///
33    /// The coordinator rejects the response unless `urls.len()` equals
34    /// [`RouteTasksEvent::handle`].
35    pub fn new(urls: Vec<Url>) -> Self {
36        Self { urls }
37    }
38}
39
40/// Optionally assigns a stage's task slots to worker URLs.
41///
42/// Handlers are evaluated in reverse registration order. Return `Some(Ok(_))` to select a
43/// complete routing response and stop dispatch, or `None` to defer to earlier handlers. If every
44/// handler returns `None`, the coordinator assigns the tasks round-robin, from a randomized worker
45/// offset. Returning `Some(Err(_))` aborts execution before tasks are submitted.
46pub trait RouteTasksHandler: Send + Sync + 'static {
47    /// Optionally assigns the tasks described by `ev` to specific worker URLs.
48    ///
49    /// Return `None` when this handler does not apply. A successful response must provide exactly
50    /// one URL for every task, in task-index order.
51    fn handle(&self, ev: RouteTasksEvent) -> Option<Result<RouteTasksEventResponse>>;
52}
53
54impl<F> RouteTasksHandler for F
55where
56    F: Send + Sync + 'static,
57    F: for<'a> Fn(RouteTasksEvent<'a>) -> Option<Result<RouteTasksEventResponse>>,
58{
59    fn handle(&self, ev: RouteTasksEvent) -> Option<Result<RouteTasksEventResponse>> {
60        self(ev)
61    }
62}
63
64impl RouteTasksHandler for Arc<dyn RouteTasksHandler> {
65    fn handle(&self, ev: RouteTasksEvent) -> Option<Result<RouteTasksEventResponse>> {
66        self.as_ref().handle(ev)
67    }
68}
69
70pub(crate) type RouteTasksHandlers = EventHandlerChain<dyn RouteTasksHandler>;
71
72impl RouteTasksHandlers {
73    pub(crate) fn handle(ev: RouteTasksEvent) -> Option<Result<RouteTasksEventResponse>> {
74        ev.task_ctx
75            .session_config()
76            .get_extension::<RouteTasksHandlers>()?
77            .find_map(|handler| handler.handle(ev.clone()))
78    }
79}
datafusion_distributed/events/route_tasks.rs

datafusion_distributed/events/
route_tasks.rs
