From a59101f8f90b525037a31c59e345c40d081ba7d7 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:11:48 -0700 Subject: [PATCH 01/11] 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 3c89d3e8314140a0baab185caa0ab1f001de87fa Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:46:51 -0700 Subject: [PATCH 02/11] 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 04584d36753cb2c73ebfab46fcd57c6dcff003b4 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 14:37:05 -0700 Subject: [PATCH 03/11] 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, From e26d61305fd0dfa210d938a4f625c5dd4251ccb2 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 9 Jul 2026 16:02:52 -0700 Subject: [PATCH 04/11] Re-exposed the worker-metrics API an out-of-crate driver needs. A transport that returns worker metrics out-of-band, rather than over the coordinator_channel return stream, needs its driver to decode the frames and file them into the executed plan's store before the per-task EXPLAIN rewrite reads it: metrics_store() with a public insert, a no-gRPC decode_task_metrics, and the frame builders collect_plan_metrics_protos and set_received_time. The metrics codec moves out from behind grpc with them, and the gRPC client drops its private copy of the decode. The test pins the frame-to-store path. Co-authored-by: Stu Hood --- src/coordinator/distributed.rs | 7 +++ src/coordinator/metrics_store.rs | 4 +- src/coordinator/mod.rs | 2 +- src/lib.rs | 11 ++-- src/protocol/grpc/mod.rs | 1 - src/protocol/grpc/worker_client.rs | 21 +------ src/protocol/grpc/worker_service.rs | 2 +- src/protocol/{grpc => }/metrics_proto.rs | 70 ++++++++++++++++++++- src/protocol/mod.rs | 4 ++ src/work_unit_feed/mod.rs | 1 + src/work_unit_feed/remote_work_unit_feed.rs | 7 +++ src/worker/impl_coordinator_channel.rs | 22 +++++++ src/worker/mod.rs | 1 + 13 files changed, 126 insertions(+), 27 deletions(-) rename src/protocol/{grpc => }/metrics_proto.rs (95%) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index a1e123f2..17fe31ab 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -64,6 +64,13 @@ impl DistributedExec { } } + /// The store where worker task metrics land at runtime, if metrics collection is enabled. + /// Exposed for a driver whose transport returns metrics out-of-band; it files decoded frames + /// here before the per-task EXPLAIN rewrite. + pub fn metrics_store(&self) -> Option> { + self.metrics_store.clone() + } + /// Enables task metrics collection from remote workers. pub fn with_metrics_collection(mut self, enabled: bool) -> Self { self.metrics_store = match enabled { diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index ed55db2c..b38d62a0 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -17,7 +17,9 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { + // Public for a driver whose transport returns worker metrics out-of-band; it files the + // decoded frames here before the per-task EXPLAIN rewrite reads them. + pub fn insert(&self, key: TaskKey, metrics: TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd..be069485 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,4 +6,4 @@ mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; -pub(crate) use metrics_store::MetricsStore; +pub use metrics_store::MetricsStore; diff --git a/src/lib.rs b/src/lib.rs index 32176a88..559f6b78 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,9 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +// `MetricsStore` is public for a driver whose transport returns worker metrics out-of-band; it +// files the decoded frames into the executed plan's store before the per-task EXPLAIN rewrite. +pub use coordinator::{DistributedExec, MetricsStore}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, @@ -51,17 +53,18 @@ pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, + WorkerToCoordinatorMsg, decode_task_metrics, get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, }; pub use work_unit_feed::{ - DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, + DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, + WorkUnitFeedProvider, set_received_time, }; pub use worker::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, - Worker, WorkerQueryContext, WorkerSessionBuilder, + Worker, WorkerQueryContext, WorkerSessionBuilder, collect_plan_metrics_protos, }; #[cfg(all(feature = "grpc", any(feature = "integration", test)))] diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs index fc367585..cc2f6576 100644 --- a/src/protocol/grpc/mod.rs +++ b/src/protocol/grpc/mod.rs @@ -1,6 +1,5 @@ mod channel_resolver; mod errors; -mod metrics_proto; mod observability; mod on_drop_stream; mod spawn_select_all; diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 93291d91..c2866827 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -1,16 +1,16 @@ use super::channel_resolver::BoxCloneSyncChannel; use super::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; -use super::metrics_proto::metrics_set_proto_to_df; 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::metrics_proto::decode_task_metrics; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHeadSpec, - SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, - WorkerChannel, WorkerToCoordinatorMsg, + SetPlanRequest, TaskKey, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, + WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -491,21 +491,6 @@ fn decode_worker_to_coordinator_msg( ) } -fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { - Ok(TaskMetrics { - pre_order_plan_metrics: task_metrics - .pre_order_plan_metrics - .into_iter() - .map(|metrics_set| metrics_set_proto_to_df(&metrics_set)) - .collect::>()?, - task_metrics: metrics_set_proto_to_df( - &task_metrics - .task_metrics - .ok_or_else(|| missing("task_metrics"))?, - )?, - }) -} - fn decode_load_info(load_info: pb::LoadInfo) -> LoadInfo { LoadInfo { partition: load_info.partition as usize, diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index f3472d43..87c1ef4b 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::metrics_proto::df_metrics_set_to_proto; use super::spawn_select_all::spawn_select_all; use crate::protocol::generated::worker as pb; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::ProducerHeadSpec; diff --git a/src/protocol/grpc/metrics_proto.rs b/src/protocol/metrics_proto.rs similarity index 95% rename from src/protocol/grpc/metrics_proto.rs rename to src/protocol/metrics_proto.rs index 654b6cee..f72e0edb 100644 --- a/src/protocol/grpc/metrics_proto.rs +++ b/src/protocol/metrics_proto.rs @@ -1,4 +1,4 @@ -use crate::protocol::generated::worker as pb; +use super::generated::worker as pb; use chrono::DateTime; use datafusion::common::internal_err; use datafusion::error::DataFusionError; @@ -55,6 +55,27 @@ pub fn metrics_set_proto_to_df( Ok(metrics_set) } +/// Decode a wire [`pb::TaskMetrics`] into the in-memory [`crate::TaskMetrics`]. The no-gRPC push +/// embedder receives metric frames as proto and files them into the plain metrics store, so it +/// needs the decode direction without the gRPC client. +pub fn decode_task_metrics( + task_metrics: pb::TaskMetrics, +) -> Result { + Ok(crate::TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .iter() + .map(metrics_set_proto_to_df) + .collect::>()?, + task_metrics: metrics_set_proto_to_df( + task_metrics + .task_metrics + .as_ref() + .ok_or_else(|| DataFusionError::Internal("Missing field 'task_metrics'".into()))?, + )?, + }) +} + /// Custom metrics are not supported in proto conversion. const CUSTOM_METRICS_NOT_SUPPORTED: &str = "custom metrics are not supported in metrics proto conversion"; @@ -1305,3 +1326,50 @@ mod tests { } } } + +#[cfg(test)] +mod decode_tests { + use super::*; + use crate::TaskKey; + use crate::coordinator::MetricsStore; + use datafusion::physical_plan::metrics::{Count, Metric, MetricValue, MetricsSet}; + + /// Pins the driver-facing decode path: a proto metrics frame (the wire form a transport + /// delivers out-of-band) decodes to the plain form and files into the store the per-task + /// EXPLAIN rewrite reads. + #[test] + fn a_metrics_frame_decodes_and_files_into_the_store() { + let count = Count::new(); + count.add(42); + let mut set = MetricsSet::new(); + set.push(Arc::new(Metric::new( + MetricValue::OutputRows(count), + Some(0), + ))); + let node_metrics = df_metrics_set_to_proto(&set).expect("encode"); + + let decoded = decode_task_metrics(pb::TaskMetrics { + pre_order_plan_metrics: vec![node_metrics], + task_metrics: Some(pb::MetricsSet::default()), + }) + .expect("decode"); + assert_eq!(decoded.pre_order_plan_metrics.len(), 1); + assert_eq!( + decoded.pre_order_plan_metrics[0].output_rows(), + Some(42), + "the counter survives the wire round trip" + ); + + let store = MetricsStore::new(); + let key = TaskKey { + query_id: uuid::Uuid::new_v4(), + stage_id: 1, + task_number: 0, + }; + store.insert(key, decoded); + assert!( + store.get(&key).is_some(), + "the driver-filed entry is readable" + ); + } +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index dd878f2d..eb7abc57 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -7,11 +7,15 @@ mod channel_resolver; // stack. pub(crate) mod generated; mod in_process; +// The metrics codec sits off gRPC for the same reason as the message types: a transport that +// delivers metrics out-of-band decodes the same frames the gRPC client does. +pub(crate) mod metrics_proto; 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 metrics_proto::decode_task_metrics; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index 1bd55ebf..ec74f52c 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -11,6 +11,7 @@ pub(crate) use remote_work_unit_feed::{ }; pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed}; +pub use remote_work_unit_feed::set_received_time; pub use work_unit::WorkUnit; pub use work_unit_feed::{WorkUnitFeed, WorkUnitFeedProto}; pub use work_unit_feed_provider::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider}; diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 18c51ffe..78ecc42a 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -99,6 +99,13 @@ pub(crate) fn set_work_unit_received_time( msg } +/// Stamps the receive time on a bare proto work unit. A transport that hands units across without +/// going through [set_work_unit_received_time] stamps each one as it crosses into the worker, so the +/// worker-side latency math has a delivery timestamp; a missing stamp reads as zero latency. +pub fn set_received_time(work_unit: &mut crate::proto::WorkUnit) { + work_unit.received_timestamp_unix_nanos = now_ns(); +} + /// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnitMsg]s coming over /// the wire from a [RemoteWorkUnitFeedRegistry]. /// diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index cf16e890..9cb5d7bc 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,5 +1,6 @@ use crate::common::TreeNodeExt; use crate::execution_plans::SamplerExec; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; use crate::worker::task_data::TaskDataMetrics; @@ -231,3 +232,24 @@ fn send_metrics_via_channel( task_metrics: task_data_metrics.to_metrics_set(), }); } + +/// Collects the per-node metrics of `plan` in pre-order traversal order, encoded as protos. A +/// push transport reuses this for its metrics frame so the per-node shape matches what the +/// coordinator channel reports. A node without metrics, or one whose metrics fail to encode, +/// contributes an empty set in place: the consumer matches entries to plan nodes by pre-order +/// position, so dropping an entry would mis-attribute every set after it. +pub fn collect_plan_metrics_protos( + plan: &Arc, + dt_ctx: DistributedTaskContext, +) -> Vec { + let mut pre_order_plan_metrics = vec![]; + let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { + pre_order_plan_metrics.push( + node.metrics() + .and_then(|m| df_metrics_set_to_proto(&m).ok()) + .unwrap_or_default(), + ); + Ok(TreeNodeRecursion::Continue) + }); + pre_order_plan_metrics +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index a0ab8959..759dee43 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -10,6 +10,7 @@ pub(crate) mod test_utils; mod worker_connection_pool; mod worker_service; +pub use impl_coordinator_channel::collect_plan_metrics_protos; pub(crate) use single_write_multi_read::SingleWriteMultiRead; pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool}; From dff98810cfc8f003eccea2e2d722bc3aef8ebb96 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 1 Jul 2026 11:34:42 -0700 Subject: [PATCH 05/11] Added a DispatchPlanSource hook for coordinator-sourced plans. An embedder whose plan nodes the coordinator's codec cannot represent, or whose serialization needs embedder-side handling the codec extension point cannot express, serializes the dispatch bytes itself. The coordinator hands it the TaskKey and the ready-to-run per-task plan it would otherwise encode; returning None falls back to the coordinator's own encode. The tests pin the contract from all three sides: consultation with nested stages already Remote, source-provided bytes running the query, and a source error failing the dispatch. Co-authored-by: Stu Hood --- src/coordinator/query_coordinator.rs | 21 ++- src/dispatch_plan_source.rs | 241 +++++++++++++++++++++++++++ src/distributed_ext.rs | 38 ++++- src/lib.rs | 2 + 4 files changed, 294 insertions(+), 8 deletions(-) create mode 100644 src/dispatch_plan_source.rs diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 887a109f..9309e541 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -11,7 +11,8 @@ use crate::{ DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_dispatch_plan_source, + get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -140,19 +141,27 @@ impl<'a> StageCoordinator<'a> { UnboundedReceiver, )> { let session_config = self.task_ctx.session_config(); - let codec = DistributedCodec::new_combined_with_user(session_config); let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?; - let plan_proto = - PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec(); - let plan_size = plan_proto.len(); - + // An embedder can serialize the dispatch bytes for this stage itself (e.g. with a codec + // the config's extension point cannot express) instead of the coordinator encoding the + // plan. Either way the bytes describe `specialized`, the ready-to-run per-task plan. let task_key = TaskKey { query_id: self.query_id, stage_id: self.stage_id, task_number: task_i, }; + let plan_proto = match get_distributed_dispatch_plan_source(session_config) + .and_then(|source| source.dispatch_plan_proto(&task_key, &specialized)) + { + Some(bytes) => bytes?, + None => { + let codec = DistributedCodec::new_combined_with_user(session_config); + PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec() + } + }; + let plan_size = plan_proto.len(); let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest { task_key, diff --git a/src/dispatch_plan_source.rs b/src/dispatch_plan_source.rs new file mode 100644 index 00000000..d7f84994 --- /dev/null +++ b/src/dispatch_plan_source.rs @@ -0,0 +1,241 @@ +use crate::TaskKey; +use datafusion::common::Result; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// Serializes the stage subplan the coordinator dispatches, instead of the coordinator encoding +/// it with its own codec. +/// +/// An embedder registers one via +/// [`crate::DistributedExt::with_distributed_dispatch_plan_source`] when the coordinator's codec +/// cannot represent its plan nodes, or when its serialization needs embedder-side handling the +/// codec extension point cannot express (the shm embedder's UDF definitions, for example). The +/// coordinator hands over `specialized`, the same ready-to-run per-task plan it would encode: +/// task-specialized, with nested stages already converted to `Remote`, so a worker executes the +/// decoded bytes as-is. `task` carries the query id, so a source registered on a session that +/// runs concurrent queries can tell them apart. +/// +/// Returning `None` for a task lets the coordinator fall back to encoding the plan itself, so a +/// source that only overrides some stages stays correct. +pub trait DispatchPlanSource: Send + Sync { + fn dispatch_plan_proto( + &self, + task: &TaskKey, + specialized: &Arc, + ) -> Option>>; +} + +#[derive(Clone)] +pub(crate) struct DispatchPlanSourceExtension(pub(crate) Arc); + +pub(crate) fn set_distributed_dispatch_plan_source( + cfg: &mut SessionConfig, + source: impl DispatchPlanSource + 'static, +) { + set_distributed_dispatch_plan_source_arc(cfg, Arc::new(source)) +} + +pub(crate) fn set_distributed_dispatch_plan_source_arc( + cfg: &mut SessionConfig, + source: Arc, +) { + cfg.set_extension(Arc::new(DispatchPlanSourceExtension(source))); +} + +/// Returns the [`DispatchPlanSource`] registered on this config, if any. +pub fn get_distributed_dispatch_plan_source( + cfg: &SessionConfig, +) -> Option> { + cfg.get_extension::() + .map(|ext| Arc::clone(&ext.0)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::InProcessChannelResolver; + use crate::{ + DistributedExt, NetworkBoundaryExt, SessionStateBuilderExt, Stage, WorkerResolver, + }; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::common::{DataFusionError, Result}; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::{ExecutionPlan, collect}; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use datafusion_proto::physical_plan::AsExecutionPlan; + use datafusion_proto::protobuf::PhysicalPlanNode; + use prost::Message; + use std::io::Write; + use std::sync::Mutex; + use url::Url; + + struct Workers(usize); + + impl WorkerResolver for Workers { + fn get_urls(&self) -> Result> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + type Calls = Arc)>>>; + + /// Records what the coordinator hands over and declines, so the coordinator's own encode + /// still runs and the query is unaffected by the recording. + #[derive(Default)] + struct RecordingSource { + calls: Calls, + } + + impl DispatchPlanSource for RecordingSource { + fn dispatch_plan_proto( + &self, + task: &TaskKey, + specialized: &Arc, + ) -> Option>> { + self.calls + .lock() + .unwrap() + .push((*task, Arc::clone(specialized))); + None + } + } + + const QUERY: &str = "SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k"; + + async fn distributed_ctx( + name: &str, + source: impl DispatchPlanSource + 'static, + ) -> Result<(SessionContext, std::path::PathBuf)> { + let path = std::env::temp_dir().join(format!("dfd_{name}_{}.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 state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(4)) + .with_distributed_planner() + .with_distributed_worker_resolver(Workers(4)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_dispatch_plan_source(source) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx = SessionContext::from(state); + ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + Ok((ctx, path)) + } + + /// The contract a serializing source relies on: it is consulted once per dispatched task, + /// and the plan it is handed is ready to run, with nested stages already converted to + /// `Remote`. + #[tokio::test] + async fn source_is_consulted_with_the_ready_to_run_plan() -> Result<()> { + let source = RecordingSource::default(); + let calls = Arc::clone(&source.calls); + let (ctx, path) = distributed_ctx("recording", source).await?; + + let physical = ctx.sql(QUERY).await?.create_physical_plan().await?; + collect(physical, ctx.task_ctx()).await?; + + let calls = calls.lock().unwrap(); + assert!( + !calls.is_empty(), + "the coordinator never consulted the source" + ); + let mut seen: datafusion::common::HashSet = Default::default(); + for (task, specialized) in calls.iter() { + assert!(seen.insert(*task), "consulted twice for {task:?}"); + specialized.apply(|node| { + if let Some(nb) = node.as_ref().as_network_boundary() { + assert!( + matches!(nb.input_stage(), Stage::Remote(_)), + "{task:?} carries a Local nested stage; the handed-over plan must be \ + ready to run" + ); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + std::fs::remove_file(&path).ok(); + Ok(()) + } + + /// Serializes with the same codec the coordinator's fallback uses, so the worker decodes + /// source-provided bytes exactly as it decodes coordinator-encoded ones. + struct SerializingSource; + + impl DispatchPlanSource for SerializingSource { + fn dispatch_plan_proto( + &self, + _task: &TaskKey, + specialized: &Arc, + ) -> Option>> { + let codec = crate::DistributedCodec::new_combined_with_user(&SessionConfig::new()); + Some( + PhysicalPlanNode::try_from_physical_plan(Arc::clone(specialized), &codec) + .map(|node| node.encode_to_vec()), + ) + } + } + + /// The bytes the source returns are what the workers run. + #[tokio::test] + async fn source_provided_bytes_run_the_query() -> Result<()> { + let (ctx, path) = distributed_ctx("serializing", SerializingSource).await?; + let got = ctx.sql(QUERY).await?.collect().await?; + let got = datafusion::arrow::util::pretty::pretty_format_batches(&got)?.to_string(); + + let serial = SessionContext::new(); + serial + .register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + let expected = serial.sql(QUERY).await?.collect().await?; + let expected = + datafusion::arrow::util::pretty::pretty_format_batches(&expected)?.to_string(); + + assert_eq!(got, expected, "source-encoded dispatch != serial"); + std::fs::remove_file(&path).ok(); + Ok(()) + } + + struct FailingSource; + + impl DispatchPlanSource for FailingSource { + fn dispatch_plan_proto( + &self, + _task: &TaskKey, + _specialized: &Arc, + ) -> Option>> { + Some(Err(DataFusionError::Internal( + "the embedder could not serialize this stage".into(), + ))) + } + } + + /// A source failure fails the dispatch instead of falling back to bytes the embedder said it + /// could not produce. + #[tokio::test] + async fn source_error_fails_the_query() -> Result<()> { + let (ctx, path) = distributed_ctx("failing", FailingSource).await?; + let result = ctx.sql(QUERY).await?.collect().await; + let err = result + .expect_err("a failing source must fail the query") + .to_string(); + assert!( + err.contains("could not serialize"), + "unexpected error: {err}" + ); + std::fs::remove_file(&path).ok(); + Ok(()) + } +} diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 9efdb420..6f71f7a4 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -2,14 +2,15 @@ use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; +use crate::dispatch_plan_source::set_distributed_dispatch_plan_source; use crate::distributed_planner::set_distributed_task_estimator; use crate::passthrough_headers::set_passthrough_headers; use crate::protocol::set_distributed_channel_resolver; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ - ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, + ChannelResolver, DispatchPlanSource, DistributedConfig, TaskEstimator, WorkUnitFeed, + WorkUnitFeedProvider, WorkerResolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -278,6 +279,16 @@ pub trait DistributedExt: Sized { resolver: T, ); + /// Registers a [DispatchPlanSource] the coordinator consults for each stage's dispatch bytes + /// instead of encoding the plan it holds. See [DispatchPlanSource] for when this is needed. + fn with_distributed_dispatch_plan_source( + self, + source: T, + ) -> Self; + + /// Same as [DistributedExt::with_distributed_dispatch_plan_source] but with an in-place mutation. + fn set_distributed_dispatch_plan_source(&mut self, source: T); + /// Adds a distributed task count estimator. [TaskEstimator]s are executed on each node /// sequentially until one returns an estimation on the number of tasks that should be /// used for the stage containing that node. @@ -634,6 +645,10 @@ impl DistributedExt for SessionConfig { set_distributed_channel_resolver(self, resolver); } + fn set_distributed_dispatch_plan_source(&mut self, source: T) { + set_distributed_dispatch_plan_source(self, source); + } + fn set_distributed_task_estimator( &mut self, estimator: T, @@ -787,6 +802,10 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + #[call(set_distributed_task_estimator)] #[expr($;self)] fn with_distributed_task_estimator(mut self, estimator: T) -> Self; @@ -889,6 +908,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1012,6 +1036,11 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1135,6 +1164,11 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] diff --git a/src/lib.rs b/src/lib.rs index 559f6b78..c5c6c8e4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ mod codec; mod common; mod config_extension_ext; mod coordinator; +mod dispatch_plan_source; mod distributed_ext; mod distributed_planner; mod execution_plans; @@ -47,6 +48,7 @@ pub use protocol::grpc; pub use protocol::generated::worker as proto; pub use codec::DistributedCodec; +pub use dispatch_plan_source::{DispatchPlanSource, get_distributed_dispatch_plan_source}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ From 22f96c9e3da16e440f1e764640fb76dc205b3d47 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 9 Jul 2026 16:02:34 -0700 Subject: [PATCH 06/11] Exposed produce-side partition routing off the network boundaries. A pull-based transport never places a produced partition: the consumer computes its own slice inside the boundary's execute. A push-based transport places every partition before any consumer asks, so it reads the boundary's consumer layout through route_partition instead of re-deriving it from node properties and drifting when the layout changes. NetworkCoalesceExec overrides it with an error: its consumers read whole per-producer-task groups, not slices. Co-authored-by: Stu Hood --- src/distributed_planner/mod.rs | 2 +- src/distributed_planner/network_boundary.rs | 130 +++++++++++++++++++- src/execution_plans/network_coalesce.rs | 10 ++ src/lib.rs | 2 +- 4 files changed, 141 insertions(+), 3 deletions(-) diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 18cfe6e0..176bd319 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -15,7 +15,7 @@ pub(crate) use inject_network_boundaries::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, inject_network_boundaries, }; pub(crate) use network_boundary::ProducerHead; -pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt}; +pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt, PartitionRoute}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use statistics::calculate_cost; pub(crate) use task_estimator::set_distributed_task_estimator; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 3b7bd3b1..05a00ad6 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -5,7 +5,7 @@ use crate::{ Stage, }; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; +use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; @@ -19,6 +19,14 @@ use datafusion_proto::protobuf::proto_error; use prost::Message; use std::sync::Arc; +/// Where a producer's output partition should be sent: which consumer task, and the local partition +/// index within that task's slice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PartitionRoute { + pub consumer_task: usize, + pub consumer_partition: usize, +} + /// This trait represents a node that introduces the necessity of a network boundary in the plan. /// The distributed planner, upon stepping into one of these, will break the plan and build a stage /// out of it. @@ -36,6 +44,31 @@ pub trait NetworkBoundary: ExecutionPlan { /// implementation have. This information is used during planning an executing for ensuring /// the head of a stage has the appropriate shape for consumption. fn producer_head(&self, consumer_tasks: usize) -> ProducerHead; + + /// Maps a producer output partition to the consumer task and the local partition within that + /// task that reads it, for the sliced layout shuffle and broadcast reads use + /// (`global = P_c * consumer_task + local`, where `P_c` is the boundary's own per-task output + /// partition count). A pull-based transport never needs this: its consumers compute their own + /// slice inside the boundary's `execute`. A push-based transport places every produced + /// partition before any consumer asks, so it reads the layout here instead of re-deriving it + /// and drifting when the layout changes. + /// + /// Boundaries whose consumers do not read that layout must override this with an error; the + /// default would silently misroute them. A zero-partition boundary is a planner bug, so it + /// errors instead of routing everything to task `0`. + fn route_partition(&self, output_partition: usize) -> Result { + let p_c = self.properties().partitioning.partition_count(); + if p_c == 0 { + return internal_err!( + "cannot route output partition {output_partition}: the boundary reports 0 \ + partitions per consumer task" + ); + } + Ok(PartitionRoute { + consumer_task: output_partition / p_c, + consumer_partition: output_partition % p_c, + }) + } } /// Defines what shape should the head node of a stage have upon getting executed. Depending @@ -155,3 +188,98 @@ impl ProducerHead { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::InProcessChannelResolver; + use crate::{DistributedExt, NetworkBoundaryExt, SessionStateBuilderExt, WorkerResolver}; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::error::DataFusionError; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use std::io::Write; + use url::Url; + + struct Workers(usize); + + impl WorkerResolver for Workers { + fn get_urls(&self) -> Result> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + /// Pins the sliced routing (`global = P_c * consumer_task + local`) on boundaries the + /// planner actually built, with `P_c` read off the boundary's own properties, so an override + /// or a change in what `properties()` reports fails here first. The data-level guarantee + /// that the slicing matches what consumers read comes from the in-process transport's + /// end-to-end test. `NetworkCoalesceExec` must refuse to route: its consumers read whole + /// per-producer-task groups, and the sliced formula would misroute them. + #[tokio::test] + async fn route_partition_matches_the_consumer_slicing() -> Result<()> { + let path = std::env::temp_dir().join(format!("dfd_routing_{}.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 state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(4)) + .with_distributed_planner() + .with_distributed_worker_resolver(Workers(4)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx = SessionContext::from(state); + ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + let physical = ctx + .sql("SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k") + .await? + .create_physical_plan() + .await?; + + let mut sliced = 0usize; + let mut grouped = 0usize; + physical.apply(|node| { + let Some(nb) = node.as_ref().as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + if node + .as_ref() + .downcast_ref::() + .is_some() + { + assert!( + nb.route_partition(0).is_err(), + "a per-task-group boundary must refuse the sliced routing" + ); + grouped += 1; + return Ok(TreeNodeRecursion::Continue); + } + let p_c = nb.properties().partitioning.partition_count(); + assert!(p_c > 0); + for consumer_task in 0..3 { + for local in 0..p_c { + let route = nb.route_partition(p_c * consumer_task + local)?; + assert_eq!(route.consumer_task, consumer_task); + assert_eq!(route.consumer_partition, local); + } + } + sliced += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert!(sliced > 0, "the plan grew no sliced-layout boundary"); + assert!(grouped > 0, "the plan grew no per-task-group boundary"); + std::fs::remove_file(&path).ok(); + Ok(()) + } +} diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 703843cc..4379910a 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -182,6 +182,16 @@ impl NetworkBoundary for NetworkCoalesceExec { fn producer_head(&self, _consumer_task_count: usize) -> ProducerHead { ProducerHead::None } + + fn route_partition( + &self, + output_partition: usize, + ) -> Result { + internal_err!( + "NetworkCoalesceExec routes by producer task group, not by output partition; \ + partition {output_partition} has no slice-layout route" + ) + } } impl DisplayAs for NetworkCoalesceExec { diff --git a/src/lib.rs b/src/lib.rs index c5c6c8e4..3f4f1a80 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,7 +23,7 @@ pub use arrow_ipc::CompressionType; pub use coordinator::{DistributedExec, MetricsStore}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ - DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, + DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, SessionStateBuilderExt, TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, }; pub use execution_plans::{ From a3416a7bf9e8e0ccdf5929ffab4a3101bbc4521f Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:14:43 -0700 Subject: [PATCH 07/11] Ported the shared-memory transport onto the channel protocol. The shm ring/DSM/mesh/framing core moves verbatim from the fork; the old WorkerTransport umbrella is gone, so the resolver hands out a channel and execute_task reads the rings per partition. The mesh is the no-gRPC transport, so it builds in both the grpc-on and grpc-off configs. Workers are passive executors: a fragment arrives with its nested stages already Remote, so the worker runs it as-is and its boundary leaves read the mesh; nothing on the worker converts or dispatches. collect_task_metrics reports an empty task-level metric set instead of a missing one, since one failed decode starves the whole store. --- Cargo.lock | 24 + Cargo.toml | 2 + src/lib.rs | 7 + src/shm/dsm.rs | 544 +++++++ src/shm/in_process.rs | 1012 +++++++++++++ src/shm/mesh.rs | 375 +++++ src/shm/mod.rs | 81 + src/shm/mpsc_ring.rs | 1550 ++++++++++++++++++++ src/shm/runtime.rs | 477 ++++++ src/shm/self_hosted.rs | 1277 ++++++++++++++++ src/shm/setup.rs | 373 +++++ src/shm/sink.rs | 49 + src/shm/transport.rs | 2935 +++++++++++++++++++++++++++++++++++++ src/work_unit_feed/mod.rs | 1 + 14 files changed, 8707 insertions(+) create mode 100644 src/shm/dsm.rs create mode 100644 src/shm/in_process.rs create mode 100644 src/shm/mesh.rs create mode 100644 src/shm/mod.rs create mode 100644 src/shm/mpsc_ring.rs create mode 100644 src/shm/runtime.rs create mode 100644 src/shm/self_hosted.rs create mode 100644 src/shm/setup.rs create mode 100644 src/shm/sink.rs create mode 100644 src/shm/transport.rs diff --git a/Cargo.lock b/Cargo.lock index 3858be13..42b58fc2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -475,6 +475,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -2192,6 +2214,7 @@ dependencies = [ "arrow-flight", "arrow-ipc", "arrow-select", + "async-stream", "async-trait", "bincode", "bytes", @@ -2206,6 +2229,7 @@ dependencies = [ "hyper-util", "insta", "itertools 0.14.0", + "log", "moka", "num-traits", "object_store", diff --git a/Cargo.toml b/Cargo.toml index 1bb4874d..5f23f5b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,8 @@ datafusion = { workspace = true, features = [ ] } datafusion-proto = { workspace = true } async-trait = "0.1.89" +async-stream = "0.3" +log = "0.4" tokio = { version = "1.48", features = ["full"] } http = "1.3.1" itertools = "0.14.0" diff --git a/src/lib.rs b/src/lib.rs index 3f4f1a80..3fe0329d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,9 @@ mod execution_plans; mod metrics; mod passthrough_headers; mod protocol; +// Not feature-gated: the shared-memory mesh is the no-gRPC transport, so it has to build in both +// the `grpc`-on and `grpc`-off configs. +pub mod shm; mod stage; mod work_unit_feed; mod worker; @@ -49,6 +52,10 @@ pub use protocol::generated::worker as proto; pub use codec::DistributedCodec; pub use dispatch_plan_source::{DispatchPlanSource, get_distributed_dispatch_plan_source}; +// The producer-side sink traits live in `shm` because only a push-based transport produces through +// them; re-exported at the crate root so `crate::PartitionSink` resolves the way the shm core spells +// it. +pub use shm::{PartitionSink, WorkerSink}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ diff --git a/src/shm/dsm.rs b/src/shm/dsm.rs new file mode 100644 index 00000000..3720da0e --- /dev/null +++ b/src/shm/dsm.rs @@ -0,0 +1,544 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mesh-multiplexed DSM layout: one MPSC inbox per receiver process. +//! +//! Each MPP query allocates a single DSM region: +//! +//! ```text +//! +--- MppDsmHeader (repr C, MAXALIGN-padded) -------------+ +//! | magic, version, n_procs, queue_bytes | +//! | plan_offset, plan_len, queues_offset, region_total | +//! +-------------------------------------------------------+ +//! | plan bytes (bincode-serialized worker fragment) | +//! +-------------------------------------------------------+ +//! | padding to MAXALIGN | +//! +-------------------------------------------------------+ +//! | DsmMpscRing inbox array: n_procs inboxes | +//! | inbox(receiver) = queues_offset | +//! | + receiver * queue_bytes | +//! +-------------------------------------------------------+ +//! ``` +//! +//! - `n_procs` = 1 leader + N parallel workers. Leader is `proc_idx = 0`; workers are +//! `proc_idx = ParallelWorkerNumber + 1`. +//! - Each process attaches as receiver to its own inbox (one MPSC ring) and as sender +//! to each of N-1 peer inboxes. Senders stamp `MppFrameHeader::sender_proc` on every +//! frame so the receiver demuxes by source. +//! - Self-loop frames (proc → itself) ride an in-proc channel installed by +//! `worker_setup`, not a DSM slot: each `DsmMpscReceiver` is owned by a single +//! receiver process, so there is no DSM slot for the self-loop pair. + +use std::ffi::c_void; +use std::mem::size_of; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +use super::AliveFlag; +use super::mesh::{align_up_maxalign_checked, aligned_queue_bytes}; +use super::mpsc_ring::{self, DsmMpscReceiver, DsmMpscRingHeader, DsmMpscSender, Wakeup}; + +/// Number of slots in each per-receiver MPSC ring. With operator-visible +/// The embedder's queue-size knob divided across `RING_SLOTS` slots, each slot holds +/// up to `(queue_bytes / RING_SLOTS) - SLOT_HEADER` bytes of payload. That's enough +/// for typical Arrow batches at the bench scales we measured. Fixed at compile time +/// for now; a future GUC could expose it if bench data points at a different sweet +/// spot. +const RING_SLOTS: u32 = 8; + +/// Floor `ring_dims_for` applies to the per-slot payload capacity. The layout minimum check +/// uses the same floor so the constructed ring always fits its inbox region. +const MIN_SLOT_CAPACITY: u32 = 64; + +/// Cache-line alignment for the per-inbox DsmMpscRing header. The ring itself is plain +/// `#[repr(C)]` (the base address is only MAXALIGN-guaranteed), but spacing the per-inbox +/// offsets to 64 keeps hot headers on separate cache lines when the base cooperates; +/// both `queues_offset` and per-inbox `queue_bytes` are aligned up to this so the +/// computed `inbox_offset(r) = queues_offset + r * queue_bytes` lands at a 64-aligned +/// address for every `r`. +const RING_ALIGN: usize = 64; + +#[inline] +fn align_up_ring(n: usize) -> Option { + let mask = RING_ALIGN - 1; + n.checked_add(mask).map(|x| x & !mask) +} + +const MPP_DSM_MAGIC: u32 = 0x4D50_5052; // "MPPR" (RPC variant) +/// Bump on any wire-incompatible change to the DSM header layout or to the inbox-offset +/// math, so attaching workers reject mismatched leaders loudly rather than reading +/// garbage. Validated in [`MppDsmHeader::validate`]. +const MPP_DSM_HEADER_VERSION: u32 = 4; + +/// Absolute cap on DSM region size. 16 GiB is two orders of magnitude beyond +/// any realistic workload; the cap fails early on a pathologically oversized +/// request rather than asking PG for ~`usize::MAX` bytes. +const MPP_DSM_MAX_BYTES: usize = 16 * 1024 * 1024 * 1024; + +/// C-repr header at offset 0 of the DSM region. +/// +/// Layout: three `u32`s + padding, then five `u64`s. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct MppDsmHeader { + pub(super) magic: u32, + pub(super) header_version: u32, + /// Total proc count. Leader is `proc_idx = 0`; workers are + /// `proc_idx = ParallelWorkerNumber + 1`. The DSM region holds `n_procs` + /// per-receiver MPSC inboxes laid out contiguously after the plan bytes. + pub n_procs: u32, + pub(super) _pad: u32, + pub(super) queue_bytes: u64, + pub(super) plan_offset: u64, + pub(super) plan_len: u64, + pub(super) queues_offset: u64, + pub region_total: u64, +} + +impl MppDsmHeader { + fn from_layout(layout: &DsmLayout) -> Self { + Self { + magic: MPP_DSM_MAGIC, + header_version: MPP_DSM_HEADER_VERSION, + n_procs: layout.n_procs, + _pad: 0, + queue_bytes: layout.queue_bytes as u64, + plan_offset: layout.plan_offset as u64, + plan_len: layout.plan_len as u64, + queues_offset: layout.queues_offset as u64, + region_total: layout.region_total as u64, + } + } + + pub(super) fn validate(&self, region_total: u64) -> Result<(), &'static str> { + if self.magic != MPP_DSM_MAGIC { + return Err("mpp: DSM header magic mismatch"); + } + if self.header_version != MPP_DSM_HEADER_VERSION { + return Err("mpp: DSM header version mismatch"); + } + if self.n_procs == 0 { + return Err("mpp: header n_procs must be > 0"); + } + if self.region_total != region_total { + return Err("mpp: DSM region_total in header disagrees with attached size"); + } + match self.plan_offset.checked_add(self.plan_len) { + None => return Err("mpp: plan_offset + plan_len overflow"), + Some(end) if end > self.queues_offset => { + return Err("mpp: plan would overlap queues area"); + } + _ => {} + } + if self.queues_offset > region_total { + return Err("mpp: queues_offset past end of region"); + } + Ok(()) + } + + /// Byte offset (relative to DSM base) of `receiver_proc`'s MPSC inbox. + /// + /// The owner attaches as receiver (`DsmMpscReceiver`); every other process attaches + /// as sender (`DsmMpscSender`) to the same inbox and stamps its identity into the + /// frame header (`MppFrameHeader::sender_proc`). + pub(super) fn inbox_offset(&self, receiver_proc: u32) -> u64 { + debug_assert!(receiver_proc < self.n_procs); + self.queues_offset + (receiver_proc as u64) * self.queue_bytes + } +} + +/// Pure-math layout for [`compute_dsm_layout`]. +#[derive(Debug, Clone, Copy)] +pub(super) struct DsmLayout { + pub n_procs: u32, + pub queue_bytes: usize, + pub plan_offset: usize, + pub plan_len: usize, + pub queues_offset: usize, + pub region_total: usize, +} + +/// Compute the DSM region size and field offsets for one MPP query. +/// +/// `n_procs` is the total proc count (1 leader + N workers); the region holds one +/// `DsmMpscRing` inbox per process. Every other process attaches as sender to that +/// inbox and the receiver demultiplexes by `MppFrameHeader::sender_proc`. +/// +/// `queue_bytes` is the per-inbox total (ring header + `RING_SLOTS` slots). +/// The embedder's operator-facing queue-size knob controls this value. +pub(super) fn compute_dsm_layout( + n_procs: u32, + queue_bytes: usize, + plan_len: usize, +) -> Result { + if n_procs < 2 { + return Err("mpp: n_procs must be >= 2 (leader + at least one worker)"); + } + // MAXALIGN-round-down first (operator-friendly), then round up to the ring's + // 64-byte alignment requirement. Doing it in this order means each per-inbox + // region is both MAXALIGN-aligned (PG DSM convention) AND cache-line aligned, + // which keeps the rings' hot header fields on separate cache lines. + let queue_bytes = aligned_queue_bytes(queue_bytes); + if queue_bytes == 0 { + return Err("mpp: queue_bytes too small after alignment"); + } + let queue_bytes = align_up_ring(queue_bytes).ok_or("mpp: queue_bytes alignment overflow")?; + // Each inbox must have room for the ring `ring_dims_for` will actually build: + // `RING_SLOTS` slots with the 64-byte capacity floor. Checking against a smaller + // ring here would let `create_at` write past the inbox region for tiny + // `queue_bytes`, overlapping the next inbox. + if queue_bytes < DsmMpscRingHeader::region_bytes(RING_SLOTS, MIN_SLOT_CAPACITY) { + return Err("mpp: queue_bytes too small for ring header + min slot capacity"); + } + let header_end = align_up_maxalign_checked(size_of::()) + .ok_or("mpp: header alignment overflow")?; + let plan_offset = header_end; + let plan_end = plan_offset + .checked_add(plan_len) + .ok_or("mpp: plan offset+len overflow")?; + // Round queues_offset up to RING_ALIGN so the first inbox starts at a 64-aligned + // address; subsequent inboxes are queue_bytes apart and queue_bytes is RING_ALIGN- + // aligned, so they all land on cache-line boundaries. + let queues_offset = align_up_ring(plan_end).ok_or("mpp: queues alignment overflow")?; + let queues_bytes = (n_procs as usize) + .checked_mul(queue_bytes) + .ok_or("mpp: queues bytes overflow")?; + let region_total = queues_offset + .checked_add(queues_bytes) + .ok_or("mpp: region total overflow")?; + if region_total > MPP_DSM_MAX_BYTES { + return Err("mpp: DSM region exceeds MPP_DSM_MAX_BYTES"); + } + Ok(DsmLayout { + n_procs, + queue_bytes, + plan_offset, + plan_len, + queues_offset, + region_total, + }) +} + +/// Derive `(ring_slots, slot_capacity)` for a per-inbox region of `queue_bytes`. +/// Total ring region size = `DsmMpscRingHeader::region_bytes(ring_slots, slot_capacity)` +/// which fits within `queue_bytes` by construction. +fn ring_dims_for(queue_bytes: usize) -> (u32, u32) { + let header = std::mem::size_of::(); + // queue_bytes >= ring_dims minimum is enforced in compute_dsm_layout; we recompute + // here without re-validating. + let slot_total_bytes = queue_bytes.saturating_sub(header); + let slot_capacity = (slot_total_bytes / RING_SLOTS as usize).max(MIN_SLOT_CAPACITY as usize); + // Cap slot_capacity at u32::MAX (DsmMpscRing's field type) to avoid casting wrap. + let slot_capacity = slot_capacity.min(u32::MAX as usize) as u32; + (RING_SLOTS, slot_capacity) +} + +/// Per-proc return from `attach_proc`: N-1 outbound senders (one per peer inbox) plus a +/// single inbound receiver (this proc's own inbox). +/// +/// The own-inbox is the multiplexed entry point: every peer attaches to it as a sender +/// (each `DsmMpscSender` increments the ring's `sender_count`) and stamps its identity +/// into `MppFrameHeader::sender_proc` on every frame. The receiver side pulls frames +/// off that single ring and routes them to per-`(sender_proc, stage_id, partition)` +/// channel buffers via [`DrainHandle`]. +pub(super) struct ProcAttach { + /// `outbound_senders[i]` writes to peer `peer_proc_for_index(this_proc, i)`'s inbox. + /// `peer_proc(i) = i if i < this_proc else i + 1` skips the self-loop entry. + pub(super) outbound_senders: Vec, + /// This process's own MPSC inbox receiver. Drained inline by `DrainHandle`. + pub(super) inbound_receiver: DsmMpscReceiver, + pub(super) alive: AliveFlag, +} + +/// Read `region_total` out of the header at the start of an initialized region. +/// +/// # Safety +/// `base` must point at the start of a region a leader wrote via [`leader_init`]. +pub(super) unsafe fn read_region_total(base: *const c_void) -> u64 { + unsafe { std::ptr::read(base as *const MppDsmHeader).region_total } +} + +/// Translate a peer index (`0..n_procs - 1`) into a process index +/// (`0..n_procs`) by skipping the self-loop slot. +#[inline] +pub(super) fn peer_proc_for_index(this_proc: u32, peer_idx: u32) -> u32 { + if peer_idx < this_proc { + peer_idx + } else { + peer_idx + 1 + } +} + +/// Initialize the DSM region as the leader (`proc_idx = 0`). Writes the MppDsmHeader, +/// copies the plan bytes, initializes the N MPSC inboxes via `DsmMpscRing::create_at`, +/// and attaches the leader as receiver to its own inbox, plus (when `attach_senders`) as +/// sender to each peer for the control plane (work-unit frames). +/// +/// `attach_senders` is a commitment: a ring latches `detached` when its sender count falls to +/// zero, so a leader that attaches and then drops its senders before a worker attached poisons +/// that worker's inbox. Attach only when the senders outlive the query. +/// +/// # Safety +/// - `coordinate` must point to the start of a DSM region of size `>= layout.region_total`. +/// - The region must be uninitialized (the leader is the first writer). +pub(super) unsafe fn leader_init( + coordinate: *mut c_void, + layout: &DsmLayout, + plan_bytes: &[u8], + wakeup: Arc, + attach_senders: bool, +) -> Result { + if coordinate.is_null() { + return Err("mpp: leader_init given null coordinate".into()); + } + if plan_bytes.len() != layout.plan_len { + return Err(format!( + "mpp: plan_bytes.len()={} != layout.plan_len={}", + plan_bytes.len(), + layout.plan_len + )); + } + + let base = coordinate as *mut u8; + + // Header. + unsafe { + std::ptr::write( + base.cast::(), + MppDsmHeader::from_layout(layout), + ); + } + // Plan bytes. + unsafe { + std::ptr::copy_nonoverlapping( + plan_bytes.as_ptr(), + base.add(layout.plan_offset), + plan_bytes.len(), + ); + } + + // Initialize the N MPSC inboxes. Workers can't do this (the region is uninitialized + // at their attach time), so the leader runs DsmMpscRing::create_at for every receiver. + let header = MppDsmHeader::from_layout(layout); + let n_procs = layout.n_procs; + let (ring_slots, slot_capacity) = ring_dims_for(layout.queue_bytes); + for r in 0..n_procs { + let off = header.inbox_offset(r) as usize; + let inbox_addr = unsafe { base.add(off) }; + unsafe { mpsc_ring::create_at(inbox_addr, ring_slots, slot_capacity) }; + } + + let attach = unsafe { attach_proc(base, &header, 0, attach_senders, wakeup) }; + Ok(attach) +} + +/// Attach to the leader-initialized DSM region as `proc_idx` (`0 = leader`, `1..N` = +/// parallel workers). Attach as receiver to this proc's own MPSC inbox, and (if +/// `attach_senders` is true) as sender to every peer's inbox. +/// +/// `attach_senders = false` is for the leader (consumer-only). If the leader attached +/// as sender it would bump every peer inbox's `sender_count` and decrement it on `Drop`; +/// if it dropped before any worker bumped, the 1 → 0 transition would flip `detached` +/// on every peer inbox and every later worker send would fail `SendError::Detached`. +/// Skipping keeps `sender_count` honest (only producers ever increment). +/// +/// # Safety +/// - `base` must point to a DSM region whose header has been validated. +/// - `header.inbox_offset(r)` must point at a ring already initialized by +/// `DsmMpscRing::create_at` (the leader does this in `leader_init`). +unsafe fn attach_proc( + base: *mut u8, + header: &MppDsmHeader, + this_proc: u32, + attach_senders: bool, + wakeup: Arc, +) -> ProcAttach { + let n_procs = header.n_procs; + let peer_count = (n_procs - 1) as usize; + let mut outbound_senders = Vec::with_capacity(if attach_senders { peer_count } else { 0 }); + + let (ring_slots, slot_capacity) = ring_dims_for(header.queue_bytes as usize); + + // One liveness flag shared by every handle this attach mints. + let alive = Arc::new(AtomicBool::new(true)); + + if attach_senders { + for peer_idx in 0..(n_procs - 1) { + let r = peer_proc_for_index(this_proc, peer_idx); + let off = header.inbox_offset(r) as usize; + let inbox_addr = unsafe { base.add(off) }; + let nn = unsafe { mpsc_ring::attach_at(inbox_addr, ring_slots, slot_capacity) } + .expect("DsmMpscRing attach_at: leader-initialized region must validate"); + outbound_senders + .push(unsafe { DsmMpscSender::new(nn, Arc::clone(&wakeup), Arc::clone(&alive)) }); + } + } + + // Inbound: this proc's own inbox. Single receiver per ring (MPSC contract). + let own_off = header.inbox_offset(this_proc) as usize; + let own_inbox_addr = unsafe { base.add(own_off) }; + let own_nn = unsafe { mpsc_ring::attach_at(own_inbox_addr, ring_slots, slot_capacity) } + .expect("DsmMpscRing attach_at: own inbox must validate"); + let inbound_receiver = unsafe { DsmMpscReceiver::new(own_nn, Arc::clone(&alive)) }; + + ProcAttach { + outbound_senders, + inbound_receiver, + alive, + } +} + +/// Attach to the leader-initialized DSM region as `proc_idx` (1-based for +/// workers: PG's `ParallelWorkerNumber + 1`). +/// +/// # Safety +/// - `coordinate` must be the DSM region pointer the leader initialized. +/// - `region_total` must match the DSM's attached size. +pub(super) unsafe fn worker_attach( + coordinate: *mut c_void, + region_total: u64, + proc_idx: u32, + wakeup: Arc, +) -> Result<(MppDsmHeader, Vec, ProcAttach), String> { + if coordinate.is_null() { + return Err("mpp: worker_attach given null coordinate".into()); + } + let base = coordinate as *mut u8; + let header = unsafe { std::ptr::read(base.cast::()) }; + header + .validate(region_total) + .map_err(|e| format!("mpp: worker DSM validate: {e}"))?; + if proc_idx == 0 { + return Err( + "mpp: worker_attach must be called with proc_idx >= 1 (proc 0 is leader)".into(), + ); + } + if proc_idx >= header.n_procs { + return Err(format!( + "mpp: proc_idx {proc_idx} not in 1..{}", + header.n_procs + )); + } + + // Copy plan bytes out of DSM so the caller has an owned buffer. + let plan_bytes = unsafe { + std::slice::from_raw_parts( + base.add(header.plan_offset as usize), + header.plan_len as usize, + ) + .to_vec() + }; + + let attach = unsafe { + attach_proc( + base, &header, proc_idx, /* attach_senders */ true, wakeup, + ) + }; + Ok((header, plan_bytes, attach)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compute_dsm_layout_works() { + let l = compute_dsm_layout(4, 64 * 1024, 1024).unwrap(); + assert_eq!(l.n_procs, 4); + // 4 procs => 4 inboxes laid out contiguously after the plan bytes. + let aligned = aligned_queue_bytes(64 * 1024); + let queues_size = 4 * aligned; + assert_eq!(l.region_total, l.queues_offset + queues_size); + } + + #[test] + fn compute_dsm_layout_rejects_zero_procs() { + assert!(compute_dsm_layout(0, 64 * 1024, 0).is_err()); + } + + #[test] + fn compute_dsm_layout_rejects_oversize() { + assert!(compute_dsm_layout(u32::MAX, 64 * 1024, 0).is_err()); + } + + #[test] + fn compute_dsm_layout_scales_linearly_in_n_procs() { + // Total queue area must grow as O(N) in proc count. Pinning the math here so a + // regression that grows queues per-pair fails this test at compile time rather + // than at the N=24 wall-time cliff in production. + let queue_bytes = 64 * 1024; + let aligned = aligned_queue_bytes(queue_bytes); + for n in [2u32, 4, 8, 16, 24] { + let l = compute_dsm_layout(n, queue_bytes, 0).unwrap(); + let expected = (n as usize) * aligned; + assert_eq!( + l.region_total - l.queues_offset, + expected, + "n={n}: expected {expected} inbox bytes ({n} inboxes × {aligned})" + ); + } + } + + #[test] + fn header_inbox_offset_is_per_receiver() { + // 4 procs => 4 inboxes, contiguous, sized by queue_bytes each. + let l = compute_dsm_layout(4, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + let aligned = h.queue_bytes; + assert_eq!(h.inbox_offset(0), h.queues_offset); + assert_eq!(h.inbox_offset(1), h.queues_offset + aligned); + assert_eq!(h.inbox_offset(2), h.queues_offset + 2 * aligned); + assert_eq!(h.inbox_offset(3), h.queues_offset + 3 * aligned); + } + + #[test] + fn header_validate_accepts_self() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + assert!(h.validate(l.region_total as u64).is_ok()); + } + + #[test] + fn header_validate_rejects_wrong_version() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let mut h = MppDsmHeader::from_layout(&l); + h.header_version = MPP_DSM_HEADER_VERSION.wrapping_sub(1); + let err = h + .validate(l.region_total as u64) + .expect_err("wrong version must fail"); + assert!(err.contains("DSM header version mismatch"), "got: {err}"); + } + + #[test] + fn header_validate_rejects_size_mismatch() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + assert!(h.validate(l.region_total as u64 + 1).is_err()); + } + + /// Pins the coupling between `compute_dsm_layout`'s minimum check and `ring_dims_for`'s + /// slot-capacity floor: a queue sized below the floored ring must be rejected, or + /// `create_at` would write past its inbox region. + #[test] + fn layout_rejects_queue_smaller_than_floored_ring() { + let floored = DsmMpscRingHeader::region_bytes(RING_SLOTS, MIN_SLOT_CAPACITY); + let too_small = DsmMpscRingHeader::region_bytes(RING_SLOTS, 1); + assert!(too_small < floored); + assert!(compute_dsm_layout(2, too_small, 0).is_err()); + assert!(compute_dsm_layout(2, floored.next_multiple_of(64), 0).is_ok()); + } +} diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs new file mode 100644 index 00000000..632edf62 --- /dev/null +++ b/src/shm/in_process.rs @@ -0,0 +1,1012 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! In-process instantiation of the shared-memory transport, plus an end-to-end test. +//! +//! A real distributed query runs through [`ShmChannelResolver`] with no Postgres and no Flight. +//! This is the rebase-safety payoff: a single process plays every role that production splits +//! across the leader (build and slice the plan), Postgres (allocate the DSM region, launch the +//! workers), and each worker (run a producer fragment and push it through the mesh). If an upstream +//! change breaks the channel-protocol contract, the boundary routing API, or the ready-to-run +//! fragment contract, the test below fails here, before any downstream embedder rebuilds. +//! +//! Workers are passive executors: they run the ready-to-run per-task plans the coordinator +//! dispatches (task-specialized, nested stages `Remote`), whose boundary leaves read the mesh +//! through the session's [`ShmChannelResolver`]. Nothing on the worker re-plans, converts, or +//! dispatches. A [`DispatchPlanSource`] on the leader session captures each dispatched plan, +//! standing in for a production source that serializes it with the embedder's codec. +//! +//! What's faithful vs. simplified relative to the Postgres path: +//! - Faithful: the DSM ring mesh, the framing, the cooperative drain, `ShmWorkerChannel::execute_task`, +//! the per-fragment routing (`collect_dispatched_stages`), `run_worker_fragment`, the leader +//! consuming via `DistributedExec::execute`, and the dispatch handoff through +//! `DispatchPlanSource`. +//! - Simplified: the dispatched subplans are captured as `Arc`s rather than serialized through +//! the `SetPlan` frames the coordinator routes over the mesh (all roles live in one address +//! space), the wakeup is a no-op (the cooperative consumer yields rather than parking), and +//! there's no cancellation source. + +use std::alloc::Layout; +use std::ffi::c_void; +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::catalog::memory::DataSourceExec; +use datafusion::common::runtime::JoinSet; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::common::{DataFusionError, HashMap, Result}; +use datafusion::config::ConfigOptions; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::{SessionStateBuilder, TaskContext}; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion::prelude::{SessionConfig, SessionContext}; + +use crate::{ + DispatchPlanSource, DistributedConfig, DistributedExec, DistributedExt, DistributedLeafExec, + DistributedTaskContext, NetworkBoundaryExt, NetworkBroadcastExec, NetworkCoalesceExec, + NetworkShuffleExec, PartitionSink, SessionStateBuilderExt, Stage, TaskEstimation, + TaskEstimator, TaskKey, WorkerSink, decode_task_metrics, +}; + +use super::mpsc_ring::Wakeup; +use super::runtime::{InProcessWorkerResolver, MppMesh, ShmChannelResolver, proc_for_task}; +use super::setup::{ + collect_task_metrics, dsm_region_bytes, leader_setup, run_worker_fragment, worker_setup, +}; +use super::transport::{ + CooperativeDrainSet, MppFrameHeader, MppPartitionSink, MppSender, NoInterrupt, +}; + +/// Per-inbox DSM ring size for the in-process mesh. Generous: the test ships a handful of tiny +/// batches, so backpressure never kicks in. Production sizes this from `paradedb.mpp_queue_size`. +const IN_PROCESS_QUEUE_BYTES: usize = 1 << 20; + +/// No-op wakeup. The cooperative consumer in `runtime::pull_partition_stream` yields rather +/// than parking, so a publish never needs to wake a blocked thread. The real cross-thread wakeup +/// extension point is covered by `mpsc_ring`'s `injected_wakeup_unparks_blocked_consumer`. +struct NoopWakeup; +impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} +} + +/// Owns the single heap buffer that stands in for the PG DSM segment. Every proc (leader + workers) +/// reads and writes the same region through raw pointers; the lock-free rings make concurrent +/// access sound. Kept alive in the harness until after all producer tasks join. +struct HeapRegion { + ptr: *mut u8, + layout: Layout, +} + +impl HeapRegion { + fn new(bytes: usize) -> Self { + // 64-byte alignment so each per-inbox ring header lands on its own cache line; the + // dsm layout aligns the offsets within the region, but only if the base is aligned too. + let layout = Layout::from_size_align(bytes, 64).expect("dsm region layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "dsm region alloc failed"); + Self { ptr, layout } + } + + fn base(&self) -> *mut c_void { + self.ptr as *mut c_void + } +} + +impl Drop for HeapRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } +} + +/// Send wrapper so the leader-init / worker-attach raw base pointer can be handed to the per-proc +/// setup. All setup runs on the harness thread, so the pointer never crosses into a spawned task; +/// only the resulting `Send` meshes and senders do. +#[derive(Clone, Copy)] +struct SharedBase(*mut c_void); +unsafe impl Send for SharedBase {} + +/// Opaque, non-sentinel receiver token. `NoopWakeup` ignores the value; this just exercises the +/// `set_receiver` path with something the producer won't skip as "no consumer registered". +fn receiver_token(proc_idx: u32) -> u64 { + proc_idx as u64 + 1 +} + +// --------------------------------------------------------------------------------------------- +// Leader-side producer-stage discovery (in-process port of pg_search's `collect_dispatched_stages`). +// The crate's gRPC path keys dispatch on resolver URLs and never decides this; the shm_mq peers are +// push-driven without URLs, so the embedder classifies each boundary's routing here. +// --------------------------------------------------------------------------------------------- + +/// Routing rule for a producer fragment's output partitions. +#[derive(Clone, Debug)] +enum FragmentRouting { + /// Every output partition goes to one destination proc (a `NetworkCoalesceExec`, or the + /// top-level gather to the leader). + Coalesce { dest_proc: u32 }, + /// Hash-partitioned mesh (`NetworkShuffleExec` / `NetworkBroadcastExec`): output partition `q` + /// goes to the consumer task the crate's `route_partition(q)` selects. + Hashed { + consumer_task: Vec, + broadcast: bool, + }, +} + +/// One producer stage's routing metadata, captured from a network boundary. The stage's plan is +/// not captured here: workers run the specialized plans the coordinator dispatches, delivered +/// through the leader session's [`CapturingPlanSource`]. +struct StageEntry { + stage_num: u32, + task_count: usize, + routing: FragmentRouting, +} + +/// The dispatched plans, keyed by `(stage_id, task_number)`. Filled by [`CapturingPlanSource`] +/// while the leader's `execute` prepares the query; read by the worker tasks. One query per +/// mesh, so the key drops the query id. +type CapturedPlans = Arc>>>; + +/// Test-harness [`DispatchPlanSource`]: records each `(stage, task)` specialized plan the +/// coordinator hands over, standing in for a production source that serializes it with the +/// embedder's codec. The returned empty bytes keep the coordinator from encoding plans nobody +/// decodes; the workers run the captured `Arc`s directly. +struct CapturingPlanSource(CapturedPlans); + +impl DispatchPlanSource for CapturingPlanSource { + fn dispatch_plan_proto( + &self, + task: &crate::TaskKey, + specialized: &Arc, + ) -> Option>> { + self.0 + .lock() + .unwrap() + .insert((task.stage_id, task.task_number), Arc::clone(specialized)); + Some(Ok(Vec::new())) + } +} + +/// Walk the distributed physical plan and collect every producer stage, once per boundary. +fn collect_dispatched_stages(root: &Arc, n_workers: u32) -> Vec { + let mut out = Vec::new(); + collect_stages(root, n_workers, /* nested = */ false, &mut out); + out +} + +fn collect_stages( + plan: &Arc, + n_workers: u32, + nested: bool, + out: &mut Vec, +) { + if let Some(nb) = plan.as_ref().as_network_boundary() { + let stage = nb.input_stage(); + let stage_id = stage.num() as u32; + let route_consumer_tasks = || { + // The crate owns the consumer-slice layout, so the push side reads it from + // `route_partition` instead of re-deriving it and drifting when the layout changes. + let n_out = stage + .local_plan() + .map_or(0, |p| p.properties().partitioning.partition_count()); + (0..n_out) + .map(|q| { + nb.route_partition(q) + .expect("route_partition") + .consumer_task as u32 + }) + .collect::>() + }; + let plan_any = plan.as_ref(); + let routing = if plan_any.is::() { + if nested { + FragmentRouting::Coalesce { + dest_proc: proc_for_task(n_workers, 0), + } + } else { + FragmentRouting::Coalesce { dest_proc: 0 } + } + } else if plan_any.is::() { + assert!(nested, "top-level NetworkShuffleExec is unsupported"); + FragmentRouting::Hashed { + consumer_task: route_consumer_tasks(), + broadcast: false, + } + } else if plan_any.is::() { + assert!(nested, "top-level NetworkBroadcastExec is unsupported"); + FragmentRouting::Hashed { + consumer_task: route_consumer_tasks(), + broadcast: true, + } + } else { + panic!("unrecognized network boundary {}", plan.name()); + }; + + let task_count = stage.task_count(); + if let Some(stage_plan) = stage.local_plan() { + out.push(StageEntry { + stage_num: stage_id, + task_count, + routing, + }); + // The boundary's children() returns [stage.plan], so descending here would double-count + // every nested stage. Recurse through the stage plan directly with nested = true. + collect_stages(stage_plan, n_workers, true, out); + } + return; + } + for child in plan.children() { + collect_stages(child, n_workers, nested, out); + } +} + +/// One producer fragment assigned to a worker proc: a single task of a producer stage. The plan +/// arrives separately, through the [`CapturedPlans`] map the leader's dispatch fills. +struct FragmentAssignment { + stage_id: u32, + task_idx: usize, + task_count: usize, + routing: FragmentRouting, +} + +/// Rebuild a plan subtree so each fragment executes its own node instances. In production every +/// proc decodes its own copy of the plan; a captured `Arc` is shared with the coordinator and +/// with sibling tasks of the same stage, and sharing execute-once state breaks (a +/// `RepartitionExec` panics when a second fragment executes a partition the first already +/// consumed; a boundary's connection pool hands each partition stream out once). A `Remote` +/// boundary is childless, so it gets a fresh node explicitly; other leaves are shared (they +/// carry no execute-once state here). +fn reinstantiate(plan: &Arc) -> Arc { + if let Some(nb) = plan.as_ref().as_network_boundary() + && matches!(nb.input_stage(), Stage::Remote(_)) + { + return nb + .with_input_stage(nb.input_stage().clone()) + .expect("with_input_stage with the same stage"); + } + let children: Vec<_> = plan.children().into_iter().map(reinstantiate).collect(); + if children.is_empty() { + Arc::clone(plan) + } else { + Arc::clone(plan) + .with_new_children(children) + .expect("with_new_children with the same arity") + } +} + +/// Wait for the leader's dispatch to capture this fragment's plan. The capture happens inside +/// the leader's `execute` (during plan preparation), which interleaves with the worker tasks on +/// the cooperative runtime. Bounded so a dispatch that never happens fails the test instead of +/// hanging it. +async fn captured_plan( + captured: &CapturedPlans, + stage_id: u32, + task_idx: usize, +) -> Arc { + for _ in 0..100_000 { + if let Some(plan) = captured + .lock() + .unwrap() + .get(&(stage_id as usize, task_idx)) + .cloned() + { + return plan; + } + tokio::task::yield_now().await; + } + panic!("no dispatched plan captured for stage {stage_id} task {task_idx}"); +} + +/// Expand the dispatched stages into the fragments `this_proc` owns under `proc_for_task`. +fn fragments_for_proc( + entries: &[StageEntry], + this_proc: u32, + n_workers: u32, +) -> Vec { + let mut out = Vec::new(); + for entry in entries { + for task_idx in 0..entry.task_count { + if proc_for_task(n_workers, task_idx as u32) != this_proc { + continue; + } + // Broadcast caps its build subtree at task 0; the other tasks would re-emit the same + // canonical replica and the consumer's select_all would over-count. + if matches!( + entry.routing, + FragmentRouting::Hashed { + broadcast: true, + .. + } + ) && task_idx != 0 + { + continue; + } + out.push(FragmentAssignment { + stage_id: entry.stage_num, + task_idx, + task_count: entry.task_count, + routing: entry.routing.clone(), + }); + } + } + out +} + +/// Build a fragment's `TaskContext`, carrying the right `DistributedTaskContext` so nested boundary +/// nodes know their `(task_index, task_count)` and the worker session's channel resolver rides +/// along for their mesh reads. +fn fragment_task_ctx( + session: &SessionContext, + task_index: usize, + task_count: usize, +) -> Arc { + let cfg = session + .state() + .config() + .clone() + .with_extension(Arc::new(DistributedTaskContext { + task_index, + task_count, + })); + Arc::new(TaskContext::default().with_session_config(cfg)) +} + +/// Run all fragments owned by one worker proc, then signal completion. Mirrors the body of +/// pg_search's `run_mpp_worker`: build a [`WorkerSink`] that routes by partition, open a +/// [`PartitionSink`] per output partition, execute each dispatched fragment as-is (it arrives +/// ready-to-run, nested stages `Remote`, boundary leaves reading the mesh), and join. +async fn run_worker_proc( + fragments: Vec, + outbound: Vec>, + mesh: Arc, + session: SessionContext, + n_workers: u32, + captured: CapturedPlans, +) -> Result<()> { + let mut routing = HashMap::new(); + for fragment in &fragments { + routing.insert(fragment.stage_id, fragment.routing.clone()); + } + // One sink serves every stage this proc produces; it owns the base outbound senders and routes + // each (stage, partition) to the destination proc's send end. + let worker_sink = ShmMqWorkerSink { + outbound, + mesh: Arc::clone(&mesh), + n_workers, + routing, + }; + + let mut prepared = Vec::with_capacity(fragments.len()); + for fragment in &fragments { + let task_ctx = fragment_task_ctx(&session, fragment.task_idx, fragment.task_count); + // Production decodes a fresh plan per fragment; the captured Arc is shared, so copy it. + let plan = + reinstantiate(&captured_plan(&captured, fragment.stage_id, fragment.task_idx).await); + let n_out = plan.output_partitioning().partition_count(); + let mut sinks: Vec> = Vec::with_capacity(n_out); + for q in 0..n_out { + sinks.push(worker_sink.open_partition(fragment.stage_id as usize, q)?); + } + prepared.push((fragment, plan, sinks, task_ctx)); + } + // The metrics frames go to the leader after the fragments finish; the clone keeps one sender + // on the leader's inbox alive past the drop below, which only delays that ring's detach + // observation, never a per-channel EOF. + let metrics_sender_base = worker_sink + .outbound + .first() + .and_then(|s| s.as_ref()) + .map(|s| s.clone_with_header(MppFrameHeader::task_metrics(0, 0, mesh.this_proc))); + // Drop the base senders so the only senders left are the per-partition clones the fragment + // futures own; otherwise the rings never observe the last-sender detach. + drop(worker_sink); + + let mut futures = Vec::with_capacity(prepared.len()); + let mut executed = Vec::with_capacity(prepared.len()); + for (fragment, plan, sinks, task_ctx) in prepared { + executed.push(( + fragment.stage_id, + fragment.task_idx, + fragment.task_count, + Arc::clone(&plan), + )); + futures.push(run_worker_fragment(plan, sinks, task_ctx)); + } + for r in futures::future::join_all(futures).await { + r?; + } + // Report per-fragment metrics, the same frames pg's workers ship after their last EOF; the + // leader files them into the executed plan's metrics store. + if let Some(base) = metrics_sender_base { + for (stage_id, task_idx, task_count, plan) in &executed { + let frame = collect_task_metrics(plan, *task_idx, *task_count); + let sender = base.clone_with_header(MppFrameHeader::task_metrics( + *stage_id, + *task_idx as u32, + mesh.this_proc, + )); + let _ = sender.send_task_metrics_best_effort(&frame).await; + } + } + Ok(()) +} + +/// Test-harness [`WorkerSink`]: routes each `(stage, partition)` to the destination proc's outbound +/// send end, the in-process analog of what pg_search builds on a real backend. Holds the base +/// senders plus the per-stage routing so `open_partition` reproduces the header + cooperative-drain +/// wiring the produce loop used to apply inline. +struct ShmMqWorkerSink { + outbound: Vec>, + mesh: Arc, + n_workers: u32, + routing: HashMap, +} + +impl WorkerSink for ShmMqWorkerSink { + fn open_partition(&self, stage: usize, partition: usize) -> Result> { + let routing = self.routing.get(&(stage as u32)).ok_or_else(|| { + DataFusionError::Internal(format!("run_worker_proc: no routing for stage {stage}")) + })?; + let dest_proc = match routing { + FragmentRouting::Coalesce { dest_proc } => *dest_proc, + FragmentRouting::Hashed { consumer_task, .. } => { + proc_for_task(self.n_workers, consumer_task[partition]) + } + }; + let base = self.outbound[dest_proc as usize].as_ref().ok_or_else(|| { + DataFusionError::Internal(format!( + "run_worker_proc: no outbound sender for dest proc {dest_proc}" + )) + })?; + let sender = base + .clone_with_header(MppFrameHeader::batch( + stage as u32, + partition as u32, + self.mesh.this_proc, + )) + .with_cooperative_drain(Arc::clone(&self.mesh) as Arc); + Ok(Box::new(MppPartitionSink::new(sender))) + } +} + +/// Splits an in-memory `DataSourceExec` leaf across tasks, the in-memory analog of the crate's +/// `FileScanConfigTaskEstimator` (which only handles file scans). Each task reads a disjoint subset +/// of the source's partitions, so a gather over the tasks reproduces the serial result exactly. +#[derive(Debug)] +struct MemShardEstimator { + n_tasks: usize, +} + +impl MemShardEstimator { + fn mem_source(plan: &Arc) -> Option<&MemorySourceConfig> { + plan.downcast_ref::()? + .data_source() + .downcast_ref::() + } +} + +impl TaskEstimator for MemShardEstimator { + fn task_estimation( + &self, + plan: &Arc, + _cfg: &ConfigOptions, + ) -> Option { + Self::mem_source(plan).map(|_| TaskEstimation::desired(self.n_tasks)) + } + + fn scale_up_leaf_node( + &self, + plan: &Arc, + task_count: usize, + _cfg: &ConfigOptions, + ) -> Result>> { + if task_count <= 1 { + return Ok(None); + } + let Some(mem) = Self::mem_source(plan) else { + return Ok(None); + }; + let parts = mem.partitions().to_vec(); + let n_part = parts.len(); + // The stored batches are unprojected; reuse the source's exact schema + projection so each + // variant's projected output schema matches the original leaf. + let unprojected_schema: SchemaRef = parts + .iter() + .flatten() + .next() + .map(|b| b.schema()) + .unwrap_or_else(|| plan.schema()); + let projection = mem.projection().clone(); + let variants = (0..task_count).map(|i| { + // Keep every variant at the original partition count (pad with empties) so + // DistributedLeafExec's same-partition-count contract holds; round-robin the + // non-empty partitions so each task reads a disjoint slice. + let per_task: Vec> = (0..n_part) + .map(|j| { + if j % task_count == i { + parts[j].clone() + } else { + Vec::new() + } + }) + .collect(); + MemorySourceConfig::try_new_exec( + &per_task, + unprojected_schema.clone(), + projection.clone(), + ) + .expect("memory variant") as Arc + }); + Ok(Some(Arc::new(DistributedLeafExec::try_new( + Arc::clone(plan), + variants, + )?))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::datasource::MemTable; + use futures::TryStreamExt; + + /// Total procs = leader (proc 0) + `N_WORKERS` producers. With round-robin `proc_for_task`, + /// worker proc `p` runs producer task `p - 1`. + const N_WORKERS: u32 = 3; + + fn table_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + ])) + } + + /// `N_WORKERS` partitions, two rows each, so the shard estimator hands one partition per task. + fn table_partitions() -> Vec> { + let schema = table_schema(); + (0..N_WORKERS as i32) + .map(|p| { + let ids = Int32Array::from(vec![p * 2, p * 2 + 1]); + let vals = Int32Array::from(vec![p * 20, p * 20 + 10]); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(vals)]) + .unwrap(); + vec![batch] + }) + .collect() + } + + fn register_table(ctx: &SessionContext) { + let table = MemTable::try_new(table_schema(), table_partitions()).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + } + + /// A worker is a consumer too (it reads shuffle inputs), so when one of its input streams drops + /// early it has to cancel that stream's producer, not just the leader. This checks the wiring + /// end-to-end: a worker proc's `cancel_stream` reaches the producing proc's inbox through the + /// control senders `worker_setup` installs, and the producer records it. The producer then ends + /// the stream cleanly, which `producer_send_ends_when_consumer_cancels_the_stream` covers. + #[test] + fn worker_consumer_cancels_its_producer() { + use crate::shm::transport::CooperativeDrainSet; + + // procs: leader 0, plus workers 1 and 2. + let boot = bootstrap_mesh(3); + let consumer = Arc::clone(&boot.workers[0].1); // proc 1 + let producer = Arc::clone(&boot.workers[1].1); // proc 2 + + assert!(!producer.stream_cancelled(7, 0)); + + // Proc 1 abandons the `(stage 7, partition 0)` stream it reads from proc 2. + consumer.cancel_stream(2, 7, 0); + + // Proc 2 drains its inbox and sees the cancel its consumer sent. + producer.try_drain_pass().unwrap(); + assert!(producer.stream_cancelled(7, 0)); + // Scoped to that one stream: a sibling partition stays live. + assert!(!producer.stream_cancelled(7, 1)); + } + + /// `dispatch_capture` is Some on the leader session only: its coordinator is the one that + /// dispatches, and the capturing source stands in for a production embedder's serializing + /// source. + fn build_session( + mesh: Arc, + dispatch_capture: Option, + ) -> SessionContext { + let config = SessionConfig::new().with_target_partitions(N_WORKERS as usize); + let mut builder = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .with_distributed_option_extension(DistributedConfig::default()) + .with_distributed_worker_resolver(InProcessWorkerResolver::new(N_WORKERS as usize)) + .with_distributed_channel_resolver(ShmChannelResolver::new(mesh)) + .with_distributed_task_estimator(MemShardEstimator { + n_tasks: N_WORKERS as usize, + }) + .with_distributed_planner(); + if let Some(captured) = dispatch_capture { + builder = builder + .with_distributed_dispatch_plan_source(CapturingPlanSource(captured)) + .with_distributed_metrics_collection(true) + .expect("with_distributed_metrics_collection"); + } + let ctx = SessionContext::new_with_state(builder.build()); + register_table(&ctx); + ctx + } + + fn new_captured_plans() -> CapturedPlans { + Arc::new(Mutex::new(HashMap::default())) + } + + /// Collect the `id` column out of a batch list, in row order. + fn ids_of(batches: &[RecordBatch]) -> Vec { + let mut out = Vec::new(); + for b in batches { + let col = b + .column(0) + .as_any() + .downcast_ref::() + .expect("id column"); + out.extend((0..col.len()).map(|i| col.value(i))); + } + out + } + + /// A real distributed query runs end-to-end through the shm_mq transport with no Postgres and + /// no Flight: producer fragments push through the heap-backed mesh while the leader gathers via + /// `DistributedExec::execute`. The gathered, ordered result must match the serial reference. + /// + /// Single-threaded runtime on purpose: it mirrors the per-process current-thread runtime each + /// PG worker runs, and it's exactly the cooperative model the transport is built for (the + /// producer send spin and the consumer pull loop interleave by yielding, not by parallelism). + #[tokio::test(flavor = "current_thread")] + async fn in_process_distributed_query_matches_serial() { + let query = "SELECT id, val FROM t ORDER BY id"; + + // Serial reference: same query, no distribution. + let serial_ctx = SessionContext::new(); + register_table(&serial_ctx); + let expected = serial_ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + let expected_ids = ids_of(&expected); + assert_eq!(expected_ids, vec![0, 1, 2, 3, 4, 5]); + + // One heap region stands in for the DSM segment; size it for n_procs = leader + workers. + let n_procs = N_WORKERS + 1; + let region_total = dsm_region_bytes(n_procs, IN_PROCESS_QUEUE_BYTES, 0).unwrap(); + let region = HeapRegion::new(region_total); + let base = SharedBase(region.base()); + let wakeup: Arc = Arc::new(NoopWakeup); + + // Leader first (it initializes the rings), then each worker attaches. No plan bytes travel + // through the region here; all roles share the producer subplans as Arcs. + let leader_mesh = unsafe { + leader_setup( + base.0, + n_procs, + IN_PROCESS_QUEUE_BYTES, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::new(NoInterrupt), + /* attach_senders */ false, + ) + } + .unwrap() + .mesh; + let mut worker_setups = Vec::new(); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + base.0, + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::new(NoInterrupt), + ) + } + .unwrap(); + worker_setups.push((proc_idx, attach.mesh, attach.outbound_senders)); + } + + // Build the distributed plan once on the leader session; producers and consumer share it. + let captured = new_captured_plans(); + let leader_ctx = build_session(Arc::clone(&leader_mesh), Some(Arc::clone(&captured))); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + assert!( + physical.is::(), + "expected a DistributedExec root, got {}", + physical.name() + ); + + let entries = collect_dispatched_stages(&physical, N_WORKERS); + // Guard against a planner change that silently collapses the query to a trivial one-task + // gather: the producer stage must actually fan across every worker, or the transport's + // multi-task routing never gets exercised. + assert!( + entries.iter().any(|e| e.task_count == N_WORKERS as usize), + "expected a producer stage with task_count = {N_WORKERS}; got {:?}", + entries.iter().map(|e| e.task_count).collect::>() + ); + + // Launch the producer fragments before the leader pulls, so the mesh has data flowing while + // the consumer drains. On the current-thread runtime the spawned tasks interleave with the + // consumer by yielding, the same cooperative model each PG worker runs. + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in worker_setups { + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh), None); + workers.spawn(run_worker_proc( + fragments, + outbound, + mesh, + session, + N_WORKERS, + Arc::clone(&captured), + )); + } + + // Leader consumer: execute the DistributedExec root, same as the production embedder; the + // network boundary nodes pull from the mesh through ShmWorkerChannel::execute_task. + let leader_task_ctx = leader_ctx.task_ctx(); + let stream = physical.execute(0, leader_task_ctx).unwrap(); + let got: Vec = stream.try_collect().await.unwrap(); + + while let Some(res) = workers.join_next().await { + res.expect("worker task panicked").expect("worker proc"); + } + + let got_ids = ids_of(&got); + assert_eq!(got_ids, expected_ids, "distributed gather != serial"); + + // The workers shipped one metrics frame per fragment over the mesh; file them into the + // executed plan's store the way the pg embedder does, and require full coverage. + let dist = physical + .downcast_ref::() + .expect("DistributedExec"); + let store = dist.metrics_store().expect("metrics collection enabled"); + let mut query_id = None; + let mut expected_reports = 0usize; + let _ = physical.apply(|node| { + if let Some(nb) = node.as_ref().as_network_boundary() { + let stage = nb.input_stage(); + query_id.get_or_insert_with(|| stage.query_id()); + expected_reports += stage.task_count(); + } + Ok(TreeNodeRecursion::Continue) + }); + let query_id = query_id.expect("a distributed plan has at least one boundary"); + let mut rx = leader_mesh + .take_task_metrics_receiver() + .expect("task metrics receiver"); + let mut inserted = 0usize; + for _ in 0..1_000 { + let _ = leader_mesh.try_drain_pass(); + while let Ok((stage_id, task_number, metrics)) = rx.try_recv() { + let metrics = decode_task_metrics(metrics).expect("decode task metrics"); + store.insert( + TaskKey { + query_id, + stage_id: stage_id as usize, + task_number: task_number as usize, + }, + metrics, + ); + inserted += 1; + } + if inserted >= expected_reports { + break; + } + tokio::task::yield_now().await; + } + assert_eq!( + inserted, expected_reports, + "every producer task reports metrics" + ); + + // `region` is declared before the meshes, so reverse drop order frees it after every + // receiver handle into it is gone. + } + + /// Mesh bootstrap shared by the tests: leader first (it initializes the rings), then each + /// worker attaches. + struct Bootstrap { + leader_mesh: Arc, + workers: Vec<(u32, Arc, Vec>)>, + // Last field on purpose: struct fields drop in declaration order, so the region + // outlives every receiver handle into it. + _region: HeapRegion, + } + + fn bootstrap_mesh(n_procs: u32) -> Bootstrap { + let region_total = dsm_region_bytes(n_procs, IN_PROCESS_QUEUE_BYTES, 0).unwrap(); + let region = HeapRegion::new(region_total); + let base = SharedBase(region.base()); + let wakeup: Arc = Arc::new(NoopWakeup); + let leader_mesh = unsafe { + leader_setup( + base.0, + n_procs, + IN_PROCESS_QUEUE_BYTES, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::new(NoInterrupt), + /* attach_senders */ false, + ) + } + .unwrap() + .mesh; + let mut workers = Vec::new(); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + base.0, + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::new(NoInterrupt), + ) + } + .unwrap(); + workers.push((proc_idx, attach.mesh, attach.outbound_senders)); + } + Bootstrap { + leader_mesh, + workers, + _region: region, + } + } + + /// A `GROUP BY` plans a nested `NetworkShuffleExec`, so this exercises hash-routed + /// worker-to-worker traffic and the self-loop sender, which the plain gather test never + /// touches. That routing is the main thing an upstream rebase can silently break. + #[tokio::test(flavor = "current_thread")] + async fn in_process_shuffle_query_matches_serial() { + let query = "SELECT val, count(*) AS c FROM t GROUP BY val ORDER BY val"; + + let serial_ctx = SessionContext::new(); + register_table(&serial_ctx); + let expected = serial_ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let boot = bootstrap_mesh(N_WORKERS + 1); + let captured = new_captured_plans(); + let leader_ctx = build_session(Arc::clone(&boot.leader_mesh), Some(Arc::clone(&captured))); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let entries = collect_dispatched_stages(&physical, N_WORKERS); + assert!( + entries + .iter() + .any(|e| matches!(e.routing, FragmentRouting::Hashed { .. })), + "expected a hash-routed producer stage; got {:?}", + entries.iter().map(|e| &e.routing).collect::>() + ); + + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in boot.workers { + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh), None); + workers.spawn(run_worker_proc( + fragments, + outbound, + mesh, + session, + N_WORKERS, + Arc::clone(&captured), + )); + } + + let leader_task_ctx = leader_ctx.task_ctx(); + let stream = physical.execute(0, leader_task_ctx).unwrap(); + let got: Vec = stream.try_collect().await.unwrap(); + + while let Some(res) = workers.join_next().await { + res.expect("worker task panicked").expect("worker proc"); + } + + use datafusion::arrow::util::pretty::pretty_format_batches; + assert_eq!( + pretty_format_batches(&expected).unwrap().to_string(), + pretty_format_batches(&got).unwrap().to_string(), + "distributed shuffle != serial" + ); + } + + /// A producer that attaches and then goes away without sending its EOFs must fail the + /// gather, not hang it: the drain fails the channels the dead receiver fed once the ring + /// detaches. + #[tokio::test(flavor = "current_thread")] + async fn producer_loss_fails_the_gather_instead_of_hanging() { + let query = "SELECT id, val FROM t ORDER BY id"; + + let boot = bootstrap_mesh(N_WORKERS + 1); + let captured = new_captured_plans(); + let leader_ctx = build_session(Arc::clone(&boot.leader_mesh), Some(Arc::clone(&captured))); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let entries = collect_dispatched_stages(&physical, N_WORKERS); + + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in boot.workers { + if proc_idx == 1 { + // Simulated crash: the proc attached (its senders exist), then dies without + // running its fragments or sending EOF. Dropping the senders is what process + // exit does. + drop(outbound); + drop(mesh); + continue; + } + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh), None); + workers.spawn(run_worker_proc( + fragments, + outbound, + mesh, + session, + N_WORKERS, + Arc::clone(&captured), + )); + } + + let leader_task_ctx = leader_ctx.task_ctx(); + let stream = physical.execute(0, leader_task_ctx).unwrap(); + let res: Result, _> = stream.try_collect().await; + + while let Some(r) = workers.join_next().await { + r.expect("worker task panicked").expect("worker proc"); + } + + let err = res + .expect_err("gather must fail when a producer goes away") + .to_string(); + assert!( + err.contains("detached before this channel's EOF"), + "unexpected error: {err}" + ); + } +} diff --git a/src/shm/mesh.rs b/src/shm/mesh.rs new file mode 100644 index 00000000..4421cd73 --- /dev/null +++ b/src/shm/mesh.rs @@ -0,0 +1,375 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! MPP mesh adapters around `DsmMpscRing` plus the alignment helpers used to size +//! the per-inbox DSM regions. + +use datafusion::common::DataFusionError; + +use super::mpsc_ring::{ + DsmMpscReceiver, DsmMpscSender, RecvOutcome as MpscRecvOutcome, SendError as MpscSendError, +}; +use super::transport::{BatchChannelReceiver, BatchChannelSender, RecvOutcome}; + +/// Postgres MAXALIGN. The DSM layout must agree between this crate (which computes the per-inbox +/// offsets) and whatever shared buffer the embedder hands it. PG's `MAXIMUM_ALIGNOF` is 8 on every +/// supported platform, so the shared-memory layout pins it. +const MAXALIGN: usize = 8; + +/// MAXALIGN-DOWN of `queue_bytes`, so every per-inbox region in the grid is the same size and +/// offset math is a simple multiply. +#[inline] +pub(super) fn aligned_queue_bytes(queue_bytes: usize) -> usize { + queue_bytes & !(MAXALIGN - 1) +} + +/// MAXALIGN-UP `n` to the next multiple of [`MAXALIGN`]. Returns `None` on overflow. Used by +/// [`super::dsm::compute_dsm_layout`] to keep section boundaries MAXALIGN-aligned. +#[inline] +pub(super) fn align_up_maxalign_checked(n: usize) -> Option { + let mask = MAXALIGN - 1; + n.checked_add(mask).map(|x| x & !mask) +} + +/// DSM MPSC ring as a `BatchChannelSender`. Multiple producer processes hold their own +/// `DsmInboxSender` clones targeting the same receiver inbox; the ring serializes them +/// via Vyukov CAS on `tail`. Because the ring's atomics are the synchronization point, +/// the `send_bytes` / `try_send_bytes` paths don't need an attach-thread `debug_assert!`. +/// +/// Detach-on-drop: `DsmMpscSender::Drop` decrements `sender_count`; the last drop flips +/// `detached` and wakes the receiver, mirroring shm_mq's "drop the last sender, receiver +/// sees detach" guarantee. +pub(super) struct DsmInboxSender { + inner: DsmMpscSender, + send_lock: tokio::sync::Mutex<()>, +} + +// Send + Sync auto-derive. Both fields are Send + Sync (`DsmMpscSender` via its unsafe +// impls; `tokio::sync::Mutex` by definition), so no manual `unsafe impl` is needed; the +// auto-derive also surfaces a compile error if a future field is `!Send` / `!Sync`. + +impl DsmInboxSender { + /// Wrap a `DsmMpscSender` for use through the `BatchChannelSender` trait. + pub(super) fn new(inner: DsmMpscSender) -> Self { + Self { + inner, + send_lock: tokio::sync::Mutex::new(()), + } + } + + fn map_send_err(err: MpscSendError) -> DataFusionError { + match err { + MpscSendError::Detached => { + DataFusionError::Execution("mpp: DSM MPSC inbox detached".into()) + } + MpscSendError::MessageTooLarge => DataFusionError::Execution( + "mpp: DSM MPSC frame exceeds the entire ring capacity \ + (raise the embedder's queue-size knob)" + .into(), + ), + MpscSendError::Full => DataFusionError::Execution( + "mpp: DSM MPSC inbox full (caller should retry via try_send_bytes)".into(), + ), + } + } +} + +impl BatchChannelSender for DsmInboxSender { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError> { + // Fallback for callers that didn't wire `with_cooperative_drain`. The real send + // path drives `try_send_bytes` through the cooperative spin in `transport.rs`; + // this loop just spins on `yield_now` and burns the backend core under a slow + // consumer. Hitting it in production means a missing `with_cooperative_drain` + // on the fragment, not a real backpressure path. + loop { + match self.inner.try_send(bytes) { + Ok(()) => return Ok(()), + Err(MpscSendError::Full) => std::thread::yield_now(), + Err(e) => return Err(Self::map_send_err(e)), + } + } + } + + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + match self.inner.try_send(bytes) { + Ok(()) => Ok(true), + Err(MpscSendError::Full) => Ok(false), + Err(e) => Err(Self::map_send_err(e)), + } + } + + fn send_lock(&self) -> &tokio::sync::Mutex<()> { + &self.send_lock + } +} + +/// DSM MPSC ring as a `BatchChannelReceiver`. The scratch `Vec` lives behind a `Mutex` so a +/// `&self` `try_recv` can hand the populated buffer back via `mem::take` without `RefCell` runtime +/// borrow tracking. +/// +/// Single-consumer comes from the call pattern: one `DsmInboxReceiver` per process, +/// owned by `DrainHandle::cooperative_receivers`, polled inline from `try_drain_pass`. +/// No two threads ever race on the same receiver; the mutex is just interior-mutability +/// boilerplate, uncontended in production. +pub(super) struct DsmInboxReceiver { + inner: DsmMpscReceiver, + /// Scratch the inner primitive's `try_recv` reuses across calls (via reserve+set_len). + /// We `mem::take` it on every Bytes outcome to hand the bytes to the caller without a + /// copy; the inner primitive re-grows from `Vec::new()` on the next call. + scratch: std::sync::Mutex>, +} + +// `DsmMpscReceiver` is deliberately `!Sync` (single-consumer invariant). We promote +// `DsmInboxReceiver` to Sync because the caller pattern guarantees only one thread +// touches it at a time, and `Mutex>` protects the scratch from shared-reference +// UB if that invariant ever slips. Send is auto-derived. +unsafe impl Sync for DsmInboxReceiver {} + +impl DsmInboxReceiver { + /// Wrap a `DsmMpscReceiver` for use through the `BatchChannelReceiver` trait. + pub(super) fn new(inner: DsmMpscReceiver) -> Self { + Self { + inner, + scratch: std::sync::Mutex::new(Vec::new()), + } + } + + /// Register the consumer's opaque wakeup token on the underlying ring. The embedder packs + /// whatever its [`super::mpsc_ring::Wakeup`] interprets (a PG embedder packs `(pgprocno, pid)`). + pub(super) fn set_receiver(&self, token: u64) { + self.inner.set_receiver(token); + } +} + +impl BatchChannelReceiver for DsmInboxReceiver { + fn try_recv(&self) -> RecvOutcome { + let mut buf = self.scratch.lock().expect("scratch mutex poisoned"); + match self.inner.try_recv(&mut buf) { + MpscRecvOutcome::Bytes => { + // Hand the populated buffer to the caller and leave an empty Vec behind. + // The inner primitive's `try_recv` does its own `reserve(len)` on the + // next call, so we don't need to pre-allocate scratch capacity here. + RecvOutcome::Bytes(std::mem::take(&mut *buf)) + } + MpscRecvOutcome::Empty => RecvOutcome::Empty, + MpscRecvOutcome::Detached => RecvOutcome::Detached, + } + } +} + +#[cfg(test)] +mod tests { + use super::super::mpsc_ring::{self, DsmMpscRingHeader, Wakeup}; + use super::*; + + struct NoopWakeup; + impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} + } + fn test_wakeup() -> std::sync::Arc { + std::sync::Arc::new(NoopWakeup) + } + + /// Allocate a fresh ring (heap, aligned) and return `(owning region, sender, receiver)`. + /// Pairs `DsmInboxSender` + `DsmInboxReceiver` over a heap-allocated ring matching the + /// alignment contract `create_at` requires. Production allocates the region inside a + /// `dsm_segment`; this helper exists so the BatchChannel trait impls can be exercised + /// without a PG backend. + /// + /// The region is returned FIRST so callers bind it first. Rust drops locals in reverse + /// declaration order, so the region (bound first) drops LAST, after the sender/receiver whose + /// `Drop` touches the ring memory (`DsmMpscSender::drop` does `sender_count.fetch_sub`). + /// Returning the region last instead frees the bytes before those `Drop`s run, a + /// use-after-free that corrupts the heap shared with parallel tests. + fn test_dsm_inbox_pair( + ring_size: u32, + slot_capacity: u32, + ) -> (AlignedTestRegion, DsmInboxSender, DsmInboxReceiver) { + let bytes = DsmMpscRingHeader::region_bytes(ring_size, slot_capacity); + let region = AlignedTestRegion::new(bytes); + let header_ptr = + unsafe { mpsc_ring::create_at(region.as_mut_ptr(), ring_size, slot_capacity) }; + let nn = std::ptr::NonNull::new(header_ptr).expect("create_at returned null"); + let alive = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); + let sender = DsmInboxSender::new(unsafe { + DsmMpscSender::new(nn, test_wakeup(), std::sync::Arc::clone(&alive)) + }); + let receiver = DsmInboxReceiver::new(unsafe { DsmMpscReceiver::new(nn, alive) }); + (region, sender, receiver) + } + + struct AlignedTestRegion { + ptr: *mut u8, + layout: std::alloc::Layout, + } + + impl AlignedTestRegion { + fn new(bytes: usize) -> Self { + let align = std::mem::align_of::(); + let layout = std::alloc::Layout::from_size_align(bytes, align).expect("layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null()); + Self { ptr, layout } + } + fn as_mut_ptr(&self) -> *mut u8 { + self.ptr + } + } + + impl Drop for AlignedTestRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } + } + + #[test] + fn aligned_queue_bytes_rounds_down_to_maxalign() { + let maxalign = MAXALIGN; + for req in [0, 1, 7, 8, 15, 16, 1024, 1025, 1_048_576, 8 * 1024 * 1024] { + let got = aligned_queue_bytes(req); + assert!(got <= req); + assert_eq!(got % maxalign, 0, "not aligned: {got} for req {req}"); + assert!(req - got < maxalign); + } + } + + #[test] + fn align_up_maxalign_rounds_up_to_maxalign() { + let maxalign = MAXALIGN; + for req in [0, 1, 7, 8, 15, 16, 1024, 1025] { + let got = align_up_maxalign_checked(req).unwrap(); + assert!(got >= req); + assert_eq!(got % maxalign, 0); + assert!(got - req < maxalign); + } + assert!(align_up_maxalign_checked(usize::MAX).is_none()); + } + + #[test] + fn dsm_inbox_batch_channel_round_trip() { + let (_region, tx, rx) = test_dsm_inbox_pair(4, 64); + // Sanity: try_send_bytes succeeds, try_recv hands back the bytes, send_lock + // returns a usable Mutex. + assert!(tx.try_send_bytes(b"hello").unwrap()); + assert!(tx.try_send_bytes(b"world").unwrap()); + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"hello"), + other => panic!("expected Bytes(hello), got {other:?}"), + } + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"world"), + other => panic!("expected Bytes(world), got {other:?}"), + } + assert!(matches!(rx.try_recv(), RecvOutcome::Empty)); + // send_lock is just exercised: the per-instance Mutex satisfies the trait. + let _guard = tx.send_lock(); + } + + #[test] + fn dsm_inbox_try_send_returns_false_when_full() { + let (_region, tx, rx) = test_dsm_inbox_pair(2, 64); + assert!(tx.try_send_bytes(b"a").unwrap()); + assert!(tx.try_send_bytes(b"b").unwrap()); + // Third send should hit Full (returns Ok(false), not Err). + assert!(!tx.try_send_bytes(b"c").unwrap()); + // Drain one, then send again succeeds. + assert!(matches!(rx.try_recv(), RecvOutcome::Bytes(_))); + assert!(tx.try_send_bytes(b"c").unwrap()); + } + + #[test] + fn dsm_inbox_multi_producer_through_trait_surface() { + use std::sync::Arc; + // Build a real Arc shared across threads to confirm + // dyn dispatch + Send/Sync impls compile and behave. + let (_region, tx, rx) = test_dsm_inbox_pair(64, 32); + let tx: Arc = Arc::new(tx); + let mut handles = Vec::new(); + const K: usize = 4; + const M: u32 = 200; + for producer_id in 0..K { + let tx = Arc::clone(&tx); + handles.push(std::thread::spawn(move || { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + let mut sent = 0u32; + while sent < M { + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send_bytes(&payload) { + Ok(true) => sent += 1, + Ok(false) => std::thread::yield_now(), + Err(e) => panic!("send failed: {e}"), + } + } + })); + } + let target = K * M as usize; + let mut seen = vec![vec![false; M as usize]; K]; + let mut got = 0usize; + while got < target { + match rx.try_recv() { + RecvOutcome::Bytes(b) => { + let producer_id = u32::from_le_bytes(b[0..4].try_into().unwrap()) as usize; + let idx = u32::from_le_bytes(b[4..8].try_into().unwrap()) as usize; + let already = std::mem::replace(&mut seen[producer_id][idx], true); + assert!(!already, "dup ({producer_id}, {idx})"); + got += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// Dropping the last `DsmInboxSender` flips `detached` and the receiver sees the + /// queued bytes followed by `Detached`. This is the structural equivalent of + /// shm_mq's "drop the last sender, receiver sees detach" guarantee, and is what + /// keeps the drain loop from wedging on a clean shutdown. + #[test] + fn dropping_last_sender_triggers_detach() { + let (_region, tx, rx) = test_dsm_inbox_pair(4, 64); + tx.try_send_bytes(b"final").unwrap(); + drop(tx); + // The queued frame is still readable. + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"final"), + other => panic!("expected Bytes(final), got {other:?}"), + } + // Then Detached. + assert!(matches!(rx.try_recv(), RecvOutcome::Detached)); + } + + #[test] + fn try_send_bytes_rejects_oversize_payload() { + // Multi-slot fragmentation lifted the per-slot ceiling; only payloads larger + // than `ring_size * (slot_capacity - SLOT_HEADER_BYTES)` are now rejected. + // ring_size=2, slot_capacity=32 -> payload_cap=16, ring-wide cap=32. 64 + // bytes still doesn't fit. + let (_region, tx, _rx) = test_dsm_inbox_pair(2, 32); + let oversize = vec![0u8; 64]; + let err = tx + .try_send_bytes(&oversize) + .expect_err("expected MessageTooLarge"); + assert!( + format!("{err}").contains("exceeds the entire ring capacity"), + "unexpected error: {err}" + ); + } +} diff --git a/src/shm/mod.rs b/src/shm/mod.rs new file mode 100644 index 00000000..bf5c95e1 --- /dev/null +++ b/src/shm/mod.rs @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared-memory transport. +//! +//! A non-gRPC [`ChannelResolver`] for co-located execution, where "workers" are tasks or +//! parallel processes sharing one machine and communicating over a shared-memory mesh rather +//! than gRPC. The transport-mechanism pieces (the MPSC ring, framing, routing, cooperative +//! drain) live here as a reusable library; an embedder supplies the two platform primitives via +//! small extension points: how to allocate the shared buffer, and how to wake a blocked consumer. +//! +//! The point of hosting it in this crate is testing: the in-process instantiation runs real +//! distributed queries through the transport in this crate's CI, so an upstream rebase that +//! breaks the channel-protocol contract fails here, before any downstream embedder rebuilds. +//! +//! [`ChannelResolver`]: crate::ChannelResolver +//! +//! Two assumptions an embedder signs up for: +//! - Execution is cooperative on a current-thread runtime: consumers spin on +//! `try_pop` + `yield_now` and producers drain their own inbound while blocked, instead of +//! parking on the `Wakeup` extension point. On a multi-thread runtime each stream burns a core while +//! idle. +//! - Inbound frames demux into unbounded per-channel buffers, so a consumer that falls behind +//! buffers the in-flight intermediate result in process memory. The rings in shared memory +//! stay bounded; the overflow lives on the consumer's heap. + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +mod dsm; +mod mesh; +mod mpsc_ring; +mod runtime; +// Deferred: the self-hosting default transport was built on the removed `WorkerTransport`/ +// `WorkerDispatch` dispatch umbrella, which the `ChannelResolver` model has no analog for; its +// no-gRPC-default role is now served by `InProcessChannelResolver` and its ring-exercising role by +// the `in_process` test, so it stays gated out until reimplemented on `coordinator_channel`. +#[cfg(any())] +mod self_hosted; +mod setup; +mod sink; +mod transport; + +// Curated public surface an embedder consumes. The embedder allocates the shared buffer and +// supplies the two extension points (`Wakeup`, `Interrupt`); everything else is built here. +pub use mpsc_ring::{NO_RECEIVER_TOKEN, Wakeup}; +pub use runtime::{InProcessWorkerResolver, MppMesh, ShmChannelResolver, proc_for_task}; +pub use setup::{ + LeaderAttach, WorkerAttach, collect_task_metrics, dsm_region_bytes, install_work_unit_channels, + leader_setup, region_total, run_worker_fragment, worker_setup, +}; +pub use sink::{PartitionSink, WorkerSink}; +pub use transport::{ + CooperativeDrainSet, Interrupt, MppFrameHeader, MppPartitionSink, MppSender, NoInterrupt, + SendBatchStats, SetPlanFrame, +}; + +/// Out-of-DSM liveness flag shared by the ring handles from one attach. The embedder flips it to +/// `false` from its dsm-detach callback while the segment is still mapped, so a handle dropped +/// afterward (e.g. by a memory-context reset) no-ops instead of dereferencing freed memory. +pub type AliveFlag = Arc; + +// In-process instantiation + the end-to-end test that runs a real distributed query through the +// transport with no Postgres. Test-only: it's how an upstream rebase that breaks the transport +// contract fails in this crate's CI. +#[cfg(test)] +mod in_process; diff --git a/src/shm/mpsc_ring.rs b/src/shm/mpsc_ring.rs new file mode 100644 index 00000000..55c7390b --- /dev/null +++ b/src/shm/mpsc_ring.rs @@ -0,0 +1,1550 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! DSM-backed MPSC ring for MPP mesh inboxes. +//! +//! PG's `shm_mq` is hard-wired SPSC (asserted in `shm_mq_set_sender`), so one inbox +//! can't be shared across N-1 peers without forking PG. This ring is the replacement: +//! a fixed-size byte-message ring sitting in a `dsm_segment`, MPSC-correct via +//! Vyukov-style per-slot sequence counters. +//! +//! Layout (contiguous bytes, `repr(C)`): +//! +//! ```text +//! +- DsmMpscRingHeader -----------------------+ +//! | magic, version, ring_size, slot_capacity | +//! | sender_count, detached, receiver_packed | +//! | (cache-line padding around head/tail) | +//! | head, tail (each on its own line) | +//! +-------------------------------------------+ +//! | Slot[0] | Slot[1] | ... | Slot[N-1] | +//! +-------------------------------------------+ +//! +//! Slot { +//! seq: AtomicU64, // Vyukov phase counter +//! len: AtomicU32, +//! data: [u8; slot_capacity - SLOT_HEADER_SIZE] +//! } +//! ``` +//! +//! Slot phase encoding (Vyukov MPMC reduced to MPSC): +//! +//! ```text +//! slot[i] in round k: seq = k * ring_size + i // empty +//! seq = k * ring_size + i + 1 // ready +//! ``` +//! +//! Producer claim at `tail = T`: read `slot[T % ring_size].seq`. `seq == T` then CAS +//! `tail: T → T+1`, winner copies payload and stores `seq = T + 1` (Release). `seq < T` +//! means ring full. `seq > T` means another producer took `T`, retry. Winner +//! `SetLatch`es the receiver. +//! +//! Consumer take at `head = H`: `slot[H % ring_size].seq == H + 1` means ready. Read +//! payload, store `seq = H + ring_size` (next round's empty marker), then `head = H+1`. +//! The single consumer owns `head` without CAS; producers contend only on `tail`. +//! +//! **Safety**: public methods on `DsmMpscSender` / `DsmMpscReceiver` are type-safe once +//! constructed. The constructors are `unsafe` because the DSM region must be correctly +//! sized and not aliased (one process calls `create_at`, everyone else `attach_at`). +//! +//! **Counter wraparound**: head/tail are `u64`, incremented by one per op. At 100M +//! ops/sec that's ~5800 years, so we ignore wrap. The seq math would break under wrap +//! (pre-wrap `seq` would exceed post-wrap `tail` and producers would spin forever); if +//! that ever matters, add a `tail < u64::MAX - margin` check and reset the ring. + +use std::ptr::NonNull; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; + +use std::sync::Arc; + +use super::AliveFlag; + +/// Wakes the ring's single consumer after a producer publishes a frame. +/// +/// The ring is transport-agnostic: the consumer registers an opaque `u64` token via +/// `DsmMpscReceiver::set_receiver`, and each publishing producer hands that token to this +/// hook. An embedder over PostgreSQL shared memory packs `(pgprocno, pid)` into the token and +/// `SetLatch`es the backend; an in-process embedder packs a registry key and unparks the +/// consumer thread. The token is stored in the ring header (shared memory), so the hook must be +/// able to resolve it from any producer. +pub trait Wakeup: Send + Sync { + fn wake(&self, token: u64); +} + +/// Sentinel token meaning "no consumer registered". Producers skip the wake when the stored +/// token equals this; `create_at` initializes the ring to it. +pub const NO_RECEIVER_TOKEN: u64 = u64::MAX; + +/// Reserved byte at the start of every slot. Sized so payload + header fits in +/// `slot_capacity` exactly. +const SLOT_HEADER_BYTES: usize = std::mem::size_of::(); + +#[repr(C)] +struct SlotHeader { + /// Vyukov sequence counter; see module docs for phase encoding. + seq: AtomicU64, + /// Bytes of payload written into THIS slot's data region; + /// `0..=slot_capacity - SLOT_HEADER_BYTES`. + len: AtomicU32, + /// Fragment metadata. Low 2 bits = [`FragmentKind`] (Complete=0 / First=1 / + /// Continue=2 / Last=3). For `First`, bits 16..32 hold the slot count of the + /// logical frame (1..=65535). For other kinds these bits are unused. + flags: AtomicU32, +} + +/// Kind bits stored in [`SlotHeader::flags`]'s low 2 bits. +/// +/// A frame fitting in `slot_capacity - SLOT_HEADER_BYTES` rides the single-slot +/// fast path as `Complete`. Anything larger spans +/// `n_slots = ceil(frame_len / per_slot_payload)` consecutive slots: `First` +/// (carrying `n_slots` in the upper bits), `Continue` for the middle, `Last` for +/// the tail. Producers grab the whole run with one +/// `tail.compare_exchange(T, T + n_slots)`, so other producers can't interleave +/// their fragments and the receiver always sees the slots in producer order. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FragmentKind { + Complete = 0, + First = 1, + Continue = 2, + Last = 3, +} + +/// Bit-mask for the kind bits in [`SlotHeader::flags`]. +const FLAGS_KIND_MASK: u32 = 0b11; +/// Shift for the `n_slots` field stored in the flags' upper half (only meaningful +/// for `First`). Limits per-frame fragmentation to 65535 slots, which is far +/// beyond any realistic ring size. +const FLAGS_NSLOTS_SHIFT: u32 = 16; +const FLAGS_NSLOTS_MAX: u32 = 0xFFFF; + +#[inline] +fn pack_flags(kind: FragmentKind, n_slots: u32) -> u32 { + debug_assert!( + kind != FragmentKind::First || (1..=FLAGS_NSLOTS_MAX).contains(&n_slots), + "First frames must carry 1..=65535 slots; got {n_slots}" + ); + (kind as u32) | (n_slots << FLAGS_NSLOTS_SHIFT) +} + +#[inline] +fn unpack_kind(flags: u32) -> Option { + match flags & FLAGS_KIND_MASK { + 0 => Some(FragmentKind::Complete), + 1 => Some(FragmentKind::First), + 2 => Some(FragmentKind::Continue), + 3 => Some(FragmentKind::Last), + _ => None, + } +} + +#[inline] +fn unpack_nslots(flags: u32) -> u32 { + flags >> FLAGS_NSLOTS_SHIFT +} + +/// Magic constant validating that an attaching process points at a `DsmMpscRingHeader` +/// rather than garbage. "MPCR" = MPSC Ring. Different value from `MppDsmHeader`'s magic +/// so a worker that picks up the wrong region fails the wrong-shape check loudly. +const MPSC_RING_MAGIC: u32 = u32::from_le_bytes(*b"MPCR"); + +/// Bump on any wire-incompatible layout change. Mirrors the discipline at +/// `MppDsmHeader::validate`. +/// +/// Wire versions: +/// - v1: `SlotHeader { seq, len, _pad }`. One frame per slot. +/// - v2: `SlotHeader { seq, len, flags }`. `flags` carries `FragmentKind` plus +/// `n_slots` on `First`, so frames bigger than +/// `slot_capacity - SLOT_HEADER_BYTES` can span N consecutive slots reserved +/// atomically by the producer. +const MPSC_RING_VERSION: u32 = 2; + +/// Assumed cache line size for false-sharing avoidance. 64 bytes covers x86_64 and arm64; +/// over-padding on smaller-cache-line targets costs a few bytes per ring, nothing more. +const CACHE_LINE: usize = 64; + +/// Ring header. Laid out first in the DSM region; slot array follows immediately after. +/// +/// Cache-line padding around `head` and `tail` isn't optional: at N=24 producer +/// contention the consumer's `head.store` and the producers' `tail.compare_exchange` +/// race on the same line, MESI-ping-ponging every claim. That's the false-sharing +/// footgun Vyukov's writeups call out; the Disruptor literature shows 5-10x throughput +/// loss from it on x86. Padding puts each hot field on its own line. +/// +/// NOT `#[repr(C, align(64))]`: PG `dsm_segment` base addresses are only +/// MAXALIGN-aligned in practice (on macOS the user-data offset can land 16-aligned +/// but not 64-aligned). Forcing `align(64)` would impose a 64-aligned destination on +/// `create_at` / `attach_at` we can't guarantee. The `_pad_*` fields below still +/// put `head`, `tail`, and the first slot on separate 64-byte regions; it's the +/// *distance* between hot fields that matters, not their absolute alignment. +#[repr(C)] +pub(super) struct DsmMpscRingHeader { + /// Magic constant; equals `MPSC_RING_MAGIC` for a valid ring. Checked in `attach_at`. + magic: u32, + /// Layout version; equals `MPSC_RING_VERSION`. Checked in `attach_at`. + version: u32, + /// Number of slots. Immutable after `create_at`. + ring_size: u32, + /// Byte capacity of each slot INCLUDING the slot header. Payload bytes per slot are + /// `slot_capacity - SLOT_HEADER_BYTES`. Immutable after `create_at`. + slot_capacity: u32, + /// Live `DsmMpscSender` count. Incremented in `DsmMpscSender::new`, decremented in + /// `DsmMpscSender::Drop`. The drop that takes the count from 1 → 0 sets `detached` + /// (with Release) and wakes the receiver, mirroring shm_mq's "drop = detach" + /// structural guarantee. + sender_count: AtomicU32, + /// Set by the consumer (or by the leader on query teardown) to tell producers to + /// fail-fast on subsequent sends. Sticky. + detached: AtomicBool, + _pad_after_detached: [u8; 3], + /// Packed `(pgprocno: i32, pid: i32)` of the registered receiver, or 0 (both + /// Opaque receiver token, set by the consumer via `set_receiver` and handed to the + /// embedder's `Wakeup` extension point on every post-publish wake. Initialized to + /// [`NO_RECEIVER_TOKEN`] (`u64::MAX`), which producers treat as "no receiver yet: + /// skip the wake". The embedder defines the token's contents (pg_search packs + /// `(pgprocno, pid)`); a single atomic keeps whatever pair it packs from being + /// observed torn mid-update. + receiver_packed: AtomicU64, + /// Padding to push `head` onto its own cache line. Header up to here uses bytes + /// 0..32; this padding fills 32..64 so `head` lands at offset 64 exactly. The + /// `header_layout_is_cache_friendly` test asserts this. + _pad_before_head: [u8; CACHE_LINE - 32], + /// Consumer's read cursor. Only the consumer writes this. Currently no producer + /// reads it (full-detection works via slot `seq`); the Release-on-store is defensive + /// for any future blocking-send variant that wants to poll consumer progress. + head: AtomicU64, + /// Padding to push `tail` onto its own cache line so the consumer's `head.store` + /// doesn't invalidate the producers' `tail` cache line. + _pad_between_head_and_tail: [u8; CACHE_LINE - 8], + /// Producers' write cursor. CAS'd to claim slot ownership for a tail value. + tail: AtomicU64, + /// Padding so the first slot doesn't share a cache line with `tail`. Producers + /// race on `tail`; the consumer's first slot read should not pull the `tail` cache + /// line into the consumer's L1 unnecessarily. + _pad_after_tail: [u8; CACHE_LINE - 8], +} + +impl DsmMpscRingHeader { + /// Bytes occupied by `ring_size` slots of `slot_capacity` each, plus the header. + pub(super) const fn region_bytes(ring_size: u32, slot_capacity: u32) -> usize { + std::mem::size_of::() + (ring_size as usize) * (slot_capacity as usize) + } +} + +// The receiver wake lives on `DsmMpscSender` (it holds the injected `Wakeup`); see +// `DsmMpscSender::wake_receiver`. + +/// Errors that `try_send` can surface to the producer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum SendError { + /// Ring is full; consumer hasn't drained enough to free a slot. + Full, + /// Receiver has detached (query teardown). Producer should stop sending. + Detached, + /// `bytes.len() + SLOT_HEADER_BYTES > slot_capacity`. The caller picked a slot + /// capacity too small for this payload; bumping `slot_capacity` at `create_at` time + /// fixes it. + MessageTooLarge, +} + +/// Outcome of a single `try_recv` call. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum RecvOutcome { + /// A frame was copied into the caller's buffer. + Bytes, + /// Ring is empty right now; try again later. + Empty, + /// All producers have detached and there's no more data to drain. + Detached, +} + +/// Caller-visible handle for the single consumer. +pub(super) struct DsmMpscReceiver { + ring: NonNull, + alive: AliveFlag, +} + +/// Caller-visible handle for any of the N-1 producers. +pub(super) struct DsmMpscSender { + ring: NonNull, + /// Hook to wake the registered consumer after a publish. Injected by the embedder so the + /// ring stays free of any process/thread-notification mechanism. + wakeup: Arc, + /// Whether this handle counts toward `sender_count`. Data senders (the producer fragments) + /// do, so the last one's drop flips `detached` and the consumer learns its producers are gone. + /// Control senders (a consumer's `Cancel` path) don't: they target a peer's inbox without being + /// one of its producers, so counting them would mask that peer's own producer-gone signal. + counts_as_data: bool, + alive: AliveFlag, +} + +// SAFETY: the ring is a `repr(C)` blob in shared memory whose atomic operations are the +// synchronization point. Both handles are stateless pointers to the same data; sending +// either across threads requires only the atomic ordering already in use. +// +// `DsmMpscReceiver` is deliberately !Sync: the type-level invariant that exactly one +// thread calls `try_recv` at a time is what makes the lock-free MPSC math correct (the +// single consumer owns `head` without CAS). `DsmMpscSender` is Sync so multiple producer +// threads can share one `Arc` and race on `tail` via CAS. +unsafe impl Send for DsmMpscReceiver {} +unsafe impl Send for DsmMpscSender {} +unsafe impl Sync for DsmMpscSender {} + +/// Initialize a freshly-allocated DSM region into a valid ring header + zeroed slot array. +/// Must be called exactly once, by the process that allocated the region (leader in our +/// case). All other processes attach via [`attach_at`] without re-initializing. +/// +/// # Safety +/// - `base` must point at the start of a region of at least +/// `DsmMpscRingHeader::region_bytes(ring_size, slot_capacity)` bytes. +/// - `ring_size >= 2` (a ring of 1 slot can't distinguish empty from full). +/// - `slot_capacity > SLOT_HEADER_BYTES` (need at least one byte of payload). +/// - The region must not be concurrently accessed by any other process or thread until +/// this returns. +pub(super) unsafe fn create_at( + base: *mut u8, + ring_size: u32, + slot_capacity: u32, +) -> *mut DsmMpscRingHeader { + debug_assert!(ring_size >= 2, "ring_size must be >= 2"); + debug_assert!( + slot_capacity as usize > SLOT_HEADER_BYTES, + "slot_capacity must leave room for at least one payload byte" + ); + // The header's natural alignment (from its largest field, AtomicU64) is 8 bytes. + // PG MAXALIGN is 8 on every supported platform, so dsm_segment user-data offsets + // computed via `align_up_maxalign_checked` always satisfy this. Tests use an + // aligned-Vec helper just to keep the assertion clean across allocators. + debug_assert!( + (base as usize).is_multiple_of(std::mem::align_of::()), + "create_at base must be aligned to {} bytes", + std::mem::align_of::() + ); + let header_ptr = base.cast::(); + // Write the immutable header fields. Use std::ptr::write so we don't construct an + // intermediate &mut that aliases the not-yet-initialized atomic fields. + unsafe { + std::ptr::write( + header_ptr, + DsmMpscRingHeader { + magic: MPSC_RING_MAGIC, + version: MPSC_RING_VERSION, + ring_size, + slot_capacity, + detached: AtomicBool::new(false), + _pad_after_detached: [0; 3], + sender_count: AtomicU32::new(0), + receiver_packed: AtomicU64::new(NO_RECEIVER_TOKEN), + _pad_before_head: [0; CACHE_LINE - 32], + head: AtomicU64::new(0), + _pad_between_head_and_tail: [0; CACHE_LINE - 8], + tail: AtomicU64::new(0), + _pad_after_tail: [0; CACHE_LINE - 8], + }, + ); + } + // Initialize slot sequences: slot[i].seq = i in round 0. + for i in 0..ring_size { + let slot = unsafe { slot_ptr(header_ptr, i, slot_capacity) }; + unsafe { + std::ptr::write( + slot, + SlotHeader { + seq: AtomicU64::new(i as u64), + len: AtomicU32::new(0), + flags: AtomicU32::new(0), + }, + ); + } + } + header_ptr +} + +/// Take an already-initialized ring header pointer and confirm its shape matches caller's +/// expectations. The caller's `expected_ring_size` / `expected_slot_capacity` must match +/// the values written at `create_at` time; mismatch is a hard error (returns null). +/// +/// # Safety +/// - `base` must point at the same region a previous `create_at` initialized. +/// - The region must not be deallocated for the lifetime of any handle returned from +/// the wrappers (`DsmMpscReceiver::new`, `DsmMpscSender::new`). +pub(super) unsafe fn attach_at( + base: *mut u8, + expected_ring_size: u32, + expected_slot_capacity: u32, +) -> Option> { + let header_ptr = base.cast::(); + let nn = NonNull::new(header_ptr)?; + let header = unsafe { nn.as_ref() }; + if header.magic != MPSC_RING_MAGIC || header.version != MPSC_RING_VERSION { + return None; + } + if header.ring_size != expected_ring_size || header.slot_capacity != expected_slot_capacity { + return None; + } + Some(nn) +} + +#[inline] +unsafe fn slot_ptr( + header: *mut DsmMpscRingHeader, + idx: u32, + slot_capacity: u32, +) -> *mut SlotHeader { + let header_bytes = std::mem::size_of::(); + let base = header.cast::(); + unsafe { base.add(header_bytes + (idx as usize) * (slot_capacity as usize)) } + .cast::() +} + +#[inline] +unsafe fn slot_data_ptr(slot: *mut SlotHeader) -> *mut u8 { + unsafe { slot.cast::().add(SLOT_HEADER_BYTES) } +} + +impl DsmMpscReceiver { + /// Wrap an already-initialized ring as the single consumer. Pairs with + /// [`DsmMpscSender::new`] on the producer side; calling code is responsible for + /// keeping exactly one `DsmMpscReceiver` per ring. + /// + /// # Safety + /// `ring` must point to a header initialized by [`create_at`] and not yet + /// deallocated. The caller guarantees no other `DsmMpscReceiver` exists for the + /// same ring (single-consumer invariant). + pub(super) unsafe fn new(ring: NonNull, alive: AliveFlag) -> Self { + Self { ring, alive } + } + + /// Register the consumer's opaque wakeup token. Producers Acquire-load it as a single + /// `u64` (so they never see a torn token) and hand it to their [`Wakeup`] after publishing. + /// The token is the embedder's to interpret: a PG embedder packs `(pgprocno, pid)`; an + /// in-process embedder packs a registry key. Must not be [`NO_RECEIVER_TOKEN`]. + pub(super) fn set_receiver(&self, token: u64) { + if !self.alive.load(Ordering::Acquire) { + return; + } + let header = unsafe { self.ring.as_ref() }; + header.receiver_packed.store(token, Ordering::Release); + } + + /// Try to read one frame into `out`. `Bytes`: `out` holds the payload. `Empty`: + /// caller should yield and retry. `Detached`: ring drained, all producers gone, + /// no more frames coming. + /// + /// Known wedge: if a producer CAS-advances `tail` then exits/crashes before + /// publishing `seq`, the consumer sees `tail > head`, the slot's `seq` stuck at + /// the prior-round empty marker, and `detached && tail <= head` never becomes + /// true. The drain returns `Empty` forever. In production PG's parallel-worker + /// death handling bounds this (worker exit fires leader ERROR, DSM tears down), + /// but a future pass should add an explicit `PGPROC` liveness check to + /// force-detach on producer death. + pub(super) fn try_recv(&self, out: &mut Vec) -> RecvOutcome { + if !self.alive.load(Ordering::Acquire) { + // Segment detached: no DSM left to read, so report drained. + return RecvOutcome::Detached; + } + let header = unsafe { self.ring.as_ref() }; + let head = header.head.load(Ordering::Relaxed); + let slot_idx = (head % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + let expected_ready = head.wrapping_add(1); + if seq != expected_ready { + // Slot not ready. Use `<=` rather than `==` so a strict invariant + // violation (tail < head, impossible under correct operation) still + // surfaces as Detached rather than wedging Empty forever. + if header.detached.load(Ordering::Acquire) + && header.tail.load(Ordering::Acquire) <= head + { + return RecvOutcome::Detached; + } + return RecvOutcome::Empty; + } + // Slot at head is ready. Inspect its kind to choose between single-slot + // fast path and multi-slot reassembly. + let flags = unsafe { (*slot).flags.load(Ordering::Relaxed) }; + let Some(kind) = unpack_kind(flags) else { + // Reserved bits set; treat as corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + }; + let payload_cap = (header.slot_capacity as usize).saturating_sub(SLOT_HEADER_BYTES); + match kind { + FragmentKind::Complete => self.recv_single_slot(header, head, slot, payload_cap, out), + FragmentKind::First => { + let n_slots = unpack_nslots(flags); + if n_slots == 0 || n_slots > header.ring_size { + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + self.recv_multi_slot(header, head, n_slots, payload_cap, out) + } + FragmentKind::Continue | FragmentKind::Last => { + // Encountering Continue/Last at `head` means a producer violated + // ascending-publish ordering or a previous reassembly didn't advance + // `head` past every fragment. Either is a contract break, not a + // recoverable condition; poison and detach. + header.detached.store(true, Ordering::Release); + RecvOutcome::Detached + } + } + } + + /// Read a single-slot `Complete` frame at `head` and advance. + fn recv_single_slot( + &self, + header: &DsmMpscRingHeader, + head: u64, + slot: *mut SlotHeader, + payload_cap: usize, + out: &mut Vec, + ) -> RecvOutcome { + let len_raw = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + // Clamp against slot's payload capacity. DSM is mapped writable by every + // attached backend, so a buggy / corrupted producer could write a garbage len. + // Without this guard, `set_len + copy_nonoverlapping` would read OOB into + // neighboring slots or other DSM contents. + if len_raw > payload_cap { + // Poison the ring rather than silently returning corrupt data. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + let len = len_raw; + out.clear(); + out.reserve(len); + let data = unsafe { slot_data_ptr(slot) }; + // copy_nonoverlapping before set_len so a hypothetical panic mid-copy doesn't + // leave `out` with logical-len > initialized-bytes. + unsafe { + std::ptr::copy_nonoverlapping(data, out.as_mut_ptr(), len); + out.set_len(len); + } + // Mark the slot empty for the next round. Round k empty marker is + // (k * ring_size) + slot_idx; head + ring_size is exactly that for next round. + let next_empty_seq = head.wrapping_add(header.ring_size as u64); + unsafe { (*slot).seq.store(next_empty_seq, Ordering::Release) }; + // Advance head AFTER publishing the slot's empty marker, so a producer racing + // to claim sees the empty slot before seeing the new head value. + header.head.store(head.wrapping_add(1), Ordering::Release); + RecvOutcome::Bytes + } + + /// Reassemble an `n_slots`-fragment frame at `head`. Returns `Empty` without + /// advancing `head` if any continuation slot isn't yet published (producer + /// mid-publish); the caller retries and the `First` slot stays put with its + /// flags intact. + fn recv_multi_slot( + &self, + header: &DsmMpscRingHeader, + head: u64, + n_slots: u32, + payload_cap: usize, + out: &mut Vec, + ) -> RecvOutcome { + // First pass: verify every fragment in the run is published. If not, bail + // with Empty so the caller can retry once the producer finishes. + let mut total_len: usize = 0; + for i in 0..n_slots { + let h = head.wrapping_add(i as u64); + let slot_idx = (h % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + if seq != h.wrapping_add(1) { + // Continuation slot not yet ready. Don't touch `head` or any slot + // metadata; producer will publish soon, caller will retry. + return RecvOutcome::Empty; + } + let expected_kind = if i == 0 { + FragmentKind::First + } else if i + 1 == n_slots { + FragmentKind::Last + } else { + FragmentKind::Continue + }; + let flags = unsafe { (*slot).flags.load(Ordering::Relaxed) }; + if unpack_kind(flags) != Some(expected_kind) { + // Run integrity check: the producer's multi-slot publish never + // interleaves with another producer's, so the kind sequence must be + // First, Continue*, Last. Anything else is corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + let len = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + if len > payload_cap || (i + 1 < n_slots && len != payload_cap) { + // Non-final fragments must be full slots (producer fills them + // first); final fragment may be partial. Anything else is + // corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + total_len = match total_len.checked_add(len) { + Some(v) => v, + None => { + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + }; + } + // Second pass: concatenate payloads, mark each slot empty for next round, + // advance head past the whole run in one Release store. + out.clear(); + out.reserve(total_len); + for i in 0..n_slots { + let h = head.wrapping_add(i as u64); + let slot_idx = (h % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let len = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + let data = unsafe { slot_data_ptr(slot) }; + unsafe { + let write_at = out.as_mut_ptr().add(out.len()); + std::ptr::copy_nonoverlapping(data, write_at, len); + out.set_len(out.len() + len); + } + // Round k empty marker for slot at position h is h + ring_size. + let next_empty_seq = h.wrapping_add(header.ring_size as u64); + unsafe { (*slot).seq.store(next_empty_seq, Ordering::Release) }; + } + header + .head + .store(head.wrapping_add(n_slots as u64), Ordering::Release); + RecvOutcome::Bytes + } +} + +impl DsmMpscSender { + /// Wrap an already-initialized ring as a producer. Multiple `DsmMpscSender` + /// handles to the same ring are legal (and the point of MPSC). Increments the + /// ring's `sender_count`; the `Drop` impl decrements and, on the last drop, + /// flips `detached` and wakes the receiver. That mirrors shm_mq's + /// "drop the sender, receiver sees detach" structural guarantee. + /// + /// # Safety + /// `ring` must point to a header initialized by [`create_at`] and not yet + /// deallocated. + pub(super) unsafe fn new( + ring: NonNull, + wakeup: Arc, + alive: AliveFlag, + ) -> Self { + let header = unsafe { ring.as_ref() }; + header.sender_count.fetch_add(1, Ordering::AcqRel); + Self { + ring, + wakeup, + counts_as_data: true, + alive, + } + } + + /// A control-plane sibling onto the same ring this handle already targets: it can publish + /// frames but stays out of `sender_count`, so it never sets or delays `detached`. Used to derive + /// a proc's `Cancel` senders from its data senders without a second attach, so a consumer can + /// reach its producer without masquerading as one of that producer's data senders. + pub(super) fn to_control(&self) -> DsmMpscSender { + DsmMpscSender { + ring: self.ring, + wakeup: Arc::clone(&self.wakeup), + counts_as_data: false, + alive: Arc::clone(&self.alive), + } + } + + /// Wake the registered consumer, if any. Reads the token the consumer stored via + /// [`DsmMpscReceiver::set_receiver`] and hands it to the injected [`Wakeup`]; skips when no + /// consumer is registered ([`NO_RECEIVER_TOKEN`]). + fn wake_receiver(&self) { + let header = unsafe { self.ring.as_ref() }; + let token = header.receiver_packed.load(Ordering::Acquire); + if token != NO_RECEIVER_TOKEN { + self.wakeup.wake(token); + } + } + + /// Push one frame onto the ring. Returns immediately: + /// - `Ok(())`: published; receiver's latch was set if installed. + /// - `Err(Full)`: no slot run available right now; caller yields + retries. + /// - `Err(Detached)`: receiver has detached; caller stops. + /// - `Err(MessageTooLarge)`: frame needs more than `ring_size` slots + /// (max writable size is `ring_size * (slot_capacity - SLOT_HEADER_BYTES)`). + /// + /// Frames up to the per-slot payload capacity take the single-slot fast path. + /// Larger frames span consecutive slots claimed atomically via one + /// `tail.compare_exchange(T, T + n_slots)`, so other producers can't + /// interleave their fragments. See [`FragmentKind`]. + pub(super) fn try_send(&self, bytes: &[u8]) -> Result<(), SendError> { + if !self.alive.load(Ordering::Acquire) { + return Err(SendError::Detached); + } + let header = unsafe { self.ring.as_ref() }; + // Control senders ignore `detached`: a `Cancel` still has to reach a producer whose own + // inbox already detached because its upstream finished while it's mid-output. + if self.counts_as_data && header.detached.load(Ordering::Acquire) { + return Err(SendError::Detached); + } + let payload_cap = (header.slot_capacity as usize).saturating_sub(SLOT_HEADER_BYTES); + if payload_cap == 0 { + return Err(SendError::MessageTooLarge); + } + // Single-slot fast path: cheaper than the multi-slot CAS dance and preserves + // the v1 hot path exactly so the no-fragmentation case takes no extra branches + // inside the claim loop. + if bytes.len() <= payload_cap { + return self.try_send_single_slot(header, bytes); + } + // Multi-slot path: fragment across `n_slots` consecutive slots. + let n_slots_usize = bytes.len().div_ceil(payload_cap); + if n_slots_usize > header.ring_size as usize || n_slots_usize > FLAGS_NSLOTS_MAX as usize { + // Frame is larger than the entire ring (or larger than the n_slots field + // can encode). Either bump `mpp_queue_size` or land a chunked-stream + // protocol that spans rounds. + return Err(SendError::MessageTooLarge); + } + let n_slots = n_slots_usize as u32; + self.try_send_multi_slot(header, bytes, n_slots, payload_cap) + } + + /// Single-slot path. Identical to the v1 layout's hot path with a `flags` write + /// added; the kind is always `Complete` here. + fn try_send_single_slot( + &self, + header: &DsmMpscRingHeader, + bytes: &[u8], + ) -> Result<(), SendError> { + loop { + // Acquire load on `tail` pairs defensively with any future blocking-send + // variant that may want to observe consumer progress via `head` (we don't + // today; full-detection rides on the slot's `seq`). Cheaper to keep the + // ordering tight than to relax-then-tighten under a future audit. + let tail = header.tail.load(Ordering::Acquire); + let slot_idx = (tail % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + // Three-way compare per Vyukov MPMC. + match seq.cmp(&tail) { + std::cmp::Ordering::Equal => { + // Slot is empty in our round. Try to claim by advancing tail. + // AcqRel on success so a subsequent producer's Acquire load of + // `tail` sees our claim and skips the slot we own. Relaxed on + // failure (we just retry the loop on a fresh tail load). + match header.tail.compare_exchange_weak( + tail, + tail.wrapping_add(1), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + // We own slot[slot_idx] for tail value `tail`. + unsafe { + (*slot).len.store(bytes.len() as u32, Ordering::Relaxed); + (*slot).flags.store( + pack_flags(FragmentKind::Complete, 0), + Ordering::Relaxed, + ); + let data = slot_data_ptr(slot); + std::ptr::copy_nonoverlapping(bytes.as_ptr(), data, bytes.len()); + // Publish: ready in round k is (k * ring_size + i + 1) = tail + 1. + (*slot).seq.store(tail.wrapping_add(1), Ordering::Release); + } + // Wake the consumer (resolves the latch via pgprocno + pid). + self.wake_receiver(); + return Ok(()); + } + Err(_) => continue, // another producer took this tail; retry + } + } + std::cmp::Ordering::Less => { + // seq < tail: the consumer hasn't reclaimed slot[slot_idx] for our + // round yet. Ring is full. + return Err(SendError::Full); + } + std::cmp::Ordering::Greater => { + // seq > tail: another producer has already claimed tail. Reload and + // retry. + continue; + } + } + } + } + + /// Multi-slot path. CAS-advance `tail` from `T` to `T + n_slots` to claim the run, + /// then publish each fragment ascending with `First` / `Continue` / `Last` flags. + /// The CAS requires all N target slots to already be in their expected empty + /// round, which generalizes v1's per-slot `seq == tail` check to a run. + /// + /// Ascending publish order isn't optional: the receiver only starts reassembling + /// once `slot[head]` (the `First`) is ready, then waits for each subsequent slot. + /// Out-of-order publish would make the receiver block on a `Continue` whose data + /// is already there but whose `seq` hasn't been stored yet, wasting one drain + /// pass per fragment. + // Fairness caveat: a multi-slot frame needs `n_slots` consecutive empty slots starting at + // `tail`, and competing single-slot producers can keep winning the CAS, so a large frame + // has no progress bound under sustained contention. The cooperative drain keeps retrying + // (no deadlock), but latency is unbounded; revisit with a reservation scheme if it shows + // up in send stats. + fn try_send_multi_slot( + &self, + header: &DsmMpscRingHeader, + bytes: &[u8], + n_slots: u32, + payload_cap: usize, + ) -> Result<(), SendError> { + loop { + let tail = header.tail.load(Ordering::Acquire); + // Verify every target slot is currently in its expected empty round. + // Slot at position (tail + i) % ring_size in round-of-(tail+i) has empty + // marker == (tail + i). If ANY is not empty, the ring is too contended / + // not drained enough for a run of this length. + let mut all_empty = true; + for i in 0..n_slots { + let t = tail.wrapping_add(i as u64); + let slot_idx = (t % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + match seq.cmp(&t) { + std::cmp::Ordering::Equal => {} + std::cmp::Ordering::Less => { + // Slot at offset i hasn't been reclaimed by the consumer for + // round-of-t. Run not available right now. + return Err(SendError::Full); + } + std::cmp::Ordering::Greater => { + // Another producer already claimed (T + i); our view of `tail` + // is stale, retry. + all_empty = false; + break; + } + } + } + if !all_empty { + continue; + } + // All N target slots are empty in our round. Atomically claim the run. + match header.tail.compare_exchange_weak( + tail, + tail.wrapping_add(n_slots as u64), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + // We own slots tail..tail+n_slots. Write each fragment and + // publish its seq in ascending order so the receiver can drain + // them as soon as the prefix is ready. + let mut offset = 0usize; + for i in 0..n_slots { + let t = tail.wrapping_add(i as u64); + let slot_idx = (t % header.ring_size as u64) as u32; + let slot = + unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let remaining = bytes.len() - offset; + let chunk_len = remaining.min(payload_cap); + let kind = if i == 0 { + FragmentKind::First + } else if i + 1 == n_slots { + FragmentKind::Last + } else { + FragmentKind::Continue + }; + let flags_word = if kind == FragmentKind::First { + pack_flags(kind, n_slots) + } else { + pack_flags(kind, 0) + }; + unsafe { + (*slot).len.store(chunk_len as u32, Ordering::Relaxed); + (*slot).flags.store(flags_word, Ordering::Relaxed); + let data = slot_data_ptr(slot); + std::ptr::copy_nonoverlapping( + bytes.as_ptr().add(offset), + data, + chunk_len, + ); + // Publish: ready in round-of-t for slot at position t is t + 1. + (*slot).seq.store(t.wrapping_add(1), Ordering::Release); + } + offset += chunk_len; + } + debug_assert_eq!(offset, bytes.len(), "multi-slot send copied wrong length"); + // Wake the consumer once for the whole run; the drain pass that + // wakes for the First slot will keep going through all our + // already-published Continue/Last fragments without sleeping. + self.wake_receiver(); + return Ok(()); + } + Err(_) => continue, + } + } + } +} + +impl Drop for DsmMpscSender { + fn drop(&mut self) { + if !self.alive.load(Ordering::Acquire) { + // Segment detached: the `sender_count.fetch_sub` below would write into freed DSM. + return; + } + if !self.counts_as_data { + // Control senders never joined `sender_count`, so there's nothing to release and no + // detach to trigger. + return; + } + let header = unsafe { self.ring.as_ref() }; + // AcqRel: decrement is observed by other producers (they don't care, but the + // Release pairs with the receiver's Acquire load on `detached` below). + let prev = header.sender_count.fetch_sub(1, Ordering::AcqRel); + if prev == 1 { + // We were the last sender. Tell the consumer. + header.detached.store(true, Ordering::Release); + self.wake_receiver(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering as O}; + + /// No-op wakeup for the busy-poll tests, which spin on `try_recv` and never rely on being + /// woken. + struct NoopWakeup; + impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} + } + fn noop_wakeup() -> Arc { + Arc::new(NoopWakeup) + } + + /// In-process wakeup: unparks the consumer thread registered under a token. Proves the + /// injected extension point carries a real cross-thread notification with no PG / `SetLatch`. + #[derive(Default)] + struct ThreadWakeup { + threads: std::sync::Mutex>, + wakes: AtomicUsize, + } + impl ThreadWakeup { + fn register(&self, token: u64, thread: std::thread::Thread) { + self.threads.lock().unwrap().insert(token, thread); + } + } + impl Wakeup for ThreadWakeup { + fn wake(&self, token: u64) { + self.wakes.fetch_add(1, O::Relaxed); + if let Some(t) = self.threads.lock().unwrap().get(&token) { + t.unpark(); + } + } + } + + /// Send + Copy wrapper so tests can pass the ring pointer into spawned threads. The + /// real production handles (DsmMpscSender / Receiver) already impl Send; this exists + /// only because constructing them is unsafe and we want the test to do it inside the + /// thread so each thread has its own handle. + #[derive(Clone, Copy)] + struct SharedRing(NonNull); + unsafe impl Send for SharedRing {} + + /// Owning aligned region for a ring. The default `Vec` allocator can't + /// promise the alignment the `create_at` write wants, so allocate via the global + /// allocator with an explicit `Layout` and free in `Drop`. Production uses PG + /// `dsm_segment` (page-aligned), so this dance is test-only. + struct AlignedRegion { + ptr: *mut u8, + layout: std::alloc::Layout, + } + impl AlignedRegion { + fn new(bytes: usize) -> Self { + let align = std::mem::align_of::(); + let layout = std::alloc::Layout::from_size_align(bytes, align).expect("invalid layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "allocator returned null"); + Self { ptr, layout } + } + fn as_mut_ptr(&self) -> *mut u8 { + self.ptr + } + } + impl Drop for AlignedRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } + } + + /// Allocate a heap region big enough for a ring with `ring_size` slots of + /// `slot_capacity` bytes each, initialize it, and return `(owner, receiver, + /// sender_template)`. The owner is returned first so callers bind it first. + /// Rust drops locals in reverse declaration order, so `_region` bound first drops + /// last, after the handles whose `Drop` impls touch the region's memory. Reverse + /// that order and `_region` frees the bytes before `tx_template`'s + /// `sender_count.fetch_sub` runs, which is undefined behavior and surfaces as a + /// stochastic crash at process teardown. + fn make_ring( + ring_size: u32, + slot_capacity: u32, + ) -> (AlignedRegion, DsmMpscReceiver, DsmMpscSender) { + let bytes = DsmMpscRingHeader::region_bytes(ring_size, slot_capacity); + let region = AlignedRegion::new(bytes); + let header_ptr = unsafe { create_at(region.as_mut_ptr(), ring_size, slot_capacity) }; + let nn = NonNull::new(header_ptr).expect("create_at returned null"); + // Unsafe: we hand ownership of the same pointer to two handles. Safe because the + // ring is the synchronization point; the handles are stateless wrappers. + let alive = Arc::new(AtomicBool::new(true)); + let receiver = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + let sender = unsafe { DsmMpscSender::new(nn, noop_wakeup(), alive) }; + (region, receiver, sender) + } + + #[test] + fn control_sender_does_not_gate_detach() { + let bytes = DsmMpscRingHeader::region_bytes(4, 64); + let region = AlignedRegion::new(bytes); + let nn = NonNull::new(unsafe { create_at(region.as_mut_ptr(), 4, 64) }).unwrap(); + let alive = Arc::new(AtomicBool::new(true)); + let rx = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + let data = unsafe { DsmMpscSender::new(nn, noop_wakeup(), alive) }; + let control = data.to_control(); + + let mut buf = Vec::new(); + // The control sender publishes without joining `sender_count`. + control.try_send(&[1, 2, 3]).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[1, 2, 3]); + // The data sender still holds the ring open. + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + + // Dropping the only data sender detaches the ring even though the control sender is held, + // so a downstream cancel never masks a producer-gone signal. + drop(data); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Detached); + + // The control sender still reaches the now-detached inbox: a `Cancel` has to land on a + // producer whose own upstream already finished. + control.try_send(&[9]).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[9]); + drop(control); + } + + #[test] + fn spsc_round_trip_under_capacity() { + let (_region, rx, tx) = make_ring(4, 64); + for i in 0..3u8 { + tx.try_send(&[i, i + 1, i + 2]).unwrap(); + } + let mut buf = Vec::new(); + for i in 0..3u8 { + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[i, i + 1, i + 2]); + } + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + #[test] + fn fills_then_full_then_drains() { + let (_region, rx, tx) = make_ring(4, 64); + for i in 0..4u32 { + tx.try_send(&i.to_le_bytes()).unwrap(); + } + // Fifth send must fail Full. + assert_eq!(tx.try_send(&999u32.to_le_bytes()), Err(SendError::Full)); + // Drain one, send one more, drain rest. + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, 0u32.to_le_bytes()); + tx.try_send(&100u32.to_le_bytes()).unwrap(); + for expected in [1u32, 2, 3, 100] { + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, expected.to_le_bytes()); + } + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + #[test] + fn message_too_large_is_rejected() { + // `MessageTooLarge` only fires when the frame would need more than `ring_size` + // slots. Per-slot oversize splits across consecutive slots instead. + let (_region, _rx, tx) = make_ring(2, 32); + let payload_cap = 32 - SLOT_HEADER_BYTES; + // payload_cap + 1 byte: now spans 2 slots, fits in ring_size=2. + let two_slot = vec![0u8; payload_cap + 1]; + tx.try_send(&two_slot).unwrap(); + // ring_size * payload_cap + 1 byte: needs 3 slots, ring only has 2. + let oversize = vec![0u8; 2 * payload_cap + 1]; + assert_eq!(tx.try_send(&oversize), Err(SendError::MessageTooLarge)); + // Exactly at the ring-wide cap is fine. (We just sent a 2-slot frame above, + // so we need to drain it first to free the slots; the receiver only exists + // in this test for that.) + } + + /// Multi-producer concurrent send: K threads each push M unique messages; consumer + /// receives K*M messages and verifies every (producer, sequence_in_producer) pair + /// appears exactly once. + #[test] + fn mpsc_no_lost_messages_under_contention() { + const K_PRODUCERS: usize = 8; + const M_PER_PRODUCER: u32 = 2000; + let (_region, rx, tx_template) = make_ring(64, 32); + // Wrap region pointer in something we can share across threads. The handles + // themselves are Send, so we clone via NonNull copy. + let ring_nn = SharedRing(tx_template.ring); + let consumed = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + // Construct the sender on this thread so the closure captures DsmMpscSender + // (Send via unsafe impl) rather than the inner NonNull (Rust 2021 disjoint + // capture would otherwise project to the NonNull field and fail Send). + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + let h = std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected send error: {e:?}"), + } + } + }); + handles.push(h); + } + // Drain on this thread until every producer's M messages have shown up. + let mut seen: Vec> = (0..K_PRODUCERS) + .map(|_| vec![false; M_PER_PRODUCER as usize]) + .collect(); + let mut buf = Vec::new(); + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while consumed.load(O::Relaxed) < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + assert_eq!(buf.len(), 8); + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize; + assert!(producer_id < K_PRODUCERS, "bad producer id {producer_id}"); + assert!( + sent_idx < M_PER_PRODUCER as usize, + "bad sent idx {sent_idx}" + ); + let already = std::mem::replace(&mut seen[producer_id][sent_idx], true); + assert!(!already, "duplicate ({producer_id}, {sent_idx})"); + consumed.fetch_add(1, O::Relaxed); + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach mid-drain"), + } + } + for h in handles { + h.join().unwrap(); + } + for (p, row) in seen.iter().enumerate() { + assert!( + row.iter().all(|&b| b), + "producer {p} has a missed message: row = {row:?}" + ); + } + } + + /// Per-producer in-order property: a single producer's messages observed by the + /// consumer must arrive in the order the producer sent them. (Cross-producer + /// ordering is not guaranteed.) + #[test] + fn mpsc_preserves_per_producer_order() { + const K_PRODUCERS: usize = 4; + const M_PER_PRODUCER: u32 = 500; + let (_region, rx, tx_template) = make_ring(32, 32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + if tx.try_send(&payload).is_ok() { + sent += 1; + } else { + std::thread::yield_now(); + } + } + })); + } + let mut last_seq: Vec = vec![-1; K_PRODUCERS]; + let mut buf = Vec::new(); + let mut total = 0usize; + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = i64::from(u32::from_le_bytes(buf[4..8].try_into().unwrap())); + assert_eq!( + sent_idx, + last_seq[producer_id] + 1, + "out-of-order from producer {producer_id}: got {sent_idx}, expected {}", + last_seq[producer_id] + 1 + ); + last_seq[producer_id] = sent_idx; + total += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// Cache-line layout regression: `head` must land on its own cache line + /// (offset == CACHE_LINE so the preceding header fields can't false-share), + /// `tail` separated from `head` by CACHE_LINE, and the first slot not sharing a + /// line with `tail`. Catches accidental padding removal or field reorder that + /// would re-introduce the false-sharing perf cliff the padding exists to avoid. + #[test] + fn header_layout_is_cache_friendly() { + use std::mem::offset_of; + let head_off = offset_of!(DsmMpscRingHeader, head); + let tail_off = offset_of!(DsmMpscRingHeader, tail); + let header_size = std::mem::size_of::(); + // head at exactly CACHE_LINE offset means the preceding 64 bytes (header + // fields + pad) live entirely in cache line 0, head + its pad live in + // cache line 1, etc. Anything other than equality means the padding + // math drifted; fix it rather than relaxing the assertion. + assert_eq!( + head_off, CACHE_LINE, + "head must land at offset {CACHE_LINE} (its own cache line); got {head_off}" + ); + assert!( + (tail_off - head_off) >= CACHE_LINE, + "head and tail must be on separate cache lines: head_off={head_off}, tail_off={tail_off}, CACHE_LINE={CACHE_LINE}" + ); + assert!( + (header_size - tail_off) >= CACHE_LINE, + "first slot must start on its own cache line: header_size={header_size}, tail_off={tail_off}, CACHE_LINE={CACHE_LINE}" + ); + // Header's natural alignment is determined by its largest field (AtomicU64, + // 8 bytes). Cache-line padding between hot fields still keeps them on separate + // absolute lines regardless of the struct's starting alignment, as long as the + // inter-field distance is >= CACHE_LINE. + assert_eq!(std::mem::align_of::(), 8); + } + + /// Magic + version validation: an attach against a region with the wrong magic or + /// version returns None rather than handing back a NonNull with garbage state. + /// Mirrors `MppDsmHeader::validate`'s discipline. + #[test] + fn attach_at_rejects_wrong_magic_and_version() { + let bytes = DsmMpscRingHeader::region_bytes(4, 64); + let region = AlignedRegion::new(bytes); + let header_ptr = unsafe { create_at(region.as_mut_ptr(), 4, 64) }; + // Sanity: correct attach succeeds. + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_some()); + // Corrupt magic: rejected. + unsafe { (*header_ptr).magic = 0xDEADBEEF }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_none()); + // Restore magic, corrupt version. + unsafe { (*header_ptr).magic = MPSC_RING_MAGIC }; + unsafe { (*header_ptr).version = MPSC_RING_VERSION.wrapping_add(1) }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_none()); + // Restore version, mismatched ring_size. + unsafe { (*header_ptr).version = MPSC_RING_VERSION }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 8, 64) }.is_none()); + } + + /// Stress test at the production-worst contention level (K=24 producers, matching + /// the largest mesh-row size the transport drives in practice). Smoke that the + /// primitive doesn't wedge or lose messages under heavy CAS contention; doesn't + /// measure perf. + #[test] + fn mpsc_stress_at_production_worst_case() { + const K_PRODUCERS: usize = 24; + const M_PER_PRODUCER: u32 = 500; + let (_region, rx, tx_template) = make_ring(64, 32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected: {e:?}"), + } + } + })); + } + let mut buf = Vec::new(); + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + let mut total = 0usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => total += 1, + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + // Multi-round invariant: after draining 24 * 500 = 12000 messages on a 64-slot + // ring, every slot's seq must have advanced into a future round. Walk slot[0]: + // we expect seq = head + ring_size = 12000 + 64 - whatever round 0 it's in. + // Loose check: every slot's seq is bounded below by ring_size (round >= 1). + let header = unsafe { ring_nn.0.as_ref() }; + for i in 0..header.ring_size { + let slot = unsafe { slot_ptr(ring_nn.0.as_ptr(), i, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(O::Acquire) }; + assert!( + seq >= header.ring_size as u64, + "slot[{i}].seq={seq} < ring_size; slot was never reused" + ); + } + } + + // ----- multi-slot fragmentation tests ----- + + /// Two-slot frame round-trip: send a payload that's exactly `2*payload_cap` + /// bytes long with distinct first-half and second-half markers, then verify the + /// receiver reassembles them in order with no truncation or duplication. + #[test] + fn multi_slot_two_fragment_round_trip() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(4, slot_cap as u32); + let mut frame = vec![0u8; 2 * payload_cap]; + for (i, b) in frame.iter_mut().enumerate() { + *b = (i & 0xFF) as u8; + } + tx.try_send(&frame).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf.len(), frame.len()); + assert_eq!(buf, frame); + // Ring fully drained. + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Frame at the multi-slot ceiling: needs exactly `ring_size` slots, succeeds; + /// one byte more, rejected with MessageTooLarge. + #[test] + fn multi_slot_max_size_succeeds_one_more_rejected() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let ring_size = 4; + let (_region, rx, tx) = make_ring(ring_size, slot_cap as u32); + let max = vec![0xABu8; ring_size as usize * payload_cap]; + tx.try_send(&max).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, max); + let too_big = vec![0u8; ring_size as usize * payload_cap + 1]; + assert_eq!(tx.try_send(&too_big), Err(SendError::MessageTooLarge)); + } + + /// Multi-slot followed by single-slot from the same producer: each frame must + /// come out of the receiver as a separate, intact byte sequence; fragments + /// don't leak into the following frame. + #[test] + fn multi_slot_then_single_slot_preserves_boundaries() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(8, slot_cap as u32); + let big = vec![0xAAu8; 3 * payload_cap]; + tx.try_send(&big).unwrap(); + let small = b"hi".to_vec(); + tx.try_send(&small).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, small); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Multi-slot frame wrapping the ring boundary. With ring_size=4, tail starts at 0: + /// send a single-slot first to advance tail to 1, drain it, then send a 3-slot + /// frame at tail=1 occupying slots 1, 2, 3, drain the small frame to advance head, + /// then send another single-slot at tail=4 (slot 0 in round 1) and a 3-slot at + /// tail=5 (slots 1, 2, 3 in round 1). Verifies wraparound math holds. + #[test] + fn multi_slot_wraparound() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(4, slot_cap as u32); + let mut buf = Vec::new(); + // Push a small frame so the next multi-slot reservation starts mid-ring. + tx.try_send(b"warm").unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, b"warm"); + // 3-slot frame at tail=1, wrapping at slot 3 -> slot 0 won't happen here + // (slots 1, 2, 3 are still in round 0). Drain it. + let big = vec![0x11u8; 3 * payload_cap]; + tx.try_send(&big).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big); + // Now tail=4, head=4. Send another small frame at slot 0 (round 1, seq=4). + tx.try_send(b"two").unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, b"two"); + // And a 3-slot at tail=5 occupying slots 1, 2, 3 in round 1 (seqs 5, 6, 7). + let big2 = vec![0x22u8; 3 * payload_cap]; + tx.try_send(&big2).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big2); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Stress: multiple producers each push a mix of single-slot and multi-slot + /// frames; consumer verifies every frame's contents and per-producer ordering. + /// Catches interleave bugs (a multi-slot run getting mixed with another + /// producer's fragments) and missing wake-ups under high contention. + #[test] + fn multi_slot_stress_mixed_with_singles() { + const K_PRODUCERS: usize = 6; + const M_PER_PRODUCER: u32 = 300; + let slot_cap = 128; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx_template) = make_ring(16, slot_cap as u32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + // Alternate single, 2-slot, 4-slot, single, ... so the consumer + // sees every fragment-kind combination. + let n_payload = match sent % 3 { + 0 => 8, // single-slot + 1 => 2 * payload_cap, // 2-slot + _ => 4 * payload_cap, // 4-slot + }; + let mut payload = vec![0u8; n_payload + 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + for (i, b) in payload[8..].iter_mut().enumerate() { + *b = ((producer_id ^ sent as usize ^ i) & 0xFF) as u8; + } + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected send error: {e:?}"), + } + } + })); + } + let mut last_seq: Vec = vec![-1; K_PRODUCERS]; + let mut buf = Vec::new(); + let mut total = 0usize; + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + assert!(buf.len() >= 8, "frame too short: {}", buf.len()); + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = i64::from(u32::from_le_bytes(buf[4..8].try_into().unwrap())); + assert!(producer_id < K_PRODUCERS, "bad producer id"); + assert_eq!( + sent_idx, + last_seq[producer_id] + 1, + "out-of-order frame from producer {producer_id}" + ); + // Verify payload bytes (catches fragment interleave / partial reads). + for (i, b) in buf[8..].iter().enumerate() { + let expected = ((producer_id ^ sent_idx as usize ^ i) & 0xFF) as u8; + assert_eq!(*b, expected, "payload mismatch at byte {i}"); + } + last_seq[producer_id] = sent_idx; + total += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// The injected `Wakeup` carries a real cross-thread notification with no PG: a consumer that + /// parks until woken is unparked by the producer's publish, and the producer actually invoked + /// the wakeup (asserted via the counter, so a silently-broken extension point fails rather than relying + /// on park timing). + #[test] + fn injected_wakeup_unparks_blocked_consumer() { + let region = AlignedRegion::new(DsmMpscRingHeader::region_bytes(8, 256)); + let nn = NonNull::new(unsafe { create_at(region.as_mut_ptr(), 8, 256) }) + .expect("create_at returned null"); + let ring = SharedRing(nn); + let wakeup = Arc::new(ThreadWakeup::default()); + let alive = Arc::new(AtomicBool::new(true)); + const TOKEN: u64 = 7; + + let receiver = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + // Register this thread as the wake target before publishing the token, so a producer that + // races ahead still finds us. + wakeup.register(TOKEN, std::thread::current()); + receiver.set_receiver(TOKEN); + + let producer = { + let wakeup = Arc::clone(&wakeup) as Arc; + let alive = Arc::clone(&alive); + std::thread::spawn(move || { + // Rebind the whole `Send` wrapper so the closure captures it (not the bare + // `NonNull` field, which edition-2024 disjoint capture would otherwise grab). + let ring = ring; + let tx = unsafe { DsmMpscSender::new(ring.0, wakeup, alive) }; + std::thread::sleep(std::time::Duration::from_millis(20)); + assert_eq!(tx.try_send(b"hello"), Ok(())); + }) + }; + + // park_timeout (not park) so a broken extension point fails via the counter assertion below instead + // of hanging the suite. park/unpark holds a token, so an unpark that beats our park still + // wakes the next park (no lost-wakeup). + let mut out = Vec::new(); + loop { + match receiver.try_recv(&mut out) { + RecvOutcome::Bytes => break, + RecvOutcome::Empty => std::thread::park_timeout(std::time::Duration::from_secs(5)), + RecvOutcome::Detached => panic!("detached before receiving the frame"), + } + } + producer.join().unwrap(); + assert_eq!(out, b"hello"); + assert!( + wakeup.wakes.load(O::Relaxed) >= 1, + "producer never invoked the injected wakeup" + ); + } +} diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs new file mode 100644 index 00000000..5f538d83 --- /dev/null +++ b/src/shm/runtime.rs @@ -0,0 +1,477 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Runtime glue between the leader's DataFusion execution and the DSM MPSC mesh. +//! +//! [`MppMesh`] is the runtime handle the leader builds at DSM-init time. It holds the +//! single [`super::transport::DrainHandle`] (the +//! `inbound_receiver`) that consolidates this proc's DSM inbox and self-loop, and gets +//! installed on the leader's `SessionConfig` extensions before plan execution. +//! +//! [`ShmChannelResolver`] implements [`ChannelResolver`], consulted by +//! `NetworkShuffleExec` / `NetworkCoalesceExec` / `NetworkBroadcastExec` at execute time. It hands +//! out a [`ShmWorkerChannel`] whose `execute_task(task_key, partition_range)` yields one stream per +//! consumer partition from the shared `inbound_receiver`. +//! +//! [`InProcessWorkerResolver`] hands the planner `n_workers` placeholder URLs. The transport +//! routes by task index, not URL, so the URLs are never dialed; the resolver exists only because +//! the planner sizes stages from the URL count. It replaces the placeholder URL the fork used to +//! substitute internally under the old `in_process_mode` flag. + +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; +use futures::stream::{self, BoxStream, StreamExt}; +use http::HeaderMap; +use url::Url; + +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedTxs; +use crate::{ + ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, + GetWorkerInfoResponse, WorkerChannel, WorkerResolver, WorkerToCoordinatorMsg, +}; + +use super::AliveFlag; +use super::transport::{CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppSender}; + +/// A proc's outbound senders to each peer inbox, shared between the mesh (for `Cancel` frames) and +/// the embedder that owns their lifetime. Indexed by destination `proc_idx`; this proc's own slot +/// is `None`. +pub type PeerSenders = Arc>>>; + +/// `task_idx → proc_idx` round-robin over the worker procs. The leader is `proc_idx = 0` +/// (consumer-only), workers are `1..n_procs` (each hosts producer fragments). +/// +/// A stage's task count is set by the DF-D task estimator chain, not by the worker proc count. +/// The transport does NOT support more tasks than producer procs: channels are keyed +/// `(sender_proc, stage, partition)`, so two tasks wrapped onto one proc would interleave on one +/// channel and the first EOF would truncate the other task's output. `ShmWorkerChannel::execute_task` +/// rejects that shape; the modulo here only keeps the function total. +#[inline] +pub fn proc_for_task(n_workers: u32, task_idx: u32) -> u32 { + 1 + (task_idx % n_workers.max(1)) +} + +/// Runtime handle the customscan populates at DSM-init time. +/// +/// Each process owns one MPSC inbox in DSM that receives frames from every peer. +/// `inbound_receiver` consolidates that inbox plus the in-proc self-loop channel (for +/// producer-and-consumer-on-same-worker fragments) into a single `DrainHandle`. Frames +/// carry `(sender_proc, stage_id, partition)` in their header so the routing registry +/// inside the handle delivers each frame to the matching consumer. +/// +/// [`MppFrameHeader`]: super::transport::MppFrameHeader +pub struct MppMesh { + /// This process's `proc_idx` (= 0 for the leader, `ParallelWorkerNumber + 1` for workers). + /// Frames addressed to this proc arrive on this proc's own inbox. + pub this_proc: u32, + /// Total proc count. Bounds the producer/consumer proc lookups in + /// [`ShmWorkerChannel::execute_task`]. + pub n_procs: u32, + /// Single cooperative inbound handle pulling every frame addressed to this proc. The + /// DSM MPSC inbox and an in-proc self-loop receiver both feed into this handle. Demux + /// to per-`(sender_proc, stage_id, partition)` channel buffers happens inside via + /// `DrainHandle::register_channel`. + pub(super) inbound_receiver: Arc, + /// Cancellation hook, injected by the embedder, checked at the transport's block points (the + /// cooperative send spin and the consumer pull loop). + interrupt: Arc, + /// This proc's outbound senders to each peer inbox, shared with the embedder that owns their + /// lifetime (it clears the `Vec` before the DSM unmaps). Used by [`Self::cancel_stream`] to ship + /// a `Cancel` frame to a producing proc when this proc abandons a stream. `None` until the + /// embedder installs it. + cancel_senders: Mutex>, + alive: AliveFlag, +} + +impl MppMesh { + /// Build a fresh mesh. + pub fn new( + this_proc: u32, + n_procs: u32, + inbound_receiver: Arc, + interrupt: Arc, + alive: AliveFlag, + ) -> Self { + Self { + this_proc, + n_procs, + inbound_receiver, + interrupt, + cancel_senders: Mutex::new(None), + alive, + } + } + + /// Mark this proc's DSM mesh detached so every ring handle's operations become no-ops. + /// The embedder calls this from its dsm-detach callback, while the segment is still mapped, + /// so handles dropped afterward (e.g. by a memory-context reset) never touch freed memory. + pub fn mark_detached(&self) { + self.alive.store(false, Ordering::Release); + } + + /// The raw liveness flag shared with every ring handle, for embedders that register a C + /// dsm-detach callback against it directly. + pub fn detached_flag(&self) -> AliveFlag { + Arc::clone(&self.alive) + } + + /// Share this proc's outbound senders so [`Self::cancel_stream`] can reach the producers. The + /// embedder keeps owning their lifetime: it passes a clone of the same `Arc` it releases before + /// the DSM unmaps, so the mesh never extends the senders past teardown. + pub fn set_cancel_senders(&self, senders: PeerSenders) { + *self.cancel_senders.lock().unwrap() = Some(senders); + } + + /// Tell the producer on `producer_proc` to stop the `(stage_id, partition)` stream: this proc's + /// consumer of it stopped reading before EOF. Ships one `Cancel` frame, leaving the rings + /// healthy for metrics and every other stream. A no-op when no senders are installed (the + /// embedder hasn't wired this proc, or teardown cleared them). + /// + /// Stream-level so any consumer can cancel its own input: every `(producer_proc, stage, + /// partition)` channel has a single consumer, so one stream's drop never cuts off a sibling's. + pub fn cancel_stream(&self, producer_proc: u32, stage_id: u32, partition: u32) { + let guard = self.cancel_senders.lock().unwrap(); + let Some(senders) = guard.as_ref() else { + return; + }; + let senders = senders.lock().unwrap(); + if let Some(Some(sender)) = senders.get(producer_proc as usize) { + sender.try_send_cancel(stage_id, partition); + } + } + + /// The single cooperative inbound handle that pulls frames from every peer (and the + /// self-loop) into per-`(sender_proc, stage_id, partition)` channel buffers. + pub(super) fn inbound_receiver(&self) -> &Arc { + &self.inbound_receiver + } + + /// Install the senders of one task's work-unit feed channels on this proc's drain, so + /// inbound `WorkUnit` frames for `(stage_id, task_number)` flow into them. Units that + /// arrived first are flushed; a `FeedEof` that already came through closes the channels + /// immediately. The drain only fills channels: something on this proc must keep draining + /// (a consumer pull loop, a producer send spin, or an explicit + /// [`CooperativeDrainSet::try_drain_pass`] pump) or a fragment blocked on its feed starves. + pub fn register_work_unit_senders( + &self, + stage_id: u32, + task_number: u32, + senders: RemoteWorkUnitFeedTxs, + ) { + self.inbound_receiver + .register_work_unit_senders(stage_id, task_number, senders); + } + + /// Take the plan delivered for `(stage_id, task_number)` as a `SetPlan` frame on this proc's + /// inbox, waiting for it if it has not arrived yet. Something on this proc must keep + /// draining (a pump, a consumer pull loop, or a cooperative send spin) or the wait starves. + pub async fn take_set_plan( + &self, + stage_id: u32, + task_number: u32, + ) -> Result { + self.inbound_receiver + .take_set_plan(stage_id, task_number) + .await + } + + /// Take the stream of `TaskMetrics` frames arriving on this proc's inbox: + /// `(stage_id, task_number, metrics)` per producer task that reported in. Meant for the + /// leader; the first caller gets it, later calls get `None`. + pub fn take_task_metrics_receiver( + &self, + ) -> Option> { + self.inbound_receiver.take_task_metrics_receiver() + } + + /// Number of worker procs (= `n_procs - 1`, since the leader is proc 0). Used as the + /// modulus in [`proc_for_task`]. + /// + /// The embedder must guarantee `n_procs >= 2` (one consumer-only leader plus at least one + /// producer worker) before constructing an `MppMesh`; the subtraction is otherwise unsound. + /// `compute_dsm_layout` enforces the same bound. Asserted in debug builds so a future misuse + /// fails loudly. + pub fn n_workers(&self) -> u32 { + debug_assert!( + self.n_procs >= 2, + "MppMesh::n_workers() called with n_procs={} (< 2); the embedder must gate on \ + n_procs >= 2", + self.n_procs + ); + self.n_procs - 1 + } + + /// Pull from the single inbound handle. Called from + /// [`super::transport::MppSender`]'s cooperative-send spin so a + /// producer stalled on a full outbound can pull inbound peer data inline. That's what + /// prevents the symmetric-send deadlock when every peer is simultaneously stalled waiting + /// for space. + pub(super) fn drain_all_inbound(&self) -> Result<(), DataFusionError> { + self.inbound_receiver.try_drain_pass() + } +} + +impl CooperativeDrainSet for MppMesh { + fn try_drain_pass(&self) -> Result<(), DataFusionError> { + self.drain_all_inbound() + } + + fn check_interrupt(&self) -> Result<(), DataFusionError> { + self.interrupt.check() + } + + fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + self.inbound_receiver.stream_cancelled(stage_id, partition) + } +} + +/// Hands out [`ShmWorkerChannel`]s over the leader's [`MppMesh`]. The shm mesh has no URLs to dial, +/// so every `get_worker_client_for_url` resolves to the same mesh; the channel routes by the +/// `task_key` it is handed at `execute_task`, not by URL. +pub struct ShmChannelResolver { + mesh: Arc, +} + +impl ShmChannelResolver { + pub fn new(mesh: Arc) -> Self { + Self { mesh } + } +} + +#[async_trait] +impl ChannelResolver for ShmChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + Ok(Box::new(ShmWorkerChannel { + mesh: Arc::clone(&self.mesh), + })) + } +} + +/// A [`WorkerChannel`] over the mesh. `execute_task((stage, task), partition_range)` translates the +/// DF-D `(stage, task)` addressing into the proc-pair grid: `proc_for_task(n_workers, task_number)` +/// selects which `sender_proc` hosts the producer-side task, and each returned stream pulls that +/// proc's `(sender_proc, stage_id, partition)` slice from the inbound drain. +struct ShmWorkerChannel { + mesh: Arc, +} + +#[async_trait] +impl WorkerChannel for ShmWorkerChannel { + /// pg_search delivers each worker's plan over DSM, not over this channel, so plan delivery is a + /// no-op (what the old `NoOpDispatch` did). Drain the inbound control stream to exhaustion so the + /// coordinator's keep-alive tail does not wedge, then complete with an empty output stream: task + /// metrics travel back over the mesh, not here. + async fn coordinator_channel( + &mut self, + _headers: HeaderMap, + mut c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { while c2w_stream.next().await.is_some() {} }); + Ok(stream::empty().boxed()) + } + + async fn execute_task( + &mut self, + _headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + _task_ctx: &Arc, + ) -> Result>>> { + MetricBuilder::new(&metrics) + .global_counter("mesh_connections_used") + .add(1); + let stage_id = u32::try_from(request.task_key.stage_id).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: stage_id={} > u32::MAX", + request.task_key.stage_id + )) + })?; + let task_number = u32::try_from(request.task_key.task_number).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: task_number={} > u32::MAX", + request.task_key.task_number + )) + })?; + let sender_proc = proc_for_task(self.mesh.n_workers(), task_number); + // More tasks than producer procs would fold two tasks onto one + // `(sender_proc, stage, partition)` channel: interleaved batches, and the first task's EOF + // truncates the second. With no input_stage here to count tasks, the equivalent guard is + // that the routed proc stays in range, as the old code also checked. + if sender_proc >= self.mesh.n_procs { + return Err(DataFusionError::Internal(format!( + "ShmWorkerChannel: sender_proc={sender_proc} >= n_procs={} \ + (stage_id={stage_id}, task_number={task_number})", + self.mesh.n_procs + ))); + } + // First line to grep when a query hangs on a channel nothing registered. + log::debug!( + "shm transport execute_task: this_proc={} stage_id={stage_id} \ + task_number={task_number} -> sender_proc={sender_proc}", + self.mesh.this_proc + ); + let mut streams = Vec::with_capacity( + request + .target_partition_end + .saturating_sub(request.target_partition_start), + ); + for partition in request.target_partition_start..request.target_partition_end { + let partition_u32 = u32::try_from(partition).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: partition={partition} > u32::MAX" + )) + })?; + streams.push(pull_partition_stream( + Arc::clone(&self.mesh), + sender_proc, + stage_id, + partition_u32, + )); + } + Ok(streams) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + Ok(GetWorkerInfoResponse { + version: String::new(), + }) + } +} + +/// Build the cooperative pull-loop stream for one `(producer_proc, stage_id, partition)` channel. +/// +/// The inbound drain runs inline on the consumer's thread. Each iteration checks for cancellation +/// (via the injected interrupt extension point), drains the receiver into the registry, pops one +/// batch to yield, then yields back to Tokio so sibling tasks (e.g. the leader's own producer +/// subplan) advance. The send side runs the same interrupt check inside `MppSender`'s retry spin. +fn pull_partition_stream( + mesh: Arc, + producer_proc: u32, + stage_id: u32, + partition: u32, +) -> BoxStream<'static, Result> { + // One drain per process, shared across all sender_procs. The channel-buffer registry keys by + // (sender_proc, stage_id, partition) so this consumer still sees only its named sender's slice + // even though the underlying inbox is shared with all peers. + let drain = Arc::clone(mesh.inbound_receiver()); + log::debug!( + "shm transport execute: register channel sender_proc={producer_proc} stage_id={stage_id} \ + partition={partition}" + ); + let buffer = drain.register_channel(producer_proc, stage_id, partition); + let stream = async_stream::stream! { + // Any consumer cancels its own input stream when it drops early. The mesh no-ops the send + // until the embedder wires this proc's outbound senders. + let mut cancel_guard = StreamCancelGuard { + mesh: Arc::clone(&mesh), + producer_proc, + stage_id, + partition, + armed: true, + }; + loop { + if let Err(e) = mesh.check_interrupt() { + yield Err(e); + return; + } + if let Err(e) = drain.try_drain_pass() { + yield Err(e); + return; + } + match buffer.try_pop() { + Some(DrainItem::Batch(batch)) => yield Ok(batch), + Some(DrainItem::Eof) => { + // Clean end: the producer already EOF'd, so there's nothing to cancel. + cancel_guard.armed = false; + break; + } + Some(DrainItem::Failed(msg)) => { + yield Err(DataFusionError::Execution(msg)); + return; + } + None => tokio::task::yield_now().await, + } + } + }; + Box::pin(stream) +} + +/// Placeholder worker resolver for the in-process MPP transport. +/// +/// The shm_mq transport routes by `target_task` (proc index), never by URL, so these URLs are +/// never dialed. The distributed planner still needs a resolver: it sizes stages and assigns +/// tasks from the URL count. `n_workers` placeholder URLs is exactly what the planner needs. This +/// replaces the placeholder URL the fork used to substitute internally under `in_process_mode`. +pub struct InProcessWorkerResolver { + n_workers: usize, +} + +impl InProcessWorkerResolver { + pub fn new(n_workers: usize) -> Self { + Self { n_workers } + } +} + +impl WorkerResolver for InProcessWorkerResolver { + fn get_urls(&self) -> Result, DataFusionError> { + (0..self.n_workers.max(1)) + .map(|i| { + Url::parse(&format!("inprocess://worker/{i}")).map_err(|e| { + DataFusionError::Internal(format!( + "InProcessWorkerResolver: invalid placeholder url: {e}" + )) + }) + }) + .collect() + } +} + +/// Cancels one `(stage_id, partition)` stream if its consumer drops before EOF, telling the +/// producer on `producer_proc` to stop. A consumer that stops pulling early (a top-N `LIMIT`, an +/// inner merge join exhausting a side, etc.) would otherwise leave that producer spinning on the +/// full inbox until the statement timeout. Disarmed on a clean EOF: there the producer already +/// finished, so there's nothing to cancel. +struct StreamCancelGuard { + mesh: Arc, + producer_proc: u32, + stage_id: u32, + partition: u32, + armed: bool, +} + +impl Drop for StreamCancelGuard { + fn drop(&mut self) { + if self.armed { + self.mesh + .cancel_stream(self.producer_proc, self.stage_id, self.partition); + } + } +} diff --git a/src/shm/self_hosted.rs b/src/shm/self_hosted.rs new file mode 100644 index 00000000..192b77a9 --- /dev/null +++ b/src/shm/self_hosted.rs @@ -0,0 +1,1277 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Deferred: `mod self_hosted` is gated out (`#[cfg(any())]`) in `shm/mod.rs`. This file still +//! targets the removed `WorkerTransport`/`WorkerDispatch` dispatch umbrella, which the +//! `ChannelResolver`/`WorkerChannel` protocol has no analog for. Its no-gRPC-default role is now +//! served by `InProcessChannelResolver`, and the in-crate ring safety net by the `in_process` test, +//! so it stays unported until reimplemented as a `ChannelResolver` driving produce over +//! `coordinator_channel`. Kept for reference; it does not compile against this branch. +//! +//! The shared-memory transport hosting its own workers, so it can serve as a session default. +//! +//! Production embedders drive the mesh themselves: they allocate the region, launch the worker +//! processes, and deliver plans out of band, so [`ShmMqWorkerTransport`]'s dispatcher is a no-op. +//! That shape cannot be a default transport: a default gets nothing but the `WorkerTransport` +//! calls. [`SelfHostedShmTransport`] closes that gap by playing the embedder itself, in-process: +//! dispatch delivers each task's plan through the worker session machinery +//! ([`Worker::set_task_plan`]: codec round-trip, work-unit feed channels, metrics back-channel) +//! and runs the producer fragments as tasks pushing through a heap-backed DSM mesh; `open` reads +//! the rings from the leader side. Every cross-stage byte moves through the same rings, framing, +//! and cooperative drain a production embedder uses. +//! +//! Per query, the harness lives from the first dispatch to the cancellation token firing (the +//! head stream dropping). The mesh is sized and built lazily at the first `open`, once every +//! stage has been dispatched and the routing is known. + +use std::alloc::Layout; +use std::ffi::c_void; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; + +use dashmap::DashMap; +use datafusion::common::instant::Instant; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, HashMap, Result, exec_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::ExecutionPlan; +use futures::StreamExt; +use futures::stream::BoxStream; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::mpsc_ring::Wakeup; +use super::runtime::{MppMesh, ShmMqWorkerTransport, proc_for_task}; +use super::setup::{dsm_region_bytes, leader_setup, worker_setup}; +use super::transport::{ + CooperativeDrainSet, Interrupt, MppFrameHeader, MppPartitionSink, MppSender, SendBatchStats, + SetPlanFrame, +}; +use crate::proto as pb; +use crate::{ + CoordinatorToWorkerMetrics, DistributedTaskContext, MetricsStore, NetworkBoundaryExt, + NetworkCoalesceExec, PartitionSink, ProducerHead, RemoteStage, TreeNodeExt, Worker, + WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerSessionBuilder, WorkerTransport, + collect_task_work_unit_feeds, encode_task_plan, execute_local_task, + get_config_extension_propagation_headers, get_distributed_cancellation_token, + get_passthrough_headers, serialize_uuid, set_distributed_worker_transport, set_sent_time, +}; + +/// Per-inbox ring size. Frames fragment across slots, so this bounds a single batch at roughly +/// the whole ring; the suite's batches sit well under it while keeping the per-query footprint +/// (`n_procs` inboxes) modest. +const SELF_HOSTED_QUEUE_BYTES: usize = 1 << 22; + +/// No-op wakeup: the cooperative consumers yield instead of parking, so a publish never has to +/// wake a blocked thread. +pub(super) struct NoopWakeup; +impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} +} + +/// Opaque, non-sentinel receiver token. The wakeup ignores the value; this just keeps the +/// producer from treating the consumer as unregistered. +pub(super) fn receiver_token(proc_idx: u32) -> u64 { + proc_idx as u64 + 1 +} + +/// Owns the heap buffer standing in for a shared-memory segment. Every proc reads and writes the +/// same region through raw pointers; the lock-free rings make the concurrent access sound. +pub(super) struct HeapRegion { + ptr: *mut u8, + layout: Layout, +} + +impl HeapRegion { + pub(super) fn new(bytes: usize) -> Self { + // 64-byte alignment so each per-inbox ring header lands on its own cache line; the + // dsm layout aligns the offsets within the region, but only if the base is aligned too. + let layout = Layout::from_size_align(bytes, 64).expect("dsm region layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "dsm region alloc failed"); + Self { ptr, layout } + } + + pub(super) fn base(&self) -> *mut c_void { + self.ptr as *mut c_void + } +} + +impl Drop for HeapRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } +} + +// The raw pointer is only dereferenced through the rings, which are designed for concurrent +// access from multiple procs. +unsafe impl Send for HeapRegion {} +unsafe impl Sync for HeapRegion {} + +/// Interrupt extension point wired to the query's cancellation token, so producers blocked in a send spin +/// and consumers in the pull loop both unwind when the head stream drops. +struct CancellationInterrupt(CancellationToken); +impl Interrupt for CancellationInterrupt { + fn check(&self) -> Result<(), DataFusionError> { + if self.0.is_cancelled() { + return Err(DataFusionError::Execution( + "mpp: query cancelled".to_string(), + )); + } + Ok(()) + } +} + +/// [`WorkerTransport`] over the shared-memory mesh, hosting its own workers in-process. +/// +/// All tasks share one [Worker] (one task registry, one session builder), like the in-memory +/// transport; what differs is the data plane: producer fragments run eagerly as background tasks +/// pushing through DSM rings, and reads pull from those rings instead of executing lazily. +/// +/// With the `flight` feature off this is the default transport. Multi-process embedders keep +/// driving [`ShmMqWorkerTransport`] directly. +#[derive(Clone)] +pub struct SelfHostedShmTransport { + worker: Worker, + queries: Arc>>, + queue_bytes: usize, +} + +impl Default for SelfHostedShmTransport { + fn default() -> Self { + Self::new(Worker::default()) + } +} + +impl SelfHostedShmTransport { + /// Builds the transport around an existing [Worker], sharing its task registry, session + /// builder, and runtime environment. + pub fn new(worker: Worker) -> Self { + Self { + worker, + queries: Arc::new(DashMap::new()), + queue_bytes: SELF_HOSTED_QUEUE_BYTES, + } + } + + /// Overrides the per-inbox ring size. Small values force multi-slot fragmentation and the + /// cooperative send spin on every query, which is how the ring mechanics get stress-tested; + /// the default is generous enough that only large batches touch them. + pub fn with_queue_bytes(mut self, queue_bytes: usize) -> Self { + self.queue_bytes = queue_bytes; + self + } + + /// Builds the transport with a custom [WorkerSessionBuilder], the same customization hook a + /// remote worker offers. + pub fn from_session_builder( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + Self::new(Worker::from_session_builder(session_builder)) + } + + /// The in-process [Worker] backing this transport. + pub fn worker(&self) -> &Worker { + &self.worker + } +} + +/// Recovers the producer head the planner already inserted at the top of a stage's shipped +/// fragment. The worker's `insert_producer_head` strips an existing head and re-adds one from the +/// [`ProducerHead`] it is given, so a `None` here would flatten a sliced stage to one partition; +/// echoing the existing head back keeps it intact. +fn producer_head_from_plan(plan: &Arc) -> ProducerHead { + use datafusion::physical_plan::repartition::RepartitionExec; + if let Some(r) = plan.downcast_ref::() { + ProducerHead::RepartitionExec { + partitioning: r.partitioning().clone(), + } + } else if let Some(b) = plan.downcast_ref::() { + ProducerHead::BroadcastExec { + output_partitions: b.properties().partitioning.partition_count(), + } + } else { + ProducerHead::None + } +} + +impl WorkerTransport for SelfHostedShmTransport { + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: std::ops::Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let Some(harness) = self + .queries + .get(&input_stage.query_id) + .map(|e| Arc::clone(e.value())) + else { + return exec_err!( + "self-hosted shm transport: no harness for query {}; stage {} was never \ + dispatched through this transport", + input_stage.query_id, + input_stage.num + ); + }; + // The first read finalizes the harness: by now every stage has been dispatched (plan + // preparation completes before the head executes), so the mesh can be sized and the + // producer drivers released. + harness.ensure_started()?; + let leader_mesh = harness.leader_mesh()?; + let inner = ShmMqWorkerTransport::new(leader_mesh).open( + input_stage, + target_partitions, + target_task, + producer_head, + ctx, + metrics, + )?; + Ok(Box::new(PinnedConnection { inner, harness })) + } + + fn dispatcher(&self) -> Box { + Box::new(SelfHostedDispatcher { + transport: self.clone(), + metrics: OnceLock::new(), + }) + } +} + +/// Keeps the query harness (and through it the heap region the rings live in) alive for as long +/// as any stream is still reading from the mesh. +struct PinnedConnection { + inner: Box, + harness: Arc, +} + +impl WorkerConnection for PinnedConnection { + fn execute( + &self, + partition: usize, + ) -> Result>> { + let stream = self.inner.execute(partition)?; + let harness = Arc::clone(&self.harness); + Ok(Box::pin(stream.map(move |item| { + let _ = &harness; // <- the region must outlive the ring receivers. + item + }))) + } +} + +/// Per-query plan-delivery state. As with the other transports, the plan-send metrics and the +/// query start timestamp live for the whole query. +struct SelfHostedDispatcher { + transport: SelfHostedShmTransport, + metrics: OnceLock, +} + +impl WorkerDispatch for SelfHostedDispatcher { + fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()> { + let WorkerDispatchRequest { + stage, + routed_urls, + task_ctx, + metrics, + metrics_store, + join_set, + .. + } = request; + let metrics = self + .metrics + .get_or_init(|| CoordinatorToWorkerMetrics::new(metrics)) + .clone(); + + let token = get_distributed_cancellation_token(task_ctx); + let harness = match self.transport.queries.entry(stage.query_id) { + dashmap::Entry::Occupied(e) => Arc::clone(e.get()), + dashmap::Entry::Vacant(e) => { + let harness = Arc::new(QueryHarness::new( + stage.query_id, + token.clone(), + self.transport.queue_bytes, + metrics_store.cloned(), + )); + e.insert(Arc::clone(&harness)); + // The token fires when the head stream drops (normal completion included), which + // is the query's end of life; drop the registry entry then. The region itself + // stays alive through the Arcs the drivers and streams hold. + let queries = Arc::clone(&self.transport.queries); + let query_id = stage.query_id; + let watched = token.clone(); + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + watched.cancelled().await; + queries.remove(&query_id); + }); + // One pump set per query: drains every proc's inbox so control frames flow even + // when no consumer or producer is actively draining (a fragment blocked on its + // feed, the metrics frames after the gather finished), and forwards the task + // metrics into the store. Lives until every driver and feed pump reported done. + join_set.spawn(run_pumps(Arc::clone(&harness), token.clone())); + // Plans travel as `SetPlan` frames like every other control message, but the + // rings only exist at finalize; this pump waits for them and ships the queue. + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(run_plan_delivery(Arc::clone(&harness), token.clone())); + harness + } + }; + + let mut encoded_tasks = Vec::with_capacity(routed_urls.len()); + for task_i in 0..routed_urls.len() { + encoded_tasks.push(encode_task_plan( + &stage.plan, + task_i, + stage.tasks, + task_ctx.session_config(), + )?); + } + harness.record_stage( + stage.num, + stage.tasks, + encoded_tasks.iter().map(|e| e.partitions).collect(), + ); + harness + .state + .lock() + .unwrap() + .coord_metrics + .get_or_insert_with(|| metrics.clone()); + harness.scan_for_child_routing(&stage.plan, stage.tasks)?; + + // The planner already baked the producer head into the shipped fragment. Re-state it so + // each worker's `insert_producer_head` re-creates the same head rather than flattening the + // stage to one partition. + let producer_head = producer_head_from_plan(&stage.plan).to_proto(task_ctx)?; + + let mut headers = get_config_extension_propagation_headers(task_ctx.session_config())?; + headers.extend(get_passthrough_headers(task_ctx.session_config())); + + for (task_i, (url, encoded)) in routed_urls.iter().zip(encoded_tasks).enumerate() { + let plan_size = encoded.plan_proto.len(); + + let task_key = pb::TaskKey { + query_id: serialize_uuid(&stage.query_id), + stage_id: stage.num as u64, + task_number: task_i as u64, + }; + let set_plan = pb::SetPlanRequest { + plan_proto: encoded.plan_proto, + task_count: stage.tasks as u64, + task_key: Some(task_key.clone()), + work_unit_feed_declarations: encoded.feed_declarations, + target_worker_url: url.to_string(), + query_start_time_ns: metrics.instantiation_time, + }; + // The plan reaches the driver as a `SetPlan` frame on its proc's inbox, the same + // wire crossing a separate-process worker would see; the driver only holds its + // address. The frames queue here because the rings are only built at finalize. + harness + .state + .lock() + .unwrap() + .pending_set_plans + .push(PendingSetPlan { + stage_num: stage.num as u32, + task_i, + frame: SetPlanFrame::from_parts(set_plan, &headers)?, + plan_size, + }); + + // Collected before spawning so the providers see the same eager `feed()` timing as + // they do under the other transports. + let feed_streams = + collect_task_work_unit_feeds(&stage.plan, task_ctx, task_i, stage.tasks)?; + let has_feeds = !feed_streams.is_empty(); + + let driver = TaskDriver { + harness: Arc::clone(&harness), + worker: self.transport.worker.clone(), + stage_num: stage.num as u32, + task_i, + n_partitions: encoded.partitions, + has_feeds, + task_key, + token: token.clone(), + producer_head: producer_head.clone(), + }; + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(driver.run()); + + // The feeds travel as `WorkUnit` frames through the leader's outbound senders, so + // the worker side reads them exactly as a separate process would. + if has_feeds { + harness.state.lock().unwrap().any_feeds = true; + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(run_leader_feed_pump( + Arc::clone(&harness), + stage.num as u32, + task_i, + feed_streams, + token.clone(), + )); + } + } + Ok(()) + } +} + +/// How a producer stage's output reaches its consumers, as parent-stage task indexes. A stage no +/// parent boundary ever claims has no entry: it is consumed by the head on the leader (proc 0). +/// +/// Built by simulating each consumer task's reads under its effective task contexts (a +/// `ChildrenIsolatorUnionExec` hands its children remapped contexts, so a boundary under one is +/// read with that inner context, not the stage-level one). `None` slots were never claimed by +/// any consumer (padding partitions); they route to the leader, where they sit buffered until +/// teardown. +enum RoutingSpec { + /// Nested `NetworkCoalesceExec`: consumers read whole producer tasks, so the destination + /// depends on the producer task only. + PerTask(Vec>), + /// Nested shuffle/broadcast: consumers read partition slices, identical across producer + /// tasks, so the destination depends on the output partition only. + PerPartition(Vec>), +} + +struct StageRec { + tasks: usize, + /// Output partitions of each task's specialized plan. Task-isolated nodes make these differ + /// from the unspecialized stage plan (and possibly from each other). + task_partitions: Vec, +} + +/// What a task driver needs to start producing: its proc's mesh and one routed sink per output +/// partition. +struct Launch { + mesh: Arc, + sinks: Vec>, + /// Send end for this task's `TaskMetrics` frame to the leader. + metrics_sender: MppSender, +} + +/// One task's plan, queued between `dispatch` and finalize, when the rings exist to carry it. +struct PendingSetPlan { + stage_num: u32, + task_i: usize, + frame: SetPlanFrame, + plan_size: usize, +} + +struct HarnessState { + stages: HashMap, + routing: HashMap, + started: bool, + /// Whether any dispatched task declared work-unit feeds; decides whether a feed pump runs. + any_feeds: bool, + /// Plans queued by `dispatch` until finalize builds the mesh; the delivery pump then ships + /// each as a `SetPlan` frame to the proc hosting its task. + pending_set_plans: Vec, + /// Dispatch-side metrics (plan send latency / bytes), recorded by the delivery pump. + coord_metrics: Option, + n_workers: u32, + leader_mesh: Option>, + /// The leader's outbound senders, the control-plane path for `WorkUnit` frames. + leader_senders: Vec>, + /// One mesh per worker proc, kept for the per-proc drain pumps. + worker_meshes: Vec>, + launches: HashMap<(u32, usize), Launch>, + /// Declared after the meshes so it would drop last either way; the harness Arcs held by + /// drivers and pinned streams are what actually keep it alive long enough. + region: Option, +} + +struct QueryHarness { + query_id: Uuid, + token: CancellationToken, + queue_bytes: usize, + metrics_store: Option>, + /// Drivers and feed pumps spawned for this query; `done` counts their exits. The pumps run + /// until the two meet, which is the deterministic "no more frames are coming" signal. + participants: AtomicUsize, + done: AtomicUsize, + state: Mutex, + ready_tx: tokio::sync::watch::Sender, + ready_rx: tokio::sync::watch::Receiver, +} + +impl QueryHarness { + fn new( + query_id: Uuid, + token: CancellationToken, + queue_bytes: usize, + metrics_store: Option>, + ) -> Self { + let (ready_tx, ready_rx) = tokio::sync::watch::channel(false); + Self { + query_id, + token, + queue_bytes, + metrics_store, + participants: AtomicUsize::new(0), + done: AtomicUsize::new(0), + state: Mutex::new(HarnessState { + stages: HashMap::new(), + routing: HashMap::new(), + started: false, + any_feeds: false, + pending_set_plans: Vec::new(), + coord_metrics: None, + n_workers: 0, + leader_mesh: None, + leader_senders: Vec::new(), + worker_meshes: Vec::new(), + launches: HashMap::new(), + region: None, + }), + ready_tx, + ready_rx, + } + } + + /// Wait until the first read finalizes the harness. + async fn ready(&self) -> Result<()> { + let mut rx = self.ready_rx.clone(); + while !*rx.borrow() { + rx.changed().await.map_err(|_| { + DataFusionError::Internal( + "self-hosted shm transport: harness dropped before start".to_string(), + ) + })?; + } + Ok(()) + } + + /// The leader's send end for one task's `WorkUnit` frames, cooperative-draining the + /// leader's own inbox so a symmetric full-ring stall cannot deadlock the feed push. + fn leader_feed_sender(&self, stage_num: u32, task_i: usize) -> Result { + let state = self.state.lock().unwrap(); + let dest_proc = proc_for_task(state.n_workers, task_i as u32); + let Some(base) = state + .leader_senders + .get(dest_proc as usize) + .and_then(|s| s.as_ref()) + else { + return internal_err!( + "self-hosted shm transport: no leader sender for proc {dest_proc}" + ); + }; + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + Ok(base + .clone_with_header(MppFrameHeader::work_unit(stage_num, task_i as u32, 0)) + .with_cooperative_drain(leader_mesh as Arc)) + } + + fn record_stage(&self, num: usize, tasks: usize, task_partitions: Vec) { + let mut state = self.state.lock().unwrap(); + state.stages.insert( + num as u32, + StageRec { + tasks, + task_partitions, + }, + ); + } + + /// Classify the routing of every child stage referenced by `plan`'s network boundaries. The + /// children were dispatched before this stage (plan preparation converts bottom-up), so their + /// records exist; stages no parent ever claims are consumed by the head on the leader. + fn scan_for_child_routing( + &self, + plan: &Arc, + consumer_tasks: usize, + ) -> Result<()> { + let mut state = self.state.lock().unwrap(); + for task_i in 0..consumer_tasks { + let d_ctx = DistributedTaskContext { + task_index: task_i, + task_count: consumer_tasks, + }; + let state = &mut *state; + let mut scan_err = Ok(()); + plan.apply_with_dt_ctx(d_ctx, |node, d_ctx| { + let Some(nb) = node.as_ref().as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + let child_num = nb.input_stage().num() as u32; + let Some(rec) = state.stages.get(&child_num) else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} referenced before dispatch" + ); + return Ok(TreeNodeRecursion::Stop); + }; + let child_tasks = rec.tasks; + let child_max_parts = rec.task_partitions.iter().copied().max().unwrap_or(0); + + if node.as_ref().is::() { + let spec = state + .routing + .entry(child_num) + .or_insert_with(|| RoutingSpec::PerTask(vec![None; child_tasks])); + let RoutingSpec::PerTask(dest) = spec else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} read through mixed \ + boundary kinds" + ); + return Ok(TreeNodeRecursion::Stop); + }; + // Mirror of the consumer's `task_group` split: contiguous groups, the first + // `extra` groups one producer task longer. + let base = child_tasks / d_ctx.task_count.max(1); + let extra = child_tasks % d_ctx.task_count.max(1); + let len = base + usize::from(d_ctx.task_index < extra); + let start = d_ctx.task_index * base + d_ctx.task_index.min(extra); + let end = (start + len).min(child_tasks); + for slot in dest[start..end].iter_mut() { + *slot = Some(task_i as u32); + } + } else { + let spec = state + .routing + .entry(child_num) + .or_insert_with(|| RoutingSpec::PerPartition(vec![None; child_max_parts])); + let RoutingSpec::PerPartition(dest) = spec else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} read through mixed \ + boundary kinds" + ); + return Ok(TreeNodeRecursion::Stop); + }; + // Shuffle and broadcast read the same partition slice per consumer context. + let p_c = nb.partitions_per_consumer_task(); + let from = (p_c * d_ctx.task_index).min(child_max_parts); + let to = (p_c * (d_ctx.task_index + 1)).min(child_max_parts); + for slot in dest[from..to].iter_mut() { + *slot = Some(task_i as u32); + } + } + Ok(TreeNodeRecursion::Continue) + })?; + scan_err?; + } + Ok(()) + } + + /// Size and build the mesh, resolve the routing, and release the waiting task drivers. Runs + /// once, on the first `open`. + fn ensure_started(&self) -> Result<()> { + let mut state = self.state.lock().unwrap(); + if state.started { + return Ok(()); + } + + let n_workers = state + .stages + .values() + .map(|s| s.tasks) + .max() + .unwrap_or(1) + .max(1) as u32; + let n_procs = n_workers + 1; + + let region_total = dsm_region_bytes(n_procs, self.queue_bytes, 0)?; + let region = HeapRegion::new(region_total); + let wakeup: Arc = Arc::new(NoopWakeup); + let interrupt: Arc = Arc::new(CancellationInterrupt(self.token.clone())); + + let leader_attach = unsafe { + leader_setup( + region.base(), + n_procs, + self.queue_bytes, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::clone(&interrupt), + // The harness holds the senders until the query ends, so attaching is safe, + // and the plan delivery pump always needs them. + true, + ) + }?; + let leader_mesh = leader_attach.mesh; + let mut worker_meshes = Vec::with_capacity(n_workers as usize); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + region.base(), + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::clone(&interrupt), + ) + }?; + worker_meshes.push((attach.mesh, attach.outbound_senders)); + } + + // Build every fragment's launch package: one routed sink per output partition, sharing + // the proc's outbound senders. The base senders drop at the end of this scope, so the + // rings observe the last-sender detach once the per-partition clones finish. + let mut launches = HashMap::new(); + for (&stage_num, rec) in state.stages.iter() { + let spec = state.routing.get(&stage_num); + for task_i in 0..rec.tasks { + let n_out = rec.task_partitions.get(task_i).copied().unwrap_or(0); + let proc = proc_for_task(n_workers, task_i as u32); + let (mesh, outbound) = &worker_meshes[(proc - 1) as usize]; + let Some(to_leader) = outbound[0].as_ref() else { + return internal_err!( + "self-hosted shm transport: no outbound sender from proc {proc} to the \ + leader" + ); + }; + // No cooperative drain on purpose: metrics frames go out after the cancellation + // token fired (it fires on normal completion), when the interrupt-checking spin + // would abort the send. + let metrics_sender = to_leader.clone_with_header(MppFrameHeader::task_metrics( + stage_num, + task_i as u32, + proc, + )); + let mut sinks: Vec> = Vec::with_capacity(n_out); + for q in 0..n_out { + let consumer = match spec { + // No parent boundary claimed this stage: the head consumes it on the + // leader. + None => None, + Some(RoutingSpec::PerTask(dest)) => dest.get(task_i).copied().flatten(), + Some(RoutingSpec::PerPartition(dest)) => dest.get(q).copied().flatten(), + }; + let dest_proc = match (spec, consumer) { + (None, _) | (_, None) => 0, + (_, Some(parent_task)) => proc_for_task(n_workers, parent_task), + }; + let Some(base) = outbound[dest_proc as usize].as_ref() else { + return internal_err!( + "self-hosted shm transport: no outbound sender from proc {proc} to \ + proc {dest_proc}" + ); + }; + let sender = base + .clone_with_header(MppFrameHeader::batch(stage_num, q as u32, proc)) + .with_cooperative_drain(Arc::clone(mesh) as Arc); + sinks.push(Box::new(MppPartitionSink::new(sender))); + } + launches.insert( + (stage_num, task_i), + Launch { + mesh: Arc::clone(mesh), + sinks, + metrics_sender, + }, + ); + } + } + + state.n_workers = n_workers; + state.leader_mesh = Some(leader_mesh); + state.leader_senders = leader_attach.outbound_senders; + state.worker_meshes = worker_meshes.iter().map(|(m, _)| Arc::clone(m)).collect(); + state.launches = launches; + state.region = Some(region); + state.started = true; + drop(state); + let _ = self.ready_tx.send(true); + Ok(()) + } + + fn leader_mesh(&self) -> Result> { + self.state + .lock() + .unwrap() + .leader_mesh + .clone() + .ok_or_else(|| { + DataFusionError::Internal( + "self-hosted shm transport: leader mesh not built".to_string(), + ) + }) + } + + /// Wait until the harness is finalized, then take this fragment's launch package. + async fn wait_launch(&self, stage_num: u32, task_i: usize) -> Result { + self.ready().await?; + let mut state = self.state.lock().unwrap(); + state.launches.remove(&(stage_num, task_i)).ok_or_else(|| { + DataFusionError::Internal(format!( + "self-hosted shm transport: no launch package for stage {stage_num} task {task_i}" + )) + }) + } +} + +/// Delivers one task's plan through the worker session machinery and produces its fragment into +/// the mesh. +struct TaskDriver { + harness: Arc, + worker: Worker, + stage_num: u32, + task_i: usize, + n_partitions: usize, + has_feeds: bool, + task_key: pb::TaskKey, + token: CancellationToken, + /// The producer head the distributed planner already baked into the shipped fragment. + /// Re-stated here so the worker's `insert_producer_head` strips and re-adds the same head + /// (a no-op) instead of flattening the stage to a single partition. + producer_head: pb::execute_task_request::ProducerHead, +} + +impl TaskDriver { + async fn run(self) -> Result<()> { + let harness = Arc::clone(&self.harness); + let result = self.run_inner().await; + harness.done.fetch_add(1, Ordering::SeqCst); + result + } + + async fn run_inner(self) -> Result<()> { + let Self { + harness, + worker, + stage_num, + task_i, + n_partitions, + has_feeds, + task_key, + token, + producer_head, + } = self; + + // The launch package arrives when the first read finalizes the harness. A query torn + // down before any read just unwinds the driver. + let launch = tokio::select! { + launch = harness.wait_launch(stage_num, task_i) => launch?, + _ = token.cancelled() => return Ok(()), + }; + + // The plan arrives as a `SetPlan` frame on this proc's inbox, shipped by the leader's + // delivery pump; the pumps drain it into the registry this take resolves from. + let set_plan_frame = tokio::select! { + frame = launch.mesh.take_set_plan(stage_num, task_i as u32) => frame?, + _ = token.cancelled() => return Ok(()), + }; + let (set_plan, headers) = set_plan_frame.into_parts()?; + + let mesh = Arc::clone(&launch.mesh); + let outcome = worker + .set_task_plan(set_plan, headers, move |mut cfg| { + // Child-stage reads inside the decoded fragment must pull from this proc's + // inbox, and its dispatcher must stay a no-op: the plans are already here. + set_distributed_worker_transport(&mut cfg, ShmMqWorkerTransport::new(mesh)); + Ok(cfg) + }) + .await?; + + // The feeds arrive as `WorkUnit` frames on this proc's inbox; install the channels so + // the drain fills them (and flushes whatever arrived first). The leader-side pump owns + // delivery and the `FeedEof` that ends the streams. + if has_feeds { + launch.mesh.register_work_unit_senders( + stage_num, + task_i as u32, + outcome.work_unit_senders, + ); + } + + let produce = async { + let request = pb::ExecuteTaskRequest { + task_key: Some(task_key), + target_partition_start: 0, + target_partition_end: n_partitions as u64, + producer_head: Some(producer_head), + }; + // Through `execute_local_task` rather than a bare `plan.execute` so the task metrics + // (added/executed/finished stamps, per-node metrics) flow exactly like a pull-based + // read would deliver them. + let (streams, _ctx) = execute_local_task(worker.task_data_entries(), request).await?; + if streams.len() != launch.sinks.len() { + return internal_err!( + "self-hosted shm transport: stage {stage_num} task {task_i} decoded into {} \ + partitions but routed {} sinks", + streams.len(), + launch.sinks.len() + ); + } + let mut futures = Vec::with_capacity(streams.len()); + for (mut stream, mut sink) in streams.into_iter().zip(launch.sinks) { + let token = token.clone(); + futures.push(async move { + let stream_result: Result<()> = async { + loop { + let batch = tokio::select! { + next = stream.next() => next, + // Head stream dropped: stop pulling so this fragment and its upstream + // scan unwind, instead of draining the input into a buffer no one reads. + _ = token.cancelled() => break, + }; + let Some(batch) = batch else { break }; + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + sink.send(&batch).await?; + // A downstream worker abandoned this stream (its mesh carries the cancel + // senders): stop pulling so the cancel cascades to this fragment's own + // producers, matching the embedder's `run_worker_fragment` loop. + if sink.cancelled() { + break; + } + } + Ok(()) + } + .await; + // EOF always, even after a failed send, so the consumer side unblocks; the + // stream error stays the primary one. + let eof_result = sink.finish().await; + stream_result.and(eof_result) + }); + } + // `join_all`, not fail-fast: cancelling sibling partitions mid-await would skip their + // EOFs and wedge the consumer's channel buffers. + let results = futures::future::join_all(futures).await; + for r in results { + r?; + } + Ok(()) + }; + let produce_res: Result<()> = produce.await; + + // The metrics receiver resolves as the last partition stream above completes, so this + // does not block on anything remote. The frame goes back over the mesh, where the + // leader-side pump forwards it into the metrics store. + if let Ok(task_metrics) = outcome.metrics_rx.await { + let _ = launch + .metrics_sender + .send_task_metrics_best_effort(&task_metrics) + .await; + } + produce_res + } +} + +/// Leader-side delivery of every dispatched plan: waits for finalize to build the rings, then +/// ships each queued plan as a `SetPlan` frame to the proc hosting its task. One pump per query, +/// counted as a participant so the drain pumps outlive it. +async fn run_plan_delivery(harness: Arc, token: CancellationToken) -> Result<()> { + let result = run_plan_delivery_inner(&harness, token).await; + harness.done.fetch_add(1, Ordering::SeqCst); + result +} + +async fn run_plan_delivery_inner( + harness: &Arc, + token: CancellationToken, +) -> Result<()> { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let (pending, senders, metrics) = { + let mut state = harness.state.lock().unwrap(); + let pending = std::mem::take(&mut state.pending_set_plans); + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + let mut senders = Vec::with_capacity(pending.len()); + for plan in &pending { + let dest_proc = proc_for_task(state.n_workers, plan.task_i as u32); + let Some(base) = state + .leader_senders + .get(dest_proc as usize) + .and_then(|s| s.as_ref()) + else { + return internal_err!( + "self-hosted shm transport: no leader sender for proc {dest_proc}" + ); + }; + senders.push( + base.clone_with_header(MppFrameHeader::set_plan( + plan.stage_num, + plan.task_i as u32, + 0, + )) + .with_cooperative_drain(Arc::clone(&leader_mesh) as Arc), + ); + } + (pending, senders, state.coord_metrics.clone()) + }; + let mut stats = SendBatchStats::default(); + for (plan, sender) in pending.into_iter().zip(senders) { + let start = Instant::now(); + sender.send_set_plan_traced(&plan.frame, &mut stats).await?; + if let Some(metrics) = &metrics { + metrics.plan_send_latency.record(&start); + metrics.plan_bytes_sent.add_bytes(plan.plan_size); + } + } + Ok(()) +} + +/// Leader-side pump for one task's work-unit feeds: drives the providers and ships each unit as +/// a `WorkUnit` frame to the proc hosting the task, then closes the task's channels with a +/// `FeedEof`. The close is unconditional, error or not: without it the fragment's feed streams +/// never end and the worker side wedges instead of surfacing this pump's error. +async fn run_leader_feed_pump( + harness: Arc, + stage_num: u32, + task_i: usize, + feed_streams: Vec>>, + token: CancellationToken, +) -> Result<()> { + let result = async { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let sender = harness.leader_feed_sender(stage_num, task_i)?; + + let pump_result = async { + let mut pumps = Vec::with_capacity(feed_streams.len()); + for mut stream in feed_streams { + let sender = &sender; + pumps.push(async move { + let mut stats = SendBatchStats::default(); + while let Some(unit) = stream.next().await { + let mut unit = unit?; + set_sent_time(&mut unit); + sender.send_work_unit_traced(&unit, &mut stats).await?; + } + Ok::<_, DataFusionError>(()) + }); + } + futures::future::try_join_all(pumps).await.map(|_| ()) + } + .await; + + let mut stats = SendBatchStats::default(); + let eof_result = sender.send_feed_eof_traced(&mut stats).await; + pump_result.and(eof_result) + } + .await; + harness.done.fetch_add(1, Ordering::SeqCst); + result +} + +/// Per-proc drain pumps plus the metrics forwarder, alive until every driver and feed pump +/// reported done. They are what moves control frames when nothing else drains: a fragment +/// blocked on its feed before producing, and the metrics frames arriving after the gather +/// already finished. +async fn run_pumps(harness: Arc, token: CancellationToken) -> Result<()> { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let (leader_mesh, worker_meshes) = { + let state = harness.state.lock().unwrap(); + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + (leader_mesh, state.worker_meshes.clone()) + }; + let metrics_rx = leader_mesh.take_task_metrics_receiver(); + + let mut meshes: Vec> = Vec::with_capacity(worker_meshes.len() + 1); + meshes.push(Arc::clone(&leader_mesh)); + meshes.extend(worker_meshes); + let pumps = meshes.into_iter().map(|mesh| { + let harness = Arc::clone(&harness); + async move { + loop { + let all_done = harness.done.load(Ordering::SeqCst) + >= harness.participants.load(Ordering::SeqCst); + // Drain errors surface through the consumers reading the same registry; + // the pump just stops contributing. + if mesh.try_drain_pass().is_err() { + break; + } + if all_done { + // The pass above ran after the last participant reported done, so every + // frame sent before that point has been routed. + break; + } + tokio::task::yield_now().await; + } + } + }); + futures::future::join_all(pumps).await; + + // Every frame is routed by now; whatever metrics arrived are in the channel. + if let (Some(mut rx), Some(store)) = (metrics_rx, harness.metrics_store.clone()) { + while let Ok((stage_id, task_number, task_metrics)) = rx.try_recv() { + store.insert( + pb::TaskKey { + query_id: serialize_uuid(&harness.query_id), + stage_id: stage_id as u64, + task_number: task_number as u64, + }, + task_metrics, + ); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; + use crate::test_utils::session_context::register_temp_parquet_table; + use crate::{DistributedConfig, DistributedExt, SessionStateBuilderExt, display_plan_ascii}; + use datafusion::arrow::array::{Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::execute_stream; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + + /// Forces the ring mechanics on every batch: with `RING_SLOTS = 8`, a 64 KiB inbox has + /// ~8 KiB slots, so the ~16 KiB frames below fragment across slots, and the ~2 MB of + /// payload wraps each ring dozens of times, exercising the cooperative send spin. + const TINY_QUEUE_BYTES: usize = 64 * 1024; + + const ROWS: usize = 2000; + + fn sample_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, false), + Field::new("val", DataType::Int32, false), + ])); + // ~1 KiB per row, unique values so the GROUP BY keeps the full volume flowing + // through the shuffle instead of compacting it away at the partial aggregate. + let strings: Vec = (0..ROWS) + .map(|i| format!("{i:06}-{}", "x".repeat(1024))) + .collect(); + let vals: Vec = (0..ROWS as i32).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(strings)), + Arc::new(Int32Array::from(vals)), + ], + ) + .unwrap() + } + + async fn run(ctx: &SessionContext) -> Result<(String, Vec)> { + // Shaped so every ring frame stays bounded by `shuffle_batch_size`. The strings cross + // the shuffle inside `max`'s partial state, which the repartition rebuilds with `take` + // into fresh per-batch arrays; the projection then reduces them to a length before the + // gather. Shipping `s` itself out of a sort or an aggregate would not work: those emit + // offset slices of their accumulated state, a sliced variable-length array ships its + // whole values buffer through arrow-ipc, and a single frame balloons to the size of the + // partition's state no matter the batch size. + let query = "SELECT val, length(max(s)) AS l FROM t GROUP BY val"; + let plan = ctx.sql(query).await?.create_physical_plan().await?; + let display = display_plan_ascii(plan.as_ref(), false); + let batches: Vec<_> = execute_stream(plan, ctx.task_ctx())?.try_collect().await?; + let mut lines: Vec = pretty_format_batches(&batches)? + .to_string() + .lines() + .map(str::to_string) + .collect(); + lines.sort(); + Ok((display, lines)) + } + + /// A high-cardinality shuffle query over rings far smaller than the data, so every + /// cross-stage byte moves through fragmented frames under send-spin backpressure. The + /// result must still match the serial reference exactly. + #[tokio::test] + async fn tiny_rings_force_fragmentation_and_backpressure() -> Result<()> { + let transport = SelfHostedShmTransport::default().with_queue_bytes(TINY_QUEUE_BYTES); + // Small producer batches keep each frame a few slots big instead of overflowing the + // whole ring (a single frame must fit within one ring). + let d_cfg = DistributedConfig { + shuffle_batch_size: 16, + ..Default::default() + }; + let mut state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_option_extension(d_cfg) + .with_distributed_planner() + .with_distributed_task_estimator(2) + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(3)) + .with_distributed_worker_transport(transport) + .build(); + state.config_mut().options_mut().execution.target_partitions = 3; + let ctx = SessionContext::from(state); + let path = + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, distributed) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + + let single = SessionContext::default(); + single + .register_parquet("t", path.to_string_lossy().as_ref(), Default::default()) + .await?; + let (_, expected) = run(&single).await?; + + assert_eq!(distributed, expected); + Ok(()) + } + + /// The per-query harness must be reclaimed when the output stream drops, or the `queries` map + /// grows one entry per query for the transport's lifetime. + #[tokio::test] + async fn query_harness_is_reclaimed_after_the_stream_drops() -> Result<()> { + let transport = SelfHostedShmTransport::default(); + let probe = transport.clone(); + let mut state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_task_estimator(2) + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(3)) + .with_distributed_worker_transport(transport) + .build(); + state.config_mut().options_mut().execution.target_partitions = 3; + let ctx = SessionContext::from(state); + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, _) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + + // The token fires as the gathered stream drops above; the registry entry is dropped on a + // task the cancelled token wakes, so poll briefly rather than assume it already ran. + for _ in 0..100 { + if probe.queries.is_empty() { + return Ok(()); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + panic!( + "queries map still holds {} entries; the harness leaked", + probe.queries.len() + ); + } +} diff --git a/src/shm/setup.rs b/src/shm/setup.rs new file mode 100644 index 00000000..d630c63c --- /dev/null +++ b/src/shm/setup.rs @@ -0,0 +1,373 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mesh construction over a caller-supplied shared buffer, the extension point between an embedder's buffer +//! allocation and the transport. +//! +//! The embedder allocates the shared region (a PG `dsm_segment`, or a heap buffer in-process), +//! sizes it with [`dsm_region_bytes`], then calls [`leader_setup`] on the proc that initializes the +//! rings and [`worker_setup`] on each producer proc. Both hand back an [`MppMesh`] the embedder +//! installs on its DataFusion session; `worker_setup` also returns the outbound senders and the +//! plan bytes copied out of the region. [`run_worker_fragment`] is the producer push loop. +//! +//! The two platform primitives the embedder supplies are the [`Wakeup`] (how to wake a blocked +//! consumer) and the [`Interrupt`] (how to check for cancellation); everything here is otherwise +//! pure Rust over the shared buffer. + +use std::ffi::c_void; +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use futures::stream::StreamExt; + +use super::dsm::{ + compute_dsm_layout, leader_init, peer_proc_for_index, read_region_total, worker_attach, +}; +use super::mesh::{DsmInboxReceiver, DsmInboxSender}; +use super::mpsc_ring::{DsmMpscSender, NO_RECEIVER_TOKEN, Wakeup}; +use super::runtime::MppMesh; +use super::transport::{ + BatchChannelSender, DrainHandle, Interrupt, MppFrameHeader, MppReceiver, MppSender, + ReceiverScope, SELF_LOOP_CAPACITY, in_proc_channel, +}; +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedRegistry; +use crate::{DistributedTaskContext, PartitionSink, collect_plan_metrics_protos}; + +/// Total bytes the shared region needs for `n_procs` inboxes plus `plan_len` plan bytes, with +/// `queue_bytes` per inbox. The embedder reserves exactly this much before [`leader_setup`]. +pub fn dsm_region_bytes(n_procs: u32, queue_bytes: usize, plan_len: usize) -> Result { + compute_dsm_layout(n_procs, queue_bytes, plan_len) + .map(|l| l.region_total) + .map_err(|e| DataFusionError::Internal(format!("mpp: dsm_region_bytes: {e}"))) +} + +/// Read `region_total` out of the header a leader wrote, so a worker that just mapped the region +/// can size its [`worker_setup`] call without knowing the header layout. +/// +/// A caller that derives the size from the header forfeits [`worker_setup`]'s size validation +/// (it would compare the header against itself). Pass the mapped size from the embedder's own +/// bookkeeping where it is available. +/// +/// # Safety +/// - `base` must point at the start of a region a leader initialized via [`leader_setup`]. +/// - `base` must be at least 8-byte aligned (the header holds `u64` fields). +pub unsafe fn region_total(base: *const c_void) -> usize { + unsafe { read_region_total(base) as usize } +} + +/// Wrap each peer-indexed `DsmMpscSender` into an outbound `MppSender` keyed by destination proc. +/// The dispatcher `clone_with_header`s these per output partition before sending, so the +/// placeholder header is never observed on the wire. Slot `this_proc` stays `None` until the +/// self-loop install. +/// +/// Returns `(data, cancel)`. `data` is the producer's output senders. `cancel` is a control-plane +/// sibling onto each peer inbox, used by [`MppMesh::cancel_stream`]: a consumer reaches its producer +/// without counting as one of that producer's data senders, so a held `Cancel` sender never masks +/// the producer-gone `detached` signal. +fn build_outbound_senders( + this_proc: u32, + total_procs: u32, + peer_senders: Vec, +) -> (Vec>, Vec>) { + let mut senders: Vec> = (0..total_procs).map(|_| None).collect(); + let mut cancel: Vec> = (0..total_procs).map(|_| None).collect(); + for (peer_idx, dsm_send) in peer_senders.into_iter().enumerate() { + let target_proc = peer_proc_for_index(this_proc, peer_idx as u32); + // A `peer_proc_for_index` regression that maps a peer onto the self slot would be + // silently overwritten by the self-loop install and only surface later as a missing + // sender at dispatch; name the bug at its source. + debug_assert!( + target_proc != this_proc, + "peer index {peer_idx} mapped to the self proc {this_proc}" + ); + debug_assert!( + target_proc < total_procs, + "peer index {peer_idx} mapped to proc {target_proc} >= total {total_procs}" + ); + let control: Arc = + Arc::new(DsmInboxSender::new(dsm_send.to_control())); + cancel[target_proc as usize] = Some(MppSender::with_header( + control, + MppFrameHeader::batch(0, 0, this_proc), + )); + let shared: Arc = Arc::new(DsmInboxSender::new(dsm_send)); + senders[target_proc as usize] = Some(MppSender::with_header( + shared, + // Stamp `sender_proc = this_proc` so a stray frame that escapes the dispatcher's + // `clone_with_header` overwrite still identifies its origin on the drain side. + MppFrameHeader::batch(0, 0, this_proc), + )); + } + (senders, cancel) +} + +/// What [`leader_setup`] hands back to the embedder. +pub struct LeaderAttach { + /// The leader's mesh, installed on its DataFusion session. + pub mesh: Arc, + /// Outbound senders keyed by destination proc index, for the control plane: work-unit + /// frames flow leader -> worker through them. Slot 0 (the leader itself) stays `None`; + /// empty unless `attach_senders` was passed. Holders must keep them alive for the whole + /// query: dropping them before a worker attaches latches that worker's inbox as detached. + pub outbound_senders: Vec>, +} + +/// Initialize the shared region as the leader (`proc 0`) and return its mesh plus its outbound +/// senders. +/// +/// Writes the region header, copies `plan_bytes` in, initializes the `n_procs` inboxes, and +/// attaches the leader as receiver to its own inbox. `receiver_token` is registered so producers +/// resolve this proc's [`Wakeup`]; `interrupt` is consulted at the transport's block points. +/// +/// # Safety +/// - `base` must point at an uninitialized region of at least `dsm_region_bytes(n_procs, +/// queue_bytes, plan_bytes.len())` bytes. +/// - `base` must be at least 8-byte (MAXALIGN) aligned; the ring headers hold atomics. +/// - The region must not be concurrently accessed until this returns. +#[allow(clippy::too_many_arguments)] // mirrors worker_setup; the args are the embedder's knobs +pub unsafe fn leader_setup( + base: *mut c_void, + n_procs: u32, + queue_bytes: usize, + plan_bytes: &[u8], + wakeup: Arc, + receiver_token: u64, + interrupt: Arc, + attach_senders: bool, +) -> Result { + if receiver_token == NO_RECEIVER_TOKEN { + return Err(DataFusionError::Internal( + "mpp: leader_setup: receiver_token is the NO_RECEIVER_TOKEN sentinel; wakeups \ + for this proc would be silently disabled" + .into(), + )); + } + let layout = compute_dsm_layout(n_procs, queue_bytes, plan_bytes.len()) + .map_err(|e| DataFusionError::Internal(format!("mpp: leader_setup compute layout: {e}")))?; + let attach = unsafe { + leader_init( + base, + &layout, + plan_bytes, + Arc::clone(&wakeup), + attach_senders, + ) + } + .map_err(DataFusionError::Internal)?; + + let inbox = DsmInboxReceiver::new(attach.inbound_receiver); + inbox.set_receiver(receiver_token); + let inbound = Arc::new(DrainHandle::cooperative( + 0, + vec![(ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox)))], + )); + // The leader hosts no producer fragments, but its senders carry the control plane: + // work-unit frames (and later dynamic filters) flow leader -> worker through them. Empty + // when the embedder did not opt in: a ring latches `detached` once its sender count hits + // zero, so senders that might drop before every worker attached must never exist. + // The leader's `Cancel` senders are wired by the embedder (it shares the same outbound senders + // it holds for plan delivery and releases before the DSM unmaps), so drop the cancel set here. + let (outbound_senders, _cancel_senders) = + build_outbound_senders(0, n_procs, attach.outbound_senders); + Ok(LeaderAttach { + mesh: Arc::new(MppMesh::new(0, n_procs, inbound, interrupt, attach.alive)), + outbound_senders, + }) +} + +/// Build one task's work-unit feed channels, install the receiving ends on `cfg` (where the +/// deserialized plan's remote feed leaves look them up), and register the sending ends on +/// `mesh`'s drain so inbound `WorkUnit` frames fill them. `feeds` lists the task's declared +/// feeds as `(feed id, partitions)`, the same pairs the plan's `WorkUnitFeedDeclaration`s carry. +/// +/// The caller must keep the proc draining (a consumer loop, a send spin, or an explicit +/// [`crate::shm::CooperativeDrainSet::try_drain_pass`] pump) while a fragment waits on its +/// feed, or the units sit in the inbox unread. +/// Build the [`TaskMetrics`] payload for one executed fragment, for embedders that run +/// fragments outside the worker task registry (pg parallel workers). Pair it with +/// [`super::transport::MppSender::send_task_metrics_best_effort`] after the fragment's streams +/// complete; the leader-side rewrite consumes the same pre-order the in-registry path produces. +/// The task-level stamps (plan added/executed/finished) carry no data on this path, but the set +/// must be present: `decode_task_metrics` rejects a missing field, and one failed decode starves +/// the whole store. +/// +/// [`TaskMetrics`]: crate::proto::TaskMetrics +pub fn collect_task_metrics( + plan: &Arc, + task_index: usize, + task_count: usize, +) -> pb::TaskMetrics { + pb::TaskMetrics { + pre_order_plan_metrics: collect_plan_metrics_protos( + plan, + DistributedTaskContext { + task_index, + task_count, + }, + ), + task_metrics: Some(pb::MetricsSet::default()), + } +} + +pub fn install_work_unit_channels( + cfg: &mut datafusion::prelude::SessionConfig, + mesh: &MppMesh, + stage_id: u32, + task_number: u32, + feeds: &[(uuid::Uuid, usize)], +) { + let mut channels = RemoteWorkUnitFeedRegistry::default(); + for (id, partitions) in feeds { + channels.add(*id, *partitions); + } + cfg.set_extension(Arc::new(channels.receivers)); + mesh.register_work_unit_senders(stage_id, task_number, channels.senders); +} + +/// What [`worker_setup`] hands back to the embedder. +pub struct WorkerAttach { + /// The worker's mesh, installed on its DataFusion session. + pub mesh: Arc, + /// Outbound senders keyed by destination proc index. The slot at `this_proc` is the in-proc + /// self-loop; every other slot writes to that peer's inbox. + pub outbound_senders: Vec>, + /// The plan bytes the leader wrote into the region, copied out for this worker. + pub plan_bytes: Vec, +} + +/// Attach to the leader-initialized region as worker `proc_idx` (`>= 1`). +/// +/// # Safety +/// - `base`/`region_total` must match the region the leader initialized via [`leader_setup`]. +/// - `base` must be at least 8-byte (MAXALIGN) aligned; the ring headers hold atomics. +pub unsafe fn worker_setup( + base: *mut c_void, + region_total: usize, + proc_idx: u32, + wakeup: Arc, + receiver_token: u64, + interrupt: Arc, +) -> Result { + if receiver_token == NO_RECEIVER_TOKEN { + return Err(DataFusionError::Internal( + "mpp: worker_setup: receiver_token is the NO_RECEIVER_TOKEN sentinel; wakeups \ + for this proc would be silently disabled" + .into(), + )); + } + let (header, plan_bytes, attach) = + unsafe { worker_attach(base, region_total as u64, proc_idx, Arc::clone(&wakeup)) } + .map_err(DataFusionError::Internal)?; + let total_procs = header.n_procs; + + let (mut outbound, cancel) = + build_outbound_senders(proc_idx, total_procs, attach.outbound_senders); + + // Self-loop in-proc channel: peer-mesh routing can land a producer and its consumer on the same + // proc, and an MPSC inbox has no slot for a proc sending to itself. The unified drain pulls from + // both the inbox and this channel. + let (self_tx, self_rx) = in_proc_channel(SELF_LOOP_CAPACITY); + let self_tx_arc: Arc = Arc::new(self_tx); + outbound[proc_idx as usize] = Some(MppSender::with_header( + self_tx_arc, + MppFrameHeader::batch(0, 0, proc_idx), + )); + + let inbox = DsmInboxReceiver::new(attach.inbound_receiver); + inbox.set_receiver(receiver_token); + let inbound = Arc::new(DrainHandle::cooperative( + proc_idx, + vec![ + (ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox))), + (ReceiverScope::SelfLoop, MppReceiver::new(Box::new(self_rx))), + ], + )); + let mesh = Arc::new(MppMesh::new( + proc_idx, + total_procs, + inbound, + interrupt, + attach.alive, + )); + // A worker consumes shuffle inputs, so it can be the consumer that stops a stream early. Give + // its mesh the control-plane cancel senders; they drop with the mesh at the end of the worker's + // run, well before the DSM unmaps, so no explicit release is needed. + mesh.set_cancel_senders(Arc::new(std::sync::Mutex::new(cancel))); + Ok(WorkerAttach { + mesh, + outbound_senders: outbound, + plan_bytes, + }) +} + +/// Run a producer fragment plan to exhaustion, pushing every output batch into the matching +/// per-partition [`PartitionSink`]. The output partition count of `plan` must equal `sinks.len()`; +/// `sinks[partition]` is the send end the caller routed for that output partition. +/// +/// Each partition's [`PartitionSink::finish`] sends a per-channel EOF when its stream ends, +/// regardless of how it ended: the shared queue is multiplexed across fragments, so dropping a sink +/// doesn't end the channel, only the EOF frame does. +pub async fn run_worker_fragment( + plan: Arc, + sinks: Vec>, + ctx: Arc, +) -> Result<()> { + let n_partitions = plan.output_partitioning().partition_count(); + if n_partitions != sinks.len() { + return Err(DataFusionError::Internal(format!( + "run_worker_fragment: plan has {n_partitions} output partitions but {} sinks", + sinks.len() + ))); + } + let mut futures = Vec::with_capacity(n_partitions); + for (partition, mut sink) in sinks.into_iter().enumerate() { + let plan = Arc::clone(&plan); + let ctx = Arc::clone(&ctx); + futures.push(async move { + let stream_result: Result<()> = async { + let mut stream = plan.execute(partition, ctx)?; + while let Some(batch) = stream.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + sink.send(&batch).await?; + // Consumer abandoned this stream. Stop pulling: dropping `stream` ends the + // upstream scan and cascades the cancel to its own producers. + if sink.cancelled() { + break; + } + } + Ok(()) + } + .await; + let eof_result = sink.finish().await; + // Surface the stream error first, then any EOF-send error, so neither disappears. + stream_result.and(eof_result) + }); + } + // `join_all`, not `try_join_all`: fail-fast would cancel sibling partitions mid-await before + // they reach `finish`, leaving the consumer's channel buffer stuck. + let results = futures::future::join_all(futures).await; + for r in results { + r?; + } + Ok(()) +} diff --git a/src/shm/sink.rs b/src/shm/sink.rs new file mode 100644 index 00000000..a6ae19d7 --- /dev/null +++ b/src/shm/sink.rs @@ -0,0 +1,49 @@ +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::Result; + +/// The producer's send end for one partition channel, symmetric to a [crate::WorkerChannel] read. +/// +/// This lives in the shm module rather than the core transport surface: only a push-based transport +/// (the shared-memory mesh) produces through sinks. Flight produces inside its gRPC worker service +/// and the in-memory transport pulls straight from the local task registry, so neither needs it. +/// +/// Contract with the produce loop: +/// - Batches arrive in `send` order and can be assumed non-empty. +/// - After a failed `send` the channel state is unspecified, but the caller still calls `finish` so +/// the consumer sees EOF; `finish` must tolerate a prior `send` error. +/// - Dropping a sink without calling `finish` does not end the channel, by design: `finish` is async +/// so Drop can't run it, and an implicit EOF would make an aborted producer look like a clean, +/// short stream. Abnormal teardown belongs to the transport, not the sink. +/// - `send` borrows the batch because transports serialize it into their own buffers; none needs +/// ownership. +#[async_trait] +pub trait PartitionSink: Send { + /// Sends one batch. Async so a blocked send can yield and let the transport make progress + /// elsewhere; a full channel must not park the calling thread. + async fn send(&mut self, batch: &RecordBatch) -> Result<()>; + /// Per-channel EOF, independent of the underlying link. Async for the same reason as `send`. + async fn finish(self: Box) -> Result<()>; + /// Whether the consumer cancelled this stream. The produce loop stops pulling its input when + /// this turns true, so a cancel doesn't just skip the send, it ends the upstream scan and drops + /// the input stream, cascading the cancel further up. Default `false` for links that don't carry + /// a cancel signal. + fn cancelled(&self) -> bool { + false + } +} + +/// The producer (write) side: opens a [PartitionSink] per output partition, symmetric to +/// [crate::WorkerChannel] (the read side). The worker's produce loop builds one and pushes each +/// output batch in. The shared-memory mesh provides the implementation; it is constructed by the +/// producer (which knows the per-partition routing), not handed out by the consume-side transport. +pub trait WorkerSink: Send + Sync { + /// Takes `stage` and `partition` separately because one sink serves every stage, unlike the + /// per-stage read connection that closes over its stage. + /// + /// `stage` is the producing stage's number and `partition` the producer task's own output + /// partition index, before routing. Several producer tasks of one stage may hold sinks for the + /// same pair, and the consumer merges them, so one `finish` is one producer task's EOF, not + /// channel completion (which stays transport-defined). + fn open_partition(&self, stage: usize, partition: usize) -> Result>; +} diff --git a/src/shm/transport.rs b/src/shm/transport.rs new file mode 100644 index 00000000..4b8fb2f6 --- /dev/null +++ b/src/shm/transport.rs @@ -0,0 +1,2935 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Transport layer for MPP shuffle. +//! +//! - [`MppFrameHeader`]: fixed 16-byte prefix tagging each wire message with +//! `(stage_id, partition)`, so one queue carries frames for many logical channels. +//! - [`encode_frame_into`] / [`decode_frame`]: Arrow IPC serialize/deserialize with +//! header prefix. Only codec entry points; tests round-trip through the same path. +//! - [`DrainBuffer`]: per-proc queue the drain writes into and the DataFusion consumer +//! reads from. Decouples consumer-side from producer-side backpressure: the drain +//! always makes forward progress on the inbound rings, so a stalled consumer can't +//! propagate backpressure to remote producers and cause a peer-mesh stall. + +use async_trait::async_trait; +use std::collections::VecDeque; +use std::sync::{Arc, Condvar, Mutex, MutexGuard}; + +use datafusion::common::{HashMap, HashSet}; +use std::time::{Duration, Instant}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::common::DataFusionError; +use prost::Message; + +use crate::common::deserialize_uuid; +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedTxs; +use crate::{WorkUnitMsg, set_received_time}; + +/// Magic bytes "MPPF" (MPP Frame) at the start of every wire message. +/// Lets receivers reject misrouted / corrupt frames before they hit Arrow IPC. +const MPP_FRAME_MAGIC: u32 = 0x4D505046; + +/// Wire-format size of [`MppFrameHeader`] in bytes. Asserted at compile time +/// below via `const _: ()`. +const MPP_FRAME_HEADER_SIZE: usize = 16; + +/// Kind of payload following [`MppFrameHeader`]. +/// +/// `Batch` is the common case. The header is followed by an Arrow IPC stream containing one +/// `RecordBatch`. `Eof` carries no payload. It signals the receiver that the named +/// `(stage_id, partition)` channel is done, even though the underlying shm_mq queue may still +/// carry frames for other channels. +/// +/// The remaining kinds are the control plane riding the same rings. For them the header's +/// `partition` field carries the task number instead: a work unit already names its +/// `(feed id, partition)` inside the payload, and metrics describe a whole task. +/// `SetPlan` (leader -> worker) carries one task's [`pb::SetPlanRequest`], the same message +/// Flight ships over its coordinator stream, plus the propagation headers that ride gRPC +/// metadata there; `WorkUnit` (leader -> worker) carries one prost-encoded work unit for +/// `(stage, task)`; `FeedEof` closes that task's feed channels (the wire stand-in for Flight's +/// stream close); `TaskMetrics` (worker -> leader) carries the task's collected metrics. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum MppFrameKind { + Batch = 0, + Eof = 1, + WorkUnit = 2, + FeedEof = 3, + TaskMetrics = 4, + SetPlan = 5, + /// Consumer -> producer: stop producing the `(stage_id, partition)` stream, the consumer + /// stopped reading it before EOF (a top-N `LIMIT`, an inner merge join exhausting a side, etc.). + /// The producer ends that one stream cleanly. Scoped to the stream, not the connection, so the + /// ring stays healthy for metrics and every other stream, the way gRPC closes one stream + /// without dropping the channel. + Cancel = 6, +} + +/// Payload of a `SetPlan` frame: the plan-delivery message a worker needs to run one task, +/// byte-compatible with what Flight sends. +/// +/// `set_plan` is the exact [`pb::SetPlanRequest`] the Flight dispatcher would put on its gRPC +/// stream. The headers are the config-extension propagation headers that travel as gRPC metadata +/// there; the ring has no metadata side channel, so they ride inside the frame as parallel +/// key/value lists (parallel lists rather than a map so repeated header names survive). +#[derive(Clone, PartialEq, prost::Message)] +pub struct SetPlanFrame { + #[prost(message, optional, tag = "1")] + pub set_plan: Option, + #[prost(string, repeated, tag = "2")] + pub header_keys: Vec, + #[prost(string, repeated, tag = "3")] + pub header_values: Vec, +} + +impl SetPlanFrame { + /// Bundle a plan-delivery message with the headers Flight would carry as gRPC metadata. + pub fn from_parts( + set_plan: pb::SetPlanRequest, + headers: &http::HeaderMap, + ) -> Result { + let mut header_keys = Vec::with_capacity(headers.len()); + let mut header_values = Vec::with_capacity(headers.len()); + for (name, value) in headers { + let value = value.to_str().map_err(|e| { + DataFusionError::Internal(format!( + "mpp: non-ASCII header {name} cannot travel in a SetPlan frame: {e}" + )) + })?; + header_keys.push(name.as_str().to_string()); + header_values.push(value.to_string()); + } + Ok(Self { + set_plan: Some(set_plan), + header_keys, + header_values, + }) + } + + /// Split back into the plan-delivery message and the propagation headers. + pub fn into_parts(self) -> Result<(pb::SetPlanRequest, http::HeaderMap), DataFusionError> { + let set_plan = self.set_plan.ok_or_else(|| { + DataFusionError::Internal("mpp: SetPlan frame carries no SetPlanRequest".to_string()) + })?; + let mut headers = http::HeaderMap::with_capacity(self.header_keys.len()); + for (key, value) in self.header_keys.iter().zip(self.header_values.iter()) { + let name = http::header::HeaderName::from_bytes(key.as_bytes()).map_err(|e| { + DataFusionError::Internal(format!("mpp: SetPlan frame header name {key:?}: {e}")) + })?; + let value = http::header::HeaderValue::from_str(value).map_err(|e| { + DataFusionError::Internal(format!("mpp: SetPlan frame header value for {key}: {e}")) + })?; + headers.append(name, value); + } + Ok((set_plan, headers)) + } +} + +/// 16-byte prefix on every transport frame. +/// +/// The fixed layout `[magic, flags, stage_id, partition]` (4×u32) is what +/// senders prepend before the Arrow IPC stream bytes and what receivers +/// parse before deciding which channel buffer the payload belongs to. +/// +/// See the `flags` bit-layout block below for the encoding of the `flags` word. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MppFrameHeader { + // Private so headers only come out of `batch()`/`eof()`: hand-built ones could bypass + // `pack_flags`'s sender bound and the reserved-bits invariant, and the consumer would + // reject them at decode, far from the producer. + pub(super) magic: u32, + pub(super) flags: u32, + pub(super) stage_id: u32, + pub(super) partition: u32, +} + +/// `flags` bit layout: +/// bits 0..8: frame kind (Batch | Eof) +/// bits 8..16: reserved (must be 0) +/// bits 16..32: sender_proc (mesh peer that wrote the frame) +const FRAME_KIND_MASK: u32 = 0x0000_00FF; +const FRAME_RESERVED_MASK: u32 = 0x0000_FF00; +const FRAME_SENDER_SHIFT: u32 = 16; +/// Maximum `sender_proc` representable in the header. Asserted at construction time so an +/// overflow becomes a hard error in the producer rather than silent flag corruption on the wire. +pub const MPP_MAX_SENDER_PROC: u32 = 0xFFFF; + +/// Spin bound for landing a control or metrics frame on a peer inbox that may be momentarily full. +/// A live peer drains its inbox within this many tries, so the frame lands well inside the bound; +/// it only runs out if the peer already exited, where the leader's wait-for-workers surfaces the +/// failure instead of hanging here. +const MAX_CONTROL_SEND_SPINS: usize = 10_000; + +const _: () = { + // shm_mq slot layout calculations depend on this being exact. + assert!(std::mem::size_of::() == MPP_FRAME_HEADER_SIZE); +}; + +#[inline] +fn pack_flags(kind: MppFrameKind, sender_proc: u32) -> u32 { + // fail_loud rather than debug_assert: in release builds the check would be compiled out and + // an out-of-range value would silently truncate to `sender_proc & 0xFFFF`. Catching the case + // where a refactor accidentally passes a task_id or partition here is the whole point. + assert!( + sender_proc <= MPP_MAX_SENDER_PROC, + "mpp: sender_proc {sender_proc} > MPP_MAX_SENDER_PROC ({MPP_MAX_SENDER_PROC})" + ); + (kind as u32) | (sender_proc << FRAME_SENDER_SHIFT) +} + +impl MppFrameHeader { + /// Build a `Batch` header for the given `(stage_id, partition)` stamped with `sender_proc`. + pub fn batch(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Batch, sender_proc), + stage_id, + partition, + } + } + + /// Build an `Eof` header for the given `(stage_id, partition)` stamped with `sender_proc`. + /// Carries no payload; receivers route it to the channel buffer's source-done counter. + pub fn eof(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Eof, sender_proc), + stage_id, + partition, + } + } + + /// Build a `WorkUnit` header addressed to `(stage_id, task_number)`. The `partition` slot + /// carries the task number; the unit's own `(feed id, partition)` ride in the payload. + pub fn work_unit(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::WorkUnit, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `FeedEof` header for `(stage_id, task_number)`: every feed of that task is done. + pub fn feed_eof(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::FeedEof, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `TaskMetrics` header for `(stage_id, task_number)`. + pub fn task_metrics(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::TaskMetrics, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `SetPlan` header for `(stage_id, task_number)`: the frame delivers that task's + /// plan to the proc hosting it. + pub fn set_plan(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::SetPlan, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `Cancel` header for the `(stage_id, partition)` stream, stamped with the consumer's + /// `sender_proc`. Carries no payload; the producer reads it as "stop sending this stream." + pub fn cancel(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Cancel, sender_proc), + stage_id, + partition, + } + } + + /// The mesh peer that wrote this frame. The drain demuxes incoming frames into the + /// per-channel buffer registry by `(sender_proc, stage_id, partition)`. + pub fn sender_proc(&self) -> u32 { + (self.flags >> FRAME_SENDER_SHIFT) & 0xFFFF + } + + /// Read the kind out of `flags`. Returns an error if the kind byte is + /// unknown or if any reserved bit (bits 8..16) is set, which catches wire-format + /// drift early. Sender_proc bits (16..32) are not validated here; readers extract + /// them with `sender_proc()`. + pub(super) fn kind(&self) -> Result { + let reserved = self.flags & FRAME_RESERVED_MASK; + if reserved != 0 { + return Err(DataFusionError::Internal(format!( + "mpp: reserved frame flag bits set ({reserved:#x})" + ))); + } + match self.flags & FRAME_KIND_MASK { + 0 => Ok(MppFrameKind::Batch), + 1 => Ok(MppFrameKind::Eof), + 2 => Ok(MppFrameKind::WorkUnit), + 3 => Ok(MppFrameKind::FeedEof), + 4 => Ok(MppFrameKind::TaskMetrics), + 5 => Ok(MppFrameKind::SetPlan), + 6 => Ok(MppFrameKind::Cancel), + other => Err(DataFusionError::Internal(format!( + "mpp: unknown frame kind {other:#x}" + ))), + } + } + + /// Serialize into the first `MPP_FRAME_HEADER_SIZE` bytes of `out`. + /// `out.len()` must be `>= MPP_FRAME_HEADER_SIZE`. + fn write_to(&self, out: &mut [u8]) { + debug_assert!(out.len() >= MPP_FRAME_HEADER_SIZE); + out[0..4].copy_from_slice(&self.magic.to_le_bytes()); + out[4..8].copy_from_slice(&self.flags.to_le_bytes()); + out[8..12].copy_from_slice(&self.stage_id.to_le_bytes()); + out[12..16].copy_from_slice(&self.partition.to_le_bytes()); + } + + /// Parse from the first `MPP_FRAME_HEADER_SIZE` bytes of `bytes`. Returns + /// `Err` if the slice is too short or the magic doesn't match. + fn parse(bytes: &[u8]) -> Result { + if bytes.len() < MPP_FRAME_HEADER_SIZE { + // No encoder in this file emits sub-header output, so a short frame means the + // shm_mq stitched together payloads from different senders. Hex-dump the bytes + // so the source is identifiable from log output without a debugger. + let hex = bytes + .iter() + .map(|b| format!("{b:02x}")) + .collect::>() + .join(" "); + return Err(DataFusionError::Internal(format!( + "mpp: frame too short for header ({} < {}); bytes = [{hex}]", + bytes.len(), + MPP_FRAME_HEADER_SIZE + ))); + } + let magic = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); + if magic != MPP_FRAME_MAGIC { + return Err(DataFusionError::Internal(format!( + "mpp: bad frame magic {magic:#x} (expected {MPP_FRAME_MAGIC:#x})" + ))); + } + Ok(Self { + magic, + flags: u32::from_le_bytes(bytes[4..8].try_into().unwrap()), + stage_id: u32::from_le_bytes(bytes[8..12].try_into().unwrap()), + partition: u32::from_le_bytes(bytes[12..16].try_into().unwrap()), + }) + } +} + +/// Serialize `batch` into `buf` with a 16-byte [`MppFrameHeader`] prefix +/// addressing it to `(stage_id, partition)`. Wire format: +/// +/// ```text +/// [ magic | flags | stage_id | partition ] [ Arrow IPC stream bytes ] +/// |---------- 16 bytes --------| |---- variable ----| +/// ``` +/// +/// `flags` encodes kind + sender_proc; see the bit-layout block near +/// `FRAME_KIND_MASK` for details. +/// +/// Caller is expected to hold `buf` alive across many encodes so the peak-sized +/// allocation amortizes (~500 KB/batch on the 25M GROUP BY bench). +fn encode_frame_into( + header: MppFrameHeader, + batch: &RecordBatch, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + header.write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + let mut writer = StreamWriter::try_new(&mut *buf, batch.schema_ref())?; + writer.write(batch)?; + writer.finish()?; + Ok(()) +} + +/// Serialize a payload-less [`MppFrameKind::Eof`] frame for `(stage_id, partition)` +/// into `buf`. The shm_mq peer reads this as a 16-byte message and routes it to +/// the channel buffer's source-done counter without touching Arrow IPC. +/// Consumed by [`MppSender::send_eof_traced`] when a producer fragment's +/// per-partition stream exhausts, so the receiver's `(stage_id, partition)` +/// channel buffer transitions to `Eof` even though the multiplexed shm_mq queue +/// stays attached for other channels. +fn encode_eof_frame_into( + stage_id: u32, + partition: u32, + sender_proc: u32, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + MppFrameHeader::eof(stage_id, partition, sender_proc) + .write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + Ok(()) +} + +/// Serialize a prost-encoded control payload (`WorkUnit` / `TaskMetrics`) behind `header`. +fn encode_prost_frame_into( + header: MppFrameHeader, + msg: &impl prost::Message, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + header.write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + msg.encode(buf) + .map_err(|e| DataFusionError::Internal(format!("mpp: prost frame encode: {e}")))?; + Ok(()) +} + +/// A decoded frame payload, routed by the drain according to its kind. +#[derive(Debug)] +enum FrameBody { + Batch(RecordBatch), + Eof, + WorkUnit(pb::WorkUnit), + FeedEof, + TaskMetrics(pb::TaskMetrics), + SetPlan(SetPlanFrame), + Cancel, +} + +/// Inverse of the frame encoders. Parses the 16-byte header and decodes the payload according +/// to the kind. Receivers branch on the body to decide routing. +fn decode_frame(bytes: &[u8]) -> Result<(MppFrameHeader, FrameBody), DataFusionError> { + let header = MppFrameHeader::parse(bytes)?; + let payload = &bytes[MPP_FRAME_HEADER_SIZE..]; + match header.kind()? { + MppFrameKind::Eof | MppFrameKind::FeedEof | MppFrameKind::Cancel => { + if bytes.len() != MPP_FRAME_HEADER_SIZE { + return Err(DataFusionError::Internal(format!( + "mpp: payload-less frame carries payload ({} > {})", + bytes.len(), + MPP_FRAME_HEADER_SIZE + ))); + } + match header.kind()? { + MppFrameKind::Eof => Ok((header, FrameBody::Eof)), + MppFrameKind::Cancel => Ok((header, FrameBody::Cancel)), + _ => Ok((header, FrameBody::FeedEof)), + } + } + MppFrameKind::WorkUnit => { + let unit = pb::WorkUnit::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: work unit decode: {e}")))?; + Ok((header, FrameBody::WorkUnit(unit))) + } + MppFrameKind::TaskMetrics => { + let metrics = pb::TaskMetrics::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: task metrics decode: {e}")))?; + Ok((header, FrameBody::TaskMetrics(metrics))) + } + MppFrameKind::SetPlan => { + let frame = SetPlanFrame::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: set-plan decode: {e}")))?; + Ok((header, FrameBody::SetPlan(frame))) + } + MppFrameKind::Batch => { + let mut reader = StreamReader::try_new(payload, None)?; + let batch = reader.next().ok_or_else(|| { + DataFusionError::Execution("mpp: empty arrow-ipc stream in decode_frame".into()) + })??; + Ok((header, FrameBody::Batch(batch))) + } + } +} + +/// Local queue between a drain (either the cooperative `try_drain_pass` or the test-only thread +/// variant) and the consumer that pops batches. +/// +/// In the cooperative path each `DrainBuffer` corresponds to one logical channel: one +/// `(stage_id, partition)` entry in the owning [`DrainHandle`]'s registry. `num_sources` is +/// always `1` there because a given drain serves a single sender_proc, which is the only producer +/// for any channel routed through it. The test-only thread path uses a single shared buffer with +/// `num_sources = N` over an N-sender setup. +/// +/// Push side: callers append deserialized batches; on source detach (or per-channel `Eof` frame) +/// [`DrainBuffer::notify_source_done`] is called. Once `sources_done >= num_sources` AND the +/// queue is empty, `try_pop` returns [`DrainItem::Eof`]. +/// +/// Pop side: cooperative consumers loop on `try_pop` + `yield_now`. The test-only `pop_front` +/// blocks on the condvar. +#[derive(Debug)] +pub(super) struct DrainBuffer { + inner: Mutex, + cond: Condvar, +} + +#[derive(Debug)] +struct DrainBufferInner { + queue: VecDeque, + num_sources: u32, + sources_done: u32, + /// Consumer-side cancel flag. When set (e.g., query cancelled or `DrainHandle` dropped), + /// `try_pop`/`pop_front` returns `Eof` even if `sources_done` hasn't reached `num_sources`. + cancelled: bool, + /// Set when the receiver feeding this channel detached (or errored) before the channel's + /// `Eof` frame arrived. Distinct from `cancelled`: cancellation is a clean teardown and + /// yields `Eof`, while a lost source must surface as an error or the consumer would treat + /// truncated output as complete. + failed: Option, +} + +/// Yielded by [`DrainBuffer::pop_front`]. +#[derive(Debug)] +pub(super) enum DrainItem { + /// A batch produced by one of the inbound shm_mqs. + Batch(RecordBatch), + /// All source queues have detached and the local queue is drained. + Eof, + /// The receiver feeding this channel went away before the channel's `Eof` frame. + Failed(String), +} + +impl DrainBuffer { + /// Create a drain buffer expecting `num_sources` inbound queues. For a + /// proc in an N-proc mesh, `num_sources == N - 1` (all peers + /// excluding self — the self-partition bypasses the buffer). + pub fn new(num_sources: u32) -> Arc { + Arc::new(Self { + inner: Mutex::new(DrainBufferInner { + queue: VecDeque::new(), + num_sources, + sources_done: 0, + cancelled: false, + failed: None, + }), + cond: Condvar::new(), + }) + } + + /// Push a freshly-received batch into the local queue. + pub fn push_batch(&self, batch: RecordBatch) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.queue.push_back(batch); + self.cond.notify_one(); + } + + /// Mark one source queue as detached. Safe to call from the drain thread + /// after observing `SHM_MQ_DETACHED` on a given inbound queue. + pub fn notify_source_done(&self) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.sources_done = guard.sources_done.saturating_add(1); + if guard.sources_done >= guard.num_sources { + self.cond.notify_all(); + } + } + + /// Mark the channel as fed by a dead receiver, unless it already completed (its `Eof` + /// arrived), was cancelled, or already failed. Consumers then see an error instead of + /// hanging on a channel nothing will ever fill. + pub fn fail_pending(&self, msg: &str) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + if guard.sources_done >= guard.num_sources || guard.cancelled || guard.failed.is_some() { + return; + } + guard.failed = Some(msg.to_string()); + self.cond.notify_all(); + } + + /// Cancel all further pushes and wake all consumers with EOF. + pub fn cancel(&self) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.cancelled = true; + self.cond.notify_all(); + } + + /// Non-blocking variant. Returns the front item, or `DrainItem::Eof` if + /// all sources have detached and the queue is drained, or `None` if more + /// data may yet arrive. Cooperative consumers loop on + /// `try_drain_pass` + `try_pop`, yielding to the executor between + /// iterations. + pub fn try_pop(&self) -> Option { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + Self::try_pop_locked(&mut guard) + } + + /// Shared body of [`try_pop`] and the test-only [`Self::pop_front`]. + /// Returns `Some(Batch)` if the queue has data, `Some(Eof)` if all + /// sources have detached or the buffer is cancelled, and `None` + /// otherwise. Lets the two entry points stay in lockstep on the + /// "buffered data wins over cancellation/EOF" invariant locked in by + /// the `drain_buffer_drains_buffered_before_eof` test. + fn try_pop_locked(guard: &mut MutexGuard<'_, DrainBufferInner>) -> Option { + if let Some(batch) = guard.queue.pop_front() { + return Some(DrainItem::Batch(batch)); + } + if let Some(msg) = &guard.failed { + return Some(DrainItem::Failed(msg.clone())); + } + if guard.cancelled || guard.sources_done >= guard.num_sources { + return Some(DrainItem::Eof); + } + None + } +} + +/// Outcome of a single non-blocking receive attempt. +#[derive(Debug)] +pub(super) enum RecvOutcome { + /// One serialized Arrow IPC message ready to decode. + Bytes(Vec), + /// No data currently available but the peer is still attached. + Empty, + /// The peer has detached; no more bytes will ever arrive on this channel. + Detached, +} + +/// Non-blocking byte channel receiver. Implementations: `DsmInboxReceiver` (production), +/// `std::sync::mpsc` (tests). Must be `Send` because the drain thread takes ownership. +pub(super) trait BatchChannelReceiver: Send + Sync { + fn try_recv(&self) -> RecvOutcome; +} + +/// Byte channel sender paired with [`BatchChannelReceiver`]. `send` blocks when +/// the channel is full. Dropping the sender signals EOF to the receiver. +/// +/// `Send` is required because unit tests and future producer-pump threads move +/// senders across thread boundaries. +pub(crate) trait BatchChannelSender: Send + Sync { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError>; + + /// Non-blocking variant. Returns `Ok(true)` on success, `Ok(false)` + /// when the channel is full (caller should retry), `Err` on detach / + /// transport error. Default falls back to the blocking send — safe + /// for in-proc channels used by tests where "full" doesn't arise. + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + self.send_bytes(bytes).map(|()| true) + } + + /// Async lock the send paths hold across the cooperative-drain spin so two tasks can't + /// interleave partial writes on the same handle. PG's `shm_mq_send` requires the same + /// `(nbytes, data)` on retry after `SHM_MQ_WOULD_BLOCK`. Multiple [`MppSender`] clones + /// multiplex onto one channel, and the spin's `yield_now().await` would otherwise let a + /// sibling task land a different payload mid-message and corrupt the queue. In-proc + /// channels return a per-instance mutex too, just to keep the call sites uniform. + fn send_lock(&self) -> &tokio::sync::Mutex<()>; +} + +/// Pluggable "drain everything inbound" hook for [`MppSender`]'s cooperative send spin. The +/// peer-mesh deadlock-breaking pattern needs the producer to pump ALL inbound queues (not just +/// one) while waiting for a full outbound, so the implementation typically delegates to +/// `MppMesh::drain_all_inbound()` which iterates every per-sender-proc drain. +pub trait CooperativeDrainSet: Send + Sync { + fn try_drain_pass(&self) -> Result<(), DataFusionError>; + + /// Checked in the send spin alongside `try_drain_pass`: returns `Err` (or aborts) if the + /// query should stop. Default is a no-op, for embedders with no external interrupt source; a + /// Postgres embedder overrides this to run `check_for_interrupts!`, which longjmps on cancel. + fn check_interrupt(&self) -> Result<(), DataFusionError> { + Ok(()) + } + + /// Whether a consumer cancelled the `(stage_id, partition)` stream (a `Cancel` frame arrived on + /// this proc's inbox). The send spin ends the producer's stream cleanly when it's set. Default + /// `false` for drains that don't carry inbound control frames (in-proc test channels). + fn stream_cancelled(&self, _stage_id: u32, _partition: u32) -> bool { + false + } +} + +/// Cancellation extension point, checked at the transport's block points (the send spin and the consumer +/// pull loop). An in-process embedder uses the default no-op or a cancellation token; a Postgres +/// embedder runs `check_for_interrupts!`, which longjmps out of the backend on cancel. +pub trait Interrupt: Send + Sync { + fn check(&self) -> Result<(), DataFusionError>; +} + +/// No-op interrupt for embedders that have no external cancellation source. +pub struct NoInterrupt; +impl Interrupt for NoInterrupt { + fn check(&self) -> Result<(), DataFusionError> { + Ok(()) + } +} + +impl CooperativeDrainSet for DrainHandle { + fn try_drain_pass(&self) -> Result<(), DataFusionError> { + DrainHandle::try_drain_pass(self) + } + + fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + DrainHandle::stream_cancelled(self, stage_id, partition) + } +} + +/// High-level sender: encodes a `RecordBatch` then pushes bytes through the underlying channel. +/// +/// With `cooperative_drain` set, `send_batch` breaks the symmetric-send deadlock on a +/// single-threaded tokio runtime by interleaving send-retries with +/// `CooperativeDrainSet::try_drain_pass` on the same mesh's inbound side. Each proc's +/// sender doing the same guarantees mutual progress: our drain pulls peer-shipped rows out of +/// our inbound queues, which frees peers' outbound-to-us send space, which lets their sends +/// un-stall. +pub struct MppSender { + /// Underlying byte channel. Held behind `Arc` so multiple `MppSender`s can share one + /// `shm_mq` queue while tagging frames with different `(stage_id, partition)` headers, which + /// is the multiplexed path's natural pattern. Clone the Arc, build a new `MppSender` with a + /// different header, both write into the same queue. + pub(super) channel: Arc, + cooperative_drain: Option>, + /// Frame header prepended to every outgoing batch. Identifies the logical + /// `(stage_id, partition)` channel the receiver demultiplexes on. Per-sender rather than + /// per-call: each partition gets its own `MppSender` via `clone_with_header`, all sharing + /// the underlying `Arc` of a single shm_mq queue. + pub(super) header: MppFrameHeader, + /// Scratch buffer reused across every `encode_frame_into` on this sender. Sized by the + /// first batch; subsequent batches clear and re-fill without reallocating. Interior + /// mutability lets the caller keep the `&self` signature (each producer fragment holds + /// its `MppSender` clones behind shared borrows for the duration of + /// `worker::run_worker_fragment`). + scratch: std::cell::RefCell>, +} + +// SAFETY: only `scratch: RefCell>` and the trait-object `Arc`s are `!Sync`. Callers +// compose `send_*_traced` futures via `tokio::spawn` / `join_all`, which makes the compiler +// require `&Self: Send` and therefore `Self: Sync`. The shared-memory model runs those futures on +// a current-thread runtime (see the module docs), so the cell is never observed from two +// threads; a multi-thread embedder would additionally be serialized by `send_lock` across +// every send path that touches `scratch`. +unsafe impl Sync for MppSender {} + +impl MppSender { + /// Construct a sender that tags every outgoing batch with `header`. Production call sites + /// clone one shared `Arc` across N senders, each with a different + /// `MppFrameHeader::batch(stage, p)`. That's the multiplexed pattern for fanning multiple + /// partitions over one shm_mq queue. + pub(super) fn with_header( + channel: Arc, + header: MppFrameHeader, + ) -> Self { + Self { + channel, + cooperative_drain: None, + header, + scratch: std::cell::RefCell::new(Vec::new()), + } + } + + /// Build a new `MppSender` that shares this sender's underlying channel + /// but tags every frame with `header` instead. Used by callers that know + /// the physical plan's output partition count and need one sender per + /// partition, all multiplexed over the same shm_mq queue. + pub fn clone_with_header(&self, header: MppFrameHeader) -> Self { + Self { + channel: Arc::clone(&self.channel), + cooperative_drain: self.cooperative_drain.as_ref().map(Arc::clone), + header, + scratch: std::cell::RefCell::new(Vec::new()), + } + } + + /// Whether the consumer of this sender's `(stage, partition)` stream cancelled it. Read by the + /// produce loop to stop pulling its input, not just skip the send. `false` without a drain + /// (in-proc test channels carry no inbound cancel). + pub(super) fn stream_cancelled(&self) -> bool { + self.cooperative_drain + .as_ref() + .is_some_and(|d| d.stream_cancelled(self.header.stage_id, self.header.partition)) + } + + /// Attach a [`CooperativeDrainSet`] so `Self::send_batch_traced`'s spin + /// can drain inbound peer traffic while waiting for outbound space. + /// Required for peer-mesh fragments where every worker is both sender and + /// consumer; without it, symmetric full-queue stalls deadlock the + /// single-threaded Tokio runtime. + pub fn with_cooperative_drain(mut self, drain: Arc) -> Self { + self.cooperative_drain = Some(drain); + self + } + + /// `send_batch` variant that accumulates per-call timings and spin counts into `stats`. + /// Callers that report at EOF (e.g. `ShuffleStream`) use this to diagnose where time + /// goes when the outbound queue is full. + /// + /// Async because the spin awaits the per-handle send lock and yields between + /// `try_send_bytes` retries; see `send_with_scratch`. + pub(super) async fn send_batch_traced( + &self, + batch: &RecordBatch, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + // Take the scratch buffer out of the `RefCell` rather than + // holding a `RefMut` across the spin below. The spin contains + // the embedder's `Interrupt::check`, which may unwind or `longjmp` through + // Rust frames; a `longjmp` does not run `Drop`, so a `RefMut` + // held across it would leave the cell perpetually borrowed and + // panic the next caller. `replace` is atomic — the cell is + // never observed in a borrowed state — and we put the buffer + // back at the end so its heap allocation survives across calls. + // If the spin longjmps anyway, the cell holds the default empty + // `Vec` and the next call simply re-allocates. + let mut scratch = self.scratch.replace(Vec::new()); + let result = self.send_with_scratch(batch, &mut scratch, stats).await; + self.scratch.replace(scratch); + result + } + + /// Send a payload-less [`MppFrameKind::Eof`] frame so the receiver's `(stage_id, partition)` + /// channel buffer transitions to `Eof` and the consumer's pull loop terminates cleanly. + /// + /// Producer fragments must call this exactly once per `(stage_id, partition)` channel after + /// the local stream exhausts. Without it the multiplexed shm_mq queue stays attached (other + /// channels still flow) and the consumer channel buffer never reaches `sources_done == 1`. The + /// receive-side [`DrainHandle::try_drain_pass`] decodes the frame and calls + /// `notify_source_done` on the matching channel buffer. + /// + /// Uses the same cooperative-spin path as [`Self::send_batch_traced`] so a full outbound + /// queue doesn't deadlock the EOF send. `stats.spin_iters` / `send_wait` capture any + /// contention. + /// + /// Symmetric-EOF safety: when every peer reaches EOF simultaneously with full outbound + /// queues, each peer's cooperative [`CooperativeDrainSet::try_drain_pass`] inside the spin + /// pulls peer-sent frames out of its own inbound queues, freeing space the peers are blocked + /// on. Progress is monotone: at least one `try_send_bytes` succeeds per spin iteration + /// somewhere in the mesh, so symmetric stalls resolve within a few iterations rather than + /// deadlocking. + pub(super) async fn send_eof_traced( + &self, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = self.send_eof_with_scratch(&mut scratch, stats).await; + self.scratch.replace(scratch); + result + } + + /// Bounded synchronous send of a `Cancel` frame for the `(stage_id, partition)` stream. The + /// consumer calls it on the sender to the producing proc when it abandons that stream. + /// Synchronous so it can run from the consumer stream's drop, where `await` isn't available. + /// + /// The bound doesn't risk a stuck producer. A live producer drains its own inbox in its send + /// spin, so a slot frees and the frame lands well inside the bound even when the inbox is + /// backed up. The bound only runs out if the producer already exited or died, and a dead worker + /// makes the leader's wait-for-workers error out rather than hang. + pub fn try_send_cancel(&self, stage_id: u32, partition: u32) { + let header = MppFrameHeader::cancel(stage_id, partition, self.header.sender_proc()); + let mut buf = [0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + for _ in 0..MAX_CONTROL_SEND_SPINS { + match self.channel.try_send_bytes(&buf) { + Ok(true) => return, + Ok(false) => std::thread::yield_now(), + Err(_) => return, // the producer's inbox is gone; nothing left to cancel + } + } + } + + async fn send_eof_with_scratch( + &self, + scratch: &mut Vec, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + encode_eof_frame_into( + self.header.stage_id, + self.header.partition, + self.header.sender_proc(), + scratch, + )?; + let Some(drain) = self.cooperative_drain.as_ref() else { + return self.channel.send_bytes(scratch); + }; + // Lock the channel before the spin so a sibling task can't interleave a different + // partial write through the shared shm_mq handle. See `BatchChannelSender::send_lock`. + let _send_guard = self.channel.send_lock().lock().await; + let mut first_try = true; + let t_wait_start = Instant::now(); + loop { + drain.check_interrupt()?; + // The consumer cancelled this stream, so end the send cleanly and let the producer + // fragment complete. Checked before the send so a cancel that landed mid-spin stops + // the next iteration. + if drain.stream_cancelled(self.header.stage_id, self.header.partition) { + return Ok(()); + } + if self.spin_try_send_bytes(scratch).await? { + if !first_try { + stats.send_wait += t_wait_start.elapsed(); + } + return Ok(()); + } + first_try = false; + stats.spin_iters += 1; + let t_drain = Instant::now(); + self.spin_try_drain_pass(drain).await?; + stats.coop_drain_in_spin += t_drain.elapsed(); + tokio::task::yield_now().await; + } + } + + /// Spin-loop helper: call `channel.try_send_bytes(scratch)`. + async fn spin_try_send_bytes(&self, scratch: &[u8]) -> Result { + self.channel.try_send_bytes(scratch) + } + + /// Spin-loop helper: call `drain.try_drain_pass()`. + async fn spin_try_drain_pass( + &self, + drain: &Arc, + ) -> Result<(), DataFusionError> { + drain.try_drain_pass() + } + + async fn send_with_scratch( + &self, + batch: &RecordBatch, + scratch: &mut Vec, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let t_enc = Instant::now(); + encode_frame_into(self.header, batch, scratch)?; + stats.encode += t_enc.elapsed(); + self.spin_send_scratch(scratch, stats).await + } + + /// Push an already-encoded frame through the channel via the cooperative-drain spin (or the + /// blocking fallback when no drain is attached). Shared by every frame kind's send path. + async fn spin_send_scratch( + &self, + scratch: &[u8], + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let Some(drain) = self.cooperative_drain.as_ref() else { + // No drain attached (unit tests, in-proc channels): fall + // back to the blocking send path. + return self.channel.send_bytes(scratch); + }; + // Lock the channel before the spin so a sibling task can't interleave a different + // partial write through the shared shm_mq handle. See `BatchChannelSender::send_lock`. + // Long-term, switching shm_mq for an async-friendly ring buffer (cf. #4184) drops the + // partial-send invariant entirely and removes the need for this lock. + // + // Latent under the current-thread runtime: today every fragment owns its own + // `Arc` (one sender per `Arc`), so the FIFO Mutex below + // is uncontended. A future multi-thread runtime that shares a sender across + // sibling fragment tasks (multi-partition fan-out) would let one task starve + // another for the duration of a large shuffle; the fix at that point is to move + // the entire spin off the compute thread. + let _send_guard = self.channel.send_lock().lock().await; + let mut first_try = true; + let t_wait_start = Instant::now(); + // The spin runs inside a tokio task on the backend thread's current-thread runtime + // (DataFusion needs one to drive `Stream`s). The deadlock we're breaking is + // *cross-proc*: two peers each blocked on a full outbound. `try_drain_pass` pulls + // peer batches off our inbound on the same OS thread, freeing their slots so their + // sends advance. `yield_now().await` between iterations hands the runtime back to + // siblings if any are ready, mostly a no-op under today's linear MPP topology. + loop { + drain.check_interrupt()?; + // The consumer cancelled this stream, so end the send cleanly and let the producer + // fragment complete. Checked before the send so a cancel that landed mid-spin stops + // the next iteration. + if drain.stream_cancelled(self.header.stage_id, self.header.partition) { + return Ok(()); + } + if self.spin_try_send_bytes(scratch).await? { + if !first_try { + stats.send_wait += t_wait_start.elapsed(); + } + return Ok(()); + } + first_try = false; + stats.spin_iters += 1; + // Would-block. Pull from our inbound so peers' outbound-to-us frees up and their + // sends to us unblock; without this, symmetric full-queue sends deadlock. Errors + // propagate so a peer detaching mid-spin doesn't leave us spinning on a closed + // mesh. + let t_drain = Instant::now(); + self.spin_try_drain_pass(drain).await?; + stats.coop_drain_in_spin += t_drain.elapsed(); + tokio::task::yield_now().await; + } + } + + /// Send one work unit for the task this sender's header names. The unit's hop stamps are the + /// caller's job; the payload travels prost-encoded, never through Arrow IPC. + pub async fn send_work_unit_traced( + &self, + unit: &pb::WorkUnit, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, unit, &mut scratch)?; + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Ship one task's plan as a `SetPlan` frame: the wire stand-in for Flight's + /// `SetPlanRequest` over its coordinator stream. + pub async fn send_set_plan_traced( + &self, + frame: &SetPlanFrame, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, frame, &mut scratch)?; + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Close the feed channels of the task this sender's header names: the wire stand-in for + /// Flight closing its coordinator stream, after which the worker-side feed streams end. + pub async fn send_feed_eof_traced( + &self, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + scratch.clear(); + scratch.resize(MPP_FRAME_HEADER_SIZE, 0); + MppFrameHeader::feed_eof( + self.header.stage_id, + self.header.partition, + self.header.sender_proc(), + ) + .write_to(&mut scratch[..MPP_FRAME_HEADER_SIZE]); + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Send the task's collected metrics without consulting the interrupt: this runs after the + /// query's cancellation token already fired (it fires on normal completion too), so the + /// cooperative spin would abort exactly when delivery matters. Best-effort like Flight's + /// metrics sends: a detached leader drops them. Retrying on a full ring is safe because the + /// receiving side keeps draining until every producer reported in. + pub async fn send_task_metrics_best_effort( + &self, + metrics: &pb::TaskMetrics, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, metrics, &mut scratch)?; + let _send_guard = self.channel.send_lock().lock().await; + // The consumer stops reading the data path early but keeps draining at teardown to + // collect these metrics, so the frame lands once the data the cancel stopped drains out + // and a slot frees. Bounded so a leader that really went away can't wedge the worker on + // a full ring. + for _ in 0..MAX_CONTROL_SEND_SPINS { + match self.channel.try_send_bytes(&scratch) { + Ok(true) => return Ok(()), + Ok(false) => tokio::task::yield_now().await, + Err(_) => return Ok(()), // receiver gone; metrics are best-effort + } + } + Ok(()) + } + .await; + self.scratch.replace(scratch); + result + } +} + +/// Per-call timing + spin metrics for [`MppSender::send_batch_traced`]. +/// All fields accumulate; callers zero or reuse as needed. +#[derive(Default, Debug, Clone)] +pub struct SendBatchStats { + /// Cumulative time spent inside `encode_frame_into` (header + Arrow IPC serialization). + pub encode: Duration, + /// Cumulative wall time in the send-retry spin after the first failed + /// `try_send_bytes`. Zero if the first try succeeded. + pub send_wait: Duration, + /// Cumulative time spent in `try_drain_pass` while spinning on a + /// full outbound. A subset of `send_wait`; the remainder is the + /// `tokio::task::yield_now()` await + the (small) cost of + /// `try_send_bytes` itself. + pub coop_drain_in_spin: Duration, + /// Count of `try_send_bytes` calls that returned `Ok(false)` (full). + pub spin_iters: u64, +} + +/// A [`crate::PartitionSink`] over one [`MppSender`]: the produce loop's per-partition send end. +/// `send` runs the cooperative-drain spin and `finish` flushes the channel EOF the same way, so a +/// non-Flight embedder (the in-process harness, pg_search's worker loop) wraps each routed sender +/// in one of these and pushes batches through the trait instead of touching `MppSender` directly. +pub struct MppPartitionSink { + sender: MppSender, + stats: SendBatchStats, +} + +impl MppPartitionSink { + pub fn new(sender: MppSender) -> Self { + Self { + sender, + stats: SendBatchStats::default(), + } + } + + /// Per-channel send counters, for an embedder that traces throughput. Read them before + /// `finish`, which consumes the sink. + pub fn stats(&self) -> &SendBatchStats { + &self.stats + } +} + +#[async_trait] +impl crate::PartitionSink for MppPartitionSink { + async fn send(&mut self, batch: &RecordBatch) -> datafusion::common::Result<()> { + self.sender.send_batch_traced(batch, &mut self.stats).await + } + + async fn finish(mut self: Box) -> datafusion::common::Result<()> { + self.sender.send_eof_traced(&mut self.stats).await + } + + fn cancelled(&self) -> bool { + self.sender.stream_cancelled() + } +} + +/// High-level receiver: pulls bytes via the underlying channel and decodes them +/// into `RecordBatch`. Used by the drain thread. +pub(super) struct MppReceiver { + channel: Box, +} + +impl MppReceiver { + pub fn new(channel: Box) -> Self { + Self { channel } + } + + pub(super) fn try_recv_batch(&self) -> RecvBatchOutcome { + match self.channel.try_recv() { + RecvOutcome::Bytes(bytes) => match decode_frame(&bytes) { + Ok((header, FrameBody::Batch(batch))) => RecvBatchOutcome::Batch { header, batch }, + Ok((header, FrameBody::Eof)) => RecvBatchOutcome::Eof { header }, + Ok((header, FrameBody::WorkUnit(unit))) => { + RecvBatchOutcome::WorkUnit { header, unit } + } + Ok((header, FrameBody::FeedEof)) => RecvBatchOutcome::FeedEof { header }, + Ok((header, FrameBody::TaskMetrics(metrics))) => { + RecvBatchOutcome::TaskMetrics { header, metrics } + } + Ok((header, FrameBody::SetPlan(frame))) => { + RecvBatchOutcome::SetPlan { header, frame } + } + Ok((header, FrameBody::Cancel)) => RecvBatchOutcome::Cancel { header }, + Err(e) => RecvBatchOutcome::Error(e), + }, + RecvOutcome::Empty => RecvBatchOutcome::Empty, + RecvOutcome::Detached => RecvBatchOutcome::Detached, + } + } +} + +/// Decoded result of an [`MppReceiver::try_recv_batch`]. Carries the +/// parsed [`MppFrameHeader`] so the drain thread can route the payload to +/// the right `(stage_id, partition)` channel buffer. +#[derive(Debug)] +pub(super) enum RecvBatchOutcome { + Batch { + header: MppFrameHeader, + batch: RecordBatch, + }, + /// A payload-less `Eof` frame for `header.(stage_id, partition)`. The + /// underlying shm_mq queue is still attached. The sender is just + /// signalling that this logical channel is done, so we can EOF + /// per-channel without dropping the whole queue. + Eof { + header: MppFrameHeader, + }, + /// One work unit for the task named by `header.(stage_id, partition=task)`. + WorkUnit { + header: MppFrameHeader, + unit: pb::WorkUnit, + }, + /// Every feed of the task named by the header is done; its channels close. + FeedEof { + header: MppFrameHeader, + }, + /// The plan for the task named by `header.(stage_id, partition=task)`. + SetPlan { + header: MppFrameHeader, + frame: SetPlanFrame, + }, + /// The collected metrics of the task named by the header. + TaskMetrics { + header: MppFrameHeader, + metrics: pb::TaskMetrics, + }, + /// Consumer abandoned the `header.(stage_id, partition)` stream; its producer stops sending it. + Cancel { + header: MppFrameHeader, + }, + Empty, + Detached, + Error(DataFusionError), +} + +/// Per-`(sender_proc, stage_id, partition)` channel buffer registry owned by a cooperative +/// [`DrainHandle`]. The handle may host several cooperative receivers (DSM MPSC inbox + self-loop +/// in-proc), each demultiplexed by the [`MppFrameHeader`] prefix into the same `map`. +/// `try_drain_pass` looks up the right channel buffer on every frame and pushes the payload into +/// it. Consumers waiting on a given key only see frames matching that key. +/// +/// Each entry is a `DrainBuffer::new(1)`: exactly one sender_proc emits frames for any given +/// channel. Per-channel EOF flows via the `Eof` frame demuxed onto the matching buffer; query- +/// teardown unblock flows via [`DrainHandle::cancel_channel_buffers`] from the handle's `Drop`. +#[derive(Default)] +struct ChannelBufferRegistry { + /// Keyed by `(sender_proc, stage_id, partition)`. The unified inbox carries frames + /// from every peer, so each `(stage, partition)` consumer gets its own per-sender + /// buffer. This preserves the implicit "one stream per sender" semantics that + /// `WorkerConnection::execute` consumers rely on. + map: HashMap<(u32, u32, u32), Arc>, + /// Scopes whose receiver detached (or errored) before draining cleanly. Channels fed by a + /// dead scope fail at registration time too, so a consumer that registers after the detach + /// does not wait on a channel nothing will ever fill. + dead_inbox: bool, + dead_self_loop: bool, +} + +/// Which frames a receiver carries, so a detach can fail exactly the channels it feeds. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(super) enum ReceiverScope { + /// The proc's DSM inbox: frames from every peer proc. + Inbox, + /// The in-proc self-loop: frames this proc sends itself. + SelfLoop, +} + +/// Per-sender-proc drain: stashes the receivers and polls them inline from the cooperative spin +/// (no background thread), demuxing each frame into a per-`(stage_id, partition)` channel buffer. +/// +/// Inline polling is the production requirement: pgrx's `check_active_thread` guard panics on any +/// pg FFI call (including `shm_mq_receive`) from a non-backend thread, so the drain work has to +/// run on the backend thread. Tests that need a true thread-backed drain use +/// [`ThreadedDrainHandle`] instead. +/// +/// On drop, the handle cancels every channel buffer so any consumer blocked on `try_pop` unblocks +/// with `Eof` — the drain can therefore never outlive its query, even on a panicked teardown. +pub struct DrainHandle { + /// Per-(stage_id, partition) channel buffer registry. Populated lazily on first frame for a + /// channel, or up-front by callers (e.g. `WorkerConnection::execute`) that need a + /// buffer to wait on before any frame arrives. + channel_buffers: Mutex, + /// Receivers owned by the handle and polled inline from `DrainGatherStream::poll_next` via + /// [`Self::try_drain_pass`]. The `Mutex` is for interior mutability: `try_drain_pass(&self)` + /// marks each slot as `None` after observing `Detached` so subsequent passes skip the dead + /// receiver. `BatchChannelReceiver: Send + Sync` makes `Vec>: Sync` + /// already, so the lock is no longer doubling as the `Sync` provider — replacing it with a + /// non-locking primitive would need either an atomic per-slot detached flag or accepting + /// that detached receivers get polled once per pass (fast-returning `Detached`). The lock + /// is uncontended in production (single backend thread) so the marginal cost is in the + /// type system, not the runtime. + coop_receivers: Mutex>>, + /// This proc's index, used to map a channel's `sender_proc` to the receiver scope that + /// feeds it (`SelfLoop` iff `sender_proc == this_proc`). + this_proc: u32, + /// Worker-side destination of `WorkUnit` frames, keyed `(stage_id, task_number)`. Frames + /// arriving before the embedder registers a task's channels buffer in `Pending`; `FeedEof` + /// drops the senders so the consuming feed streams end, the wire analog of Flight closing + /// its coordinator stream. + feed_registry: Mutex, + /// Leader-side destination of `TaskMetrics` frames: `(stage_id, task_number, metrics)`. + task_metrics_tx: tokio::sync::mpsc::UnboundedSender<(u32, u32, pb::TaskMetrics)>, + task_metrics_rx: + Mutex>>, + /// Worker-side destination of `SetPlan` frames, keyed `(stage_id, task_number)`. Same + /// pending-or-waiting shape as the feed registry: a frame arriving before the task asks + /// buffers in `Pending`; a task asking first parks a oneshot the drain fulfills. + set_plan_registry: Mutex, + /// `(stage_id, partition)` streams this proc's consumers abandoned, learned from inbound + /// `Cancel` frames. A producer blocked on a full outbound checks it in its send spin and ends + /// that stream cleanly, so a consumer that stopped reading early doesn't leave it spinning to + /// the statement timeout. + cancelled_streams: Mutex>, +} + +#[derive(Default)] +struct FeedRegistry { + map: HashMap<(u32, u32), FeedSlot>, + /// Set when the inbox scope died: feeds come from a peer proc, so a dead inbox means no + /// further units or `FeedEof` can arrive. Registered channels get the failure pushed in; + /// later registrations fail immediately. + dead: Option, +} + +enum FeedSlot { + /// Frames that arrived before the embedder registered the task's channels. + Pending { + units: Vec, + done: bool, + }, + Active(RemoteWorkUnitFeedTxs), +} + +/// Push one decoded unit into the channel its `(feed id, partition)` names. A missing channel is +/// not an error: the same tolerance the Flight worker applies to its stream (a feed the plan +/// does not declare is dropped). +fn forward_unit(senders: &RemoteWorkUnitFeedTxs, unit: pb::WorkUnit) { + let Ok(id) = deserialize_uuid(&unit.id) else { + return; + }; + let Some(tx) = senders.get(&(id, unit.partition as usize)) else { + return; + }; + // The feed channels carry the protocol's `WorkUnitMsg`, so the decoded proto frame is mapped + // into it the same way the gRPC worker's `decode_work_unit` does. + let _ = tx.send(Ok(WorkUnitMsg { + id, + partition: unit.partition as usize, + body: unit.body, + created_timestamp_unix_nanos: unit.created_timestamp_unix_nanos as usize, + sent_timestamp_unix_nanos: unit.sent_timestamp_unix_nanos as usize, + received_timestamp_unix_nanos: unit.received_timestamp_unix_nanos as usize, + processed_timestamp_unix_nanos: unit.processed_timestamp_unix_nanos as usize, + })); +} + +fn fail_feed_senders(senders: &RemoteWorkUnitFeedTxs, reason: &str) { + for tx in senders.values() { + let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); + } +} + +#[derive(Default)] +struct SetPlanRegistry { + map: HashMap<(u32, u32), SetPlanSlot>, + /// Set when the inbox scope died: plans come from the leader's proc, so a dead inbox means + /// no plan can arrive. Parked takers get the failure; later takers fail immediately. + dead: Option, +} + +enum SetPlanSlot { + /// A frame that arrived before the task asked for it. + Pending(SetPlanFrame), + /// A task that asked before its frame arrived. + Waiting(tokio::sync::oneshot::Sender>), +} + +impl DrainHandle { + /// Construct a cooperative drain handle. Channel buffers are populated lazily by + /// [`Self::try_drain_pass`] when a frame arrives, or up-front by [`Self::register_channel`] + /// when a consumer needs a buffer to wait on before any frame has come in. + pub(super) fn cooperative( + this_proc: u32, + receivers: Vec<(ReceiverScope, MppReceiver)>, + ) -> Self { + let wrapped = receivers.into_iter().map(Some).collect(); + let (task_metrics_tx, task_metrics_rx) = tokio::sync::mpsc::unbounded_channel(); + Self { + channel_buffers: Mutex::new(ChannelBufferRegistry::default()), + coop_receivers: Mutex::new(wrapped), + this_proc, + feed_registry: Mutex::new(FeedRegistry::default()), + task_metrics_tx, + task_metrics_rx: Mutex::new(Some(task_metrics_rx)), + set_plan_registry: Mutex::new(SetPlanRegistry::default()), + cancelled_streams: Mutex::new(HashSet::default()), + } + } + + /// Record a `Cancel` frame: the consumer abandoned the `(stage_id, partition)` stream, so this + /// proc's producer of it stops. Idempotent. + fn note_cancel(&self, stage_id: u32, partition: u32) { + self.cancelled_streams + .lock() + .unwrap() + .insert((stage_id, partition)); + } + + /// Whether a consumer cancelled the `(stage_id, partition)` stream. Read by the send spin to + /// end a producer's stream cleanly when the consumer stopped reading. + pub(super) fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + self.cancelled_streams + .lock() + .unwrap() + .contains(&(stage_id, partition)) + } + + /// Take the receiving end of the `TaskMetrics` frame stream. The embedder (the leader) + /// drains it into its metrics store; the first caller gets it, later calls get `None`. + pub(super) fn take_task_metrics_receiver( + &self, + ) -> Option> { + self.task_metrics_rx.lock().unwrap().take() + } + + /// Install the senders of one task's feed channels, flushing any units that arrived first. + /// If the task's `FeedEof` (or the inbox death) already came through, the senders drop (or + /// fail) immediately so the consuming streams terminate instead of waiting forever. + pub(super) fn register_work_unit_senders( + &self, + stage_id: u32, + task_number: u32, + senders: RemoteWorkUnitFeedTxs, + ) { + let mut registry = self.feed_registry.lock().unwrap(); + if let Some(reason) = ®istry.dead { + fail_feed_senders(&senders, reason); + return; + } + match registry.map.remove(&(stage_id, task_number)) { + Some(FeedSlot::Pending { units, done }) => { + for unit in units { + forward_unit(&senders, unit); + } + if !done { + registry + .map + .insert((stage_id, task_number), FeedSlot::Active(senders)); + } + } + Some(FeedSlot::Active(_)) | None => { + registry + .map + .insert((stage_id, task_number), FeedSlot::Active(senders)); + } + } + } + + fn route_work_unit(&self, stage_id: u32, task_number: u32, unit: pb::WorkUnit) { + let mut registry = self.feed_registry.lock().unwrap(); + if registry.dead.is_some() { + return; + } + match registry.map.get_mut(&(stage_id, task_number)) { + Some(FeedSlot::Active(senders)) => forward_unit(senders, unit), + Some(FeedSlot::Pending { units, .. }) => units.push(unit), + None => { + registry.map.insert( + (stage_id, task_number), + FeedSlot::Pending { + units: vec![unit], + done: false, + }, + ); + } + } + } + + fn close_feeds(&self, stage_id: u32, task_number: u32) { + let mut registry = self.feed_registry.lock().unwrap(); + match registry.map.get_mut(&(stage_id, task_number)) { + Some(FeedSlot::Active(_)) => { + // Dropping the senders is the close: the consuming streams see end-of-input. + registry.map.remove(&(stage_id, task_number)); + } + Some(FeedSlot::Pending { done, .. }) => *done = true, + None => { + registry.map.insert( + (stage_id, task_number), + FeedSlot::Pending { + units: Vec::new(), + done: true, + }, + ); + } + } + } + + /// Route one decoded `SetPlan` frame to whoever asked for `(stage_id, task_number)`, or + /// buffer it until they do. A duplicate for an already-buffered slot keeps the first frame. + fn route_set_plan(&self, stage_id: u32, task_number: u32, frame: SetPlanFrame) { + let mut registry = self.set_plan_registry.lock().unwrap(); + match registry.map.remove(&(stage_id, task_number)) { + Some(SetPlanSlot::Waiting(tx)) => { + let _ = tx.send(Ok(frame)); + } + Some(pending @ SetPlanSlot::Pending(_)) => { + log::debug!( + "mpp: duplicate SetPlan frame for stage {stage_id} task {task_number}; \ + keeping the first" + ); + registry.map.insert((stage_id, task_number), pending); + } + None => { + registry + .map + .insert((stage_id, task_number), SetPlanSlot::Pending(frame)); + } + } + } + + /// Take the plan delivered for `(stage_id, task_number)`, waiting for its `SetPlan` frame if + /// it has not arrived yet. Something on this proc must keep draining (a pump or a + /// cooperative send spin) or the wait starves; same contract as the feed channels. + pub(super) async fn take_set_plan( + &self, + stage_id: u32, + task_number: u32, + ) -> Result { + let rx = { + let mut registry = self.set_plan_registry.lock().unwrap(); + if let Some(reason) = ®istry.dead { + return Err(DataFusionError::Execution(reason.clone())); + } + match registry.map.remove(&(stage_id, task_number)) { + Some(SetPlanSlot::Pending(frame)) => return Ok(frame), + Some(SetPlanSlot::Waiting(_)) => { + return Err(DataFusionError::Internal(format!( + "mpp: two takers for the SetPlan frame of stage {stage_id} task \ + {task_number}" + ))); + } + None => { + let (tx, rx) = tokio::sync::oneshot::channel(); + registry + .map + .insert((stage_id, task_number), SetPlanSlot::Waiting(tx)); + rx + } + } + }; + rx.await.map_err(|_| { + DataFusionError::Execution( + "mpp: transport torn down before this task's plan arrived".to_string(), + ) + })? + } + + fn scope_for_sender(&self, sender_proc: u32) -> ReceiverScope { + if sender_proc == self.this_proc { + ReceiverScope::SelfLoop + } else { + ReceiverScope::Inbox + } + } + + /// Fail every registered channel fed by `scope` that has not completed yet, and remember the + /// scope as dead so later registrations fail too. Channels whose `Eof` already arrived are + /// untouched: a detach after a clean drain is the normal end of life for a ring. + fn fail_scope(&self, scope: ReceiverScope, reason: &str) { + let to_fail = { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + match scope { + ReceiverScope::Inbox => guard.dead_inbox = true, + ReceiverScope::SelfLoop => guard.dead_self_loop = true, + } + guard + .map + .iter() + .filter(|((sender_proc, _, _), _)| self.scope_for_sender(*sender_proc) == scope) + .map(|(_, buf)| buf.clone()) + .collect::>() + }; + for buf in to_fail { + buf.fail_pending(reason); + } + if scope == ReceiverScope::Inbox { + let mut registry = self.feed_registry.lock().unwrap(); + registry.dead = Some(reason.to_string()); + for (_, slot) in registry.map.drain() { + if let FeedSlot::Active(senders) = slot { + fail_feed_senders(&senders, reason); + } + } + drop(registry); + let mut plans = self.set_plan_registry.lock().unwrap(); + plans.dead = Some(reason.to_string()); + for (_, slot) in plans.map.drain() { + if let SetPlanSlot::Waiting(tx) = slot { + let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); + } + } + } + } + + /// Register (or look up) the channel buffer for `(sender_proc, stage_id, partition)`. + /// The returned `Arc` is the canonical destination for frames matching + /// that key: `try_drain_pass` pushes into the same entry on every `Batch { header, .. }` + /// whose `header.sender_proc()` / `stage_id` / `partition` matches. + pub(super) fn register_channel( + &self, + sender_proc: u32, + stage_id: u32, + partition: u32, + ) -> Arc { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + let scope_dead = match self.scope_for_sender(sender_proc) { + ReceiverScope::Inbox => guard.dead_inbox, + ReceiverScope::SelfLoop => guard.dead_self_loop, + }; + let buf = guard + .map + .entry((sender_proc, stage_id, partition)) + .or_insert_with(|| { + // num_sources stays 1: each (sender_proc, stage, partition) tuple has + // exactly one upstream (the named sender), even though the underlying + // inbox is shared across all senders. + DrainBuffer::new(1) + }) + .clone(); + drop(guard); + if scope_dead { + buf.fail_pending( + "transport receiver detached before this channel's EOF; the producer went away", + ); + } + buf + } + + /// Cancel every registered channel buffer. Called from `Drop` to unblock any consumer waiting on + /// a channel buffer when the handle goes away mid-query. + /// + /// Collects buffer handles under the registry lock, then notifies after releasing + /// it. Notifying inline would block any concurrent [`Self::register_channel`] for + /// as long as it takes to acquire `DrainBuffer::inner` N times. Fine today (single + /// backend thread), but cheap insurance against the multi-thread variant landing + /// later. + fn cancel_channel_buffers(&self) { + let to_cancel = { + let guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + guard.map.values().cloned().collect::>() + }; + for buf in to_cancel { + buf.cancel(); + } + } + + /// Pull batches from each live receiver and demux them into the per-`(stage_id, partition)` + /// channel buffer registry. Called from `DrainGatherStream::poll_next` and from + /// `MppSender::send_batch`'s cooperative spin. Drain work happens on the backend thread + /// (pgrx-safe). No-op for thread-backed handles. + /// + /// Each pass drains *every available* batch from each receiver (up to a safety cap). Pulling + /// only one batch per source per call would mean that under steady producer pressure the + /// cooperative sender's spin-loop can't keep up: we'd fall N:1 behind peers' sends and the + /// mesh would stall once any queue fills. Draining until the receiver reports `Empty` bounds + /// each pass by queue depth rather than by spin-loop iteration count. + /// + /// Returns `Ok(())` once every cooperative receiver has been pulled until `Empty` (or + /// detached). Errors propagate as `Err` so a transport-level failure surfaces at the call + /// site rather than getting silently dropped. + /// + /// Routing rules per outcome: + /// - `Batch { header, batch }`: look up (or lazily create) the + /// `(header.stage_id, header.partition)` channel buffer and push `batch`. + /// - `Eof { header }`: per-channel EOF. Resolve the channel buffer and call + /// `notify_source_done`. Other channels on the same queue keep flowing, + /// so the receiver slot stays live. + /// - `Detached` / `Error`: queue-wide shutdown. Notify every registered + /// channel buffer, mark the handle detached, and drop the slot. + pub fn try_drain_pass(&self) -> Result<(), DataFusionError> { + // Bound per-source pulls per call. The upper limit exists to give + // the caller a chance to re-try its own send between drains — + // otherwise a proc with a very fast peer could drain + // indefinitely on one source and starve its own outbound. + const MAX_BATCHES_PER_SOURCE_PER_PASS: usize = 256; + + let mut slots = self.coop_receivers.lock().unwrap(); + for slot in slots.iter_mut() { + let Some((scope, rx)) = slot.as_ref() else { + continue; + }; + let scope = *scope; + for _ in 0..MAX_BATCHES_PER_SOURCE_PER_PASS { + match rx.try_recv_batch() { + RecvBatchOutcome::Batch { header, batch } => { + let buf = self.register_channel( + header.sender_proc(), + header.stage_id, + header.partition, + ); + buf.push_batch(batch); + } + RecvBatchOutcome::Eof { header } => { + let buf = self.register_channel( + header.sender_proc(), + header.stage_id, + header.partition, + ); + buf.notify_source_done(); + // Other channels may still flow on this queue, so the receiver slot + // stays live. + } + RecvBatchOutcome::WorkUnit { header, mut unit } => { + set_received_time(&mut unit); + self.route_work_unit(header.stage_id, header.partition, unit); + } + RecvBatchOutcome::FeedEof { header } => { + self.close_feeds(header.stage_id, header.partition); + } + RecvBatchOutcome::TaskMetrics { header, metrics } => { + // The embedder may have dropped the receiver; metrics are best-effort. + let _ = + self.task_metrics_tx + .send((header.stage_id, header.partition, metrics)); + } + RecvBatchOutcome::SetPlan { header, frame } => { + self.route_set_plan(header.stage_id, header.partition, frame); + } + RecvBatchOutcome::Cancel { header } => { + self.note_cancel(header.stage_id, header.partition); + } + RecvBatchOutcome::Empty => break, + RecvBatchOutcome::Detached => { + // Only THIS receiver is dead. The drain holds multiple receivers + // (own-inbox MPSC + self-loop in-proc); one going away doesn't + // imply the others have. Fail only the channels this receiver's + // scope feeds, and only those still waiting on their `Eof`: after a + // clean drain the detach is the ring's normal end of life, but a + // channel that never got its `Eof` (producer crash, early sender + // drop) would otherwise spin on `try_pop -> None` forever. + *slot = None; + self.fail_scope( + scope, + "transport receiver detached before this channel's EOF; the \ + producer went away", + ); + break; + } + RecvBatchOutcome::Error(e) => { + // Same scoping as Detached, but the ring reported corruption (or the + // receiver poisoned itself), so even completed siblings can't be + // trusted to have been the last word. Still scope-limited: the other + // receiver is an independent transport. The error also propagates to + // this caller directly. + *slot = None; + self.fail_scope(scope, &format!("transport receiver failed: {e}")); + return Err(e); + } + } + } + } + Ok(()) + } +} + +impl Drop for DrainHandle { + fn drop(&mut self) { + // Unblock any consumer blocked on a channel buffer when the handle is torn down before EOF + // flows naturally (e.g. a query error en route to ExecEndCustomScan). + self.cancel_channel_buffers(); + } +} +/// SPSC channel pair for two use cases: +/// - Unit tests (bounded capacity, exercising backpressure). +/// - Production self-loop slots: when a worker's fragment emits a partition destined for +/// its OWN proc (e.g. peer-mesh hash routing where consumer task t lands on the same +/// worker as producer task t), the DSM layout has no self-pair inbox: a process is +/// not its own peer. The dispatcher routes those self-loops through this in-proc +/// channel, which exposes the same `BatchChannelSender` / `BatchChannelReceiver` +/// surface as the DSM ring so the drain and channel-buffer registry don't need a +/// special case. +/// +/// Production callers pass a very large `capacity` so the channel is effectively unbounded under +/// steady state. The current-thread Tokio runtime interleaves producer and consumer fragments +/// via `yield_now().await`, so backpressure would be benign anyway, but unbounded rules out any +/// chance of self-deadlock if the producer never yields. +pub(super) fn in_proc_channel(capacity: usize) -> (InProcSender, InProcReceiver) { + let (tx, rx) = std::sync::mpsc::sync_channel::>(capacity); + ( + InProcSender { + tx, + send_lock: tokio::sync::Mutex::new(()), + }, + InProcReceiver { rx: Mutex::new(rx) }, + ) +} + +pub(super) struct InProcSender { + tx: std::sync::mpsc::SyncSender>, + /// Per-instance lock so the [`BatchChannelSender::send_lock`] contract holds even when an + /// in-proc channel ends up in a code path that would otherwise need serialization. In-proc + /// `send_bytes` is already atomic (each call pushes a complete `Vec`), so the lock is + /// effectively a no-op here; keeping it uniform with `DsmInboxSender` avoids + /// special-casing the caller. + send_lock: tokio::sync::Mutex<()>, +} + +pub(super) struct InProcReceiver { + // The std::sync::mpsc receiver is !Sync; wrap in a Mutex so the drain + // thread can hold it behind a `Box` (which is + // `Send + Sync`-relaxed by design, but we only need Send for the thread + // hand-off). Tests only ever access from one thread so the Mutex is + // uncontended. + rx: Mutex>>, +} + +impl BatchChannelSender for InProcSender { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError> { + self.tx.send(bytes.to_vec()).map_err(|_| { + DataFusionError::Execution("mpp: in-proc channel detached during send".into()) + }) + } + + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + match self.tx.try_send(bytes.to_vec()) { + Ok(()) => Ok(true), + Err(std::sync::mpsc::TrySendError::Full(_)) => Ok(false), + Err(std::sync::mpsc::TrySendError::Disconnected(_)) => Err(DataFusionError::Execution( + "mpp: in-proc channel detached during try_send".into(), + )), + } + } + + fn send_lock(&self) -> &tokio::sync::Mutex<()> { + &self.send_lock + } +} + +impl BatchChannelReceiver for InProcReceiver { + fn try_recv(&self) -> RecvOutcome { + let rx = self.rx.lock().expect("InProcReceiver mutex poisoned"); + match rx.try_recv() { + Ok(bytes) => RecvOutcome::Bytes(bytes), + Err(std::sync::mpsc::TryRecvError::Empty) => RecvOutcome::Empty, + Err(std::sync::mpsc::TryRecvError::Disconnected) => RecvOutcome::Detached, + } + } +} + +/// Effectively unbounded capacity for self-loop in-proc channels. The +/// `std::sync::mpsc::sync_channel` API requires a numeric capacity; this constant picks one large +/// enough that production workloads won't reach it but small enough that a runaway producer +/// (e.g. infinite-loop bug) won't allocate billions of `Vec` before OOM. +pub(super) const SELF_LOOP_CAPACITY: usize = 1 << 20; + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int32Array, Int64Array, StringArray, UInt64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use std::sync::Arc as StdArc; + use std::thread; + + use std::thread::JoinHandle; + + impl DrainBuffer { + /// Block until a batch is available, EOF is reached, or the buffer is cancelled. + /// + /// INVARIANT: any already-buffered batch is returned *before* honoring either + /// cancellation or all-sources-done. Reordering the queue pop ahead of the cancel/eof + /// check would silently drop buffered data on an otherwise-clean shutdown; the + /// `drain_buffer_drains_buffered_before_eof` test locks this in. + fn pop_front(&self) -> DrainItem { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + loop { + if let Some(batch) = guard.queue.pop_front() { + return DrainItem::Batch(batch); + } + if let Some(msg) = &guard.failed { + return DrainItem::Failed(msg.clone()); + } + if guard.cancelled || guard.sources_done >= guard.num_sources { + return DrainItem::Eof; + } + guard = self.cond.wait(guard).expect("DrainBuffer mutex poisoned"); + } + } + + /// True if `cancel` has been called. The local `drain_loop` consults this; the + /// cooperative production path watches the flag through `notify_source_done` fan-out + /// instead. + fn is_cancelled(&self) -> bool { + self.inner + .lock() + .expect("DrainBuffer mutex poisoned") + .cancelled + } + } + + impl MppSender { + /// Construct a sender with the default `(stage_id=0, partition=0)` header. Used where + /// the header carries no actionable routing info. + fn new(channel: Arc) -> Self { + Self::with_header(channel, MppFrameHeader::batch(0, 0, 0)) + } + + /// Stats-less wrapper around `send_batch_traced`. Production call sites + /// (`ShuffleStream::process_batch`) always pass a `SendBatchStats` so per-peer + /// wall-time shows up in the EOF trace. Wraps the async send in a tiny current-thread + /// Tokio runtime so `#[test]` functions don't have to be `#[tokio::test]` and the + /// OS-thread-spawning test harnesses don't have to plumb an async runtime themselves. + fn send_batch(&self, batch: &RecordBatch) -> Result<(), DataFusionError> { + let mut stats = SendBatchStats::default(); + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .expect("test tokio runtime build"); + rt.block_on(self.send_batch_traced(batch, &mut stats)) + } + } + + /// Configuration for `spawn_drain_thread`. pgrx panics on any pg FFI call (including + /// `shm_mq_receive`) from a non-backend thread, so production never spawns a drain thread — + /// see [`DrainHandle::cooperative`] for the cooperative path. + struct DrainConfig { + /// Receivers to drain. Ownership moves into the spawned thread. + receivers: Vec, + /// Destination buffer. + buffer: Arc, + /// How long to sleep when every receiver is empty but some are still attached. Tuning: + /// small values reduce end-of-batch latency but raise CPU; 1 ms is a safe default until + /// we integrate with WaitLatch. + idle_sleep: Duration, + } + + impl DrainConfig { + fn new(receivers: Vec, buffer: Arc) -> Self { + Self { + receivers, + buffer, + idle_sleep: Duration::from_millis(1), + } + } + } + + /// Spawn the dedicated drain thread. The thread round-robins through every receiver with + /// non-blocking `try_recv`, pushes decoded batches into `buffer`, and marks each source done + /// as soon as it observes a detach or decode error. When every source is done, the thread + /// exits. + fn spawn_drain_thread(config: DrainConfig) -> JoinHandle> { + thread::spawn(move || drain_loop(config)) + } + + /// RAII wrapper: owns the drain thread's `JoinHandle` and the buffer it writes into. + /// `Drop` cancels the buffer (unblocking the consumer) and joins the thread, so the thread + /// can never outlive the test scope even on a panic. + struct ThreadedDrainHandle { + buffer: Arc, + join: Mutex>>>, + } + + impl ThreadedDrainHandle { + fn spawn(config: DrainConfig) -> Self { + let buffer = Arc::clone(&config.buffer); + let join = spawn_drain_thread(config); + Self { + buffer, + join: Mutex::new(Some(join)), + } + } + } + + impl Drop for ThreadedDrainHandle { + fn drop(&mut self) { + self.buffer.cancel(); + if let Some(join) = self.join.lock().unwrap().take() { + let _ = join.join(); + } + } + } + + /// Test-only thread-backed drain. Writes every observed frame into a single shared + /// [`DrainBuffer`] with `num_sources = N`. Per-channel `Eof` frames are treated as "this source + /// is done" rather than "this logical channel within the source is done"; sufficient for unit + /// tests that don't exercise per-channel demux. Production drains route through + /// [`DrainHandle::try_drain_pass`] (cooperative variant), which keys on the frame header. Tests + /// that need to validate production demux must use [`DrainHandle::cooperative`] and call + /// `try_drain_pass` directly. + fn drain_loop(config: DrainConfig) -> Result<(), DataFusionError> { + let DrainConfig { + receivers, + buffer, + idle_sleep, + } = config; + + let mut done = vec![false; receivers.len()]; + loop { + // Observe cancellation before each pass so a `DrainHandle::drop` with + // live peer senders tears down cleanly. Without this check, the drain + // thread would spin `try_recv` forever because no source has detached. + if buffer.is_cancelled() { + return Ok(()); + } + + let mut got_any = false; + let mut all_done = true; + for (i, rx) in receivers.iter().enumerate() { + if done[i] { + continue; + } + all_done = false; + match rx.try_recv_batch() { + RecvBatchOutcome::Batch { header: _, batch } => { + got_any = true; + buffer.push_batch(batch); + } + RecvBatchOutcome::Eof { header: _ } => { + // Per-channel Eof frame: single-channel positional design + // treats it as a source-done signal. See `try_drain_pass`. + done[i] = true; + buffer.notify_source_done(); + } + // Control frames have no place in the test-only single-buffer path. + RecvBatchOutcome::WorkUnit { .. } + | RecvBatchOutcome::FeedEof { .. } + | RecvBatchOutcome::TaskMetrics { .. } + | RecvBatchOutcome::SetPlan { .. } + | RecvBatchOutcome::Cancel { .. } => {} + RecvBatchOutcome::Empty => {} + RecvBatchOutcome::Detached => { + done[i] = true; + buffer.notify_source_done(); + } + RecvBatchOutcome::Error(e) => { + // Treat a decode error as a fatal detach for this source + // so the consumer can observe EOF and abort the query. + done[i] = true; + buffer.notify_source_done(); + return Err(e); + } + } + } + + if all_done { + return Ok(()); + } + if !got_any { + thread::sleep(idle_sleep); + } + } + } + + fn sample_batch(rows: i32) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ])); + let ids = Int32Array::from_iter_values(0..rows); + let names = StringArray::from_iter_values((0..rows).map(|i| format!("n{i}"))); + RecordBatch::try_new(schema, vec![StdArc::new(ids), StdArc::new(names)]).unwrap() + } + + #[test] + fn frame_round_trips_a_batch_with_header() { + let orig = sample_batch(64); + let header = MppFrameHeader::batch(7, 3, 0); + let mut buf = Vec::with_capacity(1024); + encode_frame_into(header, &orig, &mut buf).expect("encode_frame"); + + let (parsed, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::Batch); + let FrameBody::Batch(decoded) = body else { + panic!("Batch frame must carry a batch payload"); + }; + assert_eq!(decoded.num_rows(), 64); + assert_eq!(decoded.schema(), orig.schema()); + assert_eq!(decoded.num_columns(), orig.num_columns()); + for col in 0..orig.num_columns() { + assert_eq!(orig.column(col).as_ref(), decoded.column(col).as_ref()); + } + } + + #[test] + fn frame_round_trips_eof() { + let mut buf = Vec::new(); + encode_eof_frame_into(2, 5, 0, &mut buf).expect("encode_eof"); + assert_eq!(buf.len(), MPP_FRAME_HEADER_SIZE); + + let (header, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(header, MppFrameHeader::eof(2, 5, 0)); + assert_eq!(header.kind().unwrap(), MppFrameKind::Eof); + assert!(matches!(body, FrameBody::Eof)); + } + + #[test] + fn frame_round_trips_cancel() { + let header = MppFrameHeader::cancel(4, 2, 0); + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + + let (parsed, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::Cancel); + assert_eq!(parsed.stage_id, 4); + assert_eq!(parsed.partition, 2); + assert!(matches!(body, FrameBody::Cancel)); + } + + #[test] + fn drain_records_cancel_scoped_to_its_stream() { + let drain = DrainHandle::cooperative(0, Vec::new()); + assert!(!drain.stream_cancelled(7, 1)); + drain.note_cancel(7, 1); + assert!(drain.stream_cancelled(7, 1)); + // A cancel for one stream leaves the others alive, the way gRPC closes one stream: same + // stage but a different partition, and a different stage, both stay live. + assert!(!drain.stream_cancelled(7, 2)); + assert!(!drain.stream_cancelled(8, 1)); + } + + /// The producer-side half of the early-termination fix: a producer blocked on a full outbound + /// ends its send cleanly once a `Cancel` for its `(stage, partition)` stream lands on its inbox. + /// The test hangs if the spin doesn't observe the cancel. + #[tokio::test(flavor = "current_thread")] + async fn producer_send_ends_when_consumer_cancels_the_stream() { + // Outbound to a consumer that never drains: capacity 1, so the second send finds it full. + let (out_tx, _out_rx) = in_proc_channel(1); + // The producer's inbox, where the consumer's `Cancel` frame lands. + let (inbox_tx, inbox_rx) = in_proc_channel(4); + let drain = Arc::new(DrainHandle::cooperative( + 1, + vec![(ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox_rx)))], + )); + // Producer of the `(stage 7, partition 0)` stream. + let sender = MppSender::with_header(Arc::new(out_tx), MppFrameHeader::batch(7, 0, 1)) + .with_cooperative_drain(Arc::clone(&drain) as Arc); + + // The consumer abandons that stream (it stopped reading before EOF): a `Cancel` reaches the + // inbox. + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + MppFrameHeader::cancel(7, 0, 0).write_to(&mut buf); + inbox_tx.send_bytes(&buf).unwrap(); + + let mut stats = SendBatchStats::default(); + // First batch fills the one slot. + sender + .send_batch_traced(&sample_batch(1), &mut stats) + .await + .unwrap(); + // The second batch would block forever on the full slot. Instead the spin drains the inbox, + // sees the cancel, and ends the send cleanly. + sender + .send_batch_traced(&sample_batch(1), &mut stats) + .await + .unwrap(); + } + + #[test] + fn frame_round_trips_a_work_unit_and_feed_eof() { + let unit = pb::WorkUnit { + id: vec![7; 16], + partition: 4, + body: vec![1, 2, 3], + created_timestamp_unix_nanos: 11, + sent_timestamp_unix_nanos: 22, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + }; + let header = MppFrameHeader::work_unit(3, 1, 0); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &unit, &mut buf).expect("encode work unit"); + let (parsed, body) = decode_frame(&buf).expect("decode work unit"); + assert_eq!(parsed, header); + let FrameBody::WorkUnit(decoded) = body else { + panic!("WorkUnit frame must carry a unit"); + }; + assert_eq!(decoded, unit); + + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + MppFrameHeader::feed_eof(3, 1, 0).write_to(&mut buf); + let (parsed, body) = decode_frame(&buf).expect("decode feed eof"); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::FeedEof); + assert!(matches!(body, FrameBody::FeedEof)); + } + + #[test] + fn frame_round_trips_task_metrics() { + let metrics = pb::TaskMetrics { + pre_order_plan_metrics: vec![], + task_metrics: None, + }; + let header = MppFrameHeader::task_metrics(2, 0, 1); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &metrics, &mut buf).expect("encode metrics"); + let (parsed, body) = decode_frame(&buf).expect("decode metrics"); + assert_eq!(parsed, header); + assert!(matches!(body, FrameBody::TaskMetrics(m) if m == metrics)); + } + + #[test] + fn work_units_buffer_until_registration_and_feed_eof_closes() { + fn unit(partition: u64) -> pb::WorkUnit { + pb::WorkUnit { + id: vec![9; 16], + partition, + body: vec![], + created_timestamp_unix_nanos: 0, + sent_timestamp_unix_nanos: 0, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + } + } + let drain = DrainHandle::cooperative(1, vec![]); + + // Units arriving before registration must buffer, not drop. + drain.route_work_unit(5, 0, unit(0)); + drain.route_work_unit(5, 0, unit(0)); + + let id = crate::common::deserialize_uuid(&[9; 16]).unwrap(); + let mut channels = crate::work_unit_feed::RemoteWorkUnitFeedRegistry::default(); + channels.add(id, 1); + let mut rx = channels + .receivers + .get(&(id, 0)) + .unwrap() + .lock() + .unwrap() + .take() + .unwrap(); + drain.register_work_unit_senders(5, 0, channels.senders); + + assert!(rx.try_recv().unwrap().is_ok()); + assert!(rx.try_recv().unwrap().is_ok()); + // Still open: the producer may send more units. + assert!(rx.try_recv().is_err()); + + // FeedEof drops the senders, which ends the stream. + drain.route_work_unit(5, 0, unit(0)); + drain.close_feeds(5, 0); + assert!(rx.try_recv().unwrap().is_ok()); + assert!(matches!( + rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + )); + } + + #[test] + fn frame_round_trips_a_set_plan_with_headers() { + let set_plan = pb::SetPlanRequest { + plan_proto: vec![1, 2, 3, 4], + task_count: 2, + task_key: Some(pb::TaskKey { + query_id: vec![5; 16], + stage_id: 3, + task_number: 1, + }), + work_unit_feed_declarations: vec![], + target_worker_url: "inprocess://worker/1".to_string(), + query_start_time_ns: 42, + }; + let mut headers = http::HeaderMap::new(); + headers.insert("x-datafusion-distributed-config", "abc".parse().unwrap()); + headers.append("x-repeated", "one".parse().unwrap()); + headers.append("x-repeated", "two".parse().unwrap()); + + let frame = SetPlanFrame::from_parts(set_plan.clone(), &headers).expect("from_parts"); + let header = MppFrameHeader::set_plan(3, 1, 0); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &frame, &mut buf).expect("encode set plan"); + let (parsed, body) = decode_frame(&buf).expect("decode set plan"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::SetPlan); + let FrameBody::SetPlan(decoded) = body else { + panic!("SetPlan frame must carry a SetPlanFrame"); + }; + let (decoded_plan, decoded_headers) = decoded.into_parts().expect("into_parts"); + assert_eq!(decoded_plan, set_plan); + assert_eq!(decoded_headers, headers); + } + + fn sample_set_plan_frame(plan_proto: Vec) -> SetPlanFrame { + SetPlanFrame { + set_plan: Some(pb::SetPlanRequest { + plan_proto, + task_count: 1, + task_key: None, + work_unit_feed_declarations: vec![], + target_worker_url: String::new(), + query_start_time_ns: 0, + }), + header_keys: vec![], + header_values: vec![], + } + } + + #[tokio::test] + async fn set_plan_serves_taker_in_either_arrival_order() { + let drain = DrainHandle::cooperative(1, vec![]); + + // Frame first: the take resolves from the pending slot. + drain.route_set_plan(7, 0, sample_set_plan_frame(vec![1])); + let frame = drain.take_set_plan(7, 0).await.expect("pending take"); + assert_eq!(frame.set_plan.unwrap().plan_proto, vec![1]); + + // Taker first: the frame fulfills the parked oneshot. + let take = drain.take_set_plan(7, 1); + futures::pin_mut!(take); + assert!(futures::poll!(take.as_mut()).is_pending()); + drain.route_set_plan(7, 1, sample_set_plan_frame(vec![2])); + let frame = take.await.expect("waiting take"); + assert_eq!(frame.set_plan.unwrap().plan_proto, vec![2]); + } + + #[tokio::test] + async fn set_plan_take_fails_when_the_inbox_dies() { + let drain = DrainHandle::cooperative(1, vec![]); + let take = drain.take_set_plan(9, 0); + futures::pin_mut!(take); + assert!(futures::poll!(take.as_mut()).is_pending()); + drain.fail_scope(ReceiverScope::Inbox, "producer went away"); + let err = take.await.expect_err("dead inbox must fail the take"); + assert!(format!("{err}").contains("producer went away")); + // Later takers fail immediately. + let err = drain.take_set_plan(9, 1).await.expect_err("dead registry"); + assert!(format!("{err}").contains("producer went away")); + } + + #[test] + fn frame_rejects_short_message() { + let too_short = vec![0u8; MPP_FRAME_HEADER_SIZE - 1]; + let err = decode_frame(&too_short).expect_err("short frame must fail"); + assert!(format!("{err}").contains("too short")); + } + + #[test] + fn frame_rejects_bad_magic() { + // Explicit non-zero, non-magic prefix. Don't rely on the + // happenstance that 0u32 != MPP_FRAME_MAGIC. + let mut bad = vec![0u8; MPP_FRAME_HEADER_SIZE]; + bad[0..4].copy_from_slice(&0xCAFEBABE_u32.to_le_bytes()); + let err = decode_frame(&bad).expect_err("bad magic must fail"); + assert!(format!("{err}").contains("bad frame magic")); + bad[0..4].copy_from_slice(&0xDEADBEEF_u32.to_le_bytes()); + let err = decode_frame(&bad).expect_err("bad magic must fail"); + assert!(format!("{err}").contains("bad frame magic")); + } + + #[test] + fn frame_rejects_unknown_kind() { + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: 0x42, // unknown kind byte, no reserved bits set + stage_id: 0, + partition: 0, + }; + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + let err = decode_frame(&buf).expect_err("unknown kind must fail"); + assert!(format!("{err}").contains("unknown frame kind")); + } + + #[test] + fn frame_rejects_reserved_flag_bits() { + // Reserved range is bits 8..16. Bits 16..32 are sender_proc and must NOT trip the + // reserved check. Cover both boundaries of the reserved range. + for bit in [0x0000_0100u32, 0x0000_8000u32] { + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: bit, // kind byte 0 (Batch), reserved bit set, no sender_proc + stage_id: 0, + partition: 0, + }; + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + let err = decode_frame(&buf).expect_err(&format!("reserved bit {bit:#x} must fail")); + assert!( + format!("{err}").contains("reserved frame flag bits"), + "bit {bit:#x}: {err}" + ); + } + } + + #[test] + fn frame_kind_coexists_with_max_sender_proc() { + // Negative-space companion to frame_rejects_reserved_flag_bits: setting every bit in + // 16..32 (= max sender_proc) along with kind=Eof in bit 0 must parse cleanly without + // tripping the reserved-bits check, and sender_proc()/kind() must read both back. + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: 0xFFFF_0001, // Eof in low byte, max sender_proc in high half, reserved=0 + stage_id: 0, + partition: 0, + }; + assert_eq!(header.kind().unwrap(), MppFrameKind::Eof); + assert_eq!(header.sender_proc(), MPP_MAX_SENDER_PROC); + } + + #[test] + fn frame_sender_proc_round_trip() { + // sender_proc lives in flags bits 16..32 and shouldn't collide with kind or reserved. + for &sp in &[0u32, 1, 7, 255, 256, 1023, 65534, MPP_MAX_SENDER_PROC] { + let header = MppFrameHeader::batch(11, 5, sp); + assert_eq!(header.sender_proc(), sp, "batch round-trip sp={sp}"); + assert_eq!(header.kind().unwrap(), MppFrameKind::Batch); + + let mut buf = Vec::with_capacity(MPP_FRAME_HEADER_SIZE); + let payload = sample_batch(8); + encode_frame_into(header, &payload, &mut buf).expect("encode batch"); + let (parsed, _) = decode_frame(&buf).expect("decode batch"); + assert_eq!(parsed.sender_proc(), sp, "decoded batch sender_proc"); + + let mut eof_buf = Vec::new(); + encode_eof_frame_into(11, 5, sp, &mut eof_buf).expect("encode eof"); + let (parsed_eof, _) = decode_frame(&eof_buf).expect("decode eof"); + assert_eq!(parsed_eof.sender_proc(), sp, "decoded eof sender_proc"); + assert_eq!(parsed_eof.kind().unwrap(), MppFrameKind::Eof); + } + } + + #[test] + fn frame_eof_with_payload_is_rejected() { + let mut buf = Vec::with_capacity(32); + encode_eof_frame_into(0, 0, 0, &mut buf).expect("encode_eof"); + buf.push(0xAB); // smuggle a payload byte after the Eof header + let err = decode_frame(&buf).expect_err("Eof+payload must fail"); + assert!(format!("{err}").contains("payload-less frame carries payload")); + } + + #[test] + fn codec_round_trips_many_batch_sizes() { + let mut buf = Vec::with_capacity(1024); + for rows in [0, 1, 7, 64, 1024] { + let orig = sample_batch(rows); + encode_frame_into(MppFrameHeader::batch(0, 0, 0), &orig, &mut buf).expect("encode"); + let (_header, body) = decode_frame(&buf).expect("decode"); + let FrameBody::Batch(decoded) = body else { + panic!("Batch frame must carry a batch payload"); + }; + assert_eq!(orig.num_rows(), decoded.num_rows()); + } + } + + #[test] + fn drain_buffer_pop_returns_pushed_batches_in_order() { + let buf = DrainBuffer::new(1); + buf.push_batch(sample_batch(3)); + buf.push_batch(sample_batch(5)); + buf.notify_source_done(); + + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 3), + DrainItem::Eof => panic!("expected batch"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 5), + DrainItem::Eof => panic!("expected batch"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + matches!(buf.pop_front(), DrainItem::Eof); + } + + #[test] + fn drain_buffer_pop_blocks_until_push_then_eof() { + let buf = DrainBuffer::new(2); + let producer = StdArc::clone(&buf); + let handle = thread::spawn(move || { + thread::sleep(Duration::from_millis(20)); + producer.push_batch(sample_batch(2)); + producer.notify_source_done(); + thread::sleep(Duration::from_millis(20)); + producer.notify_source_done(); + }); + + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 2), + DrainItem::Eof => panic!("expected batch first"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + handle.join().unwrap(); + } + + #[test] + fn drain_buffer_cancel_unblocks_waiter() { + let buf = DrainBuffer::new(1); + let canceller = StdArc::clone(&buf); + let handle = thread::spawn(move || { + thread::sleep(Duration::from_millis(20)); + canceller.cancel(); + }); + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + handle.join().unwrap(); + } + + #[test] + fn in_proc_channel_round_trips_through_mpp_sender_receiver() { + let (tx, rx) = in_proc_channel(8); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + + sender.send_batch(&sample_batch(4)).unwrap(); + std::mem::drop(sender); + + match receiver.try_recv_batch() { + RecvBatchOutcome::Batch { header: _, batch } => assert_eq!(batch.num_rows(), 4), + other => panic!("expected batch, got {other:?}"), + } + assert!(matches!( + receiver.try_recv_batch(), + RecvBatchOutcome::Detached + )); + } + + #[test] + fn drain_thread_drains_single_source() { + let (tx, rx) = in_proc_channel(4); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + + let join = spawn_drain_thread(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + thread::spawn(move || { + for rows in [1, 2, 3, 4, 5] { + sender.send_batch(&sample_batch(rows)).unwrap(); + } + // Drop sender to signal EOF + }) + .join() + .unwrap(); + + let mut received = Vec::new(); + while let DrainItem::Batch(b) = buffer.pop_front() { + received.push(b.num_rows()); + } + assert_eq!(received, vec![1, 2, 3, 4, 5]); + join.join().unwrap().unwrap(); + } + + #[test] + fn drain_handle_shutdown_joins_cleanly() { + let (tx, rx) = in_proc_channel(4); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + let handle = + ThreadedDrainHandle::spawn(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + sender.send_batch(&sample_batch(2)).unwrap(); + std::mem::drop(sender); // detach + // Pop the one batch + assert!(matches!(buffer.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buffer.pop_front(), DrainItem::Eof)); + // Drop drives production teardown (cancel + join). Test passes if + // this returns without hanging. + std::mem::drop(handle); + } + + #[test] + fn drain_handle_drop_cancels_and_joins() { + // Build a drain that never detaches (we keep the sender alive), then + // drop the handle. The Drop impl must cancel the buffer and join the + // thread without hanging. + let (tx, rx) = in_proc_channel(4); + let _sender_kept_alive = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + let handle = + ThreadedDrainHandle::spawn(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + // Simulate consumer path error: drop the handle without calling + // shutdown(). The drain thread must exit before drop returns. + let start = Instant::now(); + drop(handle); + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(2), + "ThreadedDrainHandle::drop took too long: {elapsed:?}" + ); + // Consumer observes EOF because cancel was called. + assert!(matches!(buffer.pop_front(), DrainItem::Eof)); + } + + #[test] + fn drain_thread_drains_n2_mesh_100k_batches() { + // Simulates a 2-proc mesh under load. Each of two producers + // pushes 50_000 small batches through a bounded channel; the drain + // thread interleaves and the consumer reads EOF exactly after + // receiving all 100_000 batches. Exercises backpressure (bounded + // capacity = 16) without deadlock. + const PER_SOURCE: usize = 50_000; + let (tx0, rx0) = in_proc_channel(16); + let (tx1, rx1) = in_proc_channel(16); + let receivers = vec![ + MppReceiver::new(Box::new(rx0)), + MppReceiver::new(Box::new(rx1)), + ]; + let buffer = DrainBuffer::new(2); + let drain_join = spawn_drain_thread(DrainConfig::new(receivers, StdArc::clone(&buffer))); + + let tx0_send = MppSender::new(Arc::new(tx0)); + let tx1_send = MppSender::new(Arc::new(tx1)); + let batch_template = sample_batch(1); + + let p0 = { + let b = batch_template.clone(); + thread::spawn(move || { + for _ in 0..PER_SOURCE { + tx0_send.send_batch(&b).unwrap(); + } + }) + }; + let p1 = { + let b = batch_template.clone(); + thread::spawn(move || { + for _ in 0..PER_SOURCE { + tx1_send.send_batch(&b).unwrap(); + } + }) + }; + + let mut total = 0usize; + while let DrainItem::Batch(_) = buffer.pop_front() { + total += 1; + } + assert_eq!(total, 2 * PER_SOURCE); + p0.join().unwrap(); + p1.join().unwrap(); + drain_join.join().unwrap().unwrap(); + } + + #[test] + fn drain_buffer_drains_buffered_before_eof() { + // Even if all sources have finished and cancel fires, any already- + // buffered batches must be observed before Eof. + let buf = DrainBuffer::new(1); + buf.push_batch(sample_batch(1)); + buf.push_batch(sample_batch(1)); + buf.notify_source_done(); + buf.cancel(); + + assert!(matches!(buf.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buf.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + } + + // --------------------------------------------------------------------- + // Throughput microbenches. + // + // These are `#[ignore]` by default because they spin for seconds and spam stdout. Run with: + // + // cargo test --package pg_search --release \ + // postgres::customscan::mpp::transport::tests::throughput \ + // -- --ignored --nocapture + // + // They help us bound the transport layer's cost independently of DataFusion/Tantivy. All use + // the `in_proc_channel` backend (same `MppSender`/`MppReceiver` trait boundary as the shm_mq + // one), so numbers here are an optimistic ceiling. shm_mq adds the ring-buffer copy + + // cross-process notification cost on top. If these numbers are already below the row rate + // the real query needs, we know IPC encode + channel handoff is the bottleneck without + // needing CI data. + // --------------------------------------------------------------------- + + /// Row shape matching the post-Partial shuffle in + /// `aggregate_join_groupby`: a grouping key (title string) plus two + /// partial-aggregate accumulators (COUNT u64, SUM i64). + fn postagg_shape_batch(rows: usize) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("title", DataType::Utf8, false), + Field::new("count_partial", DataType::UInt64, false), + Field::new("sum_partial", DataType::Int64, false), + ])); + // Titles averaging ~30 bytes, typical for the docs dataset. + let titles = StringArray::from_iter_values( + (0..rows).map(|i| format!("file_{i:012}_title_with_some_length")), + ); + let counts = UInt64Array::from_iter_values((0..rows as u64).map(|i| i % 64 + 1)); + let sums = Int64Array::from_iter_values((0..rows as i64).map(|i| i * 1024)); + RecordBatch::try_new( + schema, + vec![StdArc::new(titles), StdArc::new(counts), StdArc::new(sums)], + ) + .unwrap() + } + + /// Row shape matching the probe-side shuffle in the same query: + /// `pages.fileId` (u64) plus `pages.sizeInBytes` (i64). + fn probe_shape_batch(rows: usize) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("fileId", DataType::UInt64, false), + Field::new("sizeInBytes", DataType::Int64, false), + ])); + let ids = + UInt64Array::from_iter_values((0..rows as u64).map(|i| i.wrapping_mul(2654435761))); + let sizes = Int64Array::from_iter_values((0..rows as i64).map(|i| i * 37)); + RecordBatch::try_new(schema, vec![StdArc::new(ids), StdArc::new(sizes)]).unwrap() + } + + fn bench_throughput( + label: &str, + make_batch: fn(usize) -> RecordBatch, + batch_rows: usize, + total_rows: usize, + ) { + let batches = total_rows.div_ceil(batch_rows); + let template = make_batch(batch_rows); + // Encode once up front so we also report pure-encode throughput + // separately. Real queries encode inside the hot path per batch. + let enc_start = Instant::now(); + let mut enc_bytes = 0usize; + let mut enc_buf = Vec::with_capacity(1024); + for _ in 0..batches { + encode_frame_into(MppFrameHeader::batch(0, 0, 0), &template, &mut enc_buf) + .expect("encode"); + enc_bytes += enc_buf.len(); + } + let enc_elapsed = enc_start.elapsed(); + + // N=2 mesh: two senders, one drain thread, one consumer. Matches + // the gb_postagg / gb_right topology in the real query. + let (tx0, rx0) = in_proc_channel(16); + let (tx1, rx1) = in_proc_channel(16); + let receivers = vec![ + MppReceiver::new(Box::new(rx0)), + MppReceiver::new(Box::new(rx1)), + ]; + let buffer = DrainBuffer::new(2); + let drain_join = spawn_drain_thread(DrainConfig::new(receivers, StdArc::clone(&buffer))); + let tx0_send = MppSender::new(Arc::new(tx0)); + let tx1_send = MppSender::new(Arc::new(tx1)); + + let per_source = batches / 2; + let round_trip_start = Instant::now(); + let p0 = { + let b = template.clone(); + thread::spawn(move || { + for _ in 0..per_source { + tx0_send.send_batch(&b).unwrap(); + } + }) + }; + let p1 = { + let b = template.clone(); + thread::spawn(move || { + for _ in 0..per_source { + tx1_send.send_batch(&b).unwrap(); + } + }) + }; + + let mut got_rows = 0usize; + let mut got_batches = 0usize; + while let DrainItem::Batch(b) = buffer.pop_front() { + got_rows += b.num_rows(); + got_batches += 1; + } + p0.join().unwrap(); + p1.join().unwrap(); + drain_join.join().unwrap().unwrap(); + let rt_elapsed = round_trip_start.elapsed(); + + let enc_mb_per_s = (enc_bytes as f64 / (1024.0 * 1024.0)) / enc_elapsed.as_secs_f64(); + let enc_rows_per_s = (batches * batch_rows) as f64 / enc_elapsed.as_secs_f64(); + let rt_rows_per_s = got_rows as f64 / rt_elapsed.as_secs_f64(); + let rt_bytes_total_mb = enc_bytes as f64 / (1024.0 * 1024.0); + let rt_mb_per_s = rt_bytes_total_mb / rt_elapsed.as_secs_f64(); + let per_batch_us = rt_elapsed.as_micros() as f64 / got_batches as f64; + + println!( + "[throughput] {label:<18} batch_rows={batch_rows:<5} batches={got_batches:<6} rows={got_rows} \ + encode_only: {enc_rows_per_s:>11.0} rows/s {enc_mb_per_s:>7.1} MB/s | \ + round_trip: {rt_rows_per_s:>11.0} rows/s {rt_mb_per_s:>7.1} MB/s ({per_batch_us:.1}us/batch)" + ); + } + + #[test] + #[ignore] + fn throughput_postagg_shape() { + // Sweeps batch size to show per-batch fixed cost vs per-row cost. + // 1.25M total rows ≈ what one proc ships through gb_postagg at + // 25M scale. 625K per proc × 2 = 1.25M. + for batch_rows in [128, 512, 2048, 8192, 32_768] { + bench_throughput("postagg", postagg_shape_batch, batch_rows, 1_250_000); + } + } + + #[test] + #[ignore] + fn throughput_probe_shape() { + // 12.5M total rows ≈ what one proc ships through gb_right at 25M. + for batch_rows in [128, 512, 2048, 8192, 32_768] { + bench_throughput("probe", probe_shape_batch, batch_rows, 12_500_000); + } + } + + // --------------------------------------------------------------------- + // Per-`(stage_id, partition)` channel buffer registry on the cooperative `DrainHandle`. + // + // Producers stamp `MppFrameHeader::batch(stage_id, partition)` on every outgoing frame, and + // the receiver-side cooperative drain demuxes by header into a channel buffer per + // `(stage_id, partition)`. These tests use the `in_proc_channel` backend to drive + // `try_drain_pass` from the test thread. That mirrors how the production path runs the drain + // inline from `DrainGatherStream::poll_next` on the backend thread. + // --------------------------------------------------------------------- + + /// Drain a `DrainHandle::cooperative` to completion: poll until every receiver returns + /// `Empty`. With the `in_proc_channel` test backend the drain observes `Detached` once the + /// producer drops its sender, so a bounded loop of `try_drain_pass` calls is enough to flush + /// everything the producer wrote. + fn drain_until_detached(handle: &DrainHandle) { + for _ in 0..64 { + handle.try_drain_pass().expect("try_drain_pass"); + } + } + + #[test] + fn drain_handle_demuxes_frames_by_header() { + // One queue carrying two channels: `(0, 0)` and `(0, 1)`. Each + // channel buffer receives only its own batches. Per-channel EOF is out of scope + // here. See `drain_handle_eof_frame_closes_one_channel` for explicit-Eof routing + // and `drain_handle_drop_cancels_registered_channel_buffers` for the + // teardown-EOF contract. + let (tx, rx) = in_proc_channel(8); + let base = MppSender::new(Arc::new(tx)); + let s00 = base.clone_with_header(MppFrameHeader::batch(0, 0, 0)); + let s01 = base.clone_with_header(MppFrameHeader::batch(0, 1, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s00.send_batch(&sample_batch(2)).unwrap(); + s01.send_batch(&sample_batch(7)).unwrap(); + s00.send_batch(&sample_batch(3)).unwrap(); + drop(s00); + drop(s01); + drop(base); + + let buf00 = handle.register_channel(0, 0, 0); + let buf01 = handle.register_channel(0, 0, 1); + + drain_until_detached(&handle); + + let mut p0_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf00.try_pop() { + p0_rows.push(b.num_rows()); + } + let mut p1_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf01.try_pop() { + p1_rows.push(b.num_rows()); + } + assert_eq!(p0_rows, vec![2, 3]); + assert_eq!(p1_rows, vec![7]); + } + + #[test] + fn drain_handle_eof_frame_closes_one_channel() { + // An `Eof` frame on `(0, 0)` closes that channel buffer while frames on + // `(0, 1)` continue to flow on the same queue. `Detached` doesn't broadcast a + // registry-wide EOF, so `(0, 1)` surfaces EOF only when the handle's `Drop` + // runs `cancel_channel_buffers`. + let (tx, rx) = in_proc_channel(8); + let tx_arc: Arc = Arc::new(tx); + let s00 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 0, 0)); + let s01 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 1, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s00.send_batch(&sample_batch(4)).unwrap(); + let mut eof_buf = Vec::new(); + encode_eof_frame_into(0, 0, 0, &mut eof_buf).unwrap(); + tx_arc.send_bytes(&eof_buf).unwrap(); + s01.send_batch(&sample_batch(6)).unwrap(); + + let buf00 = handle.register_channel(0, 0, 0); + let buf01 = handle.register_channel(0, 0, 1); + + drop(s00); + drop(s01); + drop(tx_arc); + drain_until_detached(&handle); + + match buf00.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 4), + other => panic!("expected (0,0) batch, got {other:?}"), + } + assert!(matches!(buf00.try_pop(), Some(DrainItem::Eof))); + + match buf01.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 6), + other => panic!("expected (0,1) batch, got {other:?}"), + } + assert!(buf01.try_pop().is_none()); + drop(handle); + assert!(matches!(buf01.try_pop(), Some(DrainItem::Eof))); + } + + #[test] + fn drain_handle_register_channel_is_idempotent() { + // Two calls for the same key return Arcs pointing to the same + // DrainBuffer instance. + let (_tx, rx) = in_proc_channel(8); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + let first = handle.register_channel(0, 2, 3); + let second = handle.register_channel(0, 2, 3); + assert!(Arc::ptr_eq(&first, &second)); + } + + #[test] + fn drain_handle_demuxes_frames_by_stage_id() { + // Same partition (0) for two different stage ids on the same queue. + // The registry's compound key keeps them on separate channel buffers. + let (tx, rx) = in_proc_channel(8); + let tx_arc: Arc = Arc::new(tx); + let s_stage0 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 0, 0)); + let s_stage1 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(1, 0, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s_stage0.send_batch(&sample_batch(2)).unwrap(); + s_stage1.send_batch(&sample_batch(9)).unwrap(); + s_stage0.send_batch(&sample_batch(4)).unwrap(); + drop(s_stage0); + drop(s_stage1); + drop(tx_arc); + + let buf0 = handle.register_channel(0, 0, 0); + let buf1 = handle.register_channel(0, 1, 0); + + drain_until_detached(&handle); + + let mut stage0_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf0.try_pop() { + stage0_rows.push(b.num_rows()); + } + let mut stage1_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf1.try_pop() { + stage1_rows.push(b.num_rows()); + } + assert_eq!(stage0_rows, vec![2, 4]); + assert_eq!(stage1_rows, vec![9]); + } + + #[test] + fn drain_handle_drop_cancels_registered_channel_buffers() { + // Dropping a cooperative DrainHandle must wake any consumer holding an Arc + // from `register_channel`. Otherwise a query error path that tears down the mesh would + // leave a consumer blocked on a buffer that will never see EOF. + let (_tx, rx) = in_proc_channel(8); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + let buf_a = handle.register_channel(0, 0, 0); + let buf_b = handle.register_channel(0, 7, 3); + // No data ever flows; the handle is just dropped. + drop(handle); + + assert!(matches!(buf_a.try_pop(), Some(DrainItem::Eof))); + assert!(matches!(buf_b.try_pop(), Some(DrainItem::Eof))); + } +} diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index ec74f52c..f51eb643 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -5,6 +5,7 @@ mod work_unit_feed; mod work_unit_feed_provider; mod work_unit_feed_registry; +pub(crate) use remote_work_unit_feed::RemoteWorkUnitFeedTxs; pub(crate) use remote_work_unit_feed::{ RemoteWorkUnitFeedRegistry, build_work_unit_batch_msg, set_work_unit_received_time, set_work_unit_send_time, From 925ff31f7c0f81e907e9c78719ad44c0c62d1ce7 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 1 Jul 2026 11:58:29 -0700 Subject: [PATCH 08/11] Routed dispatched plans to workers over the shm mesh. ShmWorkerChannel::coordinator_channel now ships each stage's SetPlanRequest to the worker proc that owns its task, as a SetPlan frame the worker takes with take_set_plan, so the embedder no longer routes plans out of band. --- src/shm/runtime.rs | 90 +++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 82 insertions(+), 8 deletions(-) diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs index 5f538d83..6b616905 100644 --- a/src/shm/runtime.rs +++ b/src/shm/runtime.rs @@ -37,13 +37,14 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use datafusion::arrow::array::RecordBatch; -use datafusion::common::{DataFusionError, Result}; +use datafusion::common::{DataFusionError, Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; use futures::stream::{self, BoxStream, StreamExt}; use http::HeaderMap; use url::Url; +use crate::common::serialize_uuid; use crate::proto as pb; use crate::work_unit_feed::RemoteWorkUnitFeedTxs; use crate::{ @@ -52,7 +53,10 @@ use crate::{ }; use super::AliveFlag; -use super::transport::{CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppSender}; +use super::transport::{ + CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppFrameHeader, MppSender, + SendBatchStats, SetPlanFrame, +}; /// A proc's outbound senders to each peer inbox, shared between the mesh (for `Cancel` frames) and /// the embedder that owns their lifetime. Indexed by destination `proc_idx`; this proc's own slot @@ -143,6 +147,37 @@ impl MppMesh { *self.cancel_senders.lock().unwrap() = Some(senders); } + /// Route a leader-built `SetPlan` frame to the worker proc that owns `task_number`, over the + /// leader's outbound senders (installed via [`Self::set_cancel_senders`]). This is the unified + /// dispatch path: the coordinator ships each stage's plan through here, so the embedder no + /// longer routes it out of band. Errors if the senders aren't installed or the destination + /// slot is empty; the caller treats a failure as best-effort (the worker starves on its + /// `SetPlan` wait, which the dispatch deadlock detector surfaces). + pub async fn send_set_plan( + self: &Arc, + stage_id: u32, + task_number: u32, + frame: SetPlanFrame, + ) -> Result<()> { + let dest_proc = proc_for_task(self.n_workers(), task_number); + let sender = { + let guard = self.cancel_senders.lock().unwrap(); + let Some(senders) = guard.as_ref() else { + return internal_err!( + "shm mesh: control senders not installed; cannot route SetPlan" + ); + }; + let senders = senders.lock().unwrap(); + let Some(Some(base)) = senders.get(dest_proc as usize) else { + return internal_err!("shm mesh: no control sender for proc {dest_proc}"); + }; + base.clone_with_header(MppFrameHeader::set_plan(stage_id, task_number, 0)) + .with_cooperative_drain(Arc::clone(self) as Arc) + }; + let mut stats = SendBatchStats::default(); + sender.send_set_plan_traced(&frame, &mut stats).await + } + /// Tell the producer on `producer_proc` to stop the `(stage_id, partition)` stream: this proc's /// consumer of it stopped reading before EOF. Ships one `Cancel` frame, leaving the rings /// healthy for metrics and every other stream. A no-op when no senders are installed (the @@ -281,17 +316,56 @@ struct ShmWorkerChannel { #[async_trait] impl WorkerChannel for ShmWorkerChannel { - /// pg_search delivers each worker's plan over DSM, not over this channel, so plan delivery is a - /// no-op (what the old `NoOpDispatch` did). Drain the inbound control stream to exhaustion so the - /// coordinator's keep-alive tail does not wedge, then complete with an empty output stream: task - /// metrics travel back over the mesh, not here. + /// Route each `SetPlanRequest` the coordinator ships to the worker proc that owns its task, as a + /// `SetPlan` frame the worker picks up with `take_set_plan`. The coordinator's keep-alive tail + /// emits nothing after the request, so the loop then blocks until the stream closes, draining it + /// to exhaustion. Completes with an empty output stream: task metrics travel back over the mesh, + /// not here. async fn coordinator_channel( &mut self, - _headers: HeaderMap, + headers: HeaderMap, mut c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, ) -> Result>> { + let mesh = Arc::clone(&self.mesh); #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { while c2w_stream.next().await.is_some() {} }); + tokio::spawn(async move { + while let Some(msg) = c2w_stream.next().await { + let CoordinatorToWorkerMsg::SetPlanRequest(req) = msg else { + continue; + }; + let stage_id = req.task_key.stage_id as u32; + let task_number = req.task_key.task_number as u32; + let set_plan = pb::SetPlanRequest { + task_key: Some(pb::TaskKey { + query_id: serialize_uuid(&req.task_key.query_id), + stage_id: req.task_key.stage_id as u64, + task_number: req.task_key.task_number as u64, + }), + task_count: req.task_count as u64, + plan_proto: req.plan_proto, + work_unit_feed_declarations: req + .work_unit_feed_declarations + .into_iter() + .map(|d| pb::set_plan_request::WorkUnitFeedDeclaration { + id: serialize_uuid(&d.id), + partitions: d.partitions as u64, + }) + .collect(), + target_worker_url: req.target_worker_url.to_string(), + query_start_time_ns: req.query_start_time_ns as u64, + }; + let frame = match SetPlanFrame::from_parts(set_plan, &headers) { + Ok(frame) => frame, + Err(e) => { + log::warn!("shm coordinator_channel: SetPlan frame build failed: {e}"); + continue; + } + }; + if let Err(e) = mesh.send_set_plan(stage_id, task_number, frame).await { + log::warn!("shm coordinator_channel: SetPlan route failed: {e}"); + } + } + }); Ok(stream::empty().boxed()) } From 70e3d14416460069491f3078a552586d92a1fc3b Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Tue, 14 Jul 2026 15:41:54 -0700 Subject: [PATCH 09/11] Incorporate the `paradedb/actions` rebase sync job. (#32, #33) Run CI on `target-patch-*` branches. (#35) --- .github/workflows/ci.yml | 1 + .github/workflows/sync-promote-branch.yml | 29 ++++++++++++++++++++++ .github/workflows/sync-upstream-rebase.yml | 25 +++++++++++++++++++ scripts/sync-upstream.sh | 29 ++++++++++++++++++++++ 4 files changed, 84 insertions(+) create mode 100644 .github/workflows/sync-promote-branch.yml create mode 100644 .github/workflows/sync-upstream-rebase.yml create mode 100755 scripts/sync-upstream.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5748469b..5cda50ae 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,6 +5,7 @@ on: push: branches: - main + - 'target-patch-*' env: CARGO_TERM_COLOR: always diff --git a/.github/workflows/sync-promote-branch.yml b/.github/workflows/sync-promote-branch.yml new file mode 100644 index 00000000..604fbe43 --- /dev/null +++ b/.github/workflows/sync-promote-branch.yml @@ -0,0 +1,29 @@ +# workflows/sync-promote-branch.yml +# +# Promote Target Patch Branch to Main +# Promote a resolved target-patch-* branch to main after manual approval and CI validation. + +name: Promote Target Patch Branch to Main + +on: + workflow_dispatch: + inputs: + branch_name: + description: "Branch name to promote (must match target-patch-*)" + required: true + type: string + +concurrency: + group: promote-branch + cancel-in-progress: false + +jobs: + call-promote: + uses: paradedb/actions/.github/workflows/reusable-promote.yml@v5 + with: + branch_name: ${{ inputs.branch_name }} + approvers: "philippemnoel,rebasedming,stuhood,mdashti" + github_app_client_id: ${{ vars.PARADEDB_GITHUB_APP_CLIENT_ID }} + secrets: + SLACK_WEBHOOK_URL: ${{ secrets.SLACK_GITHUB_CHANNEL_WEBHOOK_URL }} + PARADEDB_GITHUB_APP_PRIVATE_KEY: ${{ secrets.PARADEDB_GITHUB_APP_PRIVATE_KEY }} diff --git a/.github/workflows/sync-upstream-rebase.yml b/.github/workflows/sync-upstream-rebase.yml new file mode 100644 index 00000000..73a21010 --- /dev/null +++ b/.github/workflows/sync-upstream-rebase.yml @@ -0,0 +1,25 @@ +# workflows/sync-upstream-rebase.yml +# +# Upstream Rebase +# Periodically rebase the target patches onto the latest upstream commits and promote the result to main. + +name: Upstream Rebase + +on: + schedule: + - cron: "0 14,16,18,20,22 * * 1-5" # Mon-Fri 6am-2pm PST (UTC weekdays) + - cron: "0 0,2 * * 2-6" # Mon-Fri 4pm-6pm PST (next UTC day) + workflow_dispatch: # Allow manual triggering + +concurrency: + group: upstream-rebase + cancel-in-progress: false + +jobs: + call-rebase: + uses: paradedb/actions/.github/workflows/reusable-rebase.yml@v5 + with: + github_app_client_id: ${{ vars.PARADEDB_GITHUB_APP_CLIENT_ID }} + secrets: + SLACK_WEBHOOK_URL: ${{ secrets.SLACK_GITHUB_CHANNEL_WEBHOOK_URL }} + PARADEDB_GITHUB_APP_PRIVATE_KEY: ${{ secrets.PARADEDB_GITHUB_APP_PRIVATE_KEY }} diff --git a/scripts/sync-upstream.sh b/scripts/sync-upstream.sh new file mode 100755 index 00000000..afce1f42 --- /dev/null +++ b/scripts/sync-upstream.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# scripts/sync-upstream.sh +# +# Wrapper script to execute the centralized upstream sync logic. + +set -euo pipefail + +# 1. Define repository-specific configuration +export UPSTREAM_REPO="datafusion-contrib/datafusion-distributed" +export UPSTREAM_REPO_URL="https://github.com/datafusion-contrib/datafusion-distributed.git" +export TARGET_REPO="paradedb/datafusion-distributed" +export TARGET_BRANCH="main" +export UPSTREAM_BRANCH="main" + +# 2. Define the URL to the centralized script +# Using the raw content URL from the central repository +CORE_SCRIPT_URL="https://raw.githubusercontent.com/paradedb/actions/v5/scripts/sync-core.sh" + +# 3. Download and source the core logic as an API +TMP_SCRIPT=$(mktemp) +curl -fsSL "$CORE_SCRIPT_URL" -o "$TMP_SCRIPT" +# shellcheck source=/dev/null +source "$TMP_SCRIPT" +rm -f "$TMP_SCRIPT" + +# 4. Only execute the command router if run directly (not sourced) +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + sync_core_main "$@" +fi From 53f2d469e6b06d5d31092321920b51346575957d Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 19 Jun 2026 21:21:28 -0700 Subject: [PATCH 10/11] Removed the docs-deploy GitHub Pages workflow. It deploys the docs site to GitHub Pages on every push to `main`; the fork does not publish that site, so the job only fails on the fork. --- .github/workflows/docs.yml | 55 -------------------------------------- 1 file changed, 55 deletions(-) delete mode 100644 .github/workflows/docs.yml diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml deleted file mode 100644 index 78f827a0..00000000 --- a/.github/workflows/docs.yml +++ /dev/null @@ -1,55 +0,0 @@ -name: Deploy Documentation - -on: - push: - branches: - - main - workflow_dispatch: - -permissions: - contents: read - pages: write - id-token: write - -concurrency: - group: "pages" - cancel-in-progress: false - -jobs: - build: - runs-on: ubuntu-latest - steps: - - name: Checkout - uses: actions/checkout@v4 - with: - lfs: true - - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: '3.x' - - - name: Install dependencies - run: | - pip install -r docs/requirements.txt - - - name: Build documentation - run: | - cd docs - make html - - - name: Upload artifact - uses: actions/upload-pages-artifact@v4 - with: - path: docs/build/html - - deploy: - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - runs-on: ubuntu-latest - needs: build - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 From 4fd9283ea3ff369581459f53c145123a954f2874 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 15 Jul 2026 00:52:31 -0700 Subject: [PATCH 11/11] Split oversized batches so every frame fits its ring. Operators emit offset slices of their accumulated state, and arrow-ipc writes a sliced variable-length array's whole values buffer, so one frame balloons to the state's size no matter the batch size. Compacting and halving on the oversized path bounds frames by the ring instead of asking embedders to size rings for their widest state. --- src/shm/in_process.rs | 109 ++++++++++++++++- src/shm/mesh.rs | 4 + src/shm/mpsc_ring.rs | 9 ++ src/shm/self_hosted.rs | 14 +-- src/shm/transport.rs | 258 ++++++++++++++++++++++++++++++++++++++++- 5 files changed, 380 insertions(+), 14 deletions(-) diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs index 632edf62..893805ff 100644 --- a/src/shm/in_process.rs +++ b/src/shm/in_process.rs @@ -44,7 +44,7 @@ use std::alloc::Layout; use std::ffi::c_void; use std::sync::{Arc, Mutex}; -use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::catalog::memory::DataSourceExec; use datafusion::common::runtime::JoinSet; @@ -594,6 +594,42 @@ mod tests { ctx.register_table("t", Arc::new(table)).unwrap(); } + fn wide_table_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])) + } + + /// High-cardinality groups over ~1 KiB strings: each worker's partial-aggregate state runs + /// to hundreds of KiB, so its emit slices balloon past a tiny ring through arrow-ipc. + fn wide_table_partitions() -> Vec> { + let schema = wide_table_schema(); + const ROWS_PER_PART: i32 = 700; + (0..N_WORKERS as i32) + .map(|p| { + let vals = + Int32Array::from_iter_values((0..ROWS_PER_PART).map(|i| p * ROWS_PER_PART + i)); + let strings = StringArray::from_iter_values( + (0..ROWS_PER_PART).map(|i| format!("{p:02}-{i:06}-{}", "x".repeat(1024))), + ); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(vals), Arc::new(strings)]) + .unwrap(); + vec![batch] + }) + .collect() + } + + /// Swap the session's `t` for the wide-string variant. `build_session` registers the + /// standard table; the swap only matters on the leader (fragments travel as shared plan + /// Arcs, so worker sessions never re-resolve the table). + fn register_wide_table(ctx: &SessionContext) { + let _ = ctx.deregister_table("t").unwrap(); + let table = MemTable::try_new(wide_table_schema(), wide_table_partitions()).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + } + /// A worker is a consumer too (it reads shuffle inputs), so when one of its input streams drops /// early it has to cancel that stream's producer, not just the leader. This checks the wiring /// end-to-end: a worker proc's `cancel_stream` reaches the producing proc's inbox through the @@ -846,7 +882,11 @@ mod tests { } fn bootstrap_mesh(n_procs: u32) -> Bootstrap { - let region_total = dsm_region_bytes(n_procs, IN_PROCESS_QUEUE_BYTES, 0).unwrap(); + bootstrap_mesh_with_queue(n_procs, IN_PROCESS_QUEUE_BYTES) + } + + fn bootstrap_mesh_with_queue(n_procs: u32, queue_bytes: usize) -> Bootstrap { + let region_total = dsm_region_bytes(n_procs, queue_bytes, 0).unwrap(); let region = HeapRegion::new(region_total); let base = SharedBase(region.base()); let wakeup: Arc = Arc::new(NoopWakeup); @@ -854,7 +894,7 @@ mod tests { leader_setup( base.0, n_procs, - IN_PROCESS_QUEUE_BYTES, + queue_bytes, &[], Arc::clone(&wakeup), receiver_token(0), @@ -952,6 +992,69 @@ mod tests { ); } + /// An aggregate emits offset slices of its accumulated state, and arrow-ipc writes a + /// sliced variable-length array's whole values buffer, so raw frames balloon to the + /// state's size regardless of batch size. With rings far smaller than that state, the + /// send path's compact-and-split has to carry the query; the result must still match the + /// serial reference exactly. + #[tokio::test(flavor = "current_thread")] + async fn oversized_aggregate_frames_split_across_tiny_rings() { + let query = "SELECT val, max(s) AS m FROM t GROUP BY val ORDER BY val"; + + let serial_ctx = SessionContext::new(); + register_wide_table(&serial_ctx); + let expected = serial_ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + + // 64 KiB rings against ~700 KiB of per-worker aggregate state. + let boot = bootstrap_mesh_with_queue(N_WORKERS + 1, 64 * 1024); + let captured = new_captured_plans(); + let leader_ctx = build_session(Arc::clone(&boot.leader_mesh), Some(Arc::clone(&captured))); + register_wide_table(&leader_ctx); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let entries = collect_dispatched_stages(&physical, N_WORKERS); + + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in boot.workers { + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh), None); + workers.spawn(run_worker_proc( + fragments, + outbound, + mesh, + session, + N_WORKERS, + Arc::clone(&captured), + )); + } + + let leader_task_ctx = leader_ctx.task_ctx(); + let stream = physical.execute(0, leader_task_ctx).unwrap(); + let got: Vec = stream.try_collect().await.unwrap(); + + while let Some(res) = workers.join_next().await { + res.expect("worker task panicked").expect("worker proc"); + } + + use datafusion::arrow::util::pretty::pretty_format_batches; + assert_eq!( + pretty_format_batches(&expected).unwrap().to_string(), + pretty_format_batches(&got).unwrap().to_string(), + "distributed wide aggregate != serial" + ); + } + /// A producer that attaches and then goes away without sending its EOFs must fail the /// gather, not hang it: the drain fails the channels the dead receiver fed once the ring /// detaches. diff --git a/src/shm/mesh.rs b/src/shm/mesh.rs index 4421cd73..90391d53 100644 --- a/src/shm/mesh.rs +++ b/src/shm/mesh.rs @@ -115,6 +115,10 @@ impl BatchChannelSender for DsmInboxSender { fn send_lock(&self) -> &tokio::sync::Mutex<()> { &self.send_lock } + + fn max_frame_bytes(&self) -> Option { + Some(self.inner.max_frame_bytes()) + } } /// DSM MPSC ring as a `BatchChannelReceiver`. The scratch `Vec` lives behind a `Mutex` so a diff --git a/src/shm/mpsc_ring.rs b/src/shm/mpsc_ring.rs index 55c7390b..8b8f8ca9 100644 --- a/src/shm/mpsc_ring.rs +++ b/src/shm/mpsc_ring.rs @@ -663,6 +663,15 @@ impl DsmMpscSender { } } + /// Largest frame `try_send` can accept: every slot's payload capacity combined. Senders + /// consult this before encoding, so an oversized batch can be split to fit instead of + /// erroring with `MessageTooLarge`. + pub(super) fn max_frame_bytes(&self) -> usize { + let header = unsafe { self.ring.as_ref() }; + (header.slot_capacity as usize).saturating_sub(SLOT_HEADER_BYTES) + * header.ring_size as usize + } + /// Wake the registered consumer, if any. Reads the token the consumer stored via /// [`DsmMpscReceiver::set_receiver`] and hands it to the injected [`Wakeup`]; skips when no /// consumer is registered ([`NO_RECEIVER_TOKEN`]). diff --git a/src/shm/self_hosted.rs b/src/shm/self_hosted.rs index 192b77a9..c729de59 100644 --- a/src/shm/self_hosted.rs +++ b/src/shm/self_hosted.rs @@ -1175,14 +1175,12 @@ mod tests { } async fn run(ctx: &SessionContext) -> Result<(String, Vec)> { - // Shaped so every ring frame stays bounded by `shuffle_batch_size`. The strings cross - // the shuffle inside `max`'s partial state, which the repartition rebuilds with `take` - // into fresh per-batch arrays; the projection then reduces them to a length before the - // gather. Shipping `s` itself out of a sort or an aggregate would not work: those emit - // offset slices of their accumulated state, a sliced variable-length array ships its - // whole values buffer through arrow-ipc, and a single frame balloons to the size of the - // partition's state no matter the batch size. - let query = "SELECT val, length(max(s)) AS l FROM t GROUP BY val"; + // Ships `s` itself out of the aggregate on purpose: the emit is an offset slice of the + // partition's accumulated state, and arrow-ipc writes a sliced variable-length array's + // whole values buffer, so the raw frame balloons to the state's size no matter the + // batch size. The send path's compact-and-split is what keeps every frame within these + // tiny rings; this query is its end-to-end exercise. + let query = "SELECT val, max(s) AS m FROM t GROUP BY val"; let plan = ctx.sql(query).await?.create_physical_plan().await?; let display = display_plan_ascii(plan.as_ref(), false); let batches: Vec<_> = execute_stream(plan, ctx.task_ctx())?.try_collect().await?; diff --git a/src/shm/transport.rs b/src/shm/transport.rs index 4b8fb2f6..4d644d9f 100644 --- a/src/shm/transport.rs +++ b/src/shm/transport.rs @@ -33,7 +33,11 @@ use std::sync::{Arc, Condvar, Mutex, MutexGuard}; use datafusion::common::{HashMap, HashSet}; use std::time::{Duration, Instant}; -use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::{ + ArrayRef, BinaryViewArray, RecordBatch, StringViewArray, UInt64Array, +}; +use datafusion::arrow::compute::take; +use datafusion::arrow::datatypes::DataType; use datafusion::arrow::ipc::reader::StreamReader; use datafusion::arrow::ipc::writer::StreamWriter; use datafusion::common::DataFusionError; @@ -373,6 +377,45 @@ fn encode_frame_into( Ok(()) } +/// Copy `len` rows starting at `offset` into fresh, tightly-packed arrays. A plain +/// `RecordBatch::slice` keeps the original backing buffers, and arrow-ipc writes a sliced +/// variable-length array's values buffer whole, which is exactly the ballooning the oversized +/// send path exists to undo. `take` materializes only the selected rows; view arrays +/// additionally need a `gc`, because `take` copies their views but keeps referencing the +/// shared data blocks, so the encoded size wouldn't shrink no matter how few rows remain. +fn compact_rows( + batch: &RecordBatch, + offset: usize, + len: usize, +) -> Result { + let indices = UInt64Array::from_iter_values(offset as u64..(offset + len) as u64); + let columns = batch + .columns() + .iter() + .map(|column| { + let taken = take(column.as_ref(), &indices, None)?; + Ok(match taken.data_type() { + DataType::Utf8View => Arc::new( + taken + .as_any() + .downcast_ref::() + .unwrap() + .gc(), + ) as ArrayRef, + DataType::BinaryView => Arc::new( + taken + .as_any() + .downcast_ref::() + .unwrap() + .gc(), + ) as ArrayRef, + _ => taken, + }) + }) + .collect::, DataFusionError>>()?; + Ok(RecordBatch::try_new(batch.schema(), columns)?) +} + /// Serialize a payload-less [`MppFrameKind::Eof`] frame for `(stage_id, partition)` /// into `buf`. The shm_mq peer reads this as a 16-byte message and routes it to /// the channel buffer's source-done counter without touching Arrow IPC. @@ -634,6 +677,13 @@ pub(crate) trait BatchChannelSender: Send + Sync { /// sibling task land a different payload mid-message and corrupt the queue. In-proc /// channels return a per-instance mutex too, just to keep the call sites uniform. fn send_lock(&self) -> &tokio::sync::Mutex<()>; + + /// Largest single frame the channel accepts, when bounded. The batch send path splits + /// oversized batches to fit under it. `None` (the default, for in-proc channels) disables + /// splitting. + fn max_frame_bytes(&self) -> Option { + None + } } /// Pluggable "drain everything inbound" hook for [`MppSender`]'s cooperative send spin. The @@ -909,7 +959,33 @@ impl MppSender { let t_enc = Instant::now(); encode_frame_into(self.header, batch, scratch)?; stats.encode += t_enc.elapsed(); - self.spin_send_scratch(scratch, stats).await + let cap = self.channel.max_frame_bytes(); + if cap.is_none_or(|cap| scratch.len() <= cap) { + return self.spin_send_scratch(scratch, stats).await; + } + let cap = cap.expect("bounded channel checked above"); + // An over-cap frame usually means the batch carries offset slices of an upstream + // operator's accumulated state: arrow-ipc writes a sliced variable-length array's + // whole values buffer, so the frame balloons to the state's size no matter the row + // count. Compact the rows into fresh arrays and halve until every frame fits, + // depth-first left-to-right so row order survives for sorted streams. A single row + // that exceeds the ring still errors inside the channel with the raise-the-knob + // message. + let mut work = vec![(0usize, batch.num_rows())]; + while let Some((offset, len)) = work.pop() { + let compacted = compact_rows(batch, offset, len)?; + let t_enc = Instant::now(); + encode_frame_into(self.header, &compacted, scratch)?; + stats.encode += t_enc.elapsed(); + if scratch.len() > cap && len > 1 { + let half = len / 2; + work.push((offset + half, len - half)); + work.push((offset, half)); + continue; + } + self.spin_send_scratch(scratch, stats).await?; + } + Ok(()) } /// Push an already-encoded frame through the channel via the cooperative-drain spin (or the @@ -1829,7 +1905,9 @@ pub(super) const SELF_LOOP_CAPACITY: usize = 1 << 20; #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::{Int32Array, Int64Array, StringArray, UInt64Array}; + use datafusion::arrow::array::{ + Int32Array, Int64Array, StringArray, StringViewArray, UInt64Array, + }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use std::sync::Arc as StdArc; use std::thread; @@ -2413,6 +2491,180 @@ mod tests { } } + /// Bounded in-proc sink standing in for a DSM ring: rejects frames over `cap` the way + /// `try_send` does, records every accepted frame, and advertises the bound so the send + /// path's splitter engages. + struct BoundedSink { + cap: usize, + frames: Mutex>>, + send_lock: tokio::sync::Mutex<()>, + } + + impl BoundedSink { + fn new(cap: usize) -> Arc { + Arc::new(Self { + cap, + frames: Mutex::new(Vec::new()), + send_lock: tokio::sync::Mutex::new(()), + }) + } + } + + impl BatchChannelSender for BoundedSink { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError> { + if bytes.len() > self.cap { + return Err(DataFusionError::Execution( + "mpp: DSM MPSC frame exceeds the entire ring capacity \ + (raise the embedder's queue-size knob)" + .into(), + )); + } + self.frames.lock().unwrap().push(bytes.to_vec()); + Ok(()) + } + + fn send_lock(&self) -> &tokio::sync::Mutex<()> { + &self.send_lock + } + + fn max_frame_bytes(&self) -> Option { + Some(self.cap) + } + } + + /// `rows` of ~1 KiB unique strings, so slices drag a large values buffer through + /// arrow-ipc. + fn wide_batch(rows: usize) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, false), + Field::new("v", DataType::Int64, false), + ])); + let strings = StringArray::from_iter_values( + (0..rows).map(|i| format!("{i:06}-{}", "x".repeat(1024))), + ); + let vals = Int64Array::from_iter_values(0..rows as i64); + RecordBatch::try_new(schema, vec![Arc::new(strings), Arc::new(vals)]).unwrap() + } + + #[test] + fn oversized_sliced_batch_splits_to_fit() { + let cap = 64 * 1024; + let base = wide_batch(512); + // 128 rows, but the slice keeps the whole ~512 KiB values buffer alive. + let sliced = base.slice(64, 128); + + // Premise: the sliced encode really does balloon past the cap. If arrow-ipc ever + // learns to truncate sliced buffers, this assert flags the test (and the split + // path's main motivation) for review. + let mut probe = Vec::new(); + encode_frame_into(MppFrameHeader::batch(3, 7, 0), &sliced, &mut probe).expect("probe"); + assert!( + probe.len() > cap, + "premise: sliced encode ({} bytes) must exceed cap ({cap})", + probe.len() + ); + + let sink = BoundedSink::new(cap); + let sender = MppSender::with_header( + Arc::clone(&sink) as Arc, + MppFrameHeader::batch(3, 7, 0), + ); + sender.send_batch(&sliced).expect("split send"); + + let frames = sink.frames.lock().unwrap(); + assert!( + frames.len() > 1, + "expected the oversized batch to split, got {} frame(s)", + frames.len() + ); + let mut decoded = Vec::new(); + for frame in frames.iter() { + assert!( + frame.len() <= cap, + "frame of {} bytes over cap", + frame.len() + ); + let (header, body) = decode_frame(frame).expect("decode split frame"); + assert_eq!((header.stage_id, header.partition), (3, 7)); + let FrameBody::Batch(batch) = body else { + panic!("split frame must carry a batch"); + }; + decoded.push(batch); + } + let got = datafusion::arrow::compute::concat_batches(&sliced.schema(), &decoded) + .expect("concat decoded"); + // Row order and content must survive the split; compare against a compacted copy + // (`take` of every row) since the decoded side never sees the original buffers. + let want = compact_rows(&sliced, 0, sliced.num_rows()).expect("compact reference"); + assert_eq!(got, want); + } + + /// Same as above but through `Utf8View`: `take` copies the views yet keeps the shared + /// data blocks, so without the compact-time `gc` no amount of row-splitting shrinks a + /// frame and the send would error at one row. + #[test] + fn oversized_view_array_batch_splits_after_gc() { + let cap = 64 * 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + let strings = StringViewArray::from_iter_values( + (0..512).map(|i| format!("{i:06}-{}", "x".repeat(1024))), + ); + let base = RecordBatch::try_new(schema, vec![Arc::new(strings)]).unwrap(); + let sliced = base.slice(64, 128); + + let sink = BoundedSink::new(cap); + let sender = MppSender::with_header( + Arc::clone(&sink) as Arc, + MppFrameHeader::batch(1, 1, 0), + ); + sender.send_batch(&sliced).expect("view-array split send"); + + let frames = sink.frames.lock().unwrap(); + let mut rows = 0; + for frame in frames.iter() { + assert!( + frame.len() <= cap, + "frame of {} bytes over cap", + frame.len() + ); + let (_, body) = decode_frame(frame).expect("decode view frame"); + let FrameBody::Batch(batch) = body else { + panic!("expected batch frame"); + }; + rows += batch.num_rows(); + } + assert_eq!(rows, 128); + } + + #[test] + fn single_oversized_row_still_errors() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from_iter_values([ + "y".repeat(128 * 1024) + ]))], + ) + .unwrap(); + + let sink = BoundedSink::new(64 * 1024); + let sender = MppSender::with_header( + sink as Arc, + MppFrameHeader::batch(0, 0, 0), + ); + let err = sender + .send_batch(&batch) + .expect_err("one giant row can't split"); + assert!( + format!("{err}").contains("exceeds the entire ring capacity"), + "unexpected error: {err}" + ); + } + #[test] fn drain_buffer_pop_returns_pushed_batches_in_order() { let buf = DrainBuffer::new(1);