Decouple gRPC from the worker protocol and add an in-process transport

.github/workflows/ci.yml

@@ -41,6 +41,17 @@ jobs:
       - uses: ./.github/actions/setup
       - run: cargo test --features integration
 
+  # Runs the whole lib unit suite with `grpc` off, the proof that the abstraction builds and runs
+  # without gRPC. `lfs: true` because some planner unit tests load parquet testdata.
+  unit-test-no-grpc:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+      - uses: ./.github/actions/setup
+      - run: cargo test --no-default-features --lib
+
   tpch-correctness-test:
     runs-on: ubuntu-latest
     strategy:
@@ -50,7 +61,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test tpch_correctness_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_correctness_test
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -59,7 +70,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test tpch_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_plans_test
 
   tpcds-correctness-test:
     runs-on: ubuntu-latest
@@ -73,9 +84,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/tpcds/main.zip
+          path: benchmarks/testdata/tpcds/main.zip
           key: "main.zip"
-      - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }}
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }}
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -86,9 +97,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/tpcds/main.zip
+          path: benchmarks/testdata/tpcds/main.zip
           key: "main.zip"
-      - run: cargo test --features tpcds --test tpcds_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_plans_test
 
   clickbench-correctness-test:
     runs-on: ubuntu-latest
@@ -101,9 +112,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/clickbench/
+          path: benchmarks/testdata/clickbench/
           key: "data"
-      - run: cargo test --features clickbench --test clickbench_correctness_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_correctness_test
         env:
           ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }}
 
@@ -114,9 +125,9 @@ jobs:
       - uses: ./.github/actions/setup
       - uses: actions/cache@v4
         with:
-          path: testdata/clickbench/
+          path: benchmarks/testdata/clickbench/
           key: "data"
-      - run: cargo test --features clickbench --test clickbench_plans_test
+      - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_plans_test
 
   format-check:
     runs-on: ubuntu-latest

Cargo.toml

@@ -45,7 +45,7 @@ sysinfo = { version = "0.30", optional = true }
 sketches-ddsketch = { version = "0.3", features = ["use_serde"] }
 num-traits = "0.2"
 bincode = "1"
-tonic-prost = "0.14.2"
+tonic-prost = { version = "0.14.2", optional = true }
 
 # grpc-specific features
 arrow-flight = { version = "58", optional = true }
@@ -68,21 +68,17 @@ grpc = [
     "arrow-select",
     "arrow-ipc",
     "tonic",
+    "tonic-prost",
     "tower"
 ]
 
 integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"]
 
 system-metrics = ["sysinfo"]
 
-tpch = ["integration"]
-tpcds = ["integration"]
-clickbench = ["integration"]
-slow-tests = []
 sysinfo = ["dep:sysinfo"]
 
 [dev-dependencies]
-datafusion-distributed-benchmarks = { path = "benchmarks" }
 structopt = "0.3"
 insta = { version = "1.46.0", features = ["filters"] }
 parquet = "58"
@@ -92,6 +88,43 @@ hyper-util = "0.1.16"
 pretty_assertions = "1.4"
 test-case = "3.3.1"
 
+# Every example drives the gRPC transport, so a no-grpc build skips them.
+[[example]]
+name = "custom_distributed_partial_reduction_tree"
+required-features = ["integration"]
+
+[[example]]
+name = "custom_execution_plan"
+required-features = ["integration"]
+
+[[example]]
+name = "custom_worker_url_routing"
+required-features = ["integration"]
+
+[[example]]
+name = "work_unit_feed"
+required-features = ["integration"]
+
+[[example]]
+name = "in_memory_cluster"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_run"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_versioned_run"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_worker"
+required-features = ["grpc"]
+
+[[example]]
+name = "localhost_versioned_worker"
+required-features = ["grpc"]
+
 [workspace.lints.clippy]
 disallowed_types = "deny"
 disallowed-methods = "deny"

benchmarks/Cargo.toml

@@ -34,9 +34,20 @@ openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary f
 mimalloc = "0.1"
 sketches-ddsketch = "0.3"
 
+[features]
+# Gates for the dataset test suites under `tests/`. They live here rather than in the library
+# because, as a dev-dependency of the library, this crate re-enables `grpc` on every test build
+# through feature unification, which made the no-grpc config untestable.
+tpch = []
+tpcds = []
+clickbench = []
+slow-tests = []
+
 [dev-dependencies]
 criterion = "0.5"
 sysinfo = "0.30"
+pretty_assertions = "1.4"
+test-case = "3.3.1"
 
 [build-dependencies]
 built = { version = "0.8", features = ["git2", "chrono"] }

benchmarks/cdk/bin/worker.rs

@@ -1,3 +1,4 @@
+#![allow(clippy::disallowed_types)]
 use async_trait::async_trait;
 use aws_config::BehaviorVersion;
 use aws_sdk_ec2::Client as Ec2Client;

tests/clickbench_correctness_test.rs → benchmarks/tests/clickbench_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "clickbench", test))]
+#[cfg(all(feature = "clickbench", test))]
 mod tests {
     use datafusion::arrow::array::RecordBatch;
     use datafusion::common::plan_err;

tests/clickbench_plans_test.rs → benchmarks/tests/clickbench_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "clickbench", test))]
+#[cfg(all(feature = "clickbench", test))]
 mod tests {
     use datafusion::error::Result;
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;

tests/stateful_data_cleanup.rs → benchmarks/tests/stateful_data_cleanup.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion::common::instant::Instant;
     use datafusion::error::Result;

tests/tpcds_correctness_test.rs → benchmarks/tests/tpcds_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpcds", test))]
+#[cfg(all(feature = "tpcds", test))]
 mod tests {
     use datafusion::arrow::array::RecordBatch;
     use datafusion::common::plan_err;

tests/tpcds_plans_test.rs → benchmarks/tests/tpcds_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpcds", test))]
+#[cfg(all(feature = "tpcds", test))]
 mod tests {
     use datafusion::error::Result;
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;

tests/tpch_correctness_test.rs → benchmarks/tests/tpch_correctness_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion::physical_plan::execute_stream;
     use datafusion::prelude::SessionContext;

tests/tpch_plans_test.rs → benchmarks/tests/tpch_plans_test.rs

@@ -1,4 +1,4 @@
-#[cfg(all(feature = "integration", feature = "tpch", test))]
+#[cfg(all(feature = "tpch", test))]
 mod tests {
     use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context;
     use datafusion_distributed::{

console/src/main.rs

@@ -63,10 +63,10 @@ async fn run_app(
         terminal.draw(|frame| ui::render(frame, app))?;
 
         // Check for keyboard input (16ms timeout ~ 60fps responsiveness)
-        if event::poll(Duration::from_millis(16))? {
-            if let Event::Key(key) = event::read()? {
-                input::handle_key_event(app, key);
-            }
+        if event::poll(Duration::from_millis(16))?
+            && let Event::Key(key) = event::read()?
+        {
+            input::handle_key_event(app, key);
         }
 
         if app.should_quit {

console/src/worker.rs

@@ -176,32 +176,32 @@ impl WorkerConn {
 
                 // Detect completed tasks: tasks that were running but disappeared
                 for old_task in &self.tasks {
-                    if old_task.status == TaskStatus::Running as i32 {
-                        if let Some(sk) = &old_task.task_key {
-                            let key = (sk.query_id.clone(), sk.stage_id, sk.task_number);
-                            if !new_task_keys.contains(&key) {
-                                // Task disappeared — assume completed
-                                let observed_duration = self
-                                    .task_first_seen
-                                    .get(&key)
-                                    .map(|first| first.elapsed())
-                                    .unwrap_or_default();
-
-                                self.completed_tasks.push_front(CompletedTaskRecord {
-                                    query_id: sk.query_id.clone(),
-                                    stage_id: sk.stage_id,
-                                    task_number: sk.task_number,
-                                    observed_duration,
-                                });
-
-                                // Maintain bounded size
-                                while self.completed_tasks.len() > MAX_COMPLETED_TASKS {
-                                    self.completed_tasks.pop_back();
-                                }
-
-                                // Remove from first_seen tracking
-                                self.task_first_seen.remove(&key);
+                    if old_task.status == TaskStatus::Running as i32
+                        && let Some(sk) = &old_task.task_key
+                    {
+                        let key = (sk.query_id.clone(), sk.stage_id, sk.task_number);
+                        if !new_task_keys.contains(&key) {
+                            // Task disappeared — assume completed
+                            let observed_duration = self
+                                .task_first_seen
+                                .get(&key)
+                                .map(|first| first.elapsed())
+                                .unwrap_or_default();
+
+                            self.completed_tasks.push_front(CompletedTaskRecord {
+                                query_id: sk.query_id.clone(),
+                                stage_id: sk.stage_id,
+                                task_number: sk.task_number,
+                                observed_duration,
+                            });
+
+                            // Maintain bounded size
+                            while self.completed_tasks.len() > MAX_COMPLETED_TASKS {
+                                self.completed_tasks.pop_back();
                             }
+
+                            // Remove from first_seen tracking
+                            self.task_first_seen.remove(&key);
                         }
                     }
                 }

src/coordinator/metrics_store.rs

@@ -27,7 +27,7 @@ impl MetricsStore {
         self.rx.borrow().get(key).cloned()
     }
 
-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]
     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);

src/execution_plans/metrics.rs

@@ -21,7 +21,7 @@ impl MetricsWrapperExec {
         Self { inner, metrics }
     }
 
-    #[cfg(test)]
+    #[cfg(all(test, feature = "grpc"))]
     pub(crate) fn inner(&self) -> &Arc<dyn ExecutionPlan> {
         &self.inner
     }

src/execution_plans/mod.rs

@@ -8,7 +8,9 @@ mod network_coalesce;
 mod network_shuffle;
 mod sampler;
 
-#[cfg(any(test, feature = "integration"))]
+// The benchmark fixtures spin up real gRPC workers (`tonic` servers, Flight channels), so they
+// only exist when that transport is compiled in.
+#[cfg(all(feature = "grpc", any(test, feature = "integration")))]
 pub mod benchmarks;
 
 pub use broadcast::BroadcastExec;

src/lib.rs

@@ -46,14 +46,19 @@ pub mod test_utils;
 #[cfg(feature = "grpc")]
 pub use protocol::grpc;
 
+/// The worker-protocol prost message types, independent of any transport. A non-gRPC transport
+/// reaches for these to speak the same wire shape the gRPC path serializes. Unstable: this
+/// mirrors `worker.proto`, which is regenerated freely.
+pub use protocol::generated::worker as proto;
+
 pub use codec::DistributedCodec;
 pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver};
 
 pub use protocol::{
     ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest,
-    GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics,
-    WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg,
-    get_distributed_channel_resolver,
+    GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest,
+    TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel,
+    WorkerToCoordinatorMsg, get_distributed_channel_resolver,
 };
 pub use stage::{
     DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze,
@@ -66,7 +71,7 @@ pub use worker::{
     Worker, WorkerQueryContext, WorkerSessionBuilder,
 };
 
-#[cfg(any(feature = "integration", test))]
+#[cfg(all(feature = "grpc", any(feature = "integration", test)))]
 pub use execution_plans::benchmarks::{
     LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench,
     ShuffleFixture, TransportBench, TransportBenchMode, TransportFixture,

src/metrics/task_metrics_collector.rs

@@ -19,7 +19,8 @@ pub fn collect_plan_metrics(plan: &Arc<dyn ExecutionPlan>) -> Result<Vec<Metrics
     Ok(metrics)
 }
 
-#[cfg(test)]
+// These tests execute over the in-memory gRPC transport, so they need that transport compiled in.
+#[cfg(all(test, feature = "grpc"))]
 mod tests {
 
     use super::*;