Include ChannelManager in the DistributedExt trait

Author: gabotechs

src/channel_manager.rs

@@ -0,0 +1,52 @@
+use async_trait::async_trait;
+use datafusion::error::DataFusionError;
+use datafusion::prelude::SessionConfig;
+use std::sync::Arc;
+use tonic::body::BoxBody;
+use url::Url;
+
+pub(crate) fn set_channel_resolver(
+    cfg: &mut SessionConfig,
+    channel_resolver: impl ChannelResolver + Send + Sync + 'static,
+) {
+    cfg.set_extension(Arc::new(ChannelResolverExtension(Arc::new(
+        channel_resolver,
+    ))));
+}
+
+pub(crate) fn get_channel_resolver(
+    cfg: &SessionConfig,
+) -> Option<Arc<dyn ChannelResolver + Send + Sync>> {
+    cfg.get_extension::<ChannelResolverExtension>()
+        .map(|cm| cm.0.clone())
+}
+
+#[derive(Clone)]
+struct ChannelResolverExtension(Arc<dyn ChannelResolver + Send + Sync>);
+
+pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
+    http::Request<BoxBody>,
+    http::Response<BoxBody>,
+    tonic::transport::Error,
+>;
+
+/// Abstracts networking details so that users can implement their own network resolution
+/// mechanism.
+#[async_trait]
+pub trait ChannelResolver {
+    /// Gets all available worker URLs. Used during stage assignment.
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
+    /// For a given URL, get a channel for communicating to it.
+    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError>;
+}
+
+#[async_trait]
+impl ChannelResolver for Arc<dyn ChannelResolver + Send + Sync> {
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+        self.as_ref().get_urls()
+    }
+
+    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+        self.as_ref().get_channel_for_url(url).await
+    }
+}

src/channel_manager_ext.rs

@@ -1,7 +1,9 @@
+use crate::channel_manager_ext::set_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
 use crate::user_codec_ext::set_user_codec;
+use crate::ChannelResolver;
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::{SessionState, SessionStateBuilder};
@@ -122,17 +124,19 @@ pub trait DistributedExt: Sized {
         headers: &HeaderMap,
     ) -> Result<(), DataFusionError>;
 
-    /// Injects a user-defined codec that is capable of encoding/decoding custom execution nodes.
+    /// Injects a user-defined [PhysicalExtensionCodec] that is capable of encoding/decoding
+    /// custom execution nodes.
     ///
     /// Example:
     ///
     /// ```
     /// # use std::sync::Arc;
+    /// # use datafusion::common::DataFusionError;
     /// # use datafusion::execution::{SessionState, FunctionRegistry, SessionStateBuilder};
     /// # use datafusion::physical_plan::ExecutionPlan;
     /// # use datafusion::prelude::SessionConfig;
     /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-    /// # use datafusion_distributed::DistributedExt;
+    /// # use datafusion_distributed::{DistributedExt, DistributedSessionBuilderContext};
     ///
     /// #[derive(Debug)]
     /// struct CustomExecCodec;
@@ -148,11 +152,61 @@ pub trait DistributedExt: Sized {
     /// }
     ///
     /// let config = SessionConfig::new().with_user_codec(CustomExecCodec);
+    ///
+    /// async fn build_state(ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
+    ///     // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
+    ///     // know how to deserialize the CustomExtension from the gRPC metadata.
+    ///     Ok(SessionStateBuilder::new()
+    ///         .with_user_codec(CustomExecCodec)
+    ///         .build())
+    /// }
     /// ```
     fn with_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
 
     /// Same as [DistributedExt::with_user_codec] but with an in-place mutation
     fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
+
+    /// Injects a [ChannelResolver] implementation for Distributed DataFusion to resolve worker
+    /// nodes. When running in distributed mode, setting a [ChannelResolver] is required.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// # use async_trait::async_trait;
+    /// # use datafusion::common::DataFusionError;
+    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
+    /// # use datafusion::prelude::SessionConfig;
+    /// # use url::Url;
+    /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedSessionBuilderContext};
+    ///
+    /// struct CustomChannelResolver;
+    ///
+    /// #[async_trait]
+    /// impl ChannelResolver for CustomChannelResolver {
+    ///     fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+    ///         todo!()
+    ///     }
+    ///
+    ///     async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    ///         todo!()
+    ///     }
+    /// }
+    ///
+    /// let config = SessionConfig::new().with_channel_resolver(CustomChannelResolver);
+    ///
+    /// async fn build_state(ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
+    ///     // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
+    ///     // know how to deserialize the CustomExtension from the gRPC metadata.
+    ///     Ok(SessionStateBuilder::new()
+    ///         .with_channel_resolver(CustomChannelResolver)
+    ///         .build())
+    /// }
+    /// ```
+    fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T)
+        -> Self;
+
+    /// Same as [DistributedExt::with_channel_resolver] but with an in-place mutation.
+    fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
 }
 
 impl DistributedExt for SessionConfig {
@@ -174,17 +228,27 @@ impl DistributedExt for SessionConfig {
         set_user_codec(self, codec)
     }
 
+    fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T) {
+        set_channel_resolver(self, resolver)
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Result<Self, DataFusionError>;
+
             #[call(set_distributed_option_extension_from_headers)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
+
             #[call(set_user_codec)]
             #[expr($;self)]
             fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+
+            #[call(set_channel_resolver)]
+            #[expr($;self)]
+            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -196,14 +260,21 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_option_extension)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Result<Self, DataFusionError>;
+
             fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
             #[call(set_distributed_option_extension_from_headers)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
+
             fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
             #[call(set_user_codec)]
             #[expr($;self)]
             fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+
+            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_channel_resolver)]
+            #[expr($;self)]
+            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -215,14 +286,21 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_option_extension)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension<T: ConfigExtension + Default>(mut self, t: T) -> Result<Self, DataFusionError>;
+
             fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
             #[call(set_distributed_option_extension_from_headers)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
+
             fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
             #[call(set_user_codec)]
             #[expr($;self)]
             fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+
+            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_channel_resolver)]
+            #[expr($;self)]
+            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -234,14 +312,21 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_option_extension)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension<T: ConfigExtension + Default>(self, t: T) -> Result<Self, DataFusionError>;
+
             fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
             #[call(set_distributed_option_extension_from_headers)]
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
+
             fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
             #[call(set_user_codec)]
             #[expr($;self)]
             fn with_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
+
+            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_channel_resolver)]
+            #[expr($;self)]
+            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
         }
     }
 }

src/distributed_ext.rs

@@ -130,7 +130,6 @@ impl ArrowFlightEndpoint {
 
                 // Add the extensions that might be required for ExecutionPlan nodes in the plan
                 let config = state.config_mut();
-                config.set_extension(Arc::clone(&self.channel_manager));
                 config.set_extension(stage.clone());
                 config.set_extension(Arc::new(ContextGrpcMetadata(headers)));
 

src/flight_service/do_get.rs

@@ -7,4 +7,5 @@ pub(crate) use do_get::DoGet;
 pub use service::{ArrowFlightEndpoint, StageKey};
 pub use session_builder::{
     DefaultSessionBuilder, DistributedSessionBuilder, DistributedSessionBuilderContext,
+    MappedDistributedSessionBuilder, MappedDistributedSessionBuilderExt,
 };

src/flight_service/mod.rs

@@ -1,8 +1,5 @@
-use crate::channel_manager::ChannelManager;
-use crate::flight_service::session_builder::DefaultSessionBuilder;
 use crate::flight_service::DistributedSessionBuilder;
 use crate::stage::ExecutionStage;
-use crate::ChannelResolver;
 use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::{
     Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
@@ -32,29 +29,20 @@ pub struct StageKey {
 }
 
 pub struct ArrowFlightEndpoint {
-    pub(super) channel_manager: Arc<ChannelManager>,
     pub(super) runtime: Arc<RuntimeEnv>,
     #[allow(clippy::type_complexity)]
     pub(super) stages: DashMap<StageKey, Arc<OnceCell<(SessionState, Arc<ExecutionStage>)>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }
 
 impl ArrowFlightEndpoint {
-    pub fn new(channel_resolver: impl ChannelResolver + Send + Sync + 'static) -> Self {
+    pub fn new(session_builder: impl DistributedSessionBuilder + Send + Sync + 'static) -> Self {
         Self {
-            channel_manager: Arc::new(ChannelManager::new(channel_resolver)),
             runtime: Arc::new(RuntimeEnv::default()),
             stages: DashMap::new(),
-            session_builder: Arc::new(DefaultSessionBuilder),
+            session_builder: Arc::new(session_builder),
         }
     }
-
-    pub fn with_session_builder(
-        &mut self,
-        session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
-    ) {
-        self.session_builder = Arc::new(session_builder);
-    }
 }
 
 #[async_trait]