Abstract worker protocol from gRPC

.gitignore

@@ -8,7 +8,7 @@ testdata/tpcds/*
 !testdata/tpcds/README.md
 testdata/clickbench/*
 !testdata/clickbench/queries
-src/observability/gen/target/*
-src/observability/gen/Cargo.lock
-src/worker/gen/target/*
-src/worker/gen/Cargo.lock
+src/protocol/grpc/observability/gen/target/*
+src/protocol/grpc/observability/gen/Cargo.lock
+src/protocol/grpc/gen/target/*
+src/protocol/grpc/gen/Cargo.lock

Cargo.toml

@@ -23,13 +23,8 @@ datafusion = { workspace = true, features = [
   "datetime_expressions",
 ] }
 datafusion-proto = { workspace = true }
-arrow-flight = "58"
-arrow-select = "58"
-arrow-ipc = { version = "58", features = ["zstd"] }
 async-trait = "0.1.89"
 tokio = { version = "1.48", features = ["full"] }
-tonic = { version = "0.14.1", features = ["transport"] }
-tower = "0.5.2"
 http = "1.3.1"
 itertools = "0.14.0"
 futures = "0.3.31"
@@ -52,15 +47,31 @@ num-traits = "0.2"
 bincode = "1"
 tonic-prost = "0.14.2"
 
+# grpc-specific features
+arrow-flight = { version = "58", optional = true }
+arrow-select = { version = "58", optional = true }
+arrow-ipc = { version = "58", features = ["zstd"], optional = true }
+tonic = { version = "0.14.1", features = ["transport"], optional = true }
+tower = { version = "0.5.2", optional = true }
+
 # integration_tests deps
 insta = { version = "1.46.0", features = ["filters"], optional = true }
 parquet = { version = "58", optional = true }
 arrow = { version = "58", optional = true, features = ["test_utils"] }
 hyper-util = { version = "0.1.16", optional = true }
 
 [features]
+default = ["grpc"]
 avro = ["datafusion/avro"]
-integration = ["insta", "parquet", "arrow", "hyper-util"]
+grpc = [
+    "arrow-flight",
+    "arrow-select",
+    "arrow-ipc",
+    "tonic",
+    "tower"
+]
+
+integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"]
 
 system-metrics = ["sysinfo"]
 

console/src/worker.rs

@@ -1,5 +1,5 @@
 use datafusion::common::{HashMap, HashSet};
-use datafusion_distributed::{
+use datafusion_distributed::grpc::{
     GetClusterWorkersRequest, GetTaskProgressRequest, ObservabilityServiceClient, PingRequest,
     TaskProgress, TaskStatus,
 };

examples/in_memory_cluster.rs

@@ -4,9 +4,8 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion_distributed::{
-    BoxCloneSyncChannel, ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker,
-    WorkerQueryContext, WorkerResolver, WorkerServiceClient, create_worker_client,
-    display_plan_ascii,
+    ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker, WorkerChannel,
+    WorkerQueryContext, WorkerResolver, display_plan_ascii, grpc,
 };
 use futures::TryStreamExt;
 use hyper_util::rt::TokioIo;
@@ -66,7 +65,7 @@ const DUMMY_URL: &str = "http://localhost:50051";
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
 struct InMemoryChannelResolver {
-    channel: WorkerServiceClient<BoxCloneSyncChannel>,
+    channel: grpc::BoxCloneSyncChannel,
 }
 
 impl InMemoryChannelResolver {
@@ -84,7 +83,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: create_worker_client(BoxCloneSyncChannel::new(channel)),
+            channel: grpc::BoxCloneSyncChannel::new(channel),
         };
         let this_clone = this.clone();
 
@@ -110,8 +109,8 @@ impl ChannelResolver for InMemoryChannelResolver {
     async fn get_worker_client_for_url(
         &self,
         _: &url::Url,
-    ) -> Result<WorkerServiceClient<BoxCloneSyncChannel>, DataFusionError> {
-        Ok(self.channel.clone())
+    ) -> Result<Box<dyn WorkerChannel>, DataFusionError> {
+        Ok(grpc::create_worker_client(self.channel.clone()))
     }
 }
 

examples/localhost_versioned_run.rs

@@ -4,8 +4,8 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion_distributed::{
-    DefaultChannelResolver, DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt,
-    WorkerResolver, create_worker_client, display_plan_ascii,
+    DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt, WorkerResolver,
+    display_plan_ascii, grpc,
 };
 use futures::TryStreamExt;
 use std::error::Error;
@@ -40,28 +40,28 @@ struct Args {
 /// the `GetWorkerInfo` RPC. Returns `false` if the worker is unreachable, returns
 /// an error, or reports a different version.
 async fn worker_has_version(
-    channel_resolver: &DefaultChannelResolver,
+    channel_resolver: &grpc::DefaultChannelResolver,
     url: &Url,
     expected_version: &str,
 ) -> bool {
     let Ok(channel) = channel_resolver.get_channel(url).await else {
         return false;
     };
 
-    let mut client = create_worker_client(channel);
+    let mut client = grpc::create_worker_client(channel);
     let Ok(response) = client.get_worker_info(GetWorkerInfoRequest {}).await else {
         return false;
     };
 
-    response.into_inner().version == expected_version
+    response.version == expected_version
 }
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     let args = Args::from_args();
 
     let ports = if let Some(target_version) = &args.version {
-        let channel_resolver = DefaultChannelResolver::default();
+        let channel_resolver = grpc::DefaultChannelResolver::default();
         let mut compatible = Vec::new();
         for &port in &args.cluster_ports {
             let url = Url::parse(&format!("http://localhost:{port}"))?;

src/protobuf/mod.rs → src/codec/mod.rs

@@ -1,13 +1,7 @@
 mod distributed_codec;
-mod errors;
-mod producer_head;
 mod user_codec;
 
 pub use distributed_codec::DistributedCodec;
-pub(crate) use errors::{
-    datafusion_error_to_tonic_status, map_flight_to_datafusion_error,
-    tonic_status_to_datafusion_error,
-};
 pub(crate) use user_codec::{
     get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc,
 };

src/common/mod.rs

@@ -1,5 +1,4 @@
 mod children_helpers;
-mod on_drop_stream;
 mod once_lock;
 mod recursion;
 mod task_context_helpers;
@@ -8,7 +7,6 @@ mod uuid;
 mod vec;
 
 pub(crate) use children_helpers::require_one_child;
-pub(crate) use on_drop_stream::on_drop_stream;
 pub(crate) use once_lock::OnceLockResult;
 pub(crate) use recursion::TreeNodeExt;
 pub(crate) use task_context_helpers::task_ctx_with_extension;

src/common/time.rs

@@ -1,9 +1,14 @@
+use num_traits::AsPrimitive;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 /// Nanoseconds elapsed since UNIX epoch.
-pub(crate) fn now_ns() -> u64 {
+pub(crate) fn now_ns<T>() -> T
+where
+    T: 'static + Copy,
+    u128: AsPrimitive<T>,
+{
     SystemTime::now()
         .duration_since(UNIX_EPOCH)
-        .map(|duration| duration.as_nanos() as u64)
+        .map(|duration| duration.as_nanos().as_())
         .expect("SystemTime before UNIX EPOCH!")
 }

src/coordinator/distributed.rs

@@ -1,11 +1,10 @@
-use crate::DistributedConfig;
-use crate::common::{require_one_child, serialize_uuid};
+use crate::common::require_one_child;
 use crate::coordinator::metrics_store::MetricsStore;
 use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan;
 use crate::coordinator::prepare_static_plan::prepare_static_plan;
 use crate::coordinator::query_coordinator::QueryCoordinator;
 use crate::distributed_planner::NetworkBoundaryExt;
-use crate::worker::generated::worker::TaskKey;
+use crate::{DistributedConfig, TaskKey};
 use datafusion::common::internal_datafusion_err;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion::common::{Result, exec_err};
@@ -94,9 +93,9 @@ impl DistributedExec {
                 let stage = boundary.input_stage();
                 for i in 0..stage.task_count() {
                     expected_keys.push(TaskKey {
-                        query_id: serialize_uuid(&stage.query_id()),
-                        stage_id: stage.num() as u64,
-                        task_number: i as u64,
+                        query_id: stage.query_id(),
+                        stage_id: stage.num(),
+                        task_number: i,
                     });
                 }
             }

src/coordinator/metrics_store.rs

@@ -1,9 +1,8 @@
-use crate::TaskKey;
-use crate::worker::generated::worker as pb;
+use crate::{TaskKey, TaskMetrics};
 use datafusion::common::HashMap;
 use tokio::sync::watch;
 
-type MetricsMap = HashMap<TaskKey, pb::TaskMetrics>;
+type MetricsMap = HashMap<TaskKey, TaskMetrics>;
 
 /// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive.
 #[derive(Debug, Clone)]
@@ -18,20 +17,18 @@ impl MetricsStore {
         Self { tx, rx }
     }
 
-    pub(crate) fn insert(&self, key: TaskKey, metrics: pb::TaskMetrics) {
+    pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) {
         self.tx.send_modify(|map| {
             map.insert(key, metrics);
         });
     }
 
-    pub(crate) fn get(&self, key: &TaskKey) -> Option<pb::TaskMetrics> {
+    pub(crate) fn get(&self, key: &TaskKey) -> Option<TaskMetrics> {
         self.rx.borrow().get(key).cloned()
     }
 
     #[cfg(test)]
-    pub(crate) fn from_entries(
-        entries: impl IntoIterator<Item = (TaskKey, pb::TaskMetrics)>,
-    ) -> Self {
+    pub(crate) fn from_entries(entries: impl IntoIterator<Item = (TaskKey, TaskMetrics)>) -> Self {
         let map: HashMap<_, _> = entries.into_iter().collect();
         let (tx, rx) = watch::channel(map);
         Self { tx, rx }

src/coordinator/prepare_dynamic_plan.rs

@@ -8,8 +8,7 @@ use crate::distributed_planner::{
 };
 use crate::execution_plans::SamplerExec;
 use crate::stage::{LocalStage, RemoteStage};
-use crate::worker::generated::worker as pb;
-use crate::{BytesCounterMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage};
+use crate::{BytesCounterMetric, LoadInfo, NetworkBoundaryExt, NetworkCoalesceExec, Stage};
 use dashmap::DashMap;
 use datafusion::common::stats::Precision;
 use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
@@ -204,7 +203,7 @@ impl PlanReconstructor {
 
 /// Estimates the bytes per second flowing through a stage by reading sample information.
 async fn gather_runtime_statistics(
-    per_task_load_info_stream: Vec<impl Stream<Item = pb::LoadInfo> + Unpin>,
+    per_task_load_info_stream: Vec<impl Stream<Item = LoadInfo> + Unpin>,
     plan: &Arc<dyn ExecutionPlan>,
 ) -> Result<Statistics> {
     const ESTIMATED_QUERY_TIME_S: usize = 10;
@@ -237,8 +236,8 @@ async fn gather_runtime_statistics(
 
     let mut load_info_stream = futures::stream::select_all(per_task_load_info_stream);
     while let Some(load_info) = load_info_stream.next().await {
-        rows_per_second += load_info.rows_per_second as usize;
-        rows_ready += load_info.rows_ready as usize;
+        rows_per_second += load_info.rows_per_second;
+        rows_ready += load_info.rows_ready;
         per_col_bytes_per_second = element_wise_sum(
             per_col_bytes_per_second,
             &load_info.per_column_bytes_per_second,