From 335a90b0c51d75a770b95dd7062cc0f7c2778a0f Mon Sep 17 00:00:00 2001 From: Gabriel Date: Sat, 25 Jul 2026 09:54:19 +0200 Subject: [PATCH] Define worker plan rewrites in terms of events --- docs/source/advanced/03-plan-hooks.md | 55 ++++++++----- docs/source/user-guide/03-worker.md | 2 +- src/distributed_ext.rs | 72 ++++++++++++++++- src/events/common.rs | 11 +++ src/events/mod.rs | 5 ++ src/events/worker_plan_rewrite.rs | 74 ++++++++++++++++++ src/lib.rs | 3 +- src/worker/impl_coordinator_channel.rs | 11 +-- src/worker/worker_service.rs | 39 ---------- tests/worker_plan_hook.rs | 103 ++++++++++++++----------- 10 files changed, 260 insertions(+), 115 deletions(-) create mode 100644 src/events/worker_plan_rewrite.rs diff --git a/docs/source/advanced/03-plan-hooks.md b/docs/source/advanced/03-plan-hooks.md index 5ec13efd4..1de218010 100644 --- a/docs/source/advanced/03-plan-hooks.md +++ b/docs/source/advanced/03-plan-hooks.md @@ -1,36 +1,49 @@ -# Plan hooks +# Worker plan rewrite handlers -`Worker::add_on_plan_hook` registers callbacks that run after the worker session -has been built and the physical plan has been decoded, but before the task plan -is registered for execution. It's a hook for **worker-local** rewrites of the -fragment a worker is about to run. +`DistributedExt::with_distributed_worker_plan_rewrite_handler` registers handlers that run after +the worker session has been built and the physical plan has been decoded, but before the task plan +is registered for execution. It is intended for **worker-local** rewrites of the fragment a worker +is about to run. -Each hook receives the decoded `ExecutionPlan` and the per-query `SessionConfig`, -and returns the (possibly rewritten) plan: +Register handlers in the `SessionStateBuilder` used by each worker's `WorkerSessionBuilder`. +Registering one on the coordinator's session has no effect: handlers are not sent to workers with +stage plans. + +Each handler receives a `WorkerPlanRewriteEvent` and returns the possibly rewritten plan: ```rust -worker.add_on_plan_hook(|plan, session_config| { - let rule = your_physical_optimizer_rule(); - rule.optimize(plan, session_config.options()) -}); +# use datafusion::common::Result; +# use datafusion::execution::SessionState; +# use datafusion_distributed::{DistributedExt, Worker, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerQueryContext}; +async fn build_worker_session(ctx: WorkerQueryContext) -> Result { + Ok(ctx + .builder + .with_distributed_worker_plan_rewrite_handler(|event: WorkerPlanRewriteEvent<'_>| { + Ok(WorkerPlanRewriteEventResponse::new(event.plan)) + }) + .build()) +} + +let _worker = Worker::from_session_builder(build_worker_session); ``` -Hooks run in registration order — each hook sees the plan produced by the +Handlers run in registration order—each handler sees the plan produced by the previous one. -## What hooks may and may not do +## What handlers may and may not do + +Treat handlers as trusted, worker-local rewrites. Instrumentation and semantics-preserving +physical optimizer rules that retain the stage topology are appropriate uses. A handler must: -Treat hooks as trusted, worker-local rewrites. Transparent instrumentation -wrappers and semantics-preserving physical optimizer rules are appropriate uses. +- Keep the same plan topology; do not add or remove nodes or edges. +- Produce the same rows, including multiplicity, from every output partition. +- Preserve the head node's output schema, partitioning, ordering, boundedness, and emission type. -The returned plan **must preserve the contract the coordinating context planned -for**: row semantics, output schema, partitioning, and ordering requirements. Do -**not** use a hook to add or remove rows or columns, repartition the stage, or -otherwise re-plan distributed execution. If a hook returns an error, the -distributed query fails when the coordinating context tries to execute that task. +Intermediate nodes may change their output schema. If a handler returns an error, the distributed +query fails when the coordinating context tries to execute that task. ```{note} -A plan hook acts on a single worker's copy of a stage, after distribution has +A worker plan rewrite handler acts on a single worker's copy of a stage, after distribution has already been decided. It is **not** a way to change how a query is split across the cluster — for that, see [Building Custom Distributed Plans](05-custom-distributed-plans.md). diff --git a/docs/source/user-guide/03-worker.md b/docs/source/user-guide/03-worker.md index 3084cda6f..55a772dd8 100644 --- a/docs/source/user-guide/03-worker.md +++ b/docs/source/user-guide/03-worker.md @@ -109,7 +109,7 @@ object-store registrations (pass them via `Worker::from_session_builder` / A few more worker capabilities have their own pages: -- [Plan hooks](../advanced/03-plan-hooks.md) — run worker-local rewrites on each decoded plan +- [Worker plan rewrite handlers](../advanced/03-plan-hooks.md) — run worker-local rewrites on each decoded plan before it executes. - [Worker versioning](../advanced/07-worker-versioning.md) — tag workers with a version and route queries only to compatible workers during rolling deployments. diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index a61eb0a38..800d973cd 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -4,7 +4,8 @@ use crate::config_extension_ext::{ }; use crate::events::{ DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers, - ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, + ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, WorkerPlanRewriteHandler, + WorkerPlanRewriteHandlers, }; use crate::passthrough_headers::set_passthrough_headers; use crate::protocol::set_distributed_channel_resolver; @@ -690,6 +691,52 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_route_tasks_handler] but with an in-place /// mutation. fn set_distributed_route_tasks_handler(&mut self, handler: T); + + /// Registers a handler that rewrites a decoded worker stage plan before it is executed. + /// + /// Handlers registered with this method need to meet the following criteria: + /// + /// - The topology of the plan must remain the same: no nodes or edges may be added or removed. + /// - Each output partition must produce the same rows, including their multiplicity. + /// - The output schema, partitioning, ordering, boundedness, and emission type of the head + /// node must remain unchanged. Intermediate nodes may change their output schema. + /// + /// Handlers run in registration order, each receiving the plan returned by the preceding + /// handler. Returning an error aborts worker plan registration. + /// + /// Register this handler on the [`SessionStateBuilder`] used by each + /// [`crate::WorkerSessionBuilder`]. Registering it on the coordinating session has no effect: + /// the handler is not sent to workers with the stage plan. + /// + /// ```rust + /// # use datafusion::common::Result; + /// # use datafusion::execution::SessionState; + /// # use datafusion_distributed::{DistributedExt, Worker, WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerQueryContext}; + /// + /// async fn build_worker_session(ctx: WorkerQueryContext) -> Result { + /// Ok(ctx + /// .builder + /// .with_distributed_worker_plan_rewrite_handler( + /// |event: WorkerPlanRewriteEvent<'_>| { + /// Ok(WorkerPlanRewriteEventResponse::new(event.plan)) + /// }, + /// ) + /// .build()) + /// } + /// + /// let _worker = Worker::from_session_builder(build_worker_session); + /// ``` + fn with_distributed_worker_plan_rewrite_handler( + self, + handler: T, + ) -> Self; + + /// Same as [DistributedExt::with_distributed_worker_plan_rewrite_handler] but with an + /// in-place mutation. + fn set_distributed_worker_plan_rewrite_handler( + &mut self, + handler: T, + ); } /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context. @@ -866,6 +913,10 @@ impl DistributedExt for SessionConfig { RouteTasksHandlers::push_custom(self, Arc::new(h)); } + fn set_distributed_worker_plan_rewrite_handler(&mut self, h: H) { + WorkerPlanRewriteHandlers::push_custom(self, Arc::new(h)); + } + delegate! { to self { #[call(set_distributed_option_extension)] @@ -969,6 +1020,10 @@ impl DistributedExt for SessionConfig { #[call(set_distributed_route_tasks_handler)] #[expr($;self)] fn with_distributed_route_tasks_handler(mut self, h: H) -> Self; + + #[call(set_distributed_worker_plan_rewrite_handler)] + #[expr($;self)] + fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; } } } @@ -1112,6 +1167,11 @@ impl DistributedExt for SessionStateBuilder { #[call(set_distributed_route_tasks_handler)] #[expr($;self)] fn with_distributed_route_tasks_handler(mut self, h: H) -> Self; + + fn set_distributed_worker_plan_rewrite_handler(&mut self, h: H); + #[call(set_distributed_worker_plan_rewrite_handler)] + #[expr($;self)] + fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; } } } @@ -1257,6 +1317,11 @@ impl DistributedExt for SessionState { #[call(set_distributed_route_tasks_handler)] #[expr($;self)] fn with_distributed_route_tasks_handler(mut self, h: H) -> Self; + + fn set_distributed_worker_plan_rewrite_handler(&mut self, h: H); + #[call(set_distributed_worker_plan_rewrite_handler)] + #[expr($;self)] + fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; } } } @@ -1395,6 +1460,11 @@ impl DistributedExt for SessionContext { #[call(set_distributed_route_tasks_handler)] #[expr($;self)] fn with_distributed_route_tasks_handler(self, h: H) -> Self; + + fn set_distributed_worker_plan_rewrite_handler(&mut self, h: H); + #[call(set_distributed_worker_plan_rewrite_handler)] + #[expr($;self)] + fn with_distributed_worker_plan_rewrite_handler(self, h: H) -> Self; } } } diff --git a/src/events/common.rs b/src/events/common.rs index 7a056b3d7..7f715059d 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -33,6 +33,17 @@ impl EventHandlerChain { // If no user handler handled the event, use the built ins. self.builtin.iter().find_map(|handler| f(handler.as_ref())) } + + pub(super) fn try_fold( + &self, + mut value: T, + mut f: impl FnMut(T, &H) -> Result, + ) -> Result { + for handler in self.custom.iter().chain(&self.builtin) { + value = f(value, handler.as_ref())?; + } + Ok(value) + } } impl EventHandlerChain { diff --git a/src/events/mod.rs b/src/events/mod.rs index 144904e18..9eb4b0190 100644 --- a/src/events/mod.rs +++ b/src/events/mod.rs @@ -3,6 +3,7 @@ mod defaults; mod desired_task_count; mod route_tasks; mod scale_up_leaf_node; +mod worker_plan_rewrite; pub(crate) use defaults::{ file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, random_routing, @@ -19,3 +20,7 @@ pub(crate) use scale_up_leaf_node::ScaleUpLeafNodeHandlers; pub use scale_up_leaf_node::{ ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler, }; +pub(crate) use worker_plan_rewrite::WorkerPlanRewriteHandlers; +pub use worker_plan_rewrite::{ + WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler, +}; diff --git a/src/events/worker_plan_rewrite.rs b/src/events/worker_plan_rewrite.rs new file mode 100644 index 000000000..f2b59b6c8 --- /dev/null +++ b/src/events/worker_plan_rewrite.rs @@ -0,0 +1,74 @@ +use super::common::EventHandlerChain; +use datafusion::error::Result; +use datafusion::execution::config::SessionConfig; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +/// Information supplied while rewriting a decoded worker stage plan before registration. +pub struct WorkerPlanRewriteEvent<'a> { + /// The worker-local plan. Each handler receives the plan returned by the previous handler. + pub plan: Arc, + /// The configuration of the worker session that will execute the plan. + pub session_config: &'a SessionConfig, +} + +/// The worker-local plan produced by a [`WorkerPlanRewriteHandler`]. +pub struct WorkerPlanRewriteEventResponse { + /// The plan to pass to the next handler or register for execution. + pub plan: Arc, +} + +impl WorkerPlanRewriteEventResponse { + /// Returns a response containing the rewritten worker-local plan. + pub fn new(plan: Arc) -> Self { + Self { plan } + } +} + +/// Rewrites a decoded worker-local plan before it is registered for execution. +/// +/// Every registered handler runs in registration order and receives the plan returned by the +/// previous handler. Returning an error aborts plan registration. +pub trait WorkerPlanRewriteHandler: Send + Sync + 'static { + /// Returns the plan to pass to the next handler. + fn rewrite_worker_plan( + &self, + ev: WorkerPlanRewriteEvent, + ) -> Result; +} + +impl WorkerPlanRewriteHandler for F +where + F: Send + Sync + 'static, + F: for<'a> Fn(WorkerPlanRewriteEvent<'a>) -> Result, +{ + fn rewrite_worker_plan( + &self, + ev: WorkerPlanRewriteEvent, + ) -> Result { + self(ev) + } +} + +pub(crate) type WorkerPlanRewriteHandlers = EventHandlerChain; + +impl WorkerPlanRewriteHandlers { + pub(crate) fn handle(ev: WorkerPlanRewriteEvent) -> Result { + let WorkerPlanRewriteEvent { + plan, + session_config, + } = ev; + let plan = match session_config.get_extension::() { + Some(handlers) => handlers.try_fold(plan, |plan, handler| { + handler + .rewrite_worker_plan(WorkerPlanRewriteEvent { + plan, + session_config, + }) + .map(|response| response.plan) + })?, + None => plan, + }; + Ok(WorkerPlanRewriteEventResponse::new(plan)) + } +} diff --git a/src/lib.rs b/src/lib.rs index b132caf5b..8273be189 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,8 @@ pub use distributed_planner::{ pub use events::{ DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, - ScaleUpLeafNodeHandler, TaskCountAnnotation, + ScaleUpLeafNodeHandler, TaskCountAnnotation, WorkerPlanRewriteEvent, + WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler, }; pub use execution_plans::{ BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec, diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index cf16e890c..020bce948 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,5 @@ use crate::common::TreeNodeExt; +use crate::events::{WorkerPlanRewriteEvent, WorkerPlanRewriteHandlers}; use crate::execution_plans::SamplerExec; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; @@ -86,11 +87,11 @@ impl Worker { let codec = DistributedCodec::new_combined_with_user(session_state.config()); let task_ctx = session_state.task_ctx(); let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?; - let mut plan = proto_node.try_into_physical_plan(&task_ctx, &codec)?; - - for hook in self.hooks.on_plan.iter() { - plan = hook(plan, session_state.config())?; - } + let ev = WorkerPlanRewriteEvent { + plan: proto_node.try_into_physical_plan(&task_ctx, &codec)?, + session_config: session_state.config(), + }; + let plan = WorkerPlanRewriteHandlers::handle(ev)?.plan; load_info_rxs = SamplerExec::kick_off_first_sampler(Arc::clone(&plan), Arc::clone(&task_ctx))?; diff --git a/src/worker/worker_service.rs b/src/worker/worker_service.rs index f59d482ef..98e16c76e 100644 --- a/src/worker/worker_service.rs +++ b/src/worker/worker_service.rs @@ -2,8 +2,6 @@ use crate::worker::{LocalWorkerContext, SingleWriteMultiRead, WorkerSessionBuild use crate::{DefaultSessionBuilder, TaskData, TaskKey}; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::SessionConfig; use moka::future::Cache; use std::borrow::Cow; use std::sync::Arc; @@ -12,16 +10,6 @@ use url::Url; const TASK_CACHE_TTI: Duration = Duration::from_mins(10); -#[allow(clippy::type_complexity)] -type OnPlanHook = dyn Fn(Arc, &SessionConfig) -> Result, DataFusionError> - + Sync - + Send; - -#[derive(Clone, Default)] -pub(super) struct WorkerHooks { - pub(super) on_plan: Vec>, -} - pub(crate) type ResultTaskData = Result>; pub(crate) type TaskDataEntries = Cache>>; @@ -33,7 +21,6 @@ pub struct Worker { /// while allowing concurrent access to task results across multiple partition requests. pub(crate) task_data_entries: Arc, pub(super) session_builder: Arc, - pub(super) hooks: WorkerHooks, pub(crate) max_message_size: Option, pub(super) version: Cow<'static, str>, } @@ -45,7 +32,6 @@ impl Default for Worker { runtime: Arc::new(RuntimeEnv::default()), task_data_entries: Arc::new(cache), session_builder: Arc::new(DefaultSessionBuilder), - hooks: WorkerHooks::default(), max_message_size: Some(usize::MAX), version: Cow::Borrowed(""), } @@ -71,31 +57,6 @@ impl Worker { self } - /// Adds a callback for when an [ExecutionPlan] is received in the `set_plan` call. - /// - /// The callback runs after worker session construction and plan decoding, and before task - /// registration and execution. It receives the per-query [SessionConfig], so it can use - /// propagated options or config extensions when rewriting the plan. - /// - /// The callback is trusted to preserve the worker stage contract already planned by the - /// coordinator: row semantics, output schema, partitioning, and ordering requirements. It is - /// intended for transparent wrappers, such as instrumentation, or semantics-preserving physical - /// rewrites. Do not use it to add or remove rows or columns, repartition the stage, or otherwise - /// re-plan distributed execution. Returned errors are propagated as worker plan-registration - /// failures. - pub fn add_on_plan_hook( - &mut self, - hook: impl Fn( - Arc, - &SessionConfig, - ) -> Result, DataFusionError> - + Sync - + Send - + 'static, - ) { - self.hooks.on_plan.push(Arc::new(hook)); - } - /// Set the maximum message size for FlightData chunks. /// /// Defaults to `usize::MAX` to minimize chunking overhead for internal communication. diff --git a/tests/worker_plan_hook.rs b/tests/worker_plan_hook.rs index a7fc59eb2..a6d4061a5 100644 --- a/tests/worker_plan_hook.rs +++ b/tests/worker_plan_hook.rs @@ -7,12 +7,15 @@ mod tests { use datafusion::common::{HashSet, Result, assert_contains, extensions_options, internal_err}; use datafusion::config::ConfigExtension; use datafusion::error::DataFusionError; - use datafusion::execution::SessionState; + use datafusion::execution::{SessionState, SessionStateBuilder}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_configured_in_memory_context; use datafusion_distributed::test_utils::session_context::register_temp_parquet_table; - use datafusion_distributed::{DistributedExt, Worker, WorkerQueryContext, assert_snapshot}; + use datafusion_distributed::{ + DistributedExt, MappedWorkerSessionBuilderExt, WorkerPlanRewriteEvent, + WorkerPlanRewriteEventResponse, WorkerQueryContext, assert_snapshot, + }; use std::sync::Arc; use std::sync::Mutex; @@ -37,19 +40,19 @@ mod tests { } #[tokio::test] - async fn plan_hooks_receive_session_config_and_run_in_order() + async fn worker_plan_rewrite_handlers_receive_session_config_and_run_in_order() -> Result<(), Box> { let hook_calls = Arc::new(Mutex::new(HookCalls::default())); - let mut ctx = start_configured_in_memory_context(3, build_state, { + let session_builder = build_state.map({ let hook_calls = Arc::clone(&hook_calls); - move |mut worker| { - add_first_hook(&mut worker, Arc::clone(&hook_calls)); - add_second_hook(&mut worker, Arc::clone(&hook_calls)); - worker + move |mut builder| { + add_first_hook(&mut builder, Arc::clone(&hook_calls)); + add_second_hook(&mut builder, Arc::clone(&hook_calls)); + Ok(builder.build()) } - }) - .await; + }); + let mut ctx = start_configured_in_memory_context(3, session_builder, |worker| worker).await; ctx.set_distributed_option_extension(PlanHookOptions { label: HOOK_LABEL.to_string(), @@ -75,19 +78,21 @@ mod tests { #[tokio::test] async fn plan_hook_errors_propagate_to_query() -> Result<(), Box> { - let mut ctx = start_configured_in_memory_context(3, build_state, move |mut worker| { - worker.add_on_plan_hook(move |plan, session_config| { - let options = plan_hook_options(session_config)?; - if options.fail_in_hook { - return internal_err!("plan hook failed for {}", options.label); - } - - Ok(plan) - }); - - worker - }) - .await; + let session_builder = build_state.map(|builder| { + Ok(builder + .with_distributed_worker_plan_rewrite_handler( + |event: WorkerPlanRewriteEvent<'_>| { + let options = plan_hook_options(event.session_config)?; + if options.fail_in_hook { + return internal_err!("plan hook failed for {}", options.label); + } + + Ok(WorkerPlanRewriteEventResponse::new(event.plan)) + }, + ) + .build()) + }); + let mut ctx = start_configured_in_memory_context(3, session_builder, |worker| worker).await; ctx.set_distributed_option_extension(PlanHookOptions { label: HOOK_LABEL.to_string(), @@ -110,35 +115,39 @@ mod tests { second: usize, } - fn add_first_hook(worker: &mut Worker, calls: Arc>) { - worker.add_on_plan_hook(move |plan, session_config| { - let options = plan_hook_options(session_config)?; - if options.label != HOOK_LABEL { - return internal_err!("unexpected plan hook label {}", options.label); - } + fn add_first_hook(builder: &mut SessionStateBuilder, calls: Arc>) { + builder.set_distributed_worker_plan_rewrite_handler( + move |event: WorkerPlanRewriteEvent<'_>| { + let options = plan_hook_options(event.session_config)?; + if options.label != HOOK_LABEL { + return internal_err!("unexpected plan hook label {}", options.label); + } - let mut calls = calls.lock().unwrap(); - calls.pending_plan_ids.insert(plan_identity(&plan)); - calls.first += 1; - Ok(plan) - }); + let mut calls = calls.lock().unwrap(); + calls.pending_plan_ids.insert(plan_identity(&event.plan)); + calls.first += 1; + Ok(WorkerPlanRewriteEventResponse::new(event.plan)) + }, + ) } - fn add_second_hook(worker: &mut Worker, calls: Arc>) { - worker.add_on_plan_hook(move |plan, session_config| { - let options = plan_hook_options(session_config)?; - if options.label != HOOK_LABEL { - return internal_err!("unexpected plan hook label {}", options.label); - } + fn add_second_hook(builder: &mut SessionStateBuilder, calls: Arc>) { + builder.set_distributed_worker_plan_rewrite_handler( + move |event: WorkerPlanRewriteEvent<'_>| { + let options = plan_hook_options(event.session_config)?; + if options.label != HOOK_LABEL { + return internal_err!("unexpected plan hook label {}", options.label); + } - let mut calls = calls.lock().unwrap(); - if !calls.pending_plan_ids.remove(&plan_identity(&plan)) { - return internal_err!("second hook ran before first hook"); - } + let mut calls = calls.lock().unwrap(); + if !calls.pending_plan_ids.remove(&plan_identity(&event.plan)) { + return internal_err!("second hook ran before first hook"); + } - calls.second += 1; - Ok(plan) - }); + calls.second += 1; + Ok(WorkerPlanRewriteEventResponse::new(event.plan)) + }, + ) } fn plan_identity(plan: &Arc) -> usize {