Split channel resolver in two (#265)

* Split channel resolver in two

* Simplify WorkerResolverExtension and ChannelResolverExtension

* Add default builder to ArrowFlightEndpoint

* Add some docs

* Listen to clippy

* Split get_flight_client_for_url in two

* Fix conflicts

* Remove unnecessary channel resolver

* Improve WorkerResolver docs

* Use one ChannelResolver per runtime

* Improve error reporting on client connection failure

* Add a from_session_builder method for constructing an InMemoryChannelResolver

* Add ChannelResolver and WorkerResolver default implementations for Arcs

* Add test for default channel resolver

* Extend doc comment

* Fix typo

Author: gabotechs

Cargo.lock

@@ -38,6 +38,7 @@ object_store = "0.12.3"
 bytes = "1.10.1"
 pin-project = "1.1.10"
 tokio-stream = "0.1.17"
+moka = { version = "0.12", features = ["sync"] }
 
 # integration_tests deps
 insta = { version = "1.43.1", features = ["filters"], optional = true }

Cargo.toml

@@ -1,23 +1,19 @@
-use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use aws_config::BehaviorVersion;
 use aws_sdk_ec2::Client as Ec2Client;
 use axum::{Json, Router, extract::Query, http::StatusCode, routing::get};
-use dashmap::{DashMap, Entry};
 use datafusion::common::DataFusionError;
 use datafusion::common::instant::Instant;
 use datafusion::common::runtime::SpawnedTask;
-use datafusion::execution::{SessionState, SessionStateBuilder};
+use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
 use datafusion_distributed::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
-    create_flight_client, display_plan_ascii,
+    ArrowFlightEndpoint, DistributedExt, DistributedPhysicalOptimizerRule,
+    DistributedSessionBuilderContext, WorkerResolver, display_plan_ascii,
 };
 use futures::{StreamExt, TryFutureExt};
 use log::{error, info, warn};
-use object_store::ObjectStore;
 use object_store::aws::AmazonS3Builder;
 use serde::Serialize;
 use std::collections::HashMap;
@@ -27,7 +23,7 @@ use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use structopt::StructOpt;
-use tonic::transport::{Channel, Server};
+use tonic::transport::Server;
 use url::Url;
 
 #[derive(Serialize)]
@@ -44,43 +40,6 @@ struct Cmd {
     bucket: String,
 }
 
-#[derive(Clone)]
-struct BenchSessionStateBuilder {
-    s3_url: Url,
-    s3: Arc<dyn ObjectStore>,
-    channel_resolver: Ec2ChannelResolver,
-}
-
-impl BenchSessionStateBuilder {
-    fn new(s3_url: Url) -> Result<Self, Box<dyn Error>> {
-        let s3 = AmazonS3Builder::from_env()
-            .with_bucket_name(s3_url.host().unwrap().to_string())
-            .build()?;
-        Ok(Self {
-            s3_url,
-            s3: Arc::new(s3),
-            channel_resolver: Ec2ChannelResolver::new(),
-        })
-    }
-}
-
-#[async_trait]
-impl DistributedSessionBuilder for BenchSessionStateBuilder {
-    async fn build_session_state(
-        &self,
-        ctx: DistributedSessionBuilderContext,
-    ) -> Result<SessionState, DataFusionError> {
-        let state = SessionStateBuilder::new()
-            .with_default_features()
-            .with_runtime_env(ctx.runtime_env)
-            .with_object_store(&self.s3_url, Arc::clone(&self.s3))
-            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
-            .with_distributed_channel_resolver(self.channel_resolver.clone())
-            .build();
-        Ok(state)
-    }
-}
-
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     env_logger::builder()
@@ -98,15 +57,27 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     // Register S3 object store
     let s3_url = Url::parse(&format!("s3://{}", cmd.bucket))?;
-    let state_builder = BenchSessionStateBuilder::new(s3_url)?;
 
     info!("Building shared SessionContext for the whole lifetime of the HTTP listener...");
-    let state = state_builder
-        .build_session_state(Default::default())
-        .await?;
+    let s3 = Arc::new(
+        AmazonS3Builder::from_env()
+            .with_bucket_name(s3_url.host().unwrap().to_string())
+            .build()?,
+    );
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_object_store(&s3_url, Arc::clone(&s3) as _)
+        .with_distributed_worker_resolver(Ec2WorkerResolver::new())
+        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .build();
     let ctx = SessionContext::from(state);
 
-    let arrow_flight_endpoint = ArrowFlightEndpoint::try_new(state_builder.clone())?;
+    let arrow_flight_endpoint =
+        ArrowFlightEndpoint::from_session_builder(move |ctx: DistributedSessionBuilderContext| {
+            let s3 = s3.clone();
+            let s3_url = s3_url.clone();
+            async move { Ok(ctx.builder.with_object_store(&s3_url, s3).build()) }
+        });
     let http_server = axum::serve(
         listener,
         Router::new().route(
@@ -213,9 +184,8 @@ fn err(s: impl Display) -> (StatusCode, String) {
 }
 
 #[derive(Clone)]
-struct Ec2ChannelResolver {
+struct Ec2WorkerResolver {
     urls: Arc<RwLock<Vec<Url>>>,
-    channels: Arc<DashMap<Url, BoxCloneSyncChannel>>,
 }
 
 async fn background_ec2_worker_resolver(urls: Arc<RwLock<Vec<Url>>>) {
@@ -273,35 +243,17 @@ async fn background_ec2_worker_resolver(urls: Arc<RwLock<Vec<Url>>>) {
     });
 }
 
-impl Ec2ChannelResolver {
+impl Ec2WorkerResolver {
     fn new() -> Self {
         let urls = Arc::new(RwLock::new(Vec::new()));
-        let channels = Arc::new(DashMap::new());
         tokio::spawn(background_ec2_worker_resolver(urls.clone()));
-        Self { urls, channels }
+        Self { urls }
     }
 }
 
 #[async_trait]
-impl ChannelResolver for Ec2ChannelResolver {
+impl WorkerResolver for Ec2WorkerResolver {
     fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
         Ok(self.urls.read().unwrap().clone())
     }
-
-    async fn get_flight_client_for_url(
-        &self,
-        url: &Url,
-    ) -> Result<FlightServiceClient<BoxCloneSyncChannel>, DataFusionError> {
-        let channel = match self.channels.entry(url.clone()) {
-            Entry::Occupied(v) => v.get().clone(),
-            Entry::Vacant(v) => {
-                let endpoint = Channel::from_shared(url.to_string()).unwrap();
-                let channel = endpoint.connect_lazy();
-                let channel = BoxCloneSyncChannel::new(channel);
-                v.insert(channel.clone());
-                channel
-            }
-        };
-        Ok(create_flight_client(channel))
-    }
 }

benchmarks/cdk/bin/worker.rs

@@ -43,7 +43,7 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
 use datafusion_distributed::test_utils::localhost::{
-    LocalHostChannelResolver, spawn_flight_service,
+    LocalHostWorkerResolver, spawn_flight_service,
 };
 use datafusion_distributed::{
     DistributedExt, DistributedPhysicalOptimizerRule, DistributedSessionBuilder,
@@ -143,7 +143,7 @@ impl DistributedSessionBuilder for RunOpt {
             .with_default_features()
             .with_config(config)
             .with_distributed_user_codec(InMemoryCacheExecCodec)
-            .with_distributed_channel_resolver(LocalHostChannelResolver::new(self.workers.clone()))
+            .with_distributed_worker_resolver(LocalHostWorkerResolver::new(self.workers.clone()))
             .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
             .with_distributed_option_extension_from_headers::<WarmingUpMarker>(&ctx.headers)?
             .with_distributed_files_per_task(

benchmarks/src/tpch/run.rs

@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 datafusion = { version = "51" }
-datafusion-distributed = { path = "..", features = ["avro"] }
+datafusion-distributed = { path = "..", features = ["avro", "integration"] }
 datafusion-cli = { version = "51", default-features = false }
 tokio = { version = "1.46.1", features = ["full"] }
 clap = { version = "4", features = ["derive"] }

cli/Cargo.toml

@@ -17,8 +17,6 @@
 
 // File mainly copied from https://github.com/apache/datafusion/blob/main/datafusion-cli/src/main.rs
 
-use arrow_flight::flight_service_client::FlightServiceClient;
-use async_trait::async_trait;
 use clap::Parser;
 use datafusion::common::config_err;
 use datafusion::config::ConfigOptions;
@@ -35,16 +33,14 @@ use datafusion_cli::{
     print_format::PrintFormat,
     print_options::{MaxRows, PrintOptions},
 };
-use datafusion_distributed::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext, create_flight_client,
+use datafusion_distributed::test_utils::in_memory_channel_resolver::{
+    InMemoryChannelResolver, InMemoryWorkerResolver,
 };
-use hyper_util::rt::TokioIo;
+use datafusion_distributed::{DistributedExt, DistributedPhysicalOptimizerRule};
 use std::env;
 use std::path::Path;
 use std::process::ExitCode;
 use std::sync::Arc;
-use tonic::transport::{Endpoint, Server};
 
 #[derive(Debug, Parser, PartialEq)]
 #[clap(author, version, about, long_about= None)]
@@ -153,7 +149,8 @@ async fn main_inner() -> Result<()> {
         .with_config(session_config)
         .with_runtime_env(runtime_env)
         .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
-        .with_distributed_channel_resolver(InMemoryChannelResolver::new())
+        .with_distributed_worker_resolver(InMemoryWorkerResolver::new(16))
+        .with_distributed_channel_resolver(InMemoryChannelResolver::default())
         .build();
 
     // enable dynamic file query
@@ -265,69 +262,3 @@ fn parse_command(command: &str) -> Result<String, String> {
         Err("-c flag expects only non empty commands".to_string())
     }
 }
-
-const DUMMY_URL: &str = "http://localhost:50051";
-
-/// [ChannelResolver] implementation that returns gRPC clients baked by an in-memory
-/// tokio duplex rather than a TCP connection.
-#[derive(Clone)]
-struct InMemoryChannelResolver {
-    channel: FlightServiceClient<BoxCloneSyncChannel>,
-}
-
-impl InMemoryChannelResolver {
-    fn new() -> Self {
-        let (client, server) = tokio::io::duplex(1024 * 1024);
-
-        let mut client = Some(client);
-        let channel = Endpoint::try_from(DUMMY_URL)
-            .expect("Invalid dummy URL for building an endpoint. This should never happen")
-            .connect_with_connector_lazy(tower::service_fn(move |_| {
-                let client = client
-                    .take()
-                    .expect("Client taken twice. This should never happen");
-                async move { Ok::<_, std::io::Error>(TokioIo::new(client)) }
-            }));
-
-        let this = Self {
-            channel: create_flight_client(BoxCloneSyncChannel::new(channel)),
-        };
-        let this_clone = this.clone();
-
-        let endpoint =
-            ArrowFlightEndpoint::try_new(move |ctx: DistributedSessionBuilderContext| {
-                let this = this.clone();
-                async move {
-                    let builder = SessionStateBuilder::new()
-                        .with_default_features()
-                        .with_distributed_channel_resolver(this)
-                        .with_runtime_env(ctx.runtime_env.clone());
-                    Ok(builder.build())
-                }
-            })
-            .unwrap();
-
-        tokio::spawn(async move {
-            Server::builder()
-                .add_service(endpoint.into_flight_server())
-                .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
-                .await
-        });
-
-        this_clone
-    }
-}
-
-#[async_trait]
-impl ChannelResolver for InMemoryChannelResolver {
-    fn get_urls(&self) -> std::result::Result<Vec<url::Url>, DataFusionError> {
-        Ok(vec![url::Url::parse(DUMMY_URL).unwrap(); 16]) // simulate 16 workers
-    }
-
-    async fn get_flight_client_for_url(
-        &self,
-        _: &url::Url,
-    ) -> std::result::Result<FlightServiceClient<BoxCloneSyncChannel>, DataFusionError> {
-        Ok(self.channel.clone())
-    }
-}