From 72f139ddbca4b49b76832286e496650fea8a0dc8 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:11:48 -0700 Subject: [PATCH 1/3] Moved the worker-protocol messages out from behind the grpc feature. The prost message types carry no tonic dependency, so a transport that is not gRPC can speak the same wire shape without pulling in the gRPC stack. Only the tonic client and server stay gated; the generator emits those gates so a regeneration cannot drop them. tonic-prost feeds only the generated client and server, so it moves behind the feature too. Co-authored-by: Stu Hood --- Cargo.toml | 3 ++- src/lib.rs | 5 +++++ src/protocol/generated/mod.rs | 1 + src/protocol/{grpc => }/generated/worker.rs | 2 ++ src/protocol/grpc/gen/src/main.rs | 6 +++++- src/protocol/grpc/generated/mod.rs | 1 - src/protocol/grpc/metrics_proto.rs | 2 +- src/protocol/grpc/mod.rs | 1 - src/protocol/grpc/observability/gen/src/main.rs | 2 +- src/protocol/grpc/observability/generated/observability.rs | 2 +- src/protocol/grpc/observability/service.rs | 2 +- src/protocol/grpc/worker_client.rs | 4 ++-- src/protocol/grpc/worker_service.rs | 2 +- src/protocol/mod.rs | 4 ++++ 14 files changed, 26 insertions(+), 11 deletions(-) create mode 100644 src/protocol/generated/mod.rs rename src/protocol/{grpc => }/generated/worker.rs (99%) delete mode 100644 src/protocol/grpc/generated/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 79366ed3..5cbeaa0d 100644 --- a/Cargo.toml +++ b/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,6 +68,7 @@ grpc = [ "arrow-select", "arrow-ipc", "tonic", + "tonic-prost", "tower" ] diff --git a/src/lib.rs b/src/lib.rs index 777dd5c8..aa0917d7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,6 +39,11 @@ 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}; diff --git a/src/protocol/generated/mod.rs b/src/protocol/generated/mod.rs new file mode 100644 index 00000000..2c8b8399 --- /dev/null +++ b/src/protocol/generated/mod.rs @@ -0,0 +1 @@ +pub mod worker; diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/generated/worker.rs similarity index 99% rename from src/protocol/grpc/generated/worker.rs rename to src/protocol/generated/worker.rs index 290261ed..1f104a7b 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/generated/worker.rs @@ -468,6 +468,7 @@ pub struct MaxGauge { pub value: u64, } /// Generated client implementations. +#[cfg(feature = "grpc")] pub mod worker_service_client { #![allow( unused_variables, @@ -617,6 +618,7 @@ pub mod worker_service_client { } } /// Generated server implementations. +#[cfg(feature = "grpc")] pub mod worker_service_server { #![allow( unused_variables, diff --git a/src/protocol/grpc/gen/src/main.rs b/src/protocol/grpc/gen/src/main.rs index 7505aae5..cde0e4e6 100644 --- a/src/protocol/grpc/gen/src/main.rs +++ b/src/protocol/grpc/gen/src/main.rs @@ -6,7 +6,7 @@ fn main() -> Result<(), Box> { let proto_dir = repo_root.join("src/protocol/grpc"); let proto_file = proto_dir.join("worker.proto"); - let out_dir = repo_root.join("src/protocol/grpc/generated"); + let out_dir = repo_root.join("src/protocol/generated"); fs::create_dir_all(&out_dir)?; @@ -18,6 +18,10 @@ fn main() -> Result<(), Box> { tonic_prost_build::configure() .build_server(true) .build_client(true) + // The generated messages build with `grpc` off; only the tonic client and server carry + // the feature gate. Emitted here so a regeneration cannot drop the gates. + .client_mod_attribute(".", "#[cfg(feature = \"grpc\")]") + .server_mod_attribute(".", "#[cfg(feature = \"grpc\")]") .out_dir(&out_dir) .extern_path(".worker.FlightData", "::arrow_flight::FlightData") .extern_path( diff --git a/src/protocol/grpc/generated/mod.rs b/src/protocol/grpc/generated/mod.rs deleted file mode 100644 index 844c269c..00000000 --- a/src/protocol/grpc/generated/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) mod worker; diff --git a/src/protocol/grpc/metrics_proto.rs b/src/protocol/grpc/metrics_proto.rs index 6db8cda9..654b6cee 100644 --- a/src/protocol/grpc/metrics_proto.rs +++ b/src/protocol/grpc/metrics_proto.rs @@ -1,4 +1,4 @@ -use super::generated::worker as pb; +use crate::protocol::generated::worker as pb; use chrono::DateTime; use datafusion::common::internal_err; use datafusion::error::DataFusionError; diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs index e432b607..fc367585 100644 --- a/src/protocol/grpc/mod.rs +++ b/src/protocol/grpc/mod.rs @@ -1,6 +1,5 @@ mod channel_resolver; mod errors; -mod generated; mod metrics_proto; mod observability; mod on_drop_stream; diff --git a/src/protocol/grpc/observability/gen/src/main.rs b/src/protocol/grpc/observability/gen/src/main.rs index 124f0f1b..1b55fb7d 100644 --- a/src/protocol/grpc/observability/gen/src/main.rs +++ b/src/protocol/grpc/observability/gen/src/main.rs @@ -21,7 +21,7 @@ fn main() -> Result<(), Box> { .out_dir(&out_dir) .extern_path( ".observability.TaskKey", - "crate::protocol::grpc::generated::worker::TaskKey", + "crate::protocol::generated::worker::TaskKey", ) .compile_protos(&[proto_file], &[proto_dir])?; diff --git a/src/protocol/grpc/observability/generated/observability.rs b/src/protocol/grpc/observability/generated/observability.rs index 3b0a95cf..98aa592d 100644 --- a/src/protocol/grpc/observability/generated/observability.rs +++ b/src/protocol/grpc/observability/generated/observability.rs @@ -12,7 +12,7 @@ pub struct GetTaskProgressRequest {} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct TaskProgress { #[prost(message, optional, tag = "1")] - pub task_key: ::core::option::Option, + pub task_key: ::core::option::Option, #[prost(enumeration = "TaskStatus", tag = "4")] pub status: i32, #[prost(uint64, tag = "5")] diff --git a/src/protocol/grpc/observability/service.rs b/src/protocol/grpc/observability/service.rs index 27edee2e..04c2aa26 100644 --- a/src/protocol/grpc/observability/service.rs +++ b/src/protocol/grpc/observability/service.rs @@ -4,7 +4,7 @@ use super::{ }; use crate::common::serialize_uuid; use crate::grpc::{GetClusterWorkersRequest, GetClusterWorkersResponse}; -use crate::protocol::grpc::generated::worker as worker_pb; +use crate::protocol::generated::worker as worker_pb; use crate::worker::{SingleWriteMultiRead, TaskData}; use crate::{TaskKey, WorkerResolver}; use datafusion::error::DataFusionError; diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 9cc94de2..93291d91 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -1,10 +1,10 @@ use super::channel_resolver::BoxCloneSyncChannel; use super::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; -use super::generated::worker as pb; use super::metrics_proto::metrics_set_proto_to_df; use crate::common::serialize_uuid; -use crate::grpc::generated::worker::FlightAppMetadata; use crate::grpc::on_drop_stream::on_drop_stream; +use crate::protocol::generated::worker as pb; +use crate::protocol::generated::worker::FlightAppMetadata; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index fa677a10..f3472d43 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -1,7 +1,7 @@ use super::errors::{datafusion_error_to_tonic_status, map_status_to_datafusion_error}; -use super::generated::worker as pb; use super::metrics_proto::df_metrics_set_to_proto; use super::spawn_select_all::spawn_select_all; +use crate::protocol::generated::worker as pb; use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::ProducerHeadSpec; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 5c6bdd1a..3f4dd666 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -2,6 +2,10 @@ pub mod grpc; mod channel_resolver; +// The prost message types carry no tonic dependency, so a non-gRPC transport (an in-process +// worker, a shared-memory mesh) can speak the same wire shape without pulling in the whole gRPC +// stack. +pub(crate) mod generated; mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; From 703113b7eda2cebeea61ef6071697404d988713a Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:46:51 -0700 Subject: [PATCH 2/3] Made the no-gRPC build compile and its unit suite run in CI. The benchmarks crate's dev-dependency on the lib re-unified grpc into every test build, so a genuine no-gRPC test run was impossible; the dataset suites move into the benchmarks crate and the gRPC-coupled test utilities gate behind grpc. A unit-test-no-grpc job then runs the whole lib suite with the feature off. Co-authored-by: Stu Hood --- .github/workflows/ci.yml | 31 ++- Cargo.lock | 3 +- Cargo.toml | 42 ++- benchmarks/Cargo.toml | 11 + .../tests}/clickbench_correctness_test.rs | 2 +- .../tests}/clickbench_plans_test.rs | 2 +- .../tests}/stateful_data_cleanup.rs | 2 +- .../tests}/tpcds_correctness_test.rs | 2 +- .../tests}/tpcds_plans_test.rs | 2 +- .../tests}/tpch_correctness_test.rs | 2 +- .../tests}/tpch_plans_test.rs | 2 +- src/coordinator/metrics_store.rs | 2 +- src/execution_plans/metrics.rs | 2 +- src/execution_plans/mod.rs | 4 +- src/lib.rs | 2 +- src/metrics/task_metrics_collector.rs | 3 +- src/metrics/task_metrics_rewriter.rs | 3 +- src/test_utils/in_memory_channel_resolver.rs | 254 ++++++++++-------- src/test_utils/mod.rs | 1 + src/test_utils/routing.rs | 2 +- src/worker/mod.rs | 4 +- 21 files changed, 228 insertions(+), 150 deletions(-) rename {tests => benchmarks/tests}/clickbench_correctness_test.rs (99%) rename {tests => benchmarks/tests}/clickbench_plans_test.rs (99%) rename {tests => benchmarks/tests}/stateful_data_cleanup.rs (98%) rename {tests => benchmarks/tests}/tpcds_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpcds_plans_test.rs (99%) rename {tests => benchmarks/tests}/tpch_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpch_plans_test.rs (99%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3c5da3fd..5748469b 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Cargo.lock b/Cargo.lock index 7044adaa..3858be13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2199,7 +2199,6 @@ dependencies = [ "crossbeam-queue", "dashmap", "datafusion", - "datafusion-distributed-benchmarks", "datafusion-proto", "delegate", "futures", @@ -2249,11 +2248,13 @@ dependencies = [ "object_store", "openssl", "parquet", + "pretty_assertions", "reqwest", "serde", "serde_json", "structopt", "sysinfo", + "test-case", "tokio", "tonic", "tpchgen", diff --git a/Cargo.toml b/Cargo.toml index 5cbeaa0d..1bb4874d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,14 +76,9 @@ 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" @@ -93,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" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 0a03ead0..18197f42 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -33,9 +33,20 @@ aws-sdk-ec2 = "1" openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary for the remote benchmarks worker. mimalloc = "0.1" +[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"] } diff --git a/tests/clickbench_correctness_test.rs b/benchmarks/tests/clickbench_correctness_test.rs similarity index 99% rename from tests/clickbench_correctness_test.rs rename to benchmarks/tests/clickbench_correctness_test.rs index acf3daa8..25fb38f8 100644 --- a/tests/clickbench_correctness_test.rs +++ b/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; diff --git a/tests/clickbench_plans_test.rs b/benchmarks/tests/clickbench_plans_test.rs similarity index 99% rename from tests/clickbench_plans_test.rs rename to benchmarks/tests/clickbench_plans_test.rs index b9c32bb1..7291d4f7 100644 --- a/tests/clickbench_plans_test.rs +++ b/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; diff --git a/tests/stateful_data_cleanup.rs b/benchmarks/tests/stateful_data_cleanup.rs similarity index 98% rename from tests/stateful_data_cleanup.rs rename to benchmarks/tests/stateful_data_cleanup.rs index fae892d2..f977c61a 100644 --- a/tests/stateful_data_cleanup.rs +++ b/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; diff --git a/tests/tpcds_correctness_test.rs b/benchmarks/tests/tpcds_correctness_test.rs similarity index 99% rename from tests/tpcds_correctness_test.rs rename to benchmarks/tests/tpcds_correctness_test.rs index e4baeba9..8d5d8722 100644 --- a/tests/tpcds_correctness_test.rs +++ b/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; diff --git a/tests/tpcds_plans_test.rs b/benchmarks/tests/tpcds_plans_test.rs similarity index 99% rename from tests/tpcds_plans_test.rs rename to benchmarks/tests/tpcds_plans_test.rs index 412c50a2..b6363ad4 100644 --- a/tests/tpcds_plans_test.rs +++ b/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; diff --git a/tests/tpch_correctness_test.rs b/benchmarks/tests/tpch_correctness_test.rs similarity index 99% rename from tests/tpch_correctness_test.rs rename to benchmarks/tests/tpch_correctness_test.rs index 3ae089dc..7d6ab2ce 100644 --- a/tests/tpch_correctness_test.rs +++ b/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; diff --git a/tests/tpch_plans_test.rs b/benchmarks/tests/tpch_plans_test.rs similarity index 99% rename from tests/tpch_plans_test.rs rename to benchmarks/tests/tpch_plans_test.rs index 016e21ea..f625cfec 100644 --- a/tests/tpch_plans_test.rs +++ b/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::{ diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index 8ccc1f7c..ed55db2c 100644 --- a/src/coordinator/metrics_store.rs +++ b/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) -> Self { let map: HashMap<_, _> = entries.into_iter().collect(); let (tx, rx) = watch::channel(map); diff --git a/src/execution_plans/metrics.rs b/src/execution_plans/metrics.rs index 0fefdf3e..de97ed04 100644 --- a/src/execution_plans/metrics.rs +++ b/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 { &self.inner } diff --git a/src/execution_plans/mod.rs b/src/execution_plans/mod.rs index ecdcbc35..97904ef1 100644 --- a/src/execution_plans/mod.rs +++ b/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; diff --git a/src/lib.rs b/src/lib.rs index aa0917d7..33c8e7b2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -64,7 +64,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, diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 2d78998d..b613dcfc 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -19,7 +19,8 @@ pub fn collect_plan_metrics(plan: &Arc) -> Result Self { - Self::from_configured_worker(builder, |worker| worker) - } - - /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder] and worker setup. - pub fn from_configured_worker( - builder: impl WorkerSessionBuilder + Send + Sync + 'static, - configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, - ) -> 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: grpc::BoxCloneSyncChannel::new(channel), - }; - let this_clone = this.clone(); - - let endpoint = Worker::from_session_builder(builder.map(move |builder| { - let this = this.clone(); - Ok(builder.with_distributed_channel_resolver(this).build()) - })); - let endpoint = configure_worker(endpoint); - - #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { - Server::builder() - .add_service(endpoint.into_worker_server()) - .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server))) - .await - }); - - this_clone - } -} - -impl Default for InMemoryChannelResolver { - fn default() -> Self { - Self::from_session_builder(DefaultSessionBuilder) - } -} - -#[async_trait] -impl ChannelResolver for InMemoryChannelResolver { - async fn get_worker_client_for_url( - &self, - _: &url::Url, - ) -> Result, DataFusionError> { - Ok(grpc::create_worker_client(self.channel.clone())) - } -} - pub struct InMemoryWorkerResolver { n_workers: usize, } @@ -105,41 +23,139 @@ impl WorkerResolver for InMemoryWorkerResolver { } } -/// Creates a distributed session context backed by a single in-memory worker service. -/// The set of produced worker URLs is deterministic, taking the form http://worker-. -pub async fn start_in_memory_context( - num_workers: usize, - session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, -) -> SessionContext { - let channel_resolver = InMemoryChannelResolver::from_session_builder(session_builder); - let state = SessionStateBuilder::new() - .with_default_features() - .with_distributed_planner() - .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) - .with_distributed_channel_resolver(channel_resolver) - .build(); - SessionContext::from(state) -} +// The channel resolver runs a real `tonic` worker server over a tokio duplex, so it only compiles +// with the gRPC transport. The worker resolver above stays available without it, so the planner +// tests that only need worker URLs keep building in a no-gRPC config. +#[cfg(feature = "grpc")] +mod channel { + use super::InMemoryWorkerResolver; + use crate::{ + ChannelResolver, DefaultSessionBuilder, DistributedExt, MappedWorkerSessionBuilderExt, + SessionStateBuilderExt, Worker, WorkerChannel, WorkerSessionBuilder, grpc, + }; + use async_trait::async_trait; + use datafusion::common::DataFusionError; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::{SessionConfig, SessionContext}; + use hyper_util::rt::TokioIo; + use tonic::transport::{Endpoint, Server}; + + const DUMMY_URL: &str = "http://localhost:50051"; + + /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory + /// tokio duplex rather than a TCP connection. + #[derive(Clone)] + pub struct InMemoryChannelResolver { + channel: grpc::BoxCloneSyncChannel, + } + + impl InMemoryChannelResolver { + /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder]. + /// This allows you to inject your own DataFusion extensions in the in-memory worker + /// spawned by this method. + pub fn from_session_builder( + builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + Self::from_configured_worker(builder, |worker| worker) + } + + /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder] and worker setup. + pub fn from_configured_worker( + builder: impl WorkerSessionBuilder + Send + Sync + 'static, + configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, + ) -> 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: grpc::BoxCloneSyncChannel::new(channel), + }; + let this_clone = this.clone(); + + let endpoint = Worker::from_session_builder(builder.map(move |builder| { + let this = this.clone(); + Ok(builder.with_distributed_channel_resolver(this).build()) + })); + let endpoint = configure_worker(endpoint); + + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + Server::builder() + .add_service(endpoint.into_worker_server()) + .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server))) + .await + }); + + this_clone + } + } + + impl Default for InMemoryChannelResolver { + fn default() -> Self { + Self::from_session_builder(DefaultSessionBuilder) + } + } -/// Creates a distributed session context backed by a configurable in-memory worker service. -/// -/// Like [crate::test_utils::localhost::start_localhost_context], this uses tiny file-scan -/// partitions so small test datasets still cross worker boundaries. -pub async fn start_configured_in_memory_context( - num_workers: usize, - session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, - configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, -) -> SessionContext { - let channel_resolver = - InMemoryChannelResolver::from_configured_worker(session_builder, configure_worker); - let state = SessionStateBuilder::new() - .with_default_features() - .with_config(SessionConfig::new().with_target_partitions(num_workers)) - .with_distributed_planner() - .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) - .with_distributed_channel_resolver(channel_resolver) - .with_distributed_file_scan_config_bytes_per_partition(1) - .unwrap() - .build(); - SessionContext::from(state) + #[async_trait] + impl ChannelResolver for InMemoryChannelResolver { + async fn get_worker_client_for_url( + &self, + _: &url::Url, + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) + } + } + + /// Creates a distributed session context backed by a single in-memory worker service. + /// The set of produced worker URLs is deterministic, taking the form http://worker-. + pub async fn start_in_memory_context( + num_workers: usize, + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> SessionContext { + let channel_resolver = InMemoryChannelResolver::from_session_builder(session_builder); + let state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) + .with_distributed_channel_resolver(channel_resolver) + .build(); + SessionContext::from(state) + } + + /// Creates a distributed session context backed by a configurable in-memory worker service. + /// + /// Like [crate::test_utils::localhost::start_localhost_context], this uses tiny file-scan + /// partitions so small test datasets still cross worker boundaries. + pub async fn start_configured_in_memory_context( + num_workers: usize, + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, + ) -> SessionContext { + let channel_resolver = + InMemoryChannelResolver::from_configured_worker(session_builder, configure_worker); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(num_workers)) + .with_distributed_planner() + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) + .with_distributed_channel_resolver(channel_resolver) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + SessionContext::from(state) + } } + +#[cfg(feature = "grpc")] +pub use channel::{ + InMemoryChannelResolver, start_configured_in_memory_context, start_in_memory_context, +}; diff --git a/src/test_utils/mod.rs b/src/test_utils/mod.rs index dc66f977..6c84a3ff 100644 --- a/src/test_utils/mod.rs +++ b/src/test_utils/mod.rs @@ -1,5 +1,6 @@ pub mod in_memory_channel_resolver; pub mod insta; +#[cfg(feature = "grpc")] pub mod localhost; pub mod metrics; pub mod mock_exec; diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 26dc6c10..b36721fd 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -2,6 +2,7 @@ use arrow::{ array::{Int64Array, RecordBatch, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, }; +use async_trait::async_trait; use datafusion::{ catalog::{Session, TableFunctionImpl, TableProvider}, common::{ScalarValue, Statistics, internal_err, plan_err}, @@ -19,7 +20,6 @@ use datafusion_proto::{physical_plan::PhysicalExtensionCodec, protobuf::proto_er use futures::stream; use prost::Message; use std::{fmt::Formatter, sync::Arc}; -use tonic::async_trait; use url::Url; use crate::execution_plans::DistributedLeafExec; diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 6e8ce7b9..a0ab8959 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -3,7 +3,9 @@ mod impl_execute_task; mod session_builder; mod single_write_multi_read; mod task_data; -#[cfg(any(test, feature = "integration"))] +// `worker_handles` builds `tonic` servers and Flight channels for the benchmark fixtures, which +// only compile with the gRPC transport. +#[cfg(all(feature = "grpc", any(test, feature = "integration")))] pub(crate) mod test_utils; mod worker_connection_pool; mod worker_service; From 36af89baea29e3ab0f69a02b3eb755f2beed5cb6 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 14:37:05 -0700 Subject: [PATCH 3/3] Added an in-process worker transport that needs no gRPC. InProcessChannelResolver routes the three protocol methods straight to a co-located Worker, with no gRPC, no IPC, and no serialization round-trip: the reference implementation of the protocol for a co-located worker, and the first transport that exercises the abstraction with grpc off. Its end-to-end test (a distributed GROUP BY across tasks) runs under the no-gRPC CI job. Co-authored-by: Stu Hood --- src/lib.rs | 6 +- src/protocol/channel_resolver.rs | 3 + src/protocol/in_process.rs | 230 +++++++++++++++++++++++++++++++ src/protocol/mod.rs | 2 + 4 files changed, 238 insertions(+), 3 deletions(-) create mode 100644 src/protocol/in_process.rs diff --git a/src/lib.rs b/src/lib.rs index 33c8e7b2..32176a88 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -49,9 +49,9 @@ 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, diff --git a/src/protocol/channel_resolver.rs b/src/protocol/channel_resolver.rs index ded82e05..fee69fff 100644 --- a/src/protocol/channel_resolver.rs +++ b/src/protocol/channel_resolver.rs @@ -79,6 +79,9 @@ pub fn get_distributed_channel_resolver( #[cfg(not(feature = "grpc"))] { + // With `grpc` off there is no built-in default. A co-located deployment can register + // [`crate::InProcessChannelResolver`] (or another transport) via + // `with_distributed_channel_resolver`. panic!( "gRPC feature is not enabled, and no channel resolver was provided, so no default ChannelResolver can be provided" ); diff --git a/src/protocol/in_process.rs b/src/protocol/in_process.rs new file mode 100644 index 00000000..9a59085d --- /dev/null +++ b/src/protocol/in_process.rs @@ -0,0 +1,230 @@ +//! The reference implementation of the worker protocol for a co-located worker. +//! +//! It implements [`WorkerChannel`] by calling a [`Worker`] that lives in the same process, with no +//! gRPC, no IPC, and no networking. Its purpose is twofold: it lets the crate run distributed +//! queries with the `grpc` feature off (the protocol abstraction is only real if something other +//! than gRPC implements it), and it is the shape a custom co-located transport (for example a +//! shared-memory mesh spanning sibling processes) follows: implement [`ChannelResolver`] to hand +//! out a channel for a URL, then route the three protocol methods to wherever the worker runs. + +use crate::protocol::{ChannelResolver, WorkerChannel}; +use crate::{ + CoordinatorToWorkerMsg, DefaultSessionBuilder, DistributedExt, ExecuteTaskRequest, + GetWorkerInfoRequest, GetWorkerInfoResponse, MappedWorkerSessionBuilderExt, Worker, + WorkerSessionBuilder, WorkerToCoordinatorMsg, +}; +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::stream::{BoxStream, StreamExt}; +use http::HeaderMap; +use std::sync::{Arc, Weak}; +use url::Url; + +/// A [`ChannelResolver`] backed by a single co-located [`Worker`]. +/// +/// Every URL resolves to that one worker: with no network there is nothing to dial, so the URLs a +/// [`crate::WorkerResolver`] hands out only label the tasks the planner routes. One worker holds the +/// state for every task, keyed by [`crate::TaskKey`], the same way the gRPC worker does when several +/// tasks of a query land on it. +#[derive(Clone)] +pub struct InProcessChannelResolver { + worker: Arc, +} + +impl InProcessChannelResolver { + /// Builds the co-located worker from `session_builder`, registering this resolver into the + /// worker's own per-query sessions so a worker reading a downstream stage stays in process + /// rather than falling back to the gRPC default. + pub fn from_session_builder( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + // The worker and the resolver point at each other: the resolver runs tasks on the worker, + // and the worker resolves its own nested reads back through the resolver. A `Weak` on the + // worker's side breaks the cycle, so the returned `InProcessChannelResolver` owns the only + // strong reference and dropping it frees the worker. + let worker = Arc::new_cyclic(|weak: &Weak| { + let weak = weak.clone(); + Worker::from_session_builder(session_builder.map(move |builder| { + Ok(builder + .with_distributed_channel_resolver(WeakInProcessChannelResolver(weak.clone())) + .build()) + })) + }); + Self { worker } + } +} + +impl Default for InProcessChannelResolver { + fn default() -> Self { + Self::from_session_builder(DefaultSessionBuilder) + } +} + +#[async_trait] +impl ChannelResolver for InProcessChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + Ok(Box::new(InProcessWorkerChannel { + worker: Arc::clone(&self.worker), + })) + } +} + +/// The resolver a worker installs in its own sessions. It upgrades a [`Weak`] reference to the +/// co-located worker so a read of a downstream stage routes back in process instead of dialing out. +struct WeakInProcessChannelResolver(Weak); + +#[async_trait] +impl ChannelResolver for WeakInProcessChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + let worker = self + .0 + .upgrade() + .ok_or_else(|| internal_datafusion_err!("the in-process worker has been dropped"))?; + Ok(Box::new(InProcessWorkerChannel { worker })) + } +} + +/// A [`WorkerChannel`] that calls a co-located [`Worker`] directly. +struct InProcessWorkerChannel { + worker: Arc, +} + +#[async_trait] +impl WorkerChannel for InProcessWorkerChannel { + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + // The worker reads a fallible stream so a wire transport can surface its decode errors. + // Handing messages over in process has no such step, so each one is already an `Ok`. + self.worker + .coordinator_channel(headers, c2w_stream.map(Ok).boxed()) + .await + } + + async fn execute_task( + &mut self, + _headers: HeaderMap, + request: ExecuteTaskRequest, + _metrics: ExecutionPlanMetricsSet, + _task_ctx: &Arc, + ) -> Result>>> { + // Reading a partition runs the producer in place: the returned streams are the worker's own + // task output, so there is no IPC decode pass and no network metrics to record. The + // consumer's `task_ctx` is the consumer side's; the producer runs under the worker's own. + let (streams, _task_ctx) = self.worker.execute_task(request).await?; + Ok(streams.into_iter().map(|stream| stream.boxed()).collect()) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + Ok(GetWorkerInfoResponse { + version: self.worker.version().to_string(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{SessionStateBuilderExt, WorkerResolver, display_plan_ascii}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::collect; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use std::io::Write; + + /// Hands out as many placeholder URLs as workers. With one co-located worker behind the + /// transport, these only label the tasks the planner routes; nothing is dialed. + struct InProcessWorkers(usize); + + impl WorkerResolver for InProcessWorkers { + fn get_urls(&self) -> Result, DataFusionError> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + /// Drives a real distributed query end to end through the in-process transport. With the `grpc` + /// feature off this is the only transport that can run it; `cargo check --no-default-features` + /// covers the no-gRPC compile. + #[tokio::test] + async fn in_process_transport_runs_a_distributed_query() -> Result<()> { + const N: usize = 4; + + // A file scan round-trips through `datafusion-proto`, so a worker can rebuild it from the + // serialized stage plan. An in-memory table would not, hence a CSV on disk. + let path = std::env::temp_dir().join(format!("dfd_in_process_{}.csv", std::process::id())); + let mut file = + std::fs::File::create(&path).map_err(|e| DataFusionError::External(Box::new(e)))?; + writeln!(file, "k,v").unwrap(); + for i in 0..200 { + writeln!(file, "{},{}", ["a", "b", "c", "d"][i % 4], i).unwrap(); + } + drop(file); + let path = path.to_str().unwrap().to_string(); + + let query = "SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k"; + + // Single-node reference result. + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(N)); + ctx.register_csv("t", &path, CsvReadOptions::new()).await?; + let expected = collect( + ctx.sql(query).await?.create_physical_plan().await?, + ctx.task_ctx(), + ) + .await?; + let expected = pretty_format_batches(&expected)?.to_string(); + + // Distributed over the in-process transport. + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(N)) + .with_distributed_planner() + .with_distributed_worker_resolver(InProcessWorkers(N)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx_distributed = SessionContext::from(state); + ctx_distributed + .register_csv("t", &path, CsvReadOptions::new()) + .await?; + + let physical = ctx_distributed + .sql(query) + .await? + .create_physical_plan() + .await?; + let rendered = display_plan_ascii(physical.as_ref(), false); + assert!( + rendered.contains("DistributedExec"), + "plan was not distributed:\n{rendered}" + ); + assert!( + rendered.contains("NetworkShuffleExec"), + "no shuffle boundary, so the transport never carried a cross-task stream:\n{rendered}" + ); + + let actual = collect(physical, ctx_distributed.task_ctx()).await?; + let actual = pretty_format_batches(&actual)?.to_string(); + + let _ = std::fs::remove_file(&path); + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 3f4dd666..dd878f2d 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -6,10 +6,12 @@ mod channel_resolver; // worker, a shared-memory mesh) can speak the same wire shape without pulling in the whole gRPC // stack. pub(crate) mod generated; +mod in_process; mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; +pub use in_process::InProcessChannelResolver; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse,