Skip to content

Make desired task-count handlers async - #569

Open
gabotechs wants to merge 2 commits into
gabrielmusat/worker-plan-rewrite-eventsfrom
gabrielmusat/async-desired-task-count
Open

Make desired task-count handlers async#569
gabotechs wants to merge 2 commits into
gabrielmusat/worker-plan-rewrite-eventsfrom
gabrielmusat/async-desired-task-count

Conversation

@gabotechs

@gabotechs gabotechs commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack

  1. Introduce distributed planning lifecycle event handlers for replacing TaskEstimator #564
  2. Define routing in terms of events #567
  3. Define worker plan rewrites in terms of events #568
  4. Make desired task-count handlers async #569 <- you are here
  5. Add dynamic stage built events #570

Makes only desired task-count handlers asynchronous, so implementations can await metadata before returning a hint. The event still borrows SessionConfig, avoiding configuration copies.

#[async_trait]
impl DesiredTaskCountHandler for MyHandler {
    async fn handle(&self, event: DesiredTaskCountEvent<'_>)
        -> Option<DesiredTaskCountEventResponse>
    {
        Some(DesiredTaskCountEventResponse::desired(8))
    }
}

This was referenced Jul 27, 2026
Merged
Merged
Open
Open
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from 019c8da to d111ed0 Compare July 27, 2026 08:37
@gabotechs
gabotechs force-pushed the gabrielmusat/async-desired-task-count branch from 3c8bb87 to 68120e0 Compare July 27, 2026 08:45
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from d111ed0 to ca9c4f9 Compare July 27, 2026 08:53
@gabotechs
gabotechs force-pushed the gabrielmusat/async-desired-task-count branch from ca40839 to 30bbc84 Compare July 27, 2026 08:53
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from ca9c4f9 to ba2960f Compare July 27, 2026 09:05
@gabotechs
gabotechs force-pushed the gabrielmusat/async-desired-task-count branch from 30bbc84 to 9ff5732 Compare July 27, 2026 09:05
@gabotechs
gabotechs marked this pull request as ready for review July 27, 2026 09:54
@gabotechs gabotechs mentioned this pull request Jul 28, 2026
Open
@gabotechs
gabotechs force-pushed the gabrielmusat/async-desired-task-count branch 2 times, most recently from 6846880 to a6ac3d4 Compare July 28, 2026 09:20
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from ba2960f to f9fde4f Compare July 28, 2026 09:21
gabotechs added a commit that referenced this pull request Jul 28, 2026
@gabotechs
Introduce distributed planning lifecycle event handlers for replacing…
e3ac0fb 
… TaskEstimator (#564)

## Stack

1. #564
<- you are here
2. #567
3. #568
4. #569
5. #570

---

Introduces an event-driven API for reacting to distributed planning
lifecycle events, and replaces the previous TaskEstimator methods with
it.

## Motivation

- We need to extend existing methods with new capabilities without
imposing those changes on unrelated implementations. For example
#569.
- We need to add new event handlers without turning `TaskEstimator` into
a generic bucket for every planner lifecycle hook. For example
#570.
- We could be expressing current behavior in terms of just event handler
implementations, instead of inlining it in the code, for example:
#567.
- Converge into a single pattern all the user extension points, instead
of having different shape of APIs for them. For example:
#568

`TaskEstimator` was a convenient initial extension point, but its
methods run at different points in planning, observe different context,
and have different result semantics. Task-count estimation is evaluated
while traversing the plan; leaf specialization only runs after a stage's
final task count is known; routing runs later on the coordinator.
Keeping these operations in one trait couples otherwise independent
APIs.

This event-specific pattern is easier to extend without breaking API
consumers: a capability can evolve within the event that owns it instead
of requiring every `TaskEstimator` implementation to adopt a broader
trait. It is also easier to document and expose to users. Dedicated
event types give us focused API pages and code documentation, and
provide a quick index of the project’s extension points by the planner
lifecycle events where users can hook their code.

## What changed

Previously, one trait covered task-count estimation, leaf
specialization, and task routing:

```rust
trait TaskEstimator {
    fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions)
        -> Option<TaskEstimation>;
    fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize,
        cfg: &ConfigOptions) -> Result<Option<Arc<dyn ExecutionPlan>>>;
    fn route_tasks(&self, ctx: &TaskRoutingContext<'_>)
        -> Result<Option<Vec<Url>>>;
}
```

The old API registered one implementation containing every hook:

```rust
SessionStateBuilder::new()
    .with_distributed_task_estimator(MyTaskEstimator);
```

Those phases now have dedicated event handlers, which can be ordinary
functions and are registered independently:

```rust
fn desired_task_count(
    _event: DesiredTaskCountEvent,
) -> Option<DesiredTaskCountEventResponse> {
    None
}

fn scale_up_leaf_node(
    _event: ScaleUpLeafNodeEvent,
) -> Option<Result<ScaleUpLeafNodeEventResponse>> {
    None
}

fn route_tasks(
    _event: RouteTasksEvent,
) -> Option<Result<RouteTasksEventResponse>> {
    None
}

SessionStateBuilder::new()
    .with_distributed_desired_task_count_handler(desired_task_count)
    .with_distributed_scale_up_leaf_node_handler(scale_up_leaf_node)
    .with_distributed_route_tasks_handler(route_tasks);
```
@gabotechs
gabotechs force-pushed the gabrielmusat/worker-plan-rewrite-events branch from f9fde4f to 335a90b Compare July 28, 2026 15:14
@gabotechs
gabotechs force-pushed the gabrielmusat/async-desired-task-count branch from a6ac3d4 to a5e0b27 Compare July 28, 2026 15:14
gabotechs added a commit that referenced this pull request Jul 31, 2026
@gabotechs @jayshrivastava
Define routing in terms of events (#567)
f163e2b 
## Stack

1. #564
2. #567
<- you are here
3. #568
4. #569
5. #570

---

Moves routing policy into ordered `RouteTasksHandler`s. Built-ins now
cover child co-location, coordinator co-location, and randomized
fallback, while custom handlers can override them.

```rust
config.with_distributed_route_tasks_handler(|event| {
    Some(Ok(RouteTasksEventResponse::new(
        vec![worker_url.clone(); event.task_count],
    )))
});
```

---------

Co-authored-by: Jayant Shrivastava <jshrivastava03@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jayshrivastava jayshrivastava jayshrivastava approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@gabotechs @jayshrivastava