refactor: SessionConfig.extensions instead ConfigOptions.extensions

examples/custom_worker_url_routing.rs

@@ -41,8 +41,8 @@ use datafusion_distributed::test_utils::localhost::{
     LocalHostWorkerResolver, spawn_worker_service,
 };
 use datafusion_distributed::{
-    DistributedConfig, DistributedExt, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation,
-    TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
+    DistributedExt, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
+    TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -219,8 +219,9 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator {
     }
 
     fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
-        let d_cfg = DistributedConfig::from_task_context(&ctx.task_ctx)?;
-        let available_urls = d_cfg.worker_resolver().get_urls()?;
+        let available_urls =
+            datafusion_distributed::get_distributed_worker_resolver(ctx.task_ctx.session_config())?
+                .get_urls()?;
 
         let mut routed = None;
         ctx.plan.apply(|node| {

src/config_extension_ext.rs

@@ -21,6 +21,19 @@ pub(crate) fn set_distributed_option_extension<T: ConfigExtension + Default>(
     cfg.set_extension(Arc::new(propagation_ctx));
 }
 
+/// Registers `prefix` in [`ConfigExtensionPropagationContext`] so that headers for that
+/// extension are included in outgoing gRPC requests. Idempotent: safe to call multiple times.
+pub(crate) fn register_config_extension_prefix(cfg: &mut SessionConfig, prefix: &'static str) {
+    let mut ctx = cfg
+        .get_extension::<ConfigExtensionPropagationContext>()
+        .map(|existing| existing.as_ref().clone())
+        .unwrap_or_default();
+    if !ctx.prefixes.contains(&prefix) {
+        ctx.prefixes.push(prefix);
+        cfg.set_extension(Arc::new(ctx));
+    }
+}
+
 pub(crate) fn set_distributed_option_extension_from_headers<'a, T: ConfigExtension + Default>(
     cfg: &'a mut SessionConfig,
     headers: &HeaderMap,

src/coordinator/prepare_dynamic_plan.rs

@@ -140,7 +140,7 @@ pub(super) async fn prepare_dynamic_plan(
                 })
             })
         },
-        query_coordinator.session_config().options(),
+        query_coordinator.session_config(),
     )
     .await?;
 

src/coordinator/query_coordinator.rs

@@ -2,16 +2,18 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension};
 use crate::config_extension_ext::get_config_extension_propagation_headers;
 use crate::coordinator::MetricsStore;
 use crate::coordinator::latency_metric::LatencyMetric;
+use crate::distributed_planner::CombinedTaskEstimator;
 use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec};
 use crate::passthrough_headers::get_passthrough_headers;
 use crate::stage::LocalStage;
+use crate::work_unit_feed::WorkUnitFeedRegistry;
 use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time};
 use crate::{
     BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg,
-    DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig,
-    DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt,
-    SetPlanRequest, Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration,
-    WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver,
+    DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext,
+    DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage,
+    TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg,
+    get_distributed_channel_resolver, get_distributed_worker_resolver,
 };
 use datafusion::common::DataFusionError;
 use datafusion::common::instant::Instant;
@@ -261,8 +263,9 @@ impl<'a> StageCoordinator<'a> {
         tx: UnboundedSender<CoordinatorToWorkerMsg>,
     ) -> Result<()> {
         let session_config = self.task_ctx.session_config();
-        let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
-        let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
+        let wuf_registry = session_config
+            .get_extension::<WorkUnitFeedRegistry>()
+            .unwrap_or_default();
 
         let d_ctx = DistributedTaskContext {
             task_index: task_i,
@@ -333,8 +336,9 @@ impl<'a> StageCoordinator<'a> {
         task_i: usize,
     ) -> Result<(Arc<dyn ExecutionPlan>, Vec<WorkUnitFeedDeclaration>)> {
         let session_config = self.task_ctx.session_config();
-        let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
-        let wuf_registry = &d_cfg.__private_work_unit_feed_registry;
+        let wuf_registry = session_config
+            .get_extension::<WorkUnitFeedRegistry>()
+            .unwrap_or_default();
 
         let mut work_unit_feed_declarations = vec![];
         let d_ctx = DistributedTaskContext {
@@ -373,9 +377,11 @@ impl<'a> StageCoordinator<'a> {
     ///   [TaskEstimator::route_tasks] method.
     pub(super) fn routed_urls(&self) -> Result<Vec<Url>> {
         let session_config = self.task_ctx.session_config();
-        let d_cfg = DistributedConfig::from_config_options(session_config.options())?;
         let worker_resolver = get_distributed_worker_resolver(session_config)?;
-        let task_estimator = &d_cfg.__private_task_estimator;
+        let task_estimator = session_config
+            .get_extension::<CombinedTaskEstimator>()
+            .map(|a| a.as_ref().clone())
+            .unwrap_or_default();
 
         let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext {
             task_ctx: Arc::clone(self.task_ctx),

src/distributed_ext.rs

@@ -645,7 +645,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         bytes_per_partition: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.file_scan_config_bytes_per_partition = bytes_per_partition;
         Ok(())
     }
@@ -654,13 +654,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         factor: f64,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.cardinality_task_count_factor = factor;
         Ok(())
     }
 
     fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.collect_metrics = enabled;
         Ok(())
     }
@@ -669,13 +669,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         enabled: bool,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.children_isolator_unions = enabled;
         Ok(())
     }
 
     fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.broadcast_joins = enabled;
         Ok(())
     }
@@ -685,7 +685,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         compression: Option<arrow_ipc::CompressionType>,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.compression = match compression {
             Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(),
             Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(),
@@ -698,7 +698,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         batch_size: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.shuffle_batch_size = batch_size;
         Ok(())
     }
@@ -714,13 +714,13 @@ impl DistributedExt for SessionConfig {
         &mut self,
         max_tasks_per_stage: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.max_tasks_per_stage = max_tasks_per_stage;
         Ok(())
     }
 
     fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.partial_reduce = enabled;
         Ok(())
     }
@@ -729,7 +729,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         budget_bytes: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
         Ok(())
     }
@@ -747,7 +747,7 @@ impl DistributedExt for SessionConfig {
     }
 
     fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.dynamic_task_count = enabled;
         Ok(())
     }
@@ -756,7 +756,7 @@ impl DistributedExt for SessionConfig {
         &mut self,
         bytes_per_partition_per_second: usize,
     ) -> Result<(), DataFusionError> {
-        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        let d_cfg = DistributedConfig::from_session_config_mut(self)?;
         d_cfg.bytes_per_partition_per_second = bytes_per_partition_per_second;
         Ok(())
     }

src/distributed_planner/distributed_config.rs

@@ -1,13 +1,9 @@
-use crate::distributed_planner::task_estimator::CombinedTaskEstimator;
-use crate::protocol::ChannelResolverExtension;
-use crate::work_unit_feed::WorkUnitFeedRegistry;
-use crate::worker_resolver::WorkerResolverExtension;
-use crate::{TaskEstimator, WorkerResolver};
-use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err};
-use datafusion::config::{ConfigExtension, ConfigField, ConfigOptions, Visit};
+use crate::config_extension_ext::register_config_extension_prefix;
+use crate::config_extension_ext::set_distributed_option_extension;
+use datafusion::common::{DataFusionError, extensions_options, plan_err};
+use datafusion::config::{ConfigExtension, ConfigOptions};
 use datafusion::execution::TaskContext;
 use datafusion::prelude::SessionConfig;
-use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
 extensions_options! {
@@ -76,18 +72,6 @@ extensions_options! {
         /// If `dynamic_task_count` is enabled, this value is the amount of bytes/second each
         /// partition is expected to handle. Lower values will result in greater parallelism.
         pub bytes_per_partition_per_second: usize, default = 16 * 1024 * 1024
-        /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to
-        /// estimate how many tasks should be spawned for the [Stage] containing the leaf node.
-        pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default()
-        /// [ChannelResolver] implementation that tells the distributed planner information about
-        /// the available workers ready to execute distributed tasks.
-        pub(crate) __private_channel_resolver: ChannelResolverExtension, default = ChannelResolverExtension::default()
-        /// [WorkerResolver] implementation that tells the distributed planner information about
-        /// the available workers ready to execute distributed tasks.
-        pub(crate) __private_worker_resolver: WorkerResolverExtension, default = WorkerResolverExtension::not_implemented()
-        /// [WorkUnitFeedRegistry] that contains a set of getters that, applied to each node in a
-        /// plan, will return the [crate::WorkUnitFeed]s present in all nodes.
-        pub(crate) __private_work_unit_feed_registry: WorkUnitFeedRegistry, default = WorkUnitFeedRegistry::default()
     }
 }
 
@@ -100,19 +84,6 @@ fn cardinality_task_count_factor_default() -> f64 {
 }
 
 impl DistributedConfig {
-    /// Appends a [TaskEstimator] to the list. [TaskEstimator] will be executed sequentially in
-    /// order on leaf nodes, and the first one to provide a value is the one that gets to decide
-    /// how many tasks are used for that [Stage].
-    pub fn with_task_estimator(
-        mut self,
-        task_estimator: impl TaskEstimator + Send + Sync + 'static,
-    ) -> Self {
-        self.__private_task_estimator
-            .user_provided
-            .push(Arc::new(task_estimator));
-        self
-    }
-
     /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
     pub fn from_config_options(cfg: &ConfigOptions) -> Result<&Self, DataFusionError> {
         let Some(distributed_cfg) = cfg.extensions.get::<DistributedConfig>() else {
@@ -121,12 +92,25 @@ impl DistributedConfig {
         Ok(distributed_cfg)
     }
 
-    /// Gets the [DistributedConfig] from the [ConfigOptions]'s extensions.
-    pub fn from_config_options_mut(cfg: &mut ConfigOptions) -> Result<&mut Self, DataFusionError> {
-        let Some(distributed_cfg) = cfg.extensions.get_mut::<DistributedConfig>() else {
-            return plan_err!("DistributedConfig is not in ConfigOptions.extensions");
-        };
-        Ok(distributed_cfg)
+    /// Gets the [DistributedConfig] from the [SessionConfig], inserting a default if not present.
+    /// Always registers the `"distributed"` propagation prefix so coordinator-side settings are
+    /// included in gRPC headers sent to workers, regardless of call ordering.
+    pub fn from_session_config_mut(cfg: &mut SessionConfig) -> Result<&mut Self, DataFusionError> {
+        if cfg
+            .options()
+            .extensions
+            .get::<DistributedConfig>()
+            .is_none()
+        {
+            set_distributed_option_extension(cfg, DistributedConfig::default());
+        } else {
+            register_config_extension_prefix(cfg, DistributedConfig::PREFIX);
+        }
+        Ok(cfg
+            .options_mut()
+            .extensions
+            .get_mut::<DistributedConfig>()
+            .unwrap())
     }
 
     /// Gets the [DistributedConfig] from the [ConfigOptions]'s in the provided [SessionConfig].
@@ -138,85 +122,8 @@ impl DistributedConfig {
     pub fn from_task_context(ctx: &Arc<TaskContext>) -> Result<&Self, DataFusionError> {
         Self::from_session_config(ctx.session_config())
     }
-
-    /// Returns the [WorkerResolver] currently in scope for this [DistributedConfig].
-    pub fn worker_resolver(&self) -> &Arc<dyn WorkerResolver> {
-        &self.__private_worker_resolver.0
-    }
 }
 
 impl ConfigExtension for DistributedConfig {
     const PREFIX: &'static str = "distributed";
 }
-
-// FIXME: Ideally, both ChannelResolverExtension and TaskEstimators would be passed as
-//  extensions in SessionConfig's AnyMap instead of the ConfigOptions. However, we need
-//  to pass this as ConfigOptions as we need these two fields to be present during
-//  planning in the DistributedQueryPlanner, and the signature of the create_physical_plan()
-//  method there accepts a SessionState which only provides ConfigOptions.
-//  The following PR addresses this: https://github.com/apache/datafusion/pull/18168
-//  but it still has not been accepted or merged.
-//  Because of this, all the boilerplate trait implementations below are needed.
-impl ConfigField for ChannelResolverExtension {
-    fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
-        // nothing to do.
-    }
-
-    fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> {
-        not_impl_err!("Not implemented")
-    }
-}
-
-impl Debug for ChannelResolverExtension {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ChannelResolverExtension")
-    }
-}
-
-impl ConfigField for WorkerResolverExtension {
-    fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
-        // nothing to do.
-    }
-
-    fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> {
-        not_impl_err!("Not implemented")
-    }
-}
-
-impl Debug for WorkerResolverExtension {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "WorkerResolverExtension")
-    }
-}
-
-impl ConfigField for CombinedTaskEstimator {
-    fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
-        //nothing to do.
-    }
-
-    fn set(&mut self, _: &str, _: &str) -> Result<(), DataFusionError> {
-        not_impl_err!("not implemented")
-    }
-}
-
-impl Debug for CombinedTaskEstimator {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "TaskEstimators")
-    }
-}
-
-impl ConfigField for WorkUnitFeedRegistry {
-    fn visit<V: Visit>(&self, _: &mut V, _: &str, _: &'static str) {
-        //nothing to do.
-    }
-
-    fn set(&mut self, _: &str, _: &str) -> Result<(), DataFusionError> {
-        not_impl_err!("not implemented")
-    }
-}
-
-impl Debug for WorkUnitFeedRegistry {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "WorkUnitFeedRegistry")
-    }
-}