Introduce distributed planning lifecycle event handlers for replacing TaskEstimator

[gabotechs]

Stack

1. Introduce distributed planning lifecycle event handlers for replacing TaskEstimator #564 <- you are here
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
5. Add dynamic stage built events #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 Make desired task-count handlers async #569.
• We need to add new event handlers without turning TaskEstimator into a generic bucket for every planner lifecycle hook. For example Add dynamic stage built events #570.
• We could be expressing current behavior in terms of just event handler implementations, instead of inlining it in the code, for example: Define routing in terms of events #567.
• Converge into a single pattern all the user extension points, instead of having different shape of APIs for them. For example: Define worker plan rewrites in terms of events #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:

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:

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

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

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

[jayshrivastava]

This PR pretty much LGTM. I got a good sense of it while working on #571. I'm good to merge that into this one and then merge this into main.