From 4721f09b295b8953f6d6e802e004c2725abb115b Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:11:48 -0700 Subject: [PATCH 1/4] 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 b132caf5..a02710b5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,6 +46,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 d1f0fb4a..e60dc1ed 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/generated/worker.rs @@ -467,6 +467,7 @@ pub struct MaxGauge { pub value: u64, } /// Generated client implementations. +#[cfg(feature = "grpc")] pub mod worker_service_client { #![allow( unused_variables, @@ -616,6 +617,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 56612671..d342eebe 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 6fba14ee..831865a9 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 6df0e11b..02465876 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(crate) use channel_resolver::set_distributed_channel_resolver; From b0a7e4ae2a4bea63be51a188cb3099caf267e947 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:46:51 -0700 Subject: [PATCH 2/4] 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 + benchmarks/cdk/bin/worker.rs | 1 + .../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 +- console/src/main.rs | 8 +- console/src/worker.rs | 50 ++-- 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 +- 24 files changed, 258 insertions(+), 179 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 dbe6096c..12a4e62b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2199,7 +2199,6 @@ dependencies = [ "crossbeam-queue", "dashmap", "datafusion", - "datafusion-distributed-benchmarks", "datafusion-proto", "delegate", "futures", @@ -2249,12 +2248,14 @@ dependencies = [ "object_store", "openssl", "parquet", + "pretty_assertions", "reqwest", "serde", "serde_json", "sketches-ddsketch", "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 531a8aba..58736c66 100644 --- a/benchmarks/Cargo.toml +++ b/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"] } diff --git a/benchmarks/cdk/bin/worker.rs b/benchmarks/cdk/bin/worker.rs index 96a399c8..cdeb64a3 100644 --- a/benchmarks/cdk/bin/worker.rs +++ b/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; 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/console/src/main.rs b/console/src/main.rs index 731b504b..ab1171d7 100644 --- a/console/src/main.rs +++ b/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 { diff --git a/console/src/worker.rs b/console/src/worker.rs index 4a23c5b9..44c06260 100644 --- a/console/src/worker.rs +++ b/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); } } } 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 a02710b5..46dae7df 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -71,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, diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index ba6e6855..f0647b95 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 15fdabc6..d687e424 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 crate::execution_plans::DistributedLeafExec; use crate::worker::LocalWorkerContext; diff --git a/src/worker/mod.rs b/src/worker/mod.rs index b2e79eac..030ac357 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 5f41597283db146f067a781764a27d3ad94c4037 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 14:37:05 -0700 Subject: [PATCH 3/4] 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 46dae7df..f7becd63 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,9 +56,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 e781aa90..05da3069 100644 --- a/src/protocol/channel_resolver.rs +++ b/src/protocol/channel_resolver.rs @@ -70,6 +70,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 02465876..3aee1273 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(crate) use channel_resolver::set_distributed_channel_resolver; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; +pub use in_process::InProcessChannelResolver; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, From 38884704ee59617abd2c18d05062bedcf430dcc9 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Tue, 28 Jul 2026 11:58:19 -0700 Subject: [PATCH 4/4] Split `worker.proto` into `messages.proto` and `worker_service.proto`. --- src/protocol/generated/mod.rs | 4 + src/protocol/generated/worker.rs | 434 +---------------- src/protocol/generated/worker_grpc.rs | 447 ++++++++++++++++++ src/protocol/grpc/gen/src/main.rs | 16 +- src/protocol/grpc/worker_client.rs | 5 +- src/protocol/grpc/worker_service.proto | 15 + src/protocol/grpc/worker_service.rs | 12 +- .../{grpc/worker.proto => messages.proto} | 10 - 8 files changed, 485 insertions(+), 458 deletions(-) create mode 100644 src/protocol/generated/worker_grpc.rs create mode 100644 src/protocol/grpc/worker_service.proto rename src/protocol/{grpc/worker.proto => messages.proto} (94%) diff --git a/src/protocol/generated/mod.rs b/src/protocol/generated/mod.rs index 2c8b8399..a787a5fc 100644 --- a/src/protocol/generated/mod.rs +++ b/src/protocol/generated/mod.rs @@ -1 +1,5 @@ pub mod worker; + +#[cfg(feature = "grpc")] +#[allow(clippy::all)] +pub mod worker_grpc; diff --git a/src/protocol/generated/worker.rs b/src/protocol/generated/worker.rs index e60dc1ed..7e70450c 100644 --- a/src/protocol/generated/worker.rs +++ b/src/protocol/generated/worker.rs @@ -80,8 +80,8 @@ pub struct LoadInfo { /// The amount of rows that were pulled from leaf nodes while this partition was sampling data. #[prost(uint64, tag = "8")] pub rows_pulled_from_leaf: u64, - /// Whether the sampled partition stream reached end-of-stream by the time this LoadInfo was - /// captured. + /// Whether the sampled partition stream reached end-of-stream (i.e. the partition finished + /// producing all of its output) by the time this LoadInfo was captured. #[prost(bool, tag = "9")] pub reached_eos: bool, } @@ -466,433 +466,3 @@ pub struct MaxGauge { #[prost(uint64, tag = "2")] pub value: u64, } -/// Generated client implementations. -#[cfg(feature = "grpc")] -pub mod worker_service_client { - #![allow( - unused_variables, - dead_code, - missing_docs, - clippy::wildcard_imports, - clippy::let_unit_value - )] - use tonic::codegen::http::Uri; - use tonic::codegen::*; - #[derive(Debug, Clone)] - pub struct WorkerServiceClient { - inner: tonic::client::Grpc, - } - impl WorkerServiceClient { - /// Attempt to create a new client by connecting to a given endpoint. - pub async fn connect(dst: D) -> Result - where - D: TryInto, - D::Error: Into, - { - let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; - Ok(Self::new(conn)) - } - } - impl WorkerServiceClient - where - T: tonic::client::GrpcService, - T::Error: Into, - T::ResponseBody: Body + std::marker::Send + 'static, - ::Error: Into + std::marker::Send, - { - pub fn new(inner: T) -> Self { - let inner = tonic::client::Grpc::new(inner); - Self { inner } - } - pub fn with_origin(inner: T, origin: Uri) -> Self { - let inner = tonic::client::Grpc::with_origin(inner, origin); - Self { inner } - } - pub fn with_interceptor( - inner: T, - interceptor: F, - ) -> WorkerServiceClient> - where - F: tonic::service::Interceptor, - T::ResponseBody: Default, - T: tonic::codegen::Service< - http::Request, - Response = http::Response< - >::ResponseBody, - >, - >, - >>::Error: - Into + std::marker::Send + std::marker::Sync, - { - WorkerServiceClient::new(InterceptedService::new(inner, interceptor)) - } - /// Compress requests with the given encoding. - /// - /// This requires the server to support it otherwise it might respond with an - /// error. - #[must_use] - pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { - self.inner = self.inner.send_compressed(encoding); - self - } - /// Enable decompressing responses. - #[must_use] - pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { - self.inner = self.inner.accept_compressed(encoding); - self - } - /// Limits the maximum size of a decoded message. - /// - /// Default: `4MB` - #[must_use] - pub fn max_decoding_message_size(mut self, limit: usize) -> Self { - self.inner = self.inner.max_decoding_message_size(limit); - self - } - /// Limits the maximum size of an encoded message. - /// - /// Default: `usize::MAX` - #[must_use] - pub fn max_encoding_message_size(mut self, limit: usize) -> Self { - self.inner = self.inner.max_encoding_message_size(limit); - self - } - /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages - /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel - /// per task. - pub async fn coordinator_channel( - &mut self, - request: impl tonic::IntoStreamingRequest, - ) -> std::result::Result< - tonic::Response>, - tonic::Status, - > { - self.inner.ready().await.map_err(|e| { - tonic::Status::unknown(format!("Service was not ready: {}", e.into())) - })?; - let codec = tonic_prost::ProstCodec::default(); - let path = - http::uri::PathAndQuery::from_static("/worker.WorkerService/CoordinatorChannel"); - let mut req = request.into_streaming_request(); - req.extensions_mut().insert(GrpcMethod::new( - "worker.WorkerService", - "CoordinatorChannel", - )); - self.inner.streaming(req, path, codec).await - } - /// Executes the requested partition range of a subplan previously sent by the coordinator channel. - pub async fn execute_task( - &mut self, - request: impl tonic::IntoRequest, - ) -> std::result::Result< - tonic::Response>, - tonic::Status, - > { - self.inner.ready().await.map_err(|e| { - tonic::Status::unknown(format!("Service was not ready: {}", e.into())) - })?; - let codec = tonic_prost::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static("/worker.WorkerService/ExecuteTask"); - let mut req = request.into_request(); - req.extensions_mut() - .insert(GrpcMethod::new("worker.WorkerService", "ExecuteTask")); - self.inner.server_streaming(req, path, codec).await - } - /// Returns metadata about a worker. Currently only used for worker versioning. - pub async fn get_worker_info( - &mut self, - request: impl tonic::IntoRequest, - ) -> std::result::Result, tonic::Status> - { - self.inner.ready().await.map_err(|e| { - tonic::Status::unknown(format!("Service was not ready: {}", e.into())) - })?; - let codec = tonic_prost::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static("/worker.WorkerService/GetWorkerInfo"); - let mut req = request.into_request(); - req.extensions_mut() - .insert(GrpcMethod::new("worker.WorkerService", "GetWorkerInfo")); - self.inner.unary(req, path, codec).await - } - } -} -/// Generated server implementations. -#[cfg(feature = "grpc")] -pub mod worker_service_server { - #![allow( - unused_variables, - dead_code, - missing_docs, - clippy::wildcard_imports, - clippy::let_unit_value - )] - use tonic::codegen::*; - /// Generated trait containing gRPC methods that should be implemented for use with WorkerServiceServer. - #[async_trait] - pub trait WorkerService: std::marker::Send + std::marker::Sync + 'static { - /// Server streaming response type for the CoordinatorChannel method. - type CoordinatorChannelStream: tonic::codegen::tokio_stream::Stream< - Item = std::result::Result, - > + std::marker::Send - + 'static; - /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages - /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel - /// per task. - async fn coordinator_channel( - &self, - request: tonic::Request>, - ) -> std::result::Result, tonic::Status>; - /// Server streaming response type for the ExecuteTask method. - type ExecuteTaskStream: tonic::codegen::tokio_stream::Stream< - Item = std::result::Result<::arrow_flight::FlightData, tonic::Status>, - > + std::marker::Send - + 'static; - /// Executes the requested partition range of a subplan previously sent by the coordinator channel. - async fn execute_task( - &self, - request: tonic::Request, - ) -> std::result::Result, tonic::Status>; - /// Returns metadata about a worker. Currently only used for worker versioning. - async fn get_worker_info( - &self, - request: tonic::Request, - ) -> std::result::Result, tonic::Status>; - } - #[derive(Debug)] - pub struct WorkerServiceServer { - inner: Arc, - accept_compression_encodings: EnabledCompressionEncodings, - send_compression_encodings: EnabledCompressionEncodings, - max_decoding_message_size: Option, - max_encoding_message_size: Option, - } - impl WorkerServiceServer { - pub fn new(inner: T) -> Self { - Self::from_arc(Arc::new(inner)) - } - pub fn from_arc(inner: Arc) -> Self { - Self { - inner, - accept_compression_encodings: Default::default(), - send_compression_encodings: Default::default(), - max_decoding_message_size: None, - max_encoding_message_size: None, - } - } - pub fn with_interceptor(inner: T, interceptor: F) -> InterceptedService - where - F: tonic::service::Interceptor, - { - InterceptedService::new(Self::new(inner), interceptor) - } - /// Enable decompressing requests with the given encoding. - #[must_use] - pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { - self.accept_compression_encodings.enable(encoding); - self - } - /// Compress responses with the given encoding, if the client supports it. - #[must_use] - pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { - self.send_compression_encodings.enable(encoding); - self - } - /// Limits the maximum size of a decoded message. - /// - /// Default: `4MB` - #[must_use] - pub fn max_decoding_message_size(mut self, limit: usize) -> Self { - self.max_decoding_message_size = Some(limit); - self - } - /// Limits the maximum size of an encoded message. - /// - /// Default: `usize::MAX` - #[must_use] - pub fn max_encoding_message_size(mut self, limit: usize) -> Self { - self.max_encoding_message_size = Some(limit); - self - } - } - impl tonic::codegen::Service> for WorkerServiceServer - where - T: WorkerService, - B: Body + std::marker::Send + 'static, - B::Error: Into + std::marker::Send + 'static, - { - type Response = http::Response; - type Error = std::convert::Infallible; - type Future = BoxFuture; - fn poll_ready( - &mut self, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - fn call(&mut self, req: http::Request) -> Self::Future { - match req.uri().path() { - "/worker.WorkerService/CoordinatorChannel" => { - #[allow(non_camel_case_types)] - struct CoordinatorChannelSvc(pub Arc); - impl - tonic::server::StreamingService - for CoordinatorChannelSvc - { - type Response = super::WorkerToCoordinatorMsg; - type ResponseStream = T::CoordinatorChannelStream; - type Future = - BoxFuture, tonic::Status>; - fn call( - &mut self, - request: tonic::Request< - tonic::Streaming, - >, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::coordinator_channel(&inner, request).await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = CoordinatorChannelSvc(inner); - let codec = tonic_prost::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.streaming(method, req).await; - Ok(res) - }; - Box::pin(fut) - } - "/worker.WorkerService/ExecuteTask" => { - #[allow(non_camel_case_types)] - struct ExecuteTaskSvc(pub Arc); - impl - tonic::server::ServerStreamingService - for ExecuteTaskSvc - { - type Response = ::arrow_flight::FlightData; - type ResponseStream = T::ExecuteTaskStream; - type Future = - BoxFuture, tonic::Status>; - fn call( - &mut self, - request: tonic::Request, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::execute_task(&inner, request).await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = ExecuteTaskSvc(inner); - let codec = tonic_prost::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.server_streaming(method, req).await; - Ok(res) - }; - Box::pin(fut) - } - "/worker.WorkerService/GetWorkerInfo" => { - #[allow(non_camel_case_types)] - struct GetWorkerInfoSvc(pub Arc); - impl tonic::server::UnaryService - for GetWorkerInfoSvc - { - type Response = super::GetWorkerInfoResponse; - type Future = BoxFuture, tonic::Status>; - fn call( - &mut self, - request: tonic::Request, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::get_worker_info(&inner, request).await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = GetWorkerInfoSvc(inner); - let codec = tonic_prost::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.unary(method, req).await; - Ok(res) - }; - Box::pin(fut) - } - _ => Box::pin(async move { - let mut response = http::Response::new(tonic::body::Body::default()); - let headers = response.headers_mut(); - headers.insert( - tonic::Status::GRPC_STATUS, - (tonic::Code::Unimplemented as i32).into(), - ); - headers.insert( - http::header::CONTENT_TYPE, - tonic::metadata::GRPC_CONTENT_TYPE, - ); - Ok(response) - }), - } - } - } - impl Clone for WorkerServiceServer { - fn clone(&self) -> Self { - let inner = self.inner.clone(); - Self { - inner, - accept_compression_encodings: self.accept_compression_encodings, - send_compression_encodings: self.send_compression_encodings, - max_decoding_message_size: self.max_decoding_message_size, - max_encoding_message_size: self.max_encoding_message_size, - } - } - } - /// Generated gRPC service name - pub const SERVICE_NAME: &str = "worker.WorkerService"; - impl tonic::server::NamedService for WorkerServiceServer { - const NAME: &'static str = SERVICE_NAME; - } -} diff --git a/src/protocol/generated/worker_grpc.rs b/src/protocol/generated/worker_grpc.rs new file mode 100644 index 00000000..ba955c27 --- /dev/null +++ b/src/protocol/generated/worker_grpc.rs @@ -0,0 +1,447 @@ +// This file is @generated by prost-build. +/// Generated client implementations. +pub mod worker_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value + )] + use tonic::codegen::http::Uri; + use tonic::codegen::*; + #[derive(Debug, Clone)] + pub struct WorkerServiceClient { + inner: tonic::client::Grpc, + } + impl WorkerServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl WorkerServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> WorkerServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + >>::Error: + Into + std::marker::Send + std::marker::Sync, + { + WorkerServiceClient::new(InterceptedService::new(inner, interceptor)) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages + /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel + /// per task. + pub async fn coordinator_channel( + &mut self, + request: impl tonic::IntoStreamingRequest< + Message = super::super::worker::CoordinatorToWorkerMsg, + >, + ) -> std::result::Result< + tonic::Response>, + tonic::Status, + > { + self.inner.ready().await.map_err(|e| { + tonic::Status::unknown(format!("Service was not ready: {}", e.into())) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/worker_grpc.WorkerService/CoordinatorChannel", + ); + let mut req = request.into_streaming_request(); + req.extensions_mut().insert(GrpcMethod::new( + "worker_grpc.WorkerService", + "CoordinatorChannel", + )); + self.inner.streaming(req, path, codec).await + } + /// Executes the requested partition range of a subplan previously sent by the coordinator channel. + pub async fn execute_task( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response>, + tonic::Status, + > { + self.inner.ready().await.map_err(|e| { + tonic::Status::unknown(format!("Service was not ready: {}", e.into())) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = + http::uri::PathAndQuery::from_static("/worker_grpc.WorkerService/ExecuteTask"); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("worker_grpc.WorkerService", "ExecuteTask")); + self.inner.server_streaming(req, path, codec).await + } + /// Returns metadata about a worker. Currently only used for worker versioning. + pub async fn get_worker_info( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner.ready().await.map_err(|e| { + tonic::Status::unknown(format!("Service was not ready: {}", e.into())) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = + http::uri::PathAndQuery::from_static("/worker_grpc.WorkerService/GetWorkerInfo"); + let mut req = request.into_request(); + req.extensions_mut().insert(GrpcMethod::new( + "worker_grpc.WorkerService", + "GetWorkerInfo", + )); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated server implementations. +pub mod worker_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with WorkerServiceServer. + #[async_trait] + pub trait WorkerService: std::marker::Send + std::marker::Sync + 'static { + /// Server streaming response type for the CoordinatorChannel method. + type CoordinatorChannelStream: tonic::codegen::tokio_stream::Stream< + Item = std::result::Result< + super::super::worker::WorkerToCoordinatorMsg, + tonic::Status, + >, + > + std::marker::Send + + 'static; + /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages + /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel + /// per task. + async fn coordinator_channel( + &self, + request: tonic::Request>, + ) -> std::result::Result, tonic::Status>; + /// Server streaming response type for the ExecuteTask method. + type ExecuteTaskStream: tonic::codegen::tokio_stream::Stream< + Item = std::result::Result<::arrow_flight::FlightData, tonic::Status>, + > + std::marker::Send + + 'static; + /// Executes the requested partition range of a subplan previously sent by the coordinator channel. + async fn execute_task( + &self, + request: tonic::Request, + ) -> std::result::Result, tonic::Status>; + /// Returns metadata about a worker. Currently only used for worker versioning. + async fn get_worker_info( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct WorkerServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl WorkerServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor(inner: T, interceptor: F) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> for WorkerServiceServer + where + T: WorkerService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/worker_grpc.WorkerService/CoordinatorChannel" => { + #[allow(non_camel_case_types)] + struct CoordinatorChannelSvc(pub Arc); + impl + tonic::server::StreamingService< + super::super::worker::CoordinatorToWorkerMsg, + > for CoordinatorChannelSvc + { + type Response = super::super::worker::WorkerToCoordinatorMsg; + type ResponseStream = T::CoordinatorChannelStream; + type Future = + BoxFuture, tonic::Status>; + fn call( + &mut self, + request: tonic::Request< + tonic::Streaming, + >, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::coordinator_channel(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = CoordinatorChannelSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.streaming(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + "/worker_grpc.WorkerService/ExecuteTask" => { + #[allow(non_camel_case_types)] + struct ExecuteTaskSvc(pub Arc); + impl + tonic::server::ServerStreamingService< + super::super::worker::ExecuteTaskRequest, + > for ExecuteTaskSvc + { + type Response = ::arrow_flight::FlightData; + type ResponseStream = T::ExecuteTaskStream; + type Future = + BoxFuture, tonic::Status>; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::execute_task(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = ExecuteTaskSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.server_streaming(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + "/worker_grpc.WorkerService/GetWorkerInfo" => { + #[allow(non_camel_case_types)] + struct GetWorkerInfoSvc(pub Arc); + impl + tonic::server::UnaryService + for GetWorkerInfoSvc + { + type Response = super::super::worker::GetWorkerInfoResponse; + type Future = BoxFuture, tonic::Status>; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::get_worker_info(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = GetWorkerInfoSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + _ => Box::pin(async move { + let mut response = http::Response::new(tonic::body::Body::default()); + let headers = response.headers_mut(); + headers.insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers.insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }), + } + } + } + impl Clone for WorkerServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "worker_grpc.WorkerService"; + impl tonic::server::NamedService for WorkerServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} diff --git a/src/protocol/grpc/gen/src/main.rs b/src/protocol/grpc/gen/src/main.rs index cde0e4e6..9c6addbe 100644 --- a/src/protocol/grpc/gen/src/main.rs +++ b/src/protocol/grpc/gen/src/main.rs @@ -4,31 +4,29 @@ use std::fs; fn main() -> Result<(), Box> { let repo_root = env::current_dir()?; - let proto_dir = repo_root.join("src/protocol/grpc"); - let proto_file = proto_dir.join("worker.proto"); + let protocol_dir = repo_root.join("src/protocol"); + let messages_proto = protocol_dir.join("messages.proto"); + let worker_service_proto = protocol_dir.join("grpc/worker_service.proto"); let out_dir = repo_root.join("src/protocol/generated"); fs::create_dir_all(&out_dir)?; println!("Generating protobuf code..."); - println!("Proto dir: {proto_dir:?}"); - println!("Proto file: {proto_file:?}"); + println!("Protocol dir: {protocol_dir:?}"); + println!("Messages proto: {messages_proto:?}"); + println!("Worker service proto: {worker_service_proto:?}"); println!("Output dir: {out_dir:?}"); 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( ".worker.FlightDescriptor", "::arrow_flight::FlightDescriptor", ) - .compile_protos(&[proto_file], &[proto_dir])?; + .compile_protos(&[messages_proto, worker_service_proto], &[protocol_dir])?; println!("Successfully generated worker proto code"); diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index d342eebe..28ac9d38 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -5,6 +5,7 @@ use crate::common::serialize_uuid; use crate::grpc::on_drop_stream::on_drop_stream; use crate::protocol::generated::worker as pb; use crate::protocol::generated::worker::FlightAppMetadata; +use crate::protocol::generated::worker_grpc::worker_service_client::WorkerServiceClient; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, @@ -43,7 +44,7 @@ use tonic::metadata::MetadataMap; use tonic::{Request, Status}; #[async_trait] -impl WorkerChannel for pb::worker_service_client::WorkerServiceClient { +impl WorkerChannel for WorkerServiceClient { async fn coordinator_channel( &mut self, headers: HeaderMap, @@ -567,7 +568,7 @@ fn fanout(o_txs: &[UnboundedSender], err: Status) { /// ``` pub fn create_worker_client(channel: BoxCloneSyncChannel) -> Box { Box::new( - pb::worker_service_client::WorkerServiceClient::new(channel) + WorkerServiceClient::new(channel) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX), ) diff --git a/src/protocol/grpc/worker_service.proto b/src/protocol/grpc/worker_service.proto new file mode 100644 index 00000000..9ba6c327 --- /dev/null +++ b/src/protocol/grpc/worker_service.proto @@ -0,0 +1,15 @@ +syntax = "proto3"; +package worker_grpc; + +import "messages.proto"; + +service WorkerService { + // Establishes a bidirectional message stream between a coordinator and a worker, over which messages + // will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel + // per task. + rpc CoordinatorChannel(stream worker.CoordinatorToWorkerMsg) returns (stream worker.WorkerToCoordinatorMsg); + // Executes the requested partition range of a subplan previously sent by the coordinator channel. + rpc ExecuteTask(worker.ExecuteTaskRequest) returns (stream worker.FlightData); + // Returns metadata about a worker. Currently only used for worker versioning. + rpc GetWorkerInfo(worker.GetWorkerInfoRequest) returns (worker.GetWorkerInfoResponse); +} diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 831865a9..c3666d5c 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -1,10 +1,12 @@ use super::errors::{datafusion_error_to_tonic_status, map_status_to_datafusion_error}; 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; +use crate::protocol::generated::worker as pb; +use crate::protocol::generated::worker_grpc::worker_service_server::{ + WorkerService, WorkerServiceServer, +}; use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; use crate::{ CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, SetPlanRequest, @@ -59,8 +61,8 @@ impl Worker { /// /// # } /// ``` - pub fn into_worker_server(self) -> pb::worker_service_server::WorkerServiceServer { - pb::worker_service_server::WorkerServiceServer::new(self) + pub fn into_worker_server(self) -> WorkerServiceServer { + WorkerServiceServer::new(self) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX) } @@ -86,7 +88,7 @@ impl Worker { /// The methods are delegated to plan `impl Worker` implementations so that they can be implemented /// in different files. #[async_trait] -impl pb::worker_service_server::WorkerService for Worker { +impl WorkerService for Worker { type CoordinatorChannelStream = BoxStream<'static, Result>; async fn coordinator_channel( diff --git a/src/protocol/grpc/worker.proto b/src/protocol/messages.proto similarity index 94% rename from src/protocol/grpc/worker.proto rename to src/protocol/messages.proto index 24f5c987..f5127325 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/messages.proto @@ -1,16 +1,6 @@ syntax = "proto3"; package worker; -service WorkerService { - // Establishes a bidirectional message stream between a coordinator and a worker, over which messages - // will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel - // per task. - rpc CoordinatorChannel(stream CoordinatorToWorkerMsg) returns (stream WorkerToCoordinatorMsg); - // Executes the requested partition range of a subplan previously sent by the coordinator channel. - rpc ExecuteTask(ExecuteTaskRequest) returns (stream FlightData); - // Returns metadata about a worker. Currently only used for worker versioning. - rpc GetWorkerInfo(GetWorkerInfoRequest) returns (GetWorkerInfoResponse); -} message CoordinatorToWorkerMsg { oneof inner {