From 09938a7efbd8e52119ff79c6637c9417ec47f294 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Fri, 3 Jul 2026 12:17:17 +0200 Subject: [PATCH] Abstract away gRPC --- .gitignore | 8 +- Cargo.toml | 23 +- console/src/worker.rs | 2 +- examples/in_memory_cluster.rs | 13 +- examples/localhost_versioned_run.rs | 12 +- src/{protobuf => codec}/distributed_codec.rs | 0 src/{protobuf => codec}/mod.rs | 6 - src/{protobuf => codec}/user_codec.rs | 0 src/common/mod.rs | 2 - src/common/time.rs | 9 +- src/coordinator/distributed.rs | 11 +- src/coordinator/metrics_store.rs | 13 +- src/coordinator/prepare_dynamic_plan.rs | 9 +- src/coordinator/query_coordinator.rs | 140 ++- src/distributed_ext.rs | 48 +- src/distributed_planner/distributed_config.rs | 3 +- src/distributed_planner/network_boundary.rs | 60 +- src/distributed_planner/task_estimator.rs | 2 +- src/execution_plans/benchmarks/fixture.rs | 14 +- .../benchmarks/transport_bench.rs | 6 +- src/execution_plans/network_broadcast.rs | 8 +- src/execution_plans/network_coalesce.rs | 5 +- src/execution_plans/network_shuffle.rs | 8 +- src/execution_plans/sampler.rs | 35 +- src/lib.rs | 44 +- src/metrics/latency_metric.rs | 12 +- src/metrics/mod.rs | 1 - src/metrics/task_metrics_collector.rs | 4 +- src/metrics/task_metrics_rewriter.rs | 69 +- src/networking/mod.rs | 11 - src/protobuf/producer_head.rs | 63 -- src/protocol/channel_resolver.rs | 99 +++ .../grpc}/channel_resolver.rs | 127 +-- .../grpc}/errors/arrow_error.rs | 2 +- .../grpc}/errors/datafusion_error.rs | 12 +- .../grpc}/errors/io_error.rs | 0 src/{protobuf => protocol/grpc}/errors/mod.rs | 2 +- .../grpc}/errors/objectstore_error.rs | 0 .../grpc}/errors/parquet_error.rs | 0 .../grpc}/errors/parser_error.rs | 0 .../grpc}/errors/schema_error.rs | 0 src/{worker => protocol/grpc}/gen/Cargo.toml | 0 .../grpc}/gen/regen.sh | 2 +- src/{worker => protocol/grpc}/gen/src/main.rs | 4 +- .../grpc}/generated/mod.rs | 0 .../grpc}/generated/worker.rs | 0 .../grpc/metrics_proto.rs} | 18 +- src/protocol/grpc/mod.rs | 21 + .../grpc}/observability/README.md | 2 +- .../grpc}/observability/gen/Cargo.toml | 0 src/protocol/grpc/observability/gen/regen.sh | 6 + .../grpc}/observability/gen/src/main.rs | 6 +- .../grpc}/observability/generated/mod.rs | 0 .../observability/generated/observability.rs | 4 +- src/{ => protocol/grpc}/observability/mod.rs | 0 .../observability/proto/observability.proto | 3 +- .../grpc}/observability/service.rs | 20 +- .../grpc}/on_drop_stream.rs | 0 .../grpc}/spawn_select_all.rs | 0 src/{worker => protocol/grpc}/worker.proto | 0 src/protocol/grpc/worker_client.rs | 766 ++++++++++++++++ src/protocol/grpc/worker_service.rs | 422 +++++++++ src/protocol/mod.rs | 14 + src/protocol/worker_channel.rs | 191 ++++ src/stage.rs | 34 +- src/test_utils/in_memory_channel_resolver.rs | 14 +- src/test_utils/metrics.rs | 65 +- src/test_utils/plans.rs | 10 +- src/work_unit_feed/remote_work_unit_feed.rs | 106 +-- src/work_unit_feed/work_unit_feed_provider.rs | 2 +- src/worker/gen/regen.sh | 6 - src/worker/impl_coordinator_channel.rs | 129 +-- src/worker/impl_execute_task.rs | 231 +---- src/worker/mod.rs | 2 - src/worker/single_write_multi_read.rs | 1 + src/worker/task_data.rs | 89 +- src/worker/test_utils/worker_handles.rs | 8 +- src/worker/worker_connection_pool.rs | 827 ++++-------------- src/worker/worker_service.rs | 109 +-- src/{networking => }/worker_resolver.rs | 0 80 files changed, 2260 insertions(+), 1735 deletions(-) rename src/{protobuf => codec}/distributed_codec.rs (100%) rename src/{protobuf => codec}/mod.rs (55%) rename src/{protobuf => codec}/user_codec.rs (100%) delete mode 100644 src/networking/mod.rs delete mode 100644 src/protobuf/producer_head.rs create mode 100644 src/protocol/channel_resolver.rs rename src/{networking => protocol/grpc}/channel_resolver.rs (56%) rename src/{protobuf => protocol/grpc}/errors/arrow_error.rs (99%) rename src/{protobuf => protocol/grpc}/errors/datafusion_error.rs (97%) rename src/{protobuf => protocol/grpc}/errors/io_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/mod.rs (97%) rename src/{protobuf => protocol/grpc}/errors/objectstore_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/parquet_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/parser_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/schema_error.rs (100%) rename src/{worker => protocol/grpc}/gen/Cargo.toml (100%) rename src/{observability => protocol/grpc}/gen/regen.sh (57%) rename src/{worker => protocol/grpc}/gen/src/main.rs (87%) rename src/{worker => protocol/grpc}/generated/mod.rs (100%) rename src/{worker => protocol/grpc}/generated/worker.rs (100%) rename src/{metrics/proto.rs => protocol/grpc/metrics_proto.rs} (98%) create mode 100644 src/protocol/grpc/mod.rs rename src/{ => protocol/grpc}/observability/README.md (88%) rename src/{ => protocol/grpc}/observability/gen/Cargo.toml (100%) create mode 100755 src/protocol/grpc/observability/gen/regen.sh rename src/{ => protocol/grpc}/observability/gen/src/main.rs (76%) rename src/{ => protocol/grpc}/observability/generated/mod.rs (100%) rename src/{ => protocol/grpc}/observability/generated/observability.rs (99%) rename src/{ => protocol/grpc}/observability/mod.rs (100%) rename src/{ => protocol/grpc}/observability/proto/observability.proto (95%) rename src/{ => protocol/grpc}/observability/service.rs (89%) rename src/{common => protocol/grpc}/on_drop_stream.rs (100%) rename src/{worker => protocol/grpc}/spawn_select_all.rs (100%) rename src/{worker => protocol/grpc}/worker.proto (100%) create mode 100644 src/protocol/grpc/worker_client.rs create mode 100644 src/protocol/grpc/worker_service.rs create mode 100644 src/protocol/mod.rs create mode 100644 src/protocol/worker_channel.rs delete mode 100755 src/worker/gen/regen.sh rename src/{networking => }/worker_resolver.rs (100%) diff --git a/.gitignore b/.gitignore index 479b6815..70fa5beb 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,7 @@ testdata/tpcds/* !testdata/tpcds/README.md testdata/clickbench/* !testdata/clickbench/queries -src/observability/gen/target/* -src/observability/gen/Cargo.lock -src/worker/gen/target/* -src/worker/gen/Cargo.lock +src/protocol/grpc/observability/gen/target/* +src/protocol/grpc/observability/gen/Cargo.lock +src/protocol/grpc/gen/target/* +src/protocol/grpc/gen/Cargo.lock diff --git a/Cargo.toml b/Cargo.toml index fd5ceab2..79366ed3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,13 +23,8 @@ datafusion = { workspace = true, features = [ "datetime_expressions", ] } datafusion-proto = { workspace = true } -arrow-flight = "58" -arrow-select = "58" -arrow-ipc = { version = "58", features = ["zstd"] } async-trait = "0.1.89" tokio = { version = "1.48", features = ["full"] } -tonic = { version = "0.14.1", features = ["transport"] } -tower = "0.5.2" http = "1.3.1" itertools = "0.14.0" futures = "0.3.31" @@ -52,6 +47,13 @@ num-traits = "0.2" bincode = "1" tonic-prost = "0.14.2" +# grpc-specific features +arrow-flight = { version = "58", optional = true } +arrow-select = { version = "58", optional = true } +arrow-ipc = { version = "58", features = ["zstd"], optional = true } +tonic = { version = "0.14.1", features = ["transport"], optional = true } +tower = { version = "0.5.2", optional = true } + # integration_tests deps insta = { version = "1.46.0", features = ["filters"], optional = true } parquet = { version = "58", optional = true } @@ -59,8 +61,17 @@ arrow = { version = "58", optional = true, features = ["test_utils"] } hyper-util = { version = "0.1.16", optional = true } [features] +default = ["grpc"] avro = ["datafusion/avro"] -integration = ["insta", "parquet", "arrow", "hyper-util"] +grpc = [ + "arrow-flight", + "arrow-select", + "arrow-ipc", + "tonic", + "tower" +] + +integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"] system-metrics = ["sysinfo"] diff --git a/console/src/worker.rs b/console/src/worker.rs index d84fe602..4a23c5b9 100644 --- a/console/src/worker.rs +++ b/console/src/worker.rs @@ -1,5 +1,5 @@ use datafusion::common::{HashMap, HashSet}; -use datafusion_distributed::{ +use datafusion_distributed::grpc::{ GetClusterWorkersRequest, GetTaskProgressRequest, ObservabilityServiceClient, PingRequest, TaskProgress, TaskStatus, }; diff --git a/examples/in_memory_cluster.rs b/examples/in_memory_cluster.rs index 1d02737c..6a9d7ba1 100644 --- a/examples/in_memory_cluster.rs +++ b/examples/in_memory_cluster.rs @@ -4,9 +4,8 @@ use datafusion::common::DataFusionError; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_distributed::{ - BoxCloneSyncChannel, ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker, - WorkerQueryContext, WorkerResolver, WorkerServiceClient, create_worker_client, - display_plan_ascii, + ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker, WorkerChannel, + WorkerQueryContext, WorkerResolver, display_plan_ascii, grpc, }; use futures::TryStreamExt; use hyper_util::rt::TokioIo; @@ -66,7 +65,7 @@ const DUMMY_URL: &str = "http://localhost:50051"; /// tokio duplex rather than a TCP connection. #[derive(Clone)] struct InMemoryChannelResolver { - channel: WorkerServiceClient, + channel: grpc::BoxCloneSyncChannel, } impl InMemoryChannelResolver { @@ -84,7 +83,7 @@ impl InMemoryChannelResolver { })); let this = Self { - channel: create_worker_client(BoxCloneSyncChannel::new(channel)), + channel: grpc::BoxCloneSyncChannel::new(channel), }; let this_clone = this.clone(); @@ -110,8 +109,8 @@ impl ChannelResolver for InMemoryChannelResolver { async fn get_worker_client_for_url( &self, _: &url::Url, - ) -> Result, DataFusionError> { - Ok(self.channel.clone()) + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) } } diff --git a/examples/localhost_versioned_run.rs b/examples/localhost_versioned_run.rs index dc3f38c2..b10520cb 100644 --- a/examples/localhost_versioned_run.rs +++ b/examples/localhost_versioned_run.rs @@ -4,8 +4,8 @@ use datafusion::common::DataFusionError; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_distributed::{ - DefaultChannelResolver, DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt, - WorkerResolver, create_worker_client, display_plan_ascii, + DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt, WorkerResolver, + display_plan_ascii, grpc, }; use futures::TryStreamExt; use std::error::Error; @@ -40,7 +40,7 @@ struct Args { /// the `GetWorkerInfo` RPC. Returns `false` if the worker is unreachable, returns /// an error, or reports a different version. async fn worker_has_version( - channel_resolver: &DefaultChannelResolver, + channel_resolver: &grpc::DefaultChannelResolver, url: &Url, expected_version: &str, ) -> bool { @@ -48,12 +48,12 @@ async fn worker_has_version( return false; }; - let mut client = create_worker_client(channel); + let mut client = grpc::create_worker_client(channel); let Ok(response) = client.get_worker_info(GetWorkerInfoRequest {}).await else { return false; }; - response.into_inner().version == expected_version + response.version == expected_version } #[tokio::main] @@ -61,7 +61,7 @@ async fn main() -> Result<(), Box> { let args = Args::from_args(); let ports = if let Some(target_version) = &args.version { - let channel_resolver = DefaultChannelResolver::default(); + let channel_resolver = grpc::DefaultChannelResolver::default(); let mut compatible = Vec::new(); for &port in &args.cluster_ports { let url = Url::parse(&format!("http://localhost:{port}"))?; diff --git a/src/protobuf/distributed_codec.rs b/src/codec/distributed_codec.rs similarity index 100% rename from src/protobuf/distributed_codec.rs rename to src/codec/distributed_codec.rs diff --git a/src/protobuf/mod.rs b/src/codec/mod.rs similarity index 55% rename from src/protobuf/mod.rs rename to src/codec/mod.rs index 47deeb40..c4acaf1a 100644 --- a/src/protobuf/mod.rs +++ b/src/codec/mod.rs @@ -1,13 +1,7 @@ mod distributed_codec; -mod errors; -mod producer_head; mod user_codec; pub use distributed_codec::DistributedCodec; -pub(crate) use errors::{ - datafusion_error_to_tonic_status, map_flight_to_datafusion_error, - tonic_status_to_datafusion_error, -}; pub(crate) use user_codec::{ get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc, }; diff --git a/src/protobuf/user_codec.rs b/src/codec/user_codec.rs similarity index 100% rename from src/protobuf/user_codec.rs rename to src/codec/user_codec.rs diff --git a/src/common/mod.rs b/src/common/mod.rs index 18cb28a1..07463ca0 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,5 +1,4 @@ mod children_helpers; -mod on_drop_stream; mod once_lock; mod recursion; mod task_context_helpers; @@ -8,7 +7,6 @@ mod uuid; mod vec; pub(crate) use children_helpers::require_one_child; -pub(crate) use on_drop_stream::on_drop_stream; pub(crate) use once_lock::OnceLockResult; pub(crate) use recursion::TreeNodeExt; pub(crate) use task_context_helpers::task_ctx_with_extension; diff --git a/src/common/time.rs b/src/common/time.rs index 07757f32..6efb572c 100644 --- a/src/common/time.rs +++ b/src/common/time.rs @@ -1,9 +1,14 @@ +use num_traits::AsPrimitive; use std::time::{SystemTime, UNIX_EPOCH}; /// Nanoseconds elapsed since UNIX epoch. -pub(crate) fn now_ns() -> u64 { +pub(crate) fn now_ns() -> T +where + T: 'static + Copy, + u128: AsPrimitive, +{ SystemTime::now() .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_nanos() as u64) + .map(|duration| duration.as_nanos().as_()) .expect("SystemTime before UNIX EPOCH!") } diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 2c4bc611..a1e123f2 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,11 +1,10 @@ -use crate::DistributedConfig; -use crate::common::{require_one_child, serialize_uuid}; +use crate::common::require_one_child; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; use crate::distributed_planner::NetworkBoundaryExt; -use crate::worker::generated::worker::TaskKey; +use crate::{DistributedConfig, TaskKey}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{Result, exec_err}; @@ -94,9 +93,9 @@ impl DistributedExec { let stage = boundary.input_stage(); for i in 0..stage.task_count() { expected_keys.push(TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, - task_number: i as u64, + query_id: stage.query_id(), + stage_id: stage.num(), + task_number: i, }); } } diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index 4b9b0ffc..8ccc1f7c 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -1,9 +1,8 @@ -use crate::TaskKey; -use crate::worker::generated::worker as pb; +use crate::{TaskKey, TaskMetrics}; use datafusion::common::HashMap; use tokio::sync::watch; -type MetricsMap = HashMap; +type MetricsMap = HashMap; /// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive. #[derive(Debug, Clone)] @@ -18,20 +17,18 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: pb::TaskMetrics) { + pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); } - pub(crate) fn get(&self, key: &TaskKey) -> Option { + pub(crate) fn get(&self, key: &TaskKey) -> Option { self.rx.borrow().get(key).cloned() } #[cfg(test)] - pub(crate) fn from_entries( - entries: impl IntoIterator, - ) -> Self { + pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { let map: HashMap<_, _> = entries.into_iter().collect(); let (tx, rx) = watch::channel(map); Self { tx, rx } diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index cfea3d5d..463cc58b 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -8,8 +8,7 @@ use crate::distributed_planner::{ }; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; -use crate::worker::generated::worker as pb; -use crate::{BytesCounterMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; +use crate::{BytesCounterMetric, LoadInfo, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; use dashmap::DashMap; use datafusion::common::stats::Precision; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -204,7 +203,7 @@ impl PlanReconstructor { /// Estimates the bytes per second flowing through a stage by reading sample information. async fn gather_runtime_statistics( - per_task_load_info_stream: Vec + Unpin>, + per_task_load_info_stream: Vec + Unpin>, plan: &Arc, ) -> Result { const ESTIMATED_QUERY_TIME_S: usize = 10; @@ -237,8 +236,8 @@ async fn gather_runtime_statistics( let mut load_info_stream = futures::stream::select_all(per_task_load_info_stream); while let Some(load_info) = load_info_stream.next().await { - rows_per_second += load_info.rows_per_second as usize; - rows_ready += load_info.rows_ready as usize; + rows_per_second += load_info.rows_per_second; + rows_ready += load_info.rows_ready; per_col_bytes_per_second = element_wise_sum( per_col_bytes_per_second, &load_info.per_column_bytes_per_second, diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 2e0f2bf0..887a109f 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,25 +1,22 @@ -use crate::common::{TreeNodeExt, now_ns, serialize_uuid, task_ctx_with_extension}; +use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; use crate::coordinator::latency_metric::LatencyMetric; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; -use crate::protobuf::tonic_status_to_datafusion_error; use crate::stage::LocalStage; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; -use crate::worker::generated::worker as pb; -use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; use crate::{ - BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, - DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, NetworkBoundaryExt, - Stage, TaskEstimator, TaskKey, TaskRoutingContext, get_distributed_channel_resolver, - get_distributed_worker_resolver, + BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, + 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, }; +use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion::common::{DataFusionError, exec_datafusion_err}; use datafusion::common::{Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; @@ -27,8 +24,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_proto::physical_plan::AsExecutionPlan; use datafusion_proto::protobuf::PhysicalPlanNode; -use futures::{Stream, StreamExt}; -use http::Extensions; +use futures::{Stream, StreamExt, TryStreamExt}; use prost::Message; use rand::Rng; use std::ops::DerefMut; @@ -36,13 +32,11 @@ use std::sync::{Arc, Mutex}; use tokio::sync::Notify; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; -use tonic::Request; -use tonic::metadata::MetadataMap; use url::Url; use uuid::Uuid; -/// How many [crate::WorkUnit] messages are allowed to be chunked synchronously together in order to -/// send fewer bigger [crate::WorkUnit] batches over the wire, reducing the overhead of sending many +/// How many [crate::WorkUnitMsg] messages are allowed to be chunked synchronously together in order to +/// send fewer bigger [crate::WorkUnitMsg] batches over the wire, reducing the overhead of sending many /// small batches. See [StreamExt::ready_chunks] docs for more details about how chunking works. const WORK_UNIT_FEED_CHUNK_SIZE: usize = 256; @@ -142,8 +136,8 @@ impl<'a> StageCoordinator<'a> { task_i: usize, url: Url, ) -> Result<( - UnboundedSender, - UnboundedReceiver, + UnboundedSender, + UnboundedReceiver, )> { let session_config = self.task_ctx.session_config(); let codec = DistributedCodec::new_combined_with_user(session_config); @@ -155,21 +149,20 @@ impl<'a> StageCoordinator<'a> { let plan_size = plan_proto.len(); let task_key = TaskKey { - query_id: serialize_uuid(&self.query_id), - stage_id: self.stage_id as u64, - task_number: task_i as u64, - }; - let msg = pb::CoordinatorToWorkerMsg { - inner: Some(Inner::SetPlanRequest(pb::SetPlanRequest { - plan_proto, - task_count: self.task_count as u64, - task_key: Some(task_key.clone()), - work_unit_feed_declarations, - target_worker_url: url.to_string(), - query_start_time_ns: self.metrics.instantiation_time, - })), + query_id: self.query_id, + stage_id: self.stage_id, + task_number: task_i, }; + let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest { + task_key, + task_count: self.task_count, + plan_proto, + work_unit_feed_declarations, + target_worker_url: url.clone(), + query_start_time_ns: self.metrics.instantiation_time, + }); + let (coordinator_to_worker_tx, coordinator_to_worker_rx) = tokio::sync::mpsc::unbounded_channel(); let (worker_to_coordinator_tx, worker_to_coordinator_rx) = @@ -180,44 +173,33 @@ impl<'a> StageCoordinator<'a> { let mut headers = get_config_extension_propagation_headers(session_config)?; headers.extend(get_passthrough_headers(session_config)); - let request = Request::from_parts( - MetadataMap::from_headers(headers), - Extensions::default(), - futures::stream::once(async { msg }) - .chain(UnboundedReceiverStream::new(coordinator_to_worker_rx)) - .map(set_work_unit_send_time) - // Keep the request side of the channel open until the query ends: this tail emits - // no messages and only completes, once the `Notify` fires. Workers interpret this - // EOS of this stream as a query finished/aborted signal. The flow looks like this: - // 1. The query ends normally, as all Arrow RecordBatches are already streamed. - // 2. The end stream notifier guard is dropped in `DistributedExec::execute()`. - // 3. Here, `end_stream_notifier` fires and the coordinator->worker channel is - // gracefully ended. - // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`. - // 5. The metrics are send back in the worker->coordinator channel, and then that - // channel is closed. - .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))), - ); + let coordinator_to_worker_stream = futures::stream::once(async { msg }) + .chain(UnboundedReceiverStream::new(coordinator_to_worker_rx)) + .map(set_work_unit_send_time) + // Keep the request side of the channel open until the query ends: this tail emits + // no messages and only completes, once the `Notify` fires. Workers interpret this + // EOS of this stream as a query finished/aborted signal. The flow looks like this: + // 1. The query ends normally, as all Arrow RecordBatches are already streamed. + // 2. The end stream notifier guard is dropped in `DistributedExec::execute()`. + // 3. Here, `end_stream_notifier` fires and the coordinator->worker channel is + // gracefully ended. + // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`. + // 5. The metrics are send back in the worker->coordinator channel, and then that + // channel is closed. + .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))) + .boxed(); let metrics = self.metrics.clone(); self.join_set.lock().unwrap().spawn(async move { let start = Instant::now(); let mut client = channel_resolver.get_worker_client_for_url(&url).await?; - let response = client.coordinator_channel(request).await.map_err(|e| { - tonic_status_to_datafusion_error(&e).unwrap_or_else(|| { - exec_datafusion_err!("Error sending plan to worker {url}: {e}") - }) - })?; + let mut worker_to_coordinator_stream = client + .coordinator_channel(headers, coordinator_to_worker_stream) + .await?; metrics.plan_send_latency.record(&start); metrics.plan_bytes_sent.add_bytes(plan_size); - let mut worker_to_coordinator_stream = response.into_inner(); - while let Some(msg_or_err) = worker_to_coordinator_stream.next().await { - let msg = msg_or_err.map_err(|err| { - tonic_status_to_datafusion_error(&err).unwrap_or_else(|| { - exec_datafusion_err!("Unknown error on worker to coordinator stream: {err}") - }) - })?; + while let Some(msg) = worker_to_coordinator_stream.try_next().await? { if worker_to_coordinator_tx.send(msg).is_err() { break; // receiver dropped } @@ -234,12 +216,12 @@ impl<'a> StageCoordinator<'a> { pub(super) fn worker_to_coordinator_task( &mut self, task_i: usize, - mut worker_to_coordinator_rx: UnboundedReceiver, - ) -> UnboundedReceiver { + mut worker_to_coordinator_rx: UnboundedReceiver, + ) -> UnboundedReceiver { let task_key = TaskKey { - query_id: serialize_uuid(&self.query_id), - stage_id: self.stage_id as u64, - task_number: task_i as u64, + query_id: self.query_id, + stage_id: self.stage_id, + task_number: task_i, }; let task_metrics = self.metrics_store.clone(); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); @@ -250,20 +232,18 @@ impl<'a> StageCoordinator<'a> { #[allow(clippy::disallowed_methods)] tokio::spawn(async move { while let Some(msg) = worker_to_coordinator_rx.recv().await { - let Some(inner) = msg.inner else { continue }; - - match inner { - pb::worker_to_coordinator_msg::Inner::TaskMetrics(pre_order_metrics) => { + match msg { + WorkerToCoordinatorMsg::TaskMetrics(v) => { if let Some(task_metrics) = &task_metrics { - task_metrics.insert(task_key.clone(), pre_order_metrics); + task_metrics.insert(task_key, v); } } - pb::worker_to_coordinator_msg::Inner::LoadInfo(load_info) => { + WorkerToCoordinatorMsg::LoadInfo(load_info) => { if let Some(tx) = &load_info_tx_opt { let _ = tx.send(load_info); } } - pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { + WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } } @@ -278,7 +258,7 @@ impl<'a> StageCoordinator<'a> { pub(super) fn coordinator_to_worker_task( &mut self, task_i: usize, - tx: UnboundedSender, + tx: UnboundedSender, ) -> Result<()> { let session_config = self.task_ctx.session_config(); let d_cfg = DistributedConfig::from_config_options(session_config.options())?; @@ -329,12 +309,10 @@ impl<'a> StageCoordinator<'a> { Ok(TreeNodeRecursion::Continue) })?; - struct WorkUnitEosOnDrop(UnboundedSender); + struct WorkUnitEosOnDrop(UnboundedSender); impl Drop for WorkUnitEosOnDrop { fn drop(&mut self) { - let _ = self.0.send(pb::CoordinatorToWorkerMsg { - inner: Some(Inner::WorkUnitEos(true)), - }); + let _ = self.0.send(CoordinatorToWorkerMsg::WorkUnitEos); } } @@ -368,8 +346,8 @@ impl<'a> StageCoordinator<'a> { let transformed = plan.transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| { if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { work_unit_feed_declarations.push(WorkUnitFeedDeclaration { - id: serialize_uuid(&wuf.id()), - partitions: plan.properties().partitioning.partition_count() as u64, + id: wuf.id(), + partitions: plan.properties().partitioning.partition_count(), }); }; @@ -467,7 +445,7 @@ impl Drop for NotifyGuard { pub(super) struct CoordinatorToWorkerMetrics { pub(super) plan_bytes_sent: BytesCounterMetric, pub(super) plan_send_latency: Arc, - pub(super) instantiation_time: u64, + pub(super) instantiation_time: usize, } impl CoordinatorToWorkerMetrics { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index d33b18a6..9efdb420 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -1,16 +1,16 @@ +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::distributed_planner::set_distributed_task_estimator; -use crate::networking::{set_distributed_channel_resolver, set_distributed_worker_resolver}; use crate::passthrough_headers::set_passthrough_headers; -use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc}; +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, }; -use arrow_ipc::CompressionType; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; use datafusion::execution::{SessionState, SessionStateBuilder}; @@ -202,7 +202,7 @@ pub trait DistributedExt: Sized { /// # use datafusion::prelude::SessionConfig; /// # use url::Url; /// # use std::sync::Arc; - /// # use datafusion_distributed::{BoxCloneSyncChannel, WorkerResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext}; + /// # use datafusion_distributed::{WorkerResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext}; /// /// struct CustomWorkerResolver; /// @@ -238,14 +238,14 @@ pub trait DistributedExt: Sized { /// # use datafusion::prelude::SessionConfig; /// # use url::Url; /// # use std::sync::Arc; - /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext, WorkerServiceClient}; + /// # use datafusion_distributed::{ChannelResolver, DistributedExt, SessionStateBuilderExt, WorkerChannel, WorkerQueryContext, grpc}; /// /// struct CustomChannelResolver; /// /// #[async_trait] /// impl ChannelResolver for CustomChannelResolver { - /// async fn get_worker_client_for_url(&self, url: &Url) -> Result, DataFusionError> { - /// // Build a custom WorkerServiceClient wrapped with tower layers or something similar. + /// async fn get_worker_client_for_url(&self, url: &Url) -> Result, DataFusionError> { + /// // Build a custom worker client wrapped with tower layers or something similar. /// todo!() /// } /// } @@ -451,18 +451,20 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation. fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] /// The compression type to use for sending data over the wire. /// /// The default is [CompressionType::LZ4_FRAME]. fn with_distributed_compression( self, - compression: Option, + compression: Option, ) -> Result; + #[cfg(feature = "grpc")] /// Same as [DistributedExt::with_distributed_compression] but with an in-place mutation. fn set_distributed_compression( &mut self, - compression: Option, + compression: Option, ) -> Result<(), DataFusionError>; /// Overrides `datafusion.execution.batch_size` for worker-executed stages, letting users @@ -678,14 +680,15 @@ impl DistributedExt for SessionConfig { Ok(()) } + #[cfg(feature = "grpc")] fn set_distributed_compression( &mut self, - compression: Option, + compression: Option, ) -> Result<(), DataFusionError> { let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; d_cfg.compression = match compression { - Some(CompressionType::ZSTD) => "zstd".to_string(), - Some(CompressionType::LZ4_FRAME) => "lz4".to_string(), + Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(), + Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(), _ => "none".to_string(), }; Ok(()) @@ -810,7 +813,8 @@ impl DistributedExt for SessionConfig { #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; #[call(set_distributed_shuffle_batch_size)] #[expr($?;Ok(self))] @@ -915,10 +919,12 @@ impl DistributedExt for SessionStateBuilder { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] @@ -1036,10 +1042,12 @@ impl DistributedExt for SessionState { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] @@ -1157,10 +1165,12 @@ impl DistributedExt for SessionContext { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index e01b5ff4..92cc9352 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -1,6 +1,7 @@ use crate::distributed_planner::task_estimator::CombinedTaskEstimator; -use crate::networking::{ChannelResolverExtension, WorkerResolverExtension}; +use crate::protocol::ChannelResolverExtension; use crate::work_unit_feed::WorkUnitFeedRegistry; +use crate::worker_resolver::WorkerResolverExtension; use crate::{TaskEstimator, WorkerResolver}; use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err}; use datafusion::config::{ConfigExtension, ConfigField, ConfigOptions, Visit}; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index e6479e7e..3b7bd3b1 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -1,9 +1,22 @@ use crate::execution_plans::SamplerExec; -use crate::{BroadcastExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage}; +use crate::protocol::ProducerHeadSpec; +use crate::{ + BroadcastExec, DistributedCodec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, + Stage, +}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result; +use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion::prelude::SessionConfig; +use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; +use datafusion_proto::physical_plan::to_proto::serialize_partitioning; +use datafusion_proto::physical_plan::{DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::proto_error; +use prost::Message; use std::sync::Arc; /// This trait represents a node that introduces the necessity of a network boundary in the plan. @@ -61,6 +74,51 @@ impl NetworkBoundaryExt for dyn ExecutionPlan { } impl ProducerHead { + pub(crate) fn to_spec(&self, cfg: &SessionConfig) -> Result { + match self { + Self::None => Ok(ProducerHeadSpec::None), + Self::BroadcastExec { output_partitions } => Ok(ProducerHeadSpec::BroadcastExec { + output_partitions: *output_partitions, + }), + Self::RepartitionExec { partitioning } => { + let partitioning = serialize_partitioning( + partitioning, + &DistributedCodec::new_combined_with_user(cfg), + &DefaultPhysicalProtoConverter {}, + ) + .map(|v| v.encode_to_vec())?; + Ok(ProducerHeadSpec::RepartitionExec { partitioning }) + } + } + } + + pub(crate) fn from_spec( + spec: &ProducerHeadSpec, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result { + match spec { + ProducerHeadSpec::None => Ok(Self::None), + ProducerHeadSpec::BroadcastExec { output_partitions } => Ok(Self::BroadcastExec { + output_partitions: *output_partitions, + }), + ProducerHeadSpec::RepartitionExec { partitioning } => { + let proto_partitioning = protobuf::Partitioning::decode(partitioning.as_slice()) + .map_err(|e| proto_error(e.to_string()))?; + let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, &codec); + let partitioning = parse_protobuf_partitioning( + Some(&proto_partitioning), + &decode_ctx, + &schema, + &DefaultPhysicalProtoConverter {}, + )? + .ok_or_else(|| proto_error("Could not parse partitioning"))?; + Ok(Self::RepartitionExec { partitioning }) + } + } + } + /// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This /// can be called both during planning and lazily at runtime. pub(crate) fn insert(self, input: Arc) -> Result> { diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs index 1e792579..b915e6c1 100644 --- a/src/distributed_planner/task_estimator.rs +++ b/src/distributed_planner/task_estimator.rs @@ -388,9 +388,9 @@ impl TaskEstimator for CombinedTaskEstimator { #[cfg(test)] mod tests { use super::*; - use crate::networking::WorkerResolverExtension; use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::test_utils::parquet::register_parquet_tables; + use crate::worker_resolver::WorkerResolverExtension; use datafusion::error::DataFusionError; use datafusion::prelude::SessionContext; diff --git a/src/execution_plans/benchmarks/fixture.rs b/src/execution_plans/benchmarks/fixture.rs index 4423d7d0..d8b2cc88 100644 --- a/src/execution_plans/benchmarks/fixture.rs +++ b/src/execution_plans/benchmarks/fixture.rs @@ -1,5 +1,4 @@ -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; -use crate::{BoxCloneSyncChannel, ChannelResolver, create_worker_client}; +use crate::{ChannelResolver, WorkerChannel, grpc}; use arrow::datatypes::DataType::{ Boolean, Dictionary, Float64, Int32, Int64, List, Timestamp, UInt8, Utf8, }; @@ -72,18 +71,17 @@ pub(super) fn rows_for_producer( /// tokio duplex rather than a TCP connection. #[derive(Clone)] pub(super) struct InMemoryChannelsResolver { - pub channels: Vec, + pub channels: Vec, } #[async_trait::async_trait] impl ChannelResolver for InMemoryChannelsResolver { - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result> { + async fn get_worker_client_for_url(&self, url: &Url) -> Result> { let Some(port) = url.port() else { return exec_err!("Missing port in url {url}"); }; - Ok(create_worker_client(self.channels[port as usize].clone())) + Ok(grpc::create_worker_client( + self.channels[port as usize].clone(), + )) } } diff --git a/src/execution_plans/benchmarks/transport_bench.rs b/src/execution_plans/benchmarks/transport_bench.rs index a386c7bd..ad159cd5 100644 --- a/src/execution_plans/benchmarks/transport_bench.rs +++ b/src/execution_plans/benchmarks/transport_bench.rs @@ -4,9 +4,7 @@ use super::fixture::{ use crate::common::task_ctx_with_extension; use crate::stage::RemoteStage; use crate::worker::test_utils::worker_handles::{MemoryWorkerHandle, TcpWorkerHandle}; -use crate::{ - DefaultChannelResolver, DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, -}; +use crate::{DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, grpc}; use arrow::datatypes::Schema; use arrow_ipc::CompressionType; use datafusion::common::Result; @@ -233,7 +231,7 @@ impl TransportBench { bench: self.clone(), schema, task_ctx: SessionStateBuilder::new() - .with_distributed_channel_resolver(DefaultChannelResolver::default()) + .with_distributed_channel_resolver(grpc::DefaultChannelResolver::default()) .with_distributed_compression(self.compression)? .build() .task_ctx(), diff --git a/src/execution_plans/network_broadcast.rs b/src/execution_plans/network_broadcast.rs index 251ed29c..aa84ea78 100644 --- a/src/execution_plans/network_broadcast.rs +++ b/src/execution_plans/network_broadcast.rs @@ -248,16 +248,14 @@ impl ExecutionPlan for NetworkBroadcastExec { let mut streams = Vec::with_capacity(self.input_stage.task_count()); for input_task_index in 0..self.input_stage.task_count() { - let worker_connection = self.worker_connections.get_or_init_worker_connection( + streams.push(self.worker_connections.execute( remote_stage, off..(off + self.properties.partitioning.partition_count()), input_task_index, + off + partition, self.producer_head(task_context.task_count), &context, - )?; - - let stream = worker_connection.execute(off + partition)?; - streams.push(stream); + )?); } Ok(Box::pin(RecordBatchStreamAdapter::new( diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 1582d08c..703843cc 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -292,16 +292,15 @@ impl ExecutionPlan for NetworkCoalesceExec { let target_task = group.start_task + input_task_offset; - let worker_connection = self.worker_connections.get_or_init_worker_connection( + let stream = self.worker_connections.execute( remote_stage, 0..partitions_per_task, target_task, + target_partition, self.producer_head(task_context.task_count), &context, )?; - let stream = worker_connection.execute(target_partition)?; - Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), stream, diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 2a575b42..c7a645ae 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -225,16 +225,14 @@ impl ExecutionPlan for NetworkShuffleExec { let mut streams = Vec::with_capacity(remote_stage.workers.len()); for input_task_index in 0..remote_stage.workers.len() { - let worker_connection = self.worker_connections.get_or_init_worker_connection( + streams.push(self.worker_connections.execute( remote_stage, off..(off + self.properties.partitioning.partition_count()), input_task_index, + off + partition, self.producer_head(task_context.task_count), &context, - )?; - - let stream = worker_connection.execute(off + partition)?; - streams.push(stream); + )?); } Ok(Box::pin(RecordBatchStreamAdapter::new( diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs index 0e120219..836d9986 100644 --- a/src/execution_plans/sampler.rs +++ b/src/execution_plans/sampler.rs @@ -1,7 +1,6 @@ use crate::common::{require_one_child, vec_cast}; -use crate::worker::generated::worker as pb; use crate::{ - BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, MaxGaugeMetric, + BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, LoadInfo, MaxGaugeMetric, MaxLatencyMetric, P50LatencyMetric, }; use datafusion::arrow::array::Array; @@ -58,7 +57,7 @@ pub(crate) struct SamplerExecMetrics { /// the input arrived kick_off_to_fist_batch_p50: P50LatencyMetric, kick_off_to_fist_batch_max: MaxLatencyMetric, - /// Time since [SamplerExec::kick_off_first_sampler] was called until the [pb::LoadInfo] message + /// Time since [SamplerExec::kick_off_first_sampler] was called until the [LoadInfo] message /// was sent. kick_off_to_load_info_sent_p50: P50LatencyMetric, kick_off_to_load_info_sent_max: MaxLatencyMetric, @@ -141,7 +140,7 @@ impl SamplerExec { pub(crate) fn kick_off_first_sampler( plan: Arc, ctx: Arc, - ) -> Result>> { + ) -> Result>> { let mut receivers = vec![]; plan.apply(|plan| { let Some(sampler) = plan.downcast_ref::() else { @@ -197,7 +196,7 @@ impl PartitionSampler { self.metrics.kick_off_to_fist_batch_max.add_duration(delay); } - // Time since the sampler was kicked off until the pb::LoadInfo message was sent. + // Time since the sampler was kicked off until the LoadInfo message was sent. if let Some(t) = self.load_info_sent_at.get() { let delay = t.saturating_duration_since(*kick_off_at); self.metrics @@ -216,7 +215,7 @@ impl PartitionSampler { self.stream.lock().unwrap().take() } - fn kick_off(&self, ctx: Arc) -> Result> { + fn kick_off(&self, ctx: Arc) -> Result> { let _ = self.kick_off_at.set(Instant::now()); let (sampling_tx, sampling_rx) = oneshot::channel(); @@ -326,16 +325,16 @@ impl PartitionSampler { } } -/// Wraps a [pb::LoadInfo] and emits it on [Drop] through the provided [oneshot] channel. +/// Wraps a [LoadInfo] and emits it on [Drop] through the provided [oneshot] channel. /// /// Emitting on drop ensures that it's always emitted. struct LoadInfoDropHandler { omit: Arc, - load_info: pb::LoadInfo, + load_info: LoadInfo, bytes_ready_metric: BytesCounterMetric, bytes_per_second_metric: BytesCounterMetric, - sampling_tx: Option>, + sampling_tx: Option>, load_info_sent_at: Arc>, } @@ -359,16 +358,14 @@ impl LoadInfoDropHandler { } fn set_per_col_bytes_per_second(&mut self, total_bytes_per_second: usize) { - let per_col_ready: &[u64] = &self.load_info.per_column_bytes_ready; - let total_ready: u64 = per_col_ready.iter().sum(); + let per_col_ready = &self.load_info.per_column_bytes_ready; + let total_ready = per_col_ready.iter().sum::(); let per_col: Vec = if total_ready == 0 { vec![total_bytes_per_second / per_col_ready.len().max(1); per_col_ready.len()] } else { per_col_ready .iter() - .map(|&ready| { - (ready.saturating_mul(total_bytes_per_second as u64) / total_ready) as usize - }) + .map(|&ready| ready.saturating_mul(total_bytes_per_second) / total_ready) .collect() }; self.load_info.per_column_bytes_per_second = vec_cast(&per_col); @@ -377,11 +374,11 @@ impl LoadInfoDropHandler { } fn set_rows_ready(&mut self, rows_ready: usize) { - self.load_info.rows_ready = rows_ready as u64; + self.load_info.rows_ready = rows_ready; } fn set_rows_per_second(&mut self, rows_per_second: usize) { - self.load_info.rows_per_second = rows_per_second as u64; + self.load_info.rows_per_second = rows_per_second; } fn set_per_col_ndv(&mut self, per_column_ndv: Vec) { @@ -405,9 +402,9 @@ impl Drop for LoadInfoDropHandler { } } -fn zero_load_info(partition_idx: usize, n_cols: usize) -> pb::LoadInfo { - pb::LoadInfo { - partition: partition_idx as u64, +fn zero_load_info(partition_idx: usize, n_cols: usize) -> LoadInfo { + LoadInfo { + partition: partition_idx, rows_per_second: 0, rows_ready: 0, per_column_bytes_per_second: vec![0; n_cols], diff --git a/src/lib.rs b/src/lib.rs index d2a3a463..777dd5c8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,24 +1,21 @@ #![deny(clippy::all)] +mod codec; mod common; mod config_extension_ext; +mod coordinator; mod distributed_ext; +mod distributed_planner; mod execution_plans; mod metrics; mod passthrough_headers; +mod protocol; mod stage; -mod worker; - -mod distributed_planner; -mod networking; -mod observability; -mod protobuf; -pub use protobuf::DistributedCodec; -mod coordinator; -#[cfg(any(feature = "integration", test))] -pub mod test_utils; mod work_unit_feed; +mod worker; +mod worker_resolver; +#[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; pub use coordinator::DistributedExec; pub use distributed_ext::DistributedExt; @@ -36,9 +33,20 @@ pub use metrics::{ MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, rewrite_distributed_plan_with_metrics, }; -pub use networking::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver, - create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver, + +#[cfg(any(feature = "integration", test))] +pub mod test_utils; +#[cfg(feature = "grpc")] +pub use protocol::grpc; + +pub use codec::DistributedCodec; +pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; + +pub use protocol::{ + ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, + GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, + WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, @@ -46,21 +54,11 @@ pub use stage::{ pub use work_unit_feed::{ DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, }; -pub use worker::generated::worker::worker_service_client::WorkerServiceClient; -pub use worker::generated::worker::worker_service_server::WorkerServiceServer; -pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey}; pub use worker::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerQueryContext, WorkerSessionBuilder, }; -pub use observability::{ - GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, - GetTaskProgressResponse, ObservabilityService, ObservabilityServiceClient, - ObservabilityServiceImpl, ObservabilityServiceServer, PingRequest, PingResponse, TaskProgress, - TaskStatus, WorkerMetrics, -}; - #[cfg(any(feature = "integration", test))] pub use execution_plans::benchmarks::{ LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench, diff --git a/src/metrics/latency_metric.rs b/src/metrics/latency_metric.rs index b01f7adb..4d415f75 100644 --- a/src/metrics/latency_metric.rs +++ b/src/metrics/latency_metric.rs @@ -244,7 +244,7 @@ pub struct AvgLatencyMetric { } impl AvgLatencyMetric { - pub(crate) fn from_raw(nanos_sum: usize, count: usize) -> Self { + pub fn from_raw(nanos_sum: usize, count: usize) -> Self { Self { nanos_sum: Arc::new(AtomicUsize::new(nanos_sum)), count: Arc::new(AtomicUsize::new(count)), @@ -255,11 +255,11 @@ impl AvgLatencyMetric { self.nanos_sum.load(Relaxed) / self.count.load(Relaxed).max(1) } - pub(crate) fn nanos_sum(&self) -> usize { + pub fn nanos_sum(&self) -> usize { self.nanos_sum.load(Relaxed) } - pub(crate) fn count(&self) -> usize { + pub fn count(&self) -> usize { self.count.load(Relaxed) } @@ -410,7 +410,7 @@ macro_rules! percentile_latency_metric { } impl $name { - pub(crate) fn from_sketch(sketch: DDSketch) -> Self { + pub fn from_sketch(sketch: DDSketch) -> Self { Self { inner: Arc::new(Mutex::new(sketch)), } @@ -421,7 +421,7 @@ macro_rules! percentile_latency_metric { sketch.quantile($percentile).unwrap_or(None).unwrap_or(0.0) as usize } - pub(crate) fn serialize_sketch(&self) -> Result> { + pub fn serialize_sketch(&self) -> Result> { let sketch = self.inner.lock().unwrap(); bincode::serialize(&*sketch).map_err(|e| { datafusion::error::DataFusionError::Internal(format!( @@ -430,7 +430,7 @@ macro_rules! percentile_latency_metric { }) } - pub(crate) fn count(&self) -> usize { + pub fn count(&self) -> usize { let sketch = self.inner.lock().unwrap(); sketch.count() as usize } diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs index 6e88c231..fd113942 100644 --- a/src/metrics/mod.rs +++ b/src/metrics/mod.rs @@ -1,7 +1,6 @@ mod bytes_metric; mod latency_metric; mod max_gauge_metric; -pub(crate) mod proto; mod task_metrics_collector; mod task_metrics_rewriter; diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 553db1cc..2d78998d 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -186,7 +186,7 @@ mod tests { // empty metrics (e.g., custom execution plans without metrics), which is fine - we // just verify that a metrics set exists for each node. The count assertion above // ensures all nodes are included in the metrics collection. - let stage = stages.get(&(expected_task_key.stage_id as usize)).unwrap(); + let stage = stages.get(&expected_task_key.stage_id).unwrap(); let stage_plan = stage.local_plan().unwrap(); assert_eq!( actual_metrics.pre_order_plan_metrics.len(), @@ -306,7 +306,7 @@ mod tests { sent metrics via the coordinator channel." ) }); - let stage = stages.get(&(expected_task_key.stage_id as usize)).unwrap(); + let stage = stages.get(&expected_task_key.stage_id).unwrap(); let stage_plan = stage.local_plan().unwrap(); assert_eq!( actual_metrics.pre_order_plan_metrics.len(), diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 0e3d1a8f..8913ff87 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -1,13 +1,11 @@ -use crate::DistributedTaskContext; -use crate::common::{TreeNodeExt, serialize_uuid}; +use crate::common::TreeNodeExt; use crate::coordinator::{DistributedExec, MetricsStore}; use crate::distributed_planner::NetworkBoundaryExt; use crate::execution_plans::MetricsWrapperExec; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use crate::metrics::collect_plan_metrics; -use crate::metrics::proto::metrics_set_proto_to_df; use crate::stage::{LocalStage, Stage}; -use crate::worker::generated::worker::TaskKey; +use crate::{DistributedTaskContext, TaskKey}; use datafusion::common::HashMap; use datafusion::common::plan_err; use datafusion::common::tree_node::Transformed; @@ -228,9 +226,9 @@ pub fn stage_metrics_rewriter( task_count: stage.tasks, }; let task_key = TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }; let Some(task_metrics) = metrics_collection.get(&task_key) else { return internal_err!( @@ -249,8 +247,7 @@ pub fn stage_metrics_rewriter( ); } - let node_metrics_protos = task_metrics.pre_order_plan_metrics[per_task_counter].clone(); - let mut node_metrics = metrics_set_proto_to_df(&node_metrics_protos)?; + let mut node_metrics = task_metrics.pre_order_plan_metrics[per_task_counter].clone(); let rewrite_ctx = format.to_rewrite_ctx(task_id as u64); node_metrics = rewrite_ctx.maybe_rewrite_node_metics(node_metrics); @@ -285,41 +282,35 @@ pub fn stage_metrics_rewriter( #[cfg(test)] mod tests { + use crate::DistributedExt; use crate::coordinator::MetricsStore; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; - use crate::metrics::proto::{df_metrics_set_to_proto, metrics_set_proto_to_df}; + use crate::metrics::task_metrics_rewriter::MetricsWrapperExec; use crate::metrics::task_metrics_rewriter::{ annotate_metrics_set_with_task_id, stage_metrics_rewriter, }; use crate::metrics::{DistributedMetricsFormat, rewrite_distributed_plan_with_metrics}; + use crate::stage::LocalStage; use crate::test_utils::in_memory_channel_resolver::{ InMemoryChannelResolver, InMemoryWorkerResolver, }; - use crate::test_utils::metrics::make_test_metrics_set_proto_from_seed; + use crate::test_utils::metrics::make_test_metrics_set_from_seed; use crate::test_utils::plans::count_plan_nodes_up_to_network_boundary; use crate::test_utils::session_context::register_temp_parquet_table; - use crate::worker::generated::worker as pb; - use crate::{DistributedExec, SessionStateBuilderExt}; + use crate::{DistributedExec, SessionStateBuilderExt, TaskKey, TaskMetrics}; use datafusion::arrow::array::{Int32Array, StringArray}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::metrics::{Count, Label, Metric, MetricValue, MetricsSet}; - use test_case::test_case; - use datafusion::physical_plan::{ExecutionPlan, collect}; - use itertools::Itertools; - use uuid::Uuid; - - use crate::DistributedExt; - use crate::common::serialize_uuid; - use crate::metrics::task_metrics_rewriter::MetricsWrapperExec; - use crate::stage::LocalStage; - use crate::worker::generated::worker::TaskKey; - use datafusion::physical_plan::empty::EmptyExec; use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; + use itertools::Itertools; use std::sync::Arc; + use test_case::test_case; + use uuid::Uuid; async fn make_test_ctx() -> SessionContext { make_test_ctx_inner(false).await @@ -431,8 +422,10 @@ mod tests { fn metrics_set_eq(a: &MetricsSet, b: &MetricsSet) -> bool { println!("a: {a:?}"); println!("b: {b:?}"); - // Check equality by converting to proto representation. - df_metrics_set_to_proto(a).unwrap() == df_metrics_set_to_proto(b).unwrap() + a.iter().count() == b.iter().count() + && a.iter().zip(b.iter()).all(|(a, b)| { + a.value() == b.value() && a.partition() == b.partition() && a.labels() == b.labels() + }) } /// Asserts that we successfully re-write the metrics of a plan generated from the provided SQL query. @@ -458,20 +451,20 @@ mod tests { // Generate metrics for each task and store them in the map. let metrics_collection = MetricsStore::from_entries((0..stage.tasks).map(|task_id| { let task_key = TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }; let metrics = (0..count_plan_nodes_up_to_network_boundary(&plan)) .map(|node_id| { - make_test_metrics_set_proto_from_seed( + make_test_metrics_set_from_seed( (node_id * task_id) as u64, num_metrics_per_task_per_node, ) }) - .collect::>(); - let task_metrics = pb::TaskMetrics { - task_metrics: None, + .collect::>(); + let task_metrics = TaskMetrics { + task_metrics: MetricsSet::new(), pre_order_plan_metrics: metrics, }; (task_key, task_metrics) @@ -502,9 +495,9 @@ mod tests { { let expected_task_node_metrics = metrics_collection .get(&TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }) .unwrap() .pre_order_plan_metrics[node_id] @@ -514,9 +507,7 @@ mod tests { actual_task_node_metrics_set .for_each(|metric| actual_metrics_set.push(metric.clone())); - // Convert from proto to check for equality. - let mut expected_metrics_set = - metrics_set_proto_to_df(&expected_task_node_metrics).unwrap(); + let mut expected_metrics_set = expected_task_node_metrics; if format == DistributedMetricsFormat::PerTask { // Add task ids labels. We expect the actual metrics to be annotated by the diff --git a/src/networking/mod.rs b/src/networking/mod.rs deleted file mode 100644 index 6bab9ae0..00000000 --- a/src/networking/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -mod channel_resolver; -mod worker_resolver; - -pub use channel_resolver::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, - get_distributed_channel_resolver, -}; -pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; - -pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; -pub(crate) use worker_resolver::{WorkerResolverExtension, set_distributed_worker_resolver}; diff --git a/src/protobuf/producer_head.rs b/src/protobuf/producer_head.rs deleted file mode 100644 index d4cbfc42..00000000 --- a/src/protobuf/producer_head.rs +++ /dev/null @@ -1,63 +0,0 @@ -use crate::DistributedCodec; -use crate::distributed_planner::ProducerHead; -use crate::worker::generated::worker::{ - BroadcastExecHead, NoneHead, RepartitionExecHead, execute_task_request as pb, -}; -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; -use datafusion::execution::TaskContext; -use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; -use datafusion_proto::physical_plan::to_proto::serialize_partitioning; -use datafusion_proto::physical_plan::{DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext}; -use datafusion_proto::protobuf; -use datafusion_proto::protobuf::proto_error; -use prost::Message; -use std::sync::Arc; - -impl ProducerHead { - pub(crate) fn from_proto( - proto: pb::ProducerHead, - input_schema: &SchemaRef, - ctx: &Arc, - ) -> Result { - Ok(match proto { - pb::ProducerHead::None(_) => ProducerHead::None, - pb::ProducerHead::Broadcast(v) => ProducerHead::BroadcastExec { - output_partitions: v.output_partitions as usize, - }, - pb::ProducerHead::Repartition(v) => { - let proto_partitioning = protobuf::Partitioning::decode(v.partitioning.as_slice()) - .map_err(|e| proto_error(e.to_string()))?; - let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); - let decode_ctx = PhysicalPlanDecodeContext::new(ctx, &codec); - let partitioning = parse_protobuf_partitioning( - Some(&proto_partitioning), - &decode_ctx, - input_schema, - &DefaultPhysicalProtoConverter {}, - )? - .ok_or_else(|| proto_error("Could not parse partitioning"))?; - - ProducerHead::RepartitionExec { partitioning } - } - }) - } - - pub(crate) fn to_proto(&self, ctx: &Arc) -> Result { - Ok(match self { - ProducerHead::None => pb::ProducerHead::None(NoneHead {}), - ProducerHead::BroadcastExec { output_partitions } => { - pb::ProducerHead::Broadcast(BroadcastExecHead { - output_partitions: *output_partitions as u64, - }) - } - ProducerHead::RepartitionExec { partitioning } => { - let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); - let partitioning = - serialize_partitioning(partitioning, &codec, &DefaultPhysicalProtoConverter {}) - .map(|v| v.encode_to_vec())?; - pb::ProducerHead::Repartition(RepartitionExecHead { partitioning }) - } - }) - } -} diff --git a/src/protocol/channel_resolver.rs b/src/protocol/channel_resolver.rs new file mode 100644 index 00000000..ded82e05 --- /dev/null +++ b/src/protocol/channel_resolver.rs @@ -0,0 +1,99 @@ +use crate::config_extension_ext::set_distributed_option_extension; +#[cfg(feature = "grpc")] +use crate::protocol::grpc; +use crate::{DistributedConfig, WorkerChannel}; +use async_trait::async_trait; +use datafusion::common::DataFusionError; +use datafusion::execution::TaskContext; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; +use url::Url; + +/// Allows users to customize the way Worker clients are created. A common use case is to +/// wrap the client with tower layers or schedule it in an IO-specific tokio runtime. +/// +/// There is a default implementation of this trait that should be enough for the most common +/// use-cases. +/// +/// # Implementation Notes +/// - This is called per request, so implementors of this trait should make sure that +/// clients are reused across method calls instead of building a new Worker client every time. +/// +/// - When implementing `get_worker_client_for_url`, it is recommended to use the +/// [`create_worker_client`] helper function to ensure clients are configured with +/// appropriate message size limits for internal communication. This helps avoid message +/// size errors when transferring large datasets. +#[async_trait] +pub trait ChannelResolver { + /// For a given URL, get a Worker gRPC client for communicating to it. + /// + /// *WARNING*: This method is called for every gRPC request, so to not create + /// one client connection for each request, users are required to reuse generated clients. + /// It's recommended to rely on [DefaultChannelResolver] either by delegating method calls + /// to it or by copying the implementation. + /// + /// Consider using [`create_worker_client`] to create the client with appropriate + /// default message size limits. + async fn get_worker_client_for_url( + &self, + url: &Url, + ) -> Result, DataFusionError>; +} + +pub(crate) fn set_distributed_channel_resolver( + cfg: &mut SessionConfig, + channel_resolver: impl ChannelResolver + Send + Sync + 'static, +) { + let opts = cfg.options_mut(); + let channel_resolver = ChannelResolverExtension(Some(Arc::new(channel_resolver))); + if let Some(distributed_cfg) = opts.extensions.get_mut::() { + distributed_cfg.__private_channel_resolver = channel_resolver; + } else { + set_distributed_option_extension( + cfg, + DistributedConfig { + __private_channel_resolver: channel_resolver, + ..Default::default() + }, + ) + } +} + +pub fn get_distributed_channel_resolver( + task_ctx: &TaskContext, +) -> Arc { + let opts = task_ctx.session_config().options(); + if let Some(distributed_cfg) = opts.extensions.get::() + && let Some(cr) = &distributed_cfg.__private_channel_resolver.0 + { + return Arc::clone(cr); + } + + #[cfg(feature = "grpc")] + { + let runtime_addr = Arc::as_ptr(&task_ctx.runtime_env()) as usize; + grpc::DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME.get_with(runtime_addr, || { + Arc::new(grpc::DefaultChannelResolver::default()) + }) + } + + #[cfg(not(feature = "grpc"))] + { + panic!( + "gRPC feature is not enabled, and no channel resolver was provided, so no default ChannelResolver can be provided" + ); + } +} + +#[async_trait] +impl ChannelResolver for Arc { + async fn get_worker_client_for_url( + &self, + url: &Url, + ) -> Result, DataFusionError> { + self.as_ref().get_worker_client_for_url(url).await + } +} + +#[derive(Clone, Default)] +pub(crate) struct ChannelResolverExtension(Option>); diff --git a/src/networking/channel_resolver.rs b/src/protocol/grpc/channel_resolver.rs similarity index 56% rename from src/networking/channel_resolver.rs rename to src/protocol/grpc/channel_resolver.rs index d15c8bef..b545e21b 100644 --- a/src/networking/channel_resolver.rs +++ b/src/protocol/grpc/channel_resolver.rs @@ -1,10 +1,7 @@ -use crate::DistributedConfig; -use crate::config_extension_ext::set_distributed_option_extension; -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; +use crate::protocol::grpc::worker_client::create_worker_client; +use crate::{ChannelResolver, WorkerChannel}; use async_trait::async_trait; use datafusion::common::{DataFusionError, config_datafusion_err, exec_datafusion_err}; -use datafusion::execution::TaskContext; -use datafusion::prelude::SessionConfig; use futures::FutureExt; use futures::future::Shared; use std::sync::{Arc, LazyLock}; @@ -15,56 +12,6 @@ use tonic::transport::Channel; use tower::ServiceExt; use url::Url; -/// Allows users to customize the way Worker clients are created. A common use case is to -/// wrap the client with tower layers or schedule it in an IO-specific tokio runtime. -/// -/// There is a default implementation of this trait that should be enough for the most common -/// use-cases. -/// -/// # Implementation Notes -/// - This is called per gRPC request, so implementors of this trait should make sure that -/// clients are reused across method calls instead of building a new Worker client every time. -/// -/// - When implementing `get_worker_client_for_url`, it is recommended to use the -/// [`create_worker_client`] helper function to ensure clients are configured with -/// appropriate message size limits for internal communication. This helps avoid message -/// size errors when transferring large datasets. -#[async_trait] -pub trait ChannelResolver { - /// For a given URL, get a Worker gRPC client for communicating to it. - /// - /// *WARNING*: This method is called for every gRPC request, so to not create - /// one client connection for each request, users are required to reuse generated clients. - /// It's recommended to rely on [DefaultChannelResolver] either by delegating method calls - /// to it or by copying the implementation. - /// - /// Consider using [`create_worker_client`] to create the client with appropriate - /// default message size limits. - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result, DataFusionError>; -} - -pub(crate) fn set_distributed_channel_resolver( - cfg: &mut SessionConfig, - channel_resolver: impl ChannelResolver + Send + Sync + 'static, -) { - let opts = cfg.options_mut(); - let channel_resolver = ChannelResolverExtension(Some(Arc::new(channel_resolver))); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg.__private_channel_resolver = channel_resolver; - } else { - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_channel_resolver: channel_resolver, - ..Default::default() - }, - ) - } -} - // Unlike TaskContext, a DataFusion RuntimeEnv does not allow to introduce user-defined extensions. // For the default implementation of the ChannelResolvers, we cannot inject one DefaultChannelResolver // per TaskContext, as this holds reference to Tonic channels that must outlive a single TaskContext. @@ -72,27 +19,13 @@ pub(crate) fn set_distributed_channel_resolver( // The Tonic channels need to be established and reused under a whole RuntimeEnv scope, not a single // TaskContext, which forces us to put the default implementation in a static global variable that // stores and reuses tonic channels per RuntimeEnv's pointer address. -static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< +pub(crate) static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< moka::sync::Cache< /* Arc pointer address */ usize, /* ChannelResolver that reuses built channels */ Arc, >, > = LazyLock::new(|| moka::sync::Cache::builder().max_capacity(256).build()); -pub fn get_distributed_channel_resolver( - task_ctx: &TaskContext, -) -> Arc { - let opts = task_ctx.session_config().options(); - if let Some(distributed_cfg) = opts.extensions.get::() - && let Some(cr) = &distributed_cfg.__private_channel_resolver.0 - { - return Arc::clone(cr); - } - let runtime_addr = Arc::as_ptr(&task_ctx.runtime_env()) as usize; - DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME - .get_with(runtime_addr, || Arc::new(DefaultChannelResolver::default())) -} - pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService< http::Request, http::Response, @@ -101,9 +34,6 @@ pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService< type ChannelCacheValue = Shared>>; -#[derive(Clone, Default)] -pub(crate) struct ChannelResolverExtension(Option>); - /// Default implementation of a [ChannelResolver] that connects to the workers given the URL once /// and stores the connection instance in a TTI cache. /// @@ -159,8 +89,8 @@ impl DefaultChannelResolver { })?; Ok(BoxCloneSyncChannel::new(channel)) } - .boxed() - .shared() + .boxed() + .shared() }); channel.await.map_err(|err| { @@ -175,56 +105,11 @@ impl ChannelResolver for DefaultChannelResolver { async fn get_worker_client_for_url( &self, url: &Url, - ) -> Result, DataFusionError> { + ) -> Result, DataFusionError> { self.get_channel(url).await.map(create_worker_client) } } -#[async_trait] -impl ChannelResolver for Arc { - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result, DataFusionError> { - self.as_ref().get_worker_client_for_url(url).await - } -} - -/// Creates a [`WorkerServiceClient`] with high default message size limits. -/// -/// This is a convenience function that wraps [`WorkerServiceClient::new`] and configures -/// it with `max_decoding_message_size(usize::MAX)` and `max_encoding_message_size(usize::MAX)` -/// to avoid message size limitations for internal communication. -/// -/// Users implementing custom [`ChannelResolver`]s should use this function in their -/// `get_worker_client_for_url` implementations to ensure consistent behavior with built-in -/// implementations. -/// -/// # Example -/// -/// ```rust,ignore -/// use datafusion_distributed::{create_worker_client, BoxCloneSyncChannel, ChannelResolver}; -/// /// use tonic::transport::Channel; -/// -/// #[async_trait] -/// impl ChannelResolver for MyResolver { -/// async fn get_worker_client_for_url( -/// &self, -/// url: &Url, -/// ) -> Result, DataFusionError> { -/// let channel = Channel::from_shared(url.to_string())?.connect().await?; -/// Ok(create_worker_client(BoxCloneSyncChannel::new(channel))) -/// } -/// } -/// ``` -pub fn create_worker_client( - channel: BoxCloneSyncChannel, -) -> WorkerServiceClient { - WorkerServiceClient::new(channel) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/protobuf/errors/arrow_error.rs b/src/protocol/grpc/errors/arrow_error.rs similarity index 99% rename from src/protobuf/errors/arrow_error.rs rename to src/protocol/grpc/errors/arrow_error.rs index a8630bba..b78a950c 100644 --- a/src/protobuf/errors/arrow_error.rs +++ b/src/protocol/grpc/errors/arrow_error.rs @@ -1,6 +1,6 @@ use datafusion::arrow::error::ArrowError; -use crate::protobuf::errors::io_error::IoErrorProto; +use super::io_error::IoErrorProto; #[derive(Clone, PartialEq, ::prost::Message)] pub struct ArrowErrorProto { diff --git a/src/protobuf/errors/datafusion_error.rs b/src/protocol/grpc/errors/datafusion_error.rs similarity index 97% rename from src/protobuf/errors/datafusion_error.rs rename to src/protocol/grpc/errors/datafusion_error.rs index 1f5caede..71d41730 100644 --- a/src/protobuf/errors/datafusion_error.rs +++ b/src/protocol/grpc/errors/datafusion_error.rs @@ -1,9 +1,9 @@ -use crate::protobuf::errors::arrow_error::ArrowErrorProto; -use crate::protobuf::errors::io_error::IoErrorProto; -use crate::protobuf::errors::objectstore_error::ObjectStoreErrorProto; -use crate::protobuf::errors::parquet_error::ParquetErrorProto; -use crate::protobuf::errors::parser_error::ParserErrorProto; -use crate::protobuf::errors::schema_error::SchemaErrorProto; +use super::arrow_error::ArrowErrorProto; +use super::io_error::IoErrorProto; +use super::objectstore_error::ObjectStoreErrorProto; +use super::parquet_error::ParquetErrorProto; +use super::parser_error::ParserErrorProto; +use super::schema_error::SchemaErrorProto; use datafusion::common::{DataFusionError, Diagnostic}; use datafusion::logical_expr::sqlparser::parser::ParserError; use std::error::Error; diff --git a/src/protobuf/errors/io_error.rs b/src/protocol/grpc/errors/io_error.rs similarity index 100% rename from src/protobuf/errors/io_error.rs rename to src/protocol/grpc/errors/io_error.rs diff --git a/src/protobuf/errors/mod.rs b/src/protocol/grpc/errors/mod.rs similarity index 97% rename from src/protobuf/errors/mod.rs rename to src/protocol/grpc/errors/mod.rs index 0393964e..b940d7c5 100644 --- a/src/protobuf/errors/mod.rs +++ b/src/protocol/grpc/errors/mod.rs @@ -6,7 +6,7 @@ use datafusion::error::DataFusionError; use prost::Message; use std::borrow::Borrow; -use crate::protobuf::errors::datafusion_error::DataFusionErrorProto; +use super::errors::datafusion_error::DataFusionErrorProto; mod arrow_error; mod datafusion_error; diff --git a/src/protobuf/errors/objectstore_error.rs b/src/protocol/grpc/errors/objectstore_error.rs similarity index 100% rename from src/protobuf/errors/objectstore_error.rs rename to src/protocol/grpc/errors/objectstore_error.rs diff --git a/src/protobuf/errors/parquet_error.rs b/src/protocol/grpc/errors/parquet_error.rs similarity index 100% rename from src/protobuf/errors/parquet_error.rs rename to src/protocol/grpc/errors/parquet_error.rs diff --git a/src/protobuf/errors/parser_error.rs b/src/protocol/grpc/errors/parser_error.rs similarity index 100% rename from src/protobuf/errors/parser_error.rs rename to src/protocol/grpc/errors/parser_error.rs diff --git a/src/protobuf/errors/schema_error.rs b/src/protocol/grpc/errors/schema_error.rs similarity index 100% rename from src/protobuf/errors/schema_error.rs rename to src/protocol/grpc/errors/schema_error.rs diff --git a/src/worker/gen/Cargo.toml b/src/protocol/grpc/gen/Cargo.toml similarity index 100% rename from src/worker/gen/Cargo.toml rename to src/protocol/grpc/gen/Cargo.toml diff --git a/src/observability/gen/regen.sh b/src/protocol/grpc/gen/regen.sh similarity index 57% rename from src/observability/gen/regen.sh rename to src/protocol/grpc/gen/regen.sh index d5d48ef8..1a1d85ea 100755 --- a/src/observability/gen/regen.sh +++ b/src/protocol/grpc/gen/regen.sh @@ -3,4 +3,4 @@ set -e repo_root=$(git rev-parse --show-toplevel) -cd "$repo_root" && cargo run --manifest-path src/observability/gen/Cargo.toml +cd "$repo_root" && cargo run --manifest-path src/protocol/grpc/gen/Cargo.toml diff --git a/src/worker/gen/src/main.rs b/src/protocol/grpc/gen/src/main.rs similarity index 87% rename from src/worker/gen/src/main.rs rename to src/protocol/grpc/gen/src/main.rs index 8994d54d..7505aae5 100644 --- a/src/worker/gen/src/main.rs +++ b/src/protocol/grpc/gen/src/main.rs @@ -4,9 +4,9 @@ use std::fs; fn main() -> Result<(), Box> { let repo_root = env::current_dir()?; - let proto_dir = repo_root.join("src/worker"); + let proto_dir = repo_root.join("src/protocol/grpc"); let proto_file = proto_dir.join("worker.proto"); - let out_dir = repo_root.join("src/worker/generated"); + let out_dir = repo_root.join("src/protocol/grpc/generated"); fs::create_dir_all(&out_dir)?; diff --git a/src/worker/generated/mod.rs b/src/protocol/grpc/generated/mod.rs similarity index 100% rename from src/worker/generated/mod.rs rename to src/protocol/grpc/generated/mod.rs diff --git a/src/worker/generated/worker.rs b/src/protocol/grpc/generated/worker.rs similarity index 100% rename from src/worker/generated/worker.rs rename to src/protocol/grpc/generated/worker.rs diff --git a/src/metrics/proto.rs b/src/protocol/grpc/metrics_proto.rs similarity index 98% rename from src/metrics/proto.rs rename to src/protocol/grpc/metrics_proto.rs index c114a41c..6db8cda9 100644 --- a/src/metrics/proto.rs +++ b/src/protocol/grpc/metrics_proto.rs @@ -1,3 +1,4 @@ +use super::generated::worker as pb; use chrono::DateTime; use datafusion::common::internal_err; use datafusion::error::DataFusionError; @@ -8,15 +9,12 @@ use sketches_ddsketch::DDSketch; use std::borrow::Cow; use std::sync::Arc; -use super::bytes_metric::BytesCounterMetric; -use super::latency_metric::{ - AvgLatencyMetric, FirstLatencyMetric, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, - P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, +use crate::{ + AvgLatencyMetric, BytesCounterMetric, FirstLatencyMetric, MaxGaugeMetric, MaxLatencyMetric, + MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, }; -use crate::MaxGaugeMetric; -use crate::worker::generated::worker as pb; -/// df_metrics_set_to_proto converts a [datafusion::physical_plan::metrics::MetricsSet] to a [pb::MetricsSet]. +/// df_metrics_set_to_proto converts a [MetricsSet] to a [pb::MetricsSet]. /// Custom metrics are filtered out, but any other errors are returned. /// TODO(#140): Support custom metrics. pub fn df_metrics_set_to_proto( @@ -44,7 +42,7 @@ pub fn df_metrics_set_to_proto( Ok(pb::MetricsSet { metrics }) } -/// metrics_set_proto_to_df converts a [pb::MetricsSet] to a [datafusion::physical_plan::metrics::MetricsSet]. +/// metrics_set_proto_to_df converts a [pb::MetricsSet] to a [MetricsSet]. pub fn metrics_set_proto_to_df( metrics_set_proto: &pb::MetricsSet, ) -> Result { @@ -64,7 +62,7 @@ const CUSTOM_METRICS_NOT_SUPPORTED: &str = /// New DataFusion metrics that are not yet supported in proto conversion. const UNSUPPORTED_METRICS: &str = "metric type not supported in proto conversion"; -/// df_metric_to_proto converts a `datafusion::physical_plan::metrics::Metric` to a `pb::Metric`. It does not consume the Arc. +/// df_metric_to_proto converts a `Metric` to a `pb::Metric`. It does not consume the Arc. pub fn df_metric_to_proto(metric: Arc) -> Result { let partition = metric.partition().map(|p| p as u64); let labels = metric @@ -284,7 +282,7 @@ pub fn df_metric_to_proto(metric: Arc) -> Result Result, DataFusionError> { let partition = metric.partition.map(|p| p as usize); let labels = metric diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs new file mode 100644 index 00000000..e432b607 --- /dev/null +++ b/src/protocol/grpc/mod.rs @@ -0,0 +1,21 @@ +mod channel_resolver; +mod errors; +mod generated; +mod metrics_proto; +mod observability; +mod on_drop_stream; +mod spawn_select_all; +mod worker_client; +mod worker_service; + +// TODO: this should not be exposed. +pub(crate) use channel_resolver::DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME; + +pub use channel_resolver::{BoxCloneSyncChannel, DefaultChannelResolver}; +pub use observability::{ + GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, + GetTaskProgressResponse, ObservabilityService, ObservabilityServiceClient, + ObservabilityServiceImpl, ObservabilityServiceServer, PingRequest, PingResponse, TaskProgress, + TaskStatus, WorkerMetrics, +}; +pub use worker_client::create_worker_client; diff --git a/src/observability/README.md b/src/protocol/grpc/observability/README.md similarity index 88% rename from src/observability/README.md rename to src/protocol/grpc/observability/README.md index 73dd5389..94d9b840 100644 --- a/src/observability/README.md +++ b/src/protocol/grpc/observability/README.md @@ -8,5 +8,5 @@ for tonic/gRPC services. In the root of the datafusion-distribued repo, run: ```bash -./src/observability/gen/regen.sh +./src/protocol/grpc/observability/gen/regen.sh ``` diff --git a/src/observability/gen/Cargo.toml b/src/protocol/grpc/observability/gen/Cargo.toml similarity index 100% rename from src/observability/gen/Cargo.toml rename to src/protocol/grpc/observability/gen/Cargo.toml diff --git a/src/protocol/grpc/observability/gen/regen.sh b/src/protocol/grpc/observability/gen/regen.sh new file mode 100755 index 00000000..c95505eb --- /dev/null +++ b/src/protocol/grpc/observability/gen/regen.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash + +set -e + +repo_root=$(git rev-parse --show-toplevel) +cd "$repo_root" && cargo run --manifest-path src/protocol/grpc/observability/gen/Cargo.toml diff --git a/src/observability/gen/src/main.rs b/src/protocol/grpc/observability/gen/src/main.rs similarity index 76% rename from src/observability/gen/src/main.rs rename to src/protocol/grpc/observability/gen/src/main.rs index d414fc53..124f0f1b 100644 --- a/src/observability/gen/src/main.rs +++ b/src/protocol/grpc/observability/gen/src/main.rs @@ -4,9 +4,9 @@ use std::fs; fn main() -> Result<(), Box> { let repo_root = env::current_dir()?; - let proto_dir = repo_root.join("src/observability/proto"); + let proto_dir = repo_root.join("src/protocol/grpc/observability/proto"); let proto_file = proto_dir.join("observability.proto"); - let out_dir = repo_root.join("src/observability/generated"); + let out_dir = repo_root.join("src/protocol/grpc/observability/generated"); fs::create_dir_all(&out_dir)?; @@ -21,7 +21,7 @@ fn main() -> Result<(), Box> { .out_dir(&out_dir) .extern_path( ".observability.TaskKey", - "crate::worker::generated::worker::TaskKey", + "crate::protocol::grpc::generated::worker::TaskKey", ) .compile_protos(&[proto_file], &[proto_dir])?; diff --git a/src/observability/generated/mod.rs b/src/protocol/grpc/observability/generated/mod.rs similarity index 100% rename from src/observability/generated/mod.rs rename to src/protocol/grpc/observability/generated/mod.rs diff --git a/src/observability/generated/observability.rs b/src/protocol/grpc/observability/generated/observability.rs similarity index 99% rename from src/observability/generated/observability.rs rename to src/protocol/grpc/observability/generated/observability.rs index ee40102b..3b0a95cf 100644 --- a/src/observability/generated/observability.rs +++ b/src/protocol/grpc/observability/generated/observability.rs @@ -12,9 +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, - #[prost(uint64, tag = "2")] - pub total_partitions: u64, + pub task_key: ::core::option::Option, #[prost(enumeration = "TaskStatus", tag = "4")] pub status: i32, #[prost(uint64, tag = "5")] diff --git a/src/observability/mod.rs b/src/protocol/grpc/observability/mod.rs similarity index 100% rename from src/observability/mod.rs rename to src/protocol/grpc/observability/mod.rs diff --git a/src/observability/proto/observability.proto b/src/protocol/grpc/observability/proto/observability.proto similarity index 95% rename from src/observability/proto/observability.proto rename to src/protocol/grpc/observability/proto/observability.proto index f403aa85..8bc76b5a 100644 --- a/src/observability/proto/observability.proto +++ b/src/protocol/grpc/observability/proto/observability.proto @@ -24,8 +24,7 @@ message TaskKey { // Progress information for a single task message TaskProgress { TaskKey task_key = 1; - uint64 total_partitions = 2; - reserved 3; + reserved 2, 3; TaskStatus status = 4; uint64 output_rows = 5; } diff --git a/src/observability/service.rs b/src/protocol/grpc/observability/service.rs similarity index 89% rename from src/observability/service.rs rename to src/protocol/grpc/observability/service.rs index f49c4270..27edee2e 100644 --- a/src/observability/service.rs +++ b/src/protocol/grpc/observability/service.rs @@ -2,9 +2,11 @@ use super::{ GetTaskProgressResponse, ObservabilityService, TaskProgress, TaskStatus, WorkerMetrics, generated::observability::{GetTaskProgressRequest, PingRequest, PingResponse}, }; -use crate::worker::generated::worker::TaskKey; +use crate::common::serialize_uuid; +use crate::grpc::{GetClusterWorkersRequest, GetClusterWorkersResponse}; +use crate::protocol::grpc::generated::worker as worker_pb; use crate::worker::{SingleWriteMultiRead, TaskData}; -use crate::{GetClusterWorkersRequest, GetClusterWorkersResponse, WorkerResolver}; +use crate::{TaskKey, WorkerResolver}; use datafusion::error::DataFusionError; use datafusion::physical_plan::ExecutionPlan; use moka::future::Cache; @@ -93,12 +95,10 @@ impl ObservabilityService for ObservabilityServiceImpl { // Only include initialized tasks if let Some(Ok(task_data)) = task_data_cell.read_now() { - let total_partitions = task_data.total_partitions() as u64; let output_rows = output_rows_from_plan(&task_data.base_plan); tasks.push(TaskProgress { - task_key: Some((*internal_key).clone()), - total_partitions, + task_key: Some(task_key_to_proto(&internal_key)), status: TaskStatus::Running as i32, output_rows, }); @@ -136,7 +136,7 @@ impl ObservabilityServiceImpl { } #[cfg(feature = "system-metrics")] - return *self.system.borrow(); + *self.system.borrow() } } @@ -144,3 +144,11 @@ impl ObservabilityServiceImpl { fn output_rows_from_plan(plan: &Arc) -> u64 { plan.metrics().and_then(|m| m.output_rows()).unwrap_or(0) as u64 } + +fn task_key_to_proto(task_key: &TaskKey) -> worker_pb::TaskKey { + worker_pb::TaskKey { + query_id: serialize_uuid(&task_key.query_id), + stage_id: task_key.stage_id as u64, + task_number: task_key.task_number as u64, + } +} diff --git a/src/common/on_drop_stream.rs b/src/protocol/grpc/on_drop_stream.rs similarity index 100% rename from src/common/on_drop_stream.rs rename to src/protocol/grpc/on_drop_stream.rs diff --git a/src/worker/spawn_select_all.rs b/src/protocol/grpc/spawn_select_all.rs similarity index 100% rename from src/worker/spawn_select_all.rs rename to src/protocol/grpc/spawn_select_all.rs diff --git a/src/worker/worker.proto b/src/protocol/grpc/worker.proto similarity index 100% rename from src/worker/worker.proto rename to src/protocol/grpc/worker.proto diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs new file mode 100644 index 00000000..9cc94de2 --- /dev/null +++ b/src/protocol/grpc/worker_client.rs @@ -0,0 +1,766 @@ +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::{ + BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, + FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, + MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHeadSpec, + SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, + WorkerChannel, WorkerToCoordinatorMsg, +}; +use arrow_flight::FlightData; +use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::error::FlightError; +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::instant::Instant; +use datafusion::common::runtime::SpawnedTask; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::physical_expr_common::metrics::{Count, MetricBuilder, MetricValue, Time}; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::stream::BoxStream; +use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; +use http::{Extensions, HeaderMap}; +use pin_project::{pin_project, pinned_drop}; +use prost::Message; +use std::borrow::Cow; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::sync::Notify; +use tokio::sync::mpsc::UnboundedSender; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tokio_util::sync::CancellationToken; +use tonic::metadata::MetadataMap; +use tonic::{Request, Status}; + +#[async_trait] +impl WorkerChannel for pb::worker_service_client::WorkerServiceClient { + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + let input_stream = c2w_stream.map(encode_coordinator_to_worker_msg); + + let output_stream = self + .coordinator_channel(Request::from_parts( + MetadataMap::from_headers(headers), + Extensions::default(), + input_stream, + )) + .boxed() + .await + .map_err(map_status_to_datafusion_error)? + .into_inner() + .map_err(map_status_to_datafusion_error) + .map(|msg| decode_worker_to_coordinator_msg(msg?)) + .boxed(); + + Ok(output_stream) + } + + async fn execute_task( + &mut self, + headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + ctx: &Arc, + ) -> Result>>> { + let d_cfg = DistributedConfig::from_session_config(ctx.session_config())?; + let buffer_budget_bytes = d_cfg.worker_connection_buffer_budget_bytes; + + // We are retaining record batches in memory until they are consumed, so we need to account + // for them in the memory pool. + let memory_reservation = + Arc::new(MemoryConsumer::new("WorkerConnection").register(ctx.memory_pool())); + let memory_reservation_clone = Arc::clone(&memory_reservation); + + // Track the maximum memory used to buffer recieved messages. + let mut curr_max_mem = 0; + let max_mem_used = MetricBuilder::new(&metrics).global_gauge("max_mem_used"); + // Track the total encoded size of all recieved messages. + let bytes_transferred = MetricBuilder::new(&metrics).bytes_counter("bytes_transferred"); + let msg_count = MetricBuilder::new(&metrics).global_counter("msg_count"); + // Track end-to-end network latency distribution for messages that actually arrive. + let mut latency_metrics = NetworkLatencyMetrics::new(&metrics); + // Track the total CPU time spent in polling messages over the network + decoding them. + let elapsed_compute = Time::new(); + let elapsed_compute_clone = elapsed_compute.clone(); + MetricBuilder::new(&metrics).build(MetricValue::ElapsedCompute(elapsed_compute.clone())); + + let target_partition_range = request.target_partition_start..request.target_partition_end; + let request = pb::ExecuteTaskRequest { + task_key: Some(encode_task_key(request.task_key)), + target_partition_start: request.target_partition_start as u64, + target_partition_end: request.target_partition_end as u64, + producer_head: Some(encode_producer_head_spec(request.producer_head_spec)), + }; + let metadata = MetadataMap::from_headers(headers); + + // The senders and receivers are unbounded queues used for multiplexing the record + // batches sent through the single gRPC stream into one stream per partition. They + // are unbounded to avoid head-of-line blocking: a single bounded queue could block + // the demux task and starve all sibling partitions even though they have capacity, + // which deadlocks queries with cross-partition dependencies. + // Total memory is bounded globally below via `mem_available_notify`. + let mut per_partition_tx = Vec::with_capacity(target_partition_range.len()); + let mut per_partition_rx = Vec::with_capacity(target_partition_range.len()); + for _partition in target_partition_range.clone() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); + per_partition_tx.push(tx); + per_partition_rx.push(rx); + } + + let mem_available_notify = Arc::new(Notify::new()); + let mem_available_notify_for_task = Arc::clone(&mem_available_notify); + + let first_poll_notify = Arc::new(Notify::new()); + let first_poll_notify_for_task = Arc::clone(&first_poll_notify); + + // Cancellation token allows us to stop the background task promptly when all partition + // streams are dropped (e.g., when the query is cancelled). + let cancel_token = CancellationToken::new(); + let cancel = cancel_token.clone(); + + let mut self_clone = self.clone(); + let request_for_task = request.clone(); + let metadata_for_task = metadata.clone(); + + // This task will pull data from all the partitions in `target_partition_range`, and will + // fan them out to the appropriate `per_partition_rx` based on the "partition" declared + // in each individual record batch flight metadata. + let task = SpawnedTask::spawn(async move { + tokio::select! { + biased; + _ = cancel.cancelled() => { + // If all SendableRecordBatchStreams canceled before any poll, we need to + // anyway trigger the task execution and cancel it immediately so that the + // cancellation is propagated also in the remote worker. Otherwise, it might + // hang forever waiting for someone to execute it. + let _ = self_clone.execute_task(Request::from_parts( + metadata_for_task, + Extensions::default(), + request_for_task, + )).await; + return + }, + _ = first_poll_notify_for_task.notified() => {} + } + + let request = Request::from_parts( + metadata_for_task, + Extensions::default(), + request_for_task, + ); + let mut interleaved_stream = match self_clone.execute_task(request).await { + Ok(v) => v.into_inner(), + Err(err) => return fanout(&per_partition_tx, err), + }; + + loop { + // Backpressure gate. Per-partition channels are unbounded, so we cap + // total in-flight buffered bytes here by pausing the gRPC pull when + // consumers haven't drained enough. This propagates flow control all + // the way back to the worker without coupling sibling partitions. + // We always allow a message through when reservation == 0 to avoid + // livelock if a single message is larger than the budget. + while memory_reservation.size() >= buffer_budget_bytes { + tokio::select! { + biased; + _ = cancel.cancelled() => return, + _ = mem_available_notify_for_task.notified() => {} + } + } + + // Check for cancellation while waiting for the next message. + let flight_data = tokio::select! { + biased; + _ = cancel.cancelled() => return, + msg = interleaved_stream.next() => { + match msg { + Some(Ok(v)) => v, + Some(Err(err)) => return fanout(&per_partition_tx, err), + None => return, // Stream exhausted + } + } + }; + + // Earliest time at which the msg was received. + let msg_received_time = SystemTime::now(); + + let flight_metadata = match FlightAppMetadata::decode(flight_data.app_metadata.as_ref()) { + Ok(v) => v, + Err(err) => { + return fanout(&per_partition_tx, Status::internal(err.to_string())); + } + }; + + // Update the running latency tracker. + let sent_time = UNIX_EPOCH + Duration::from_nanos(flight_metadata.created_timestamp_unix_nanos); + if flight_metadata.created_timestamp_unix_nanos > 0 + && let Ok(delta) = msg_received_time.duration_since(sent_time) { + latency_metrics.add_duration(delta); + } + + let partition = flight_metadata.partition as usize; + // the `per_partition_tx` variable is using a normal `Vec` for storing the + // channel transmitters, so we need to subtract the `target_partition_range.start` + // to the `partition` in order to offset it to the appropriate index. + let Some(sender_i) = partition.checked_sub(target_partition_range.start) else { + let msg = format!( + "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" + ); + return fanout(&per_partition_tx, Status::internal(msg)); + }; + + let Some(o_tx) = per_partition_tx.get(sender_i) else { + let msg = format!( + "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" + ); + return fanout(&per_partition_tx, Status::internal(msg)); + }; + + // We need to send the memory reservation in the same tuple as the actual message + // so that it gets dropped as soon as the message leaves the queue. Dropping the + // memory reservation means releasing the memory from the pool for that specific + // message + let size = flight_data.encoded_len(); + memory_reservation.grow(size); + + // Update memory related metrics. + msg_count.add(1); + bytes_transferred.add_bytes(size); + let curr_mem = memory_reservation.size(); + if curr_mem > curr_max_mem { + curr_max_mem = curr_mem; + max_mem_used.set(curr_max_mem); + } + + if o_tx.send(Ok((flight_data, flight_metadata))).is_err() { + // The receiver for this partition was dropped (e.g. a hash join partition + // completed early without consuming its probe side). Don't exit: other + // partitions multiplexed over the same gRPC stream still need their data. + // Undo the memory reservation that was grown for this dropped batch. + memory_reservation.shrink(size); + continue; + }; + } + }.with_elapsed_compute(elapsed_compute)); + + let task = Arc::new(task); + let not_consumed_streams = Arc::new(AtomicUsize::new(per_partition_rx.len())); + + let mut result = Vec::with_capacity(per_partition_rx.len()); + for partition_receiver in per_partition_rx { + let task = Arc::clone(&task); + let cancel_token = cancel_token.clone(); + + let first_poll_notify = Arc::clone(&first_poll_notify); + let stream = async move { + first_poll_notify.notify_one(); + UnboundedReceiverStream::new(partition_receiver) + } + .flatten_stream(); + + let stream = stream.map_err(|err| FlightError::Tonic(Box::new(err))); + let reservation = Arc::clone(&memory_reservation_clone); + let mem_available_notify = Arc::clone(&mem_available_notify); + let stream = stream.map_ok(move |(data, _meta)| { + reservation.shrink(data.encoded_len()); + // Wake the demux task in case it is blocked on the byte budget. + mem_available_notify.notify_one(); + let _ = &task; // <- keep the task that polls data from the network alive. + data + }); + let stream = FlightRecordBatchStream::new_from_flight_data(stream); + let stream = stream.map_err(map_flight_to_datafusion_error); + let stream = stream.with_elapsed_compute(elapsed_compute_clone.clone()); + + // When the stream is dropped, cancel the background task to ensure prompt cleanup. + let not_consumed_streams = Arc::clone(¬_consumed_streams); + result.push( + on_drop_stream(stream, move || { + let remaining_streams = not_consumed_streams.fetch_sub(1, Ordering::SeqCst) - 1; + if remaining_streams == 0 { + cancel_token.cancel(); + } + }) + .boxed(), + ); + } + + Ok(result) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + let response = self + .get_worker_info(pb::GetWorkerInfoRequest {}) + .await + .map_err(map_status_to_datafusion_error)? + .into_inner(); + Ok(GetWorkerInfoResponse { + version: response.version, + }) + } +} + +type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; + +struct NetworkLatencyMetrics { + metrics: ExecutionPlanMetricsSet, + values: Option, +} + +impl NetworkLatencyMetrics { + fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + metrics: metrics.clone(), + values: None, + } + } + + fn add_duration(&mut self, duration: Duration) { + self.values + .get_or_insert_with(|| NetworkLatencyMetricValues::new(&self.metrics)) + .add_duration(duration); + } +} + +struct NetworkLatencyMetricValues { + min_latency: MinLatencyMetric, + max_latency: MaxLatencyMetric, + p50_latency: P50LatencyMetric, + p95_latency: P95LatencyMetric, + first_latency: FirstLatencyMetric, + sum_latency: Time, + latency_count: Count, +} + +impl NetworkLatencyMetricValues { + fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + let min_latency = MetricBuilder::new(metrics).min_latency("network_latency_min"); + let max_latency = MetricBuilder::new(metrics).max_latency("network_latency_max"); + let p50_latency = MetricBuilder::new(metrics).p50_latency("network_latency_p50"); + let p95_latency = MetricBuilder::new(metrics).p95_latency("network_latency_p95"); + let first_latency = MetricBuilder::new(metrics).first_latency("network_latency_first"); + let sum_latency = Time::new(); + MetricBuilder::new(metrics).build(MetricValue::Time { + name: Cow::Borrowed("network_latency_sum"), + time: sum_latency.clone(), + }); + let latency_count = MetricBuilder::new(metrics).counter("network_latency_count", 0); + + Self { + min_latency, + max_latency, + p50_latency, + p95_latency, + first_latency, + sum_latency, + latency_count, + } + } + + fn add_duration(&self, duration: Duration) { + self.min_latency.add_duration(duration); + self.max_latency.add_duration(duration); + self.p50_latency.add_duration(duration); + self.p95_latency.add_duration(duration); + self.first_latency.add_duration(duration); + self.sum_latency.add_duration(duration); + self.latency_count.add(1); + } +} + +pub(super) fn encode_producer_head_spec( + head: ProducerHeadSpec, +) -> pb::execute_task_request::ProducerHead { + match head { + ProducerHeadSpec::None => pb::execute_task_request::ProducerHead::None(pb::NoneHead {}), + ProducerHeadSpec::BroadcastExec { output_partitions } => { + pb::execute_task_request::ProducerHead::Broadcast(pb::BroadcastExecHead { + output_partitions: output_partitions as u64, + }) + } + ProducerHeadSpec::RepartitionExec { partitioning } => { + pb::execute_task_request::ProducerHead::Repartition(pb::RepartitionExecHead { + partitioning, + }) + } + } +} + +fn encode_coordinator_to_worker_msg(msg: CoordinatorToWorkerMsg) -> pb::CoordinatorToWorkerMsg { + pb::CoordinatorToWorkerMsg { + inner: Some(match msg { + CoordinatorToWorkerMsg::SetPlanRequest(request) => { + pb::coordinator_to_worker_msg::Inner::SetPlanRequest(encode_set_plan_request( + request, + )) + } + CoordinatorToWorkerMsg::WorkUnitBatch(batch) => { + pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(encode_work_unit_batch(batch)) + } + CoordinatorToWorkerMsg::WorkUnitEos => { + pb::coordinator_to_worker_msg::Inner::WorkUnitEos(true) + } + }), + } +} + +fn encode_set_plan_request(request: SetPlanRequest) -> pb::SetPlanRequest { + pb::SetPlanRequest { + task_key: Some(encode_task_key(request.task_key)), + task_count: request.task_count as u64, + plan_proto: request.plan_proto, + work_unit_feed_declarations: request + .work_unit_feed_declarations + .into_iter() + .map(encode_work_unit_feed_declaration) + .collect(), + target_worker_url: request.target_worker_url.to_string(), + query_start_time_ns: request.query_start_time_ns as u64, + } +} + +fn encode_work_unit_batch(batch: WorkUnitBatch) -> pb::WorkUnitBatch { + pb::WorkUnitBatch { + batch: batch.batch.into_iter().map(encode_work_unit).collect(), + } +} + +fn encode_work_unit(work_unit: WorkUnitMsg) -> pb::WorkUnit { + pb::WorkUnit { + id: serialize_uuid(&work_unit.id), + partition: work_unit.partition as u64, + body: work_unit.body, + created_timestamp_unix_nanos: work_unit.created_timestamp_unix_nanos as u64, + sent_timestamp_unix_nanos: work_unit.sent_timestamp_unix_nanos as u64, + received_timestamp_unix_nanos: work_unit.received_timestamp_unix_nanos as u64, + processed_timestamp_unix_nanos: work_unit.processed_timestamp_unix_nanos as u64, + } +} + +fn encode_work_unit_feed_declaration( + declaration: WorkUnitFeedDeclaration, +) -> pb::set_plan_request::WorkUnitFeedDeclaration { + pb::set_plan_request::WorkUnitFeedDeclaration { + id: serialize_uuid(&declaration.id), + partitions: declaration.partitions as u64, + } +} + +fn encode_task_key(task_key: TaskKey) -> pb::TaskKey { + pb::TaskKey { + query_id: serialize_uuid(&task_key.query_id), + stage_id: task_key.stage_id as u64, + task_number: task_key.task_number as u64, + } +} + +fn decode_worker_to_coordinator_msg( + msg: pb::WorkerToCoordinatorMsg, +) -> Result { + Ok( + match msg + .inner + .ok_or_else(|| missing("WorkerToCoordinatorMsg.inner"))? + { + pb::worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics) => { + WorkerToCoordinatorMsg::TaskMetrics(decode_task_metrics(task_metrics)?) + } + pb::worker_to_coordinator_msg::Inner::LoadInfo(load_info) => { + WorkerToCoordinatorMsg::LoadInfo(decode_load_info(load_info)) + } + pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { + WorkerToCoordinatorMsg::LoadInfoEos + } + }, + ) +} + +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, + rows_ready: load_info.rows_ready as usize, + rows_per_second: load_info.rows_per_second as usize, + per_column_bytes_ready: load_info + .per_column_bytes_ready + .into_iter() + .map(|bytes| bytes as usize) + .collect(), + per_column_bytes_per_second: load_info + .per_column_bytes_per_second + .into_iter() + .map(|bytes| bytes as usize) + .collect(), + per_column_ndv_percentage: load_info.per_column_ndv_percentage, + per_column_null_percentage: load_info.per_column_null_percentage, + } +} + +fn missing(field: &'static str) -> DataFusionError { + DataFusionError::Internal(format!("Missing field '{field}'")) +} + +fn fanout(o_txs: &[UnboundedSender], err: Status) { + for o_tx in o_txs { + let _ = o_tx.send(Err(err.clone())); + } +} + +/// Creates a [`WorkerServiceClient`] with high default message size limits. +/// +/// This is a convenience function that wraps [`WorkerServiceClient::new`] and configures +/// it with `max_decoding_message_size(usize::MAX)` and `max_encoding_message_size(usize::MAX)` +/// to avoid message size limitations for internal communication. +/// +/// Users implementing custom [`ChannelResolver`]s should use this function in their +/// `get_worker_client_for_url` implementations to ensure consistent behavior with built-in +/// implementations. +/// +/// # Example +/// +/// ```rust +/// # use datafusion::common::DataFusionError; +/// # use datafusion::error::Result; +/// # use tonic::transport::Channel; +/// # use url::Url; +/// # use datafusion_distributed::{ChannelResolver, WorkerChannel, grpc}; +/// +/// struct MyResolver; +/// +/// #[async_trait::async_trait] +/// impl ChannelResolver for MyResolver { +/// async fn get_worker_client_for_url(&self, url: &Url) -> Result> { +/// let channel = Channel::from_shared(url.to_string()) +/// .map_err(|err| DataFusionError::External(Box::new(err)))? +/// .connect() +/// .await +/// .map_err(|err| DataFusionError::External(Box::new(err)))?; +/// Ok(grpc::create_worker_client(grpc::BoxCloneSyncChannel::new(channel))) +/// } +/// } +/// ``` +pub fn create_worker_client(channel: BoxCloneSyncChannel) -> Box { + Box::new( + pb::worker_service_client::WorkerServiceClient::new(channel) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX), + ) +} + +trait ElapsedComputeFutureExt: Future + Sized { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture; +} + +trait ElapsedComputeStreamExt: Stream + Sized { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream; +} + +impl> ElapsedComputeFutureExt for F { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture { + ElapsedComputeFuture { + inner: self, + curr: Duration::default(), + elapsed_compute, + } + } +} + +impl> ElapsedComputeStreamExt for S { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { + ElapsedComputeStream { + inner: self, + curr: Duration::default(), + elapsed_compute, + } + } +} + +#[pin_project(PinnedDrop)] +struct ElapsedComputeStream { + #[pin] + inner: T, + curr: Duration, + elapsed_compute: Time, +} + +/// Drop implementation that ensures that any accumulated time is properly dumped to the metric +/// in case the stream gets dropped before completion. +#[pinned_drop] +impl PinnedDrop for ElapsedComputeStream { + fn drop(self: Pin<&mut Self>) { + if self.curr > Duration::default() { + let self_projected = self.project(); + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + } + } +} + +impl> Stream for ElapsedComputeStream { + type Item = O; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let self_projected = self.project(); + let start = Instant::now(); + let result = self_projected.inner.poll_next(cx); + *self_projected.curr += start.elapsed(); + if result.is_ready() { + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + *self_projected.curr = Duration::default(); + } + result + } +} + +#[pin_project(PinnedDrop)] +struct ElapsedComputeFuture { + #[pin] + inner: T, + curr: Duration, + elapsed_compute: Time, +} + +/// Drop implementation that ensures that any accumulated time is properly dumped to the metric +/// in case the future gets dropped before completion. +#[pinned_drop] +impl PinnedDrop for ElapsedComputeFuture { + fn drop(self: Pin<&mut Self>) { + if self.curr > Duration::default() { + let self_projected = self.project(); + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + } + } +} + +impl> Future for ElapsedComputeFuture { + type Output = O; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let self_projected = self.project(); + let start = Instant::now(); + let result = self_projected.inner.poll(cx); + *self_projected.curr += start.elapsed(); + if result.is_ready() { + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + *self_projected.curr = Duration::default(); + } + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + use futures::stream::unfold; + + #[tokio::test] + async fn elapsed_compute_future() { + async fn cheap() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + + async fn expensive() { + let mut _count = 0f64; + for i in 0..100000 { + tokio::task::yield_now().await; + _count /= i as f64 + } + } + + let cheap_time = Time::new(); + cheap().with_elapsed_compute(cheap_time.clone()).await; + println!("cheap future: {}", cheap_time.value()); + + let expensive_time = Time::new(); + expensive() + .with_elapsed_compute(expensive_time.clone()) + .await; + println!("expensive future: {}", expensive_time.value()); + + assert!(expensive_time.value() > cheap_time.value()); + } + + #[tokio::test] + async fn elapsed_compute_stream() { + fn cheap() -> impl Stream { + unfold(0i64, |state| async move { + if state < 10 { + tokio::time::sleep(Duration::from_micros(10)).await; + Some((state, state + 1)) + } else { + None + } + }) + } + + fn expensive() -> impl Stream { + unfold(0i64, |state| async move { + if state < 10 { + // Simulate expensive computation + let mut _count = 0f64; + for i in 1..100000 { + _count += (i as f64).sqrt(); + } + tokio::task::yield_now().await; + Some((state, state + 1)) + } else { + None + } + }) + } + + let cheap_time = Time::new(); + cheap() + .with_elapsed_compute(cheap_time.clone()) + .collect::>() + .await; + println!("cheap future: {}", cheap_time.value()); + + let expensive_time = Time::new(); + expensive() + .with_elapsed_compute(expensive_time.clone()) + .collect::>() + .await; + println!("expensive future: {}", expensive_time.value()); + + assert!(expensive_time.value() > cheap_time.value()); + } +} diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs new file mode 100644 index 00000000..fa677a10 --- /dev/null +++ b/src/protocol/grpc/worker_service.rs @@ -0,0 +1,422 @@ +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::common::{deserialize_uuid, now_ns}; +use crate::protocol::ProducerHeadSpec; +use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; +use crate::{ + CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, SetPlanRequest, + TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, Worker, + WorkerResolver, WorkerToCoordinatorMsg, +}; + +use arrow_flight::FlightData; +use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; +use arrow_flight::error::FlightError; +use arrow_select::dictionary::garbage_collect_any_dictionary; +use async_trait::async_trait; +use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; +use datafusion::arrow::ipc::CompressionType; +use datafusion::arrow::ipc::writer::IpcWriteOptions; +use datafusion::common::DataFusionError; +use datafusion::execution::SendableRecordBatchStream; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +use prost::Message; +use std::sync::Arc; +use tonic::{Request, Response, Status, Streaming}; +use url::Url; + +const RECORD_BATCH_BUFFER_SIZE: usize = 2; + +impl Worker { + /// Converts this [Worker] into a [`WorkerServiceServer`] with high default message size limits. + /// + /// This is a convenience method that wraps the endpoint in a [`WorkerServiceServer`] and + /// configures it with `max_decoding_message_size(usize::MAX)` and + /// `max_encoding_message_size(usize::MAX)` to avoid message size limitations for internal + /// communication. + /// + /// You can further customize the returned server by chaining additional tonic methods. + /// + /// # Example + /// + /// ``` + /// # use datafusion_distributed::Worker; + /// # use tonic::transport::Server; + /// # use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + /// # async fn f() { + /// + /// let worker = Worker::default(); + /// let server = worker.into_worker_server(); + /// + /// Server::builder() + /// .add_service(Worker::default().into_worker_server()) + /// .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)) + /// .await; + /// + /// # } + /// ``` + pub fn into_worker_server(self) -> pb::worker_service_server::WorkerServiceServer { + pb::worker_service_server::WorkerServiceServer::new(self) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX) + } + + /// Creates an [`ObservabilityServiceServer`] that exposes task progress and cluster + /// worker discovery via the provided [`WorkerResolver`]. + /// + /// The returned server is meant to be added to the same [`tonic::transport::Server`] as the + /// Flight service — gRPC multiplexes both services on a single port. + pub fn with_observability_service( + &self, + worker_resolver: Arc, + ) -> ObservabilityServiceServer { + ObservabilityServiceServer::new(ObservabilityServiceImpl::new( + self.task_data_entries.clone(), + worker_resolver, + )) + } +} + +/// Implementation of the `worker.proto` specification based on the generated Rust stubs. +/// +/// The methods are delegated to plan `impl Worker` implementations so that they can be implemented +/// in different files. +#[async_trait] +impl pb::worker_service_server::WorkerService for Worker { + type CoordinatorChannelStream = BoxStream<'static, Result>; + + async fn coordinator_channel( + &self, + request: Request>, + ) -> Result, Status> { + let (metadata, _ext, body) = request.into_parts(); + + let input_stream = body + .map_err(map_status_to_datafusion_error) + .map(move |msg| { + decode_coordinator_to_worker_msg(msg?).map_err(map_status_to_datafusion_error) + }) + .boxed(); + + let output_stream = self + .coordinator_channel(metadata.into_headers(), input_stream) + .await + .map_err(datafusion_error_to_tonic_status)? + .map(|msg| match msg { + Ok(msg) => encode_worker_to_coordinator_msg(msg), + Err(err) => Err(datafusion_error_to_tonic_status(err)), + }) + .boxed(); + + Ok(Response::new(output_stream)) + } + + type ExecuteTaskStream = BoxStream<'static, Result>; + + async fn execute_task( + &self, + request: Request, + ) -> Result, Status> { + let body = request.into_inner(); + let request = decode_execute_task_request(body).await?; + let partition_range = request.target_partition_start..request.target_partition_end; + + let (arrow_streams, task_ctx) = self + .execute_task(request) + .await + .map_err(datafusion_error_to_tonic_status)?; + + let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options()) + .map_err(datafusion_error_to_tonic_status)?; + + let compression = match d_cfg.compression.as_str() { + "lz4" => Some(CompressionType::LZ4_FRAME), + "zstd" => Some(CompressionType::ZSTD), + "none" => None, + v => Err(Status::invalid_argument(format!( + "Unknown compression type {v}" + )))?, + }; + let mut flight_streams = Vec::with_capacity(arrow_streams.len()); + for (partition, arrow_stream) in partition_range.zip(arrow_streams) { + let flight_stream = + build_flight_data_stream(arrow_stream, compression)?.map(move |msg| { + // For each FlightData produced by this stream, mark it with the appropriate + // partition. This stream will be merged with several others from other partitions, + // so marking it with the original partition allows it to be deconstructed into + // the original per-partition streams in later steps. + let flight_data = pb::FlightAppMetadata { + partition: partition as u64, + created_timestamp_unix_nanos: now_ns::(), + }; + msg.map(|v| v.with_app_metadata(flight_data.encode_to_vec())) + }); + + flight_streams.push(flight_stream); + } + + // Merge all the per-partition streams into one. Each message in the stream is marked with + // the original partition, so they can be reconstructed at the other side of the boundary. + let memory_pool = Arc::clone(&task_ctx.runtime_env().memory_pool); + let stream = spawn_select_all(flight_streams, memory_pool, RECORD_BATCH_BUFFER_SIZE); + + Ok(Response::new(Box::pin(stream.map_err(|err| match err { + FlightError::Tonic(status) => *status, + _ => Status::internal(format!("Error during flight stream: {err}")), + })))) + } + + async fn get_worker_info( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(pb::GetWorkerInfoResponse { + version: self.version().to_string(), + })) + } +} + +fn decode_coordinator_to_worker_msg( + msg: pb::CoordinatorToWorkerMsg, +) -> Result { + Ok( + match msg + .inner + .ok_or_else(missing("CoordinatorToWorkerMsg.inner"))? + { + pb::coordinator_to_worker_msg::Inner::SetPlanRequest(request) => { + CoordinatorToWorkerMsg::SetPlanRequest(decode_set_plan_request(request)?) + } + pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(batch) => { + CoordinatorToWorkerMsg::WorkUnitBatch(decode_work_unit_batch(batch)?) + } + pb::coordinator_to_worker_msg::Inner::WorkUnitEos(_) => { + CoordinatorToWorkerMsg::WorkUnitEos + } + }, + ) +} + +fn decode_set_plan_request(request: pb::SetPlanRequest) -> Result { + Ok(SetPlanRequest { + task_key: decode_task_key(request.task_key.ok_or_else(missing("task_key"))?)?, + task_count: request.task_count as usize, + plan_proto: request.plan_proto, + work_unit_feed_declarations: request + .work_unit_feed_declarations + .into_iter() + .map(decode_work_unit_feed_declaration) + .collect::>()?, + target_worker_url: parse_url(&request.target_worker_url, "target_worker_url")?, + query_start_time_ns: request.query_start_time_ns as usize, + }) +} + +async fn decode_execute_task_request( + request: pb::ExecuteTaskRequest, +) -> Result { + Ok(ExecuteTaskRequest { + task_key: decode_task_key(request.task_key.ok_or_else(missing("task_key"))?)?, + target_partition_start: request.target_partition_start as usize, + target_partition_end: request.target_partition_end as usize, + producer_head_spec: decode_producer_head_spec( + request.producer_head.ok_or_else(missing("producer_head"))?, + ), + }) +} + +pub(super) fn decode_producer_head_spec( + proto: pb::execute_task_request::ProducerHead, +) -> ProducerHeadSpec { + match proto { + pb::execute_task_request::ProducerHead::None(_) => ProducerHeadSpec::None, + pb::execute_task_request::ProducerHead::Broadcast(v) => ProducerHeadSpec::BroadcastExec { + output_partitions: v.output_partitions as usize, + }, + pb::execute_task_request::ProducerHead::Repartition(v) => { + ProducerHeadSpec::RepartitionExec { + partitioning: v.partitioning, + } + } + } +} + +fn encode_worker_to_coordinator_msg( + msg: WorkerToCoordinatorMsg, +) -> Result { + Ok(pb::WorkerToCoordinatorMsg { + inner: Some(match msg { + WorkerToCoordinatorMsg::TaskMetrics(task_metrics) => { + pb::worker_to_coordinator_msg::Inner::TaskMetrics(encode_task_metrics( + task_metrics, + )?) + } + WorkerToCoordinatorMsg::LoadInfo(load_info) => { + pb::worker_to_coordinator_msg::Inner::LoadInfo(encode_load_info(load_info)) + } + WorkerToCoordinatorMsg::LoadInfoEos => { + pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) + } + }), + }) +} + +fn encode_task_metrics(task_metrics: TaskMetrics) -> Result { + Ok(pb::TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .into_iter() + .map(|metrics_set| { + df_metrics_set_to_proto(&metrics_set).map_err(datafusion_error_to_tonic_status) + }) + .collect::>()?, + task_metrics: Some( + df_metrics_set_to_proto(&task_metrics.task_metrics) + .map_err(datafusion_error_to_tonic_status)?, + ), + }) +} + +fn encode_load_info(load_info: LoadInfo) -> pb::LoadInfo { + pb::LoadInfo { + partition: load_info.partition as u64, + rows_ready: load_info.rows_ready as u64, + rows_per_second: load_info.rows_per_second as u64, + per_column_bytes_ready: load_info + .per_column_bytes_ready + .into_iter() + .map(|bytes| bytes as u64) + .collect(), + per_column_bytes_per_second: load_info + .per_column_bytes_per_second + .into_iter() + .map(|bytes| bytes as u64) + .collect(), + per_column_ndv_percentage: load_info.per_column_ndv_percentage, + per_column_null_percentage: load_info.per_column_null_percentage, + } +} + +fn decode_work_unit_batch(batch: pb::WorkUnitBatch) -> Result { + Ok(WorkUnitBatch { + batch: batch + .batch + .into_iter() + .map(decode_work_unit) + .collect::>()?, + }) +} + +fn decode_work_unit(work_unit: pb::WorkUnit) -> Result { + Ok(WorkUnitMsg { + id: deserialize_uuid(&work_unit.id).map_err(datafusion_error_to_tonic_status)?, + partition: work_unit.partition as usize, + body: work_unit.body, + created_timestamp_unix_nanos: work_unit.created_timestamp_unix_nanos as usize, + sent_timestamp_unix_nanos: work_unit.sent_timestamp_unix_nanos as usize, + received_timestamp_unix_nanos: work_unit.received_timestamp_unix_nanos as usize, + processed_timestamp_unix_nanos: work_unit.processed_timestamp_unix_nanos as usize, + }) +} + +fn decode_work_unit_feed_declaration( + declaration: pb::set_plan_request::WorkUnitFeedDeclaration, +) -> Result { + Ok(WorkUnitFeedDeclaration { + id: deserialize_uuid(&declaration.id).map_err(datafusion_error_to_tonic_status)?, + partitions: declaration.partitions as usize, + }) +} + +fn decode_task_key(task_key: pb::TaskKey) -> Result { + Ok(TaskKey { + query_id: deserialize_uuid(&task_key.query_id).map_err(datafusion_error_to_tonic_status)?, + stage_id: task_key.stage_id as usize, + task_number: task_key.task_number as usize, + }) +} + +fn parse_url(value: &str, field: &'static str) -> Result { + Url::parse(value) + .map_err(|err| Status::invalid_argument(format!("Invalid field '{field}': {err}"))) +} + +fn missing(field: &'static str) -> impl FnOnce() -> Status { + move || Status::invalid_argument(format!("Missing field '{field}'")) +} + +fn build_flight_data_stream( + stream: SendableRecordBatchStream, + compression_type: Option, +) -> datafusion::common::Result { + let stream = FlightDataEncoderBuilder::new() + .with_options( + IpcWriteOptions::default() + .try_with_compression(compression_type) + .map_err(|err| Status::internal(err.to_string()))?, + ) + .with_schema(stream.schema()) + // This tells the encoder to send dictionaries across the wire as-is. + // The alternative (`DictionaryHandling::Hydrate`) would expand the dictionaries + // into their value types, which can potentially blow up the size of the data transfer. + // The main reason to use `DictionaryHandling::Hydrate` is for compatibility with clients + // that do not support dictionaries, but since we are using the same server/client on both + // sides, we can safely use `DictionaryHandling::Resend`. + // Note that we do garbage collection of unused dictionary values above, so we are not sending + // unused dictionary values over the wire. + .with_dictionary_handling(DictionaryHandling::Resend) + // Set max flight data size to unlimited. + // This requires servers and clients to also be configured to handle unlimited sizes. + // Using unlimited sizes avoids splitting RecordBatches into multiple FlightData messages, + // which could add significant overhead for large RecordBatches. + // The only reason to split them really is if the client/server are configured with a message size limit, + // which mainly makes sense in a public network scenario where you want to avoid DoS attacks. + // Since all of our Arrow Flight communication happens within trusted data plane networks, + // we can safely use unlimited sizes here. + .with_max_flight_data_size(usize::MAX) + .build( + stream + // Apply garbage collection of dictionary and view arrays before sending over the network + .and_then(|rb| std::future::ready(garbage_collect_arrays(rb))) + .map_err(|err| FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(err)))), + ); + Ok(stream) +} + +/// Garbage collects values sub-arrays. +/// +/// We apply this before sending RecordBatches over the network to avoid sending +/// values that are not referenced by any dictionary keys or buffers that are not used. +/// +/// Unused values can arise from operations such as filtering, where +/// some keys may no longer be referenced in the filtered result. +fn garbage_collect_arrays( + batch: RecordBatch, +) -> datafusion::common::Result { + let (schema, arrays, row_count) = batch.into_parts(); + + let arrays = arrays + .into_iter() + .map(|array| { + if let Some(array) = array.as_any_dictionary_opt() { + garbage_collect_any_dictionary(array) + } else if let Some(array) = array.as_string_view_opt() { + Ok(Arc::new(array.gc()) as Arc) + } else if let Some(array) = array.as_binary_view_opt() { + Ok(Arc::new(array.gc()) as Arc) + } else { + Ok(array) + } + }) + .collect::, _>>()?; + + Ok(RecordBatch::try_new_with_options( + schema, + arrays, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?) +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs new file mode 100644 index 00000000..5c6bdd1a --- /dev/null +++ b/src/protocol/mod.rs @@ -0,0 +1,14 @@ +#[cfg(feature = "grpc")] +pub mod grpc; + +mod channel_resolver; +mod worker_channel; + +pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; +pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; + +pub use worker_channel::{ + CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, + LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, +}; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs new file mode 100644 index 00000000..1baa23fe --- /dev/null +++ b/src/protocol/worker_channel.rs @@ -0,0 +1,191 @@ +use async_trait::async_trait; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use futures::stream::BoxStream; +use http::HeaderMap; +use std::sync::Arc; +use url::Url; +use uuid::Uuid; + +/// Abstraction over the specific transport protocol implementation. +/// +/// WARNING: The API in this trait is unstable, and it's subject to change as more things get properly +/// decoupled from details like protobuf serialization and http headers. +#[async_trait] +pub trait WorkerChannel: Send + Sync { + /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages + /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel + /// per task. + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>>; + + /// Executes the requested partition range of a subplan previously sent by the coordinator channel. + async fn execute_task( + &mut self, + headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + task_ctx: &Arc, + ) -> Result>>>; + + /// Returns metadata about a worker. Currently only used for worker versioning. + async fn get_worker_info( + &mut self, + request: GetWorkerInfoRequest, + ) -> Result; +} + +pub enum CoordinatorToWorkerMsg { + /// Sends a subplan to a worker so that a future ExecuteTask call can actually execute it. + /// The plan is identified by a TaskKey. + SetPlanRequest(SetPlanRequest), + /// A batch of messages from a work unit feed belonging to different partitions from one node from the plan set in + /// set_plan_request. A work unit feed is a per-partition stream of information that tells the node what should + /// be executed within a partition, for example, a stream of file addresses that should be read. + WorkUnitBatch(WorkUnitBatch), + /// Signals an EOS for WorkUnits. After this message is received, no more WorkUnits will be sent. + WorkUnitEos, +} + +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +pub struct TaskKey { + /// Our query id. + pub query_id: Uuid, + /// Our stage id. + pub stage_id: usize, + /// The task number within the stage. + pub task_number: usize, +} + +pub struct WorkUnitFeedDeclaration { + /// Unique identifier of the node to which work unit feeds are expected to be streamed. + pub id: Uuid, + /// The amount of partitions expected to be streamed. + pub partitions: usize, +} + +pub struct SetPlanRequest { + /// The unique identifier of the task to which the subplan belongs to. + pub task_key: TaskKey, + /// The amount of tasks that share the same subplan. Necessary for building the DistributedTaskContext during execution. + pub task_count: usize, + /// The subplan the worker is expected to execute. + // TODO: this still forces implementations to pass a serialized plan. In-memory implementations + // might want to omit the serde step, so there should be a way to pass here a normal plan, and + // pass the serializer/deserialized separately instead of being coupled to protobuf serialization + pub plan_proto: Vec, + /// Information about all the work unit feeds that will be streamed from coordinator to worker. + /// This information is needed here because at the moment of setting the plan, all the appropriate + /// channels for the incoming work unit feeds need to be constructed. + /// + /// If no WorkUnitFeedExec nodes are present in the plan, this should be empty. + pub work_unit_feed_declarations: Vec, + /// The worker URL to which this message will go. The receiving worker will use this information + /// to identify itself, and avoid further calls in case it needs to call itself for executing tasks. + pub target_worker_url: Url, + /// Unix nanos when the query started as reported by the coordinator. Used for collecting temporal metrics + /// relative to when the query was fired in the coordinator. + pub query_start_time_ns: usize, +} + +pub struct WorkUnitBatch { + /// A batch of WorkUnits. + pub batch: Vec, +} + +pub struct WorkUnitMsg { + /// Identifier of the node to which this work unit feed belongs to. + pub id: Uuid, + /// The partition index within the node to which the work unit feed belongs to. + pub partition: usize, + /// Arbitrary user-defined data (e.g., a file address) necessary during execution. + pub body: Vec, + /// Unix timestamp in nanoseconds at which this message was created in the coordinator. + pub created_timestamp_unix_nanos: usize, + /// Unix timestamp in nanoseconds at which this message was sent by the coordinator. + pub sent_timestamp_unix_nanos: usize, + /// Unix timestamp in nanoseconds at which this message was received by a worker. + pub received_timestamp_unix_nanos: usize, + /// Unix timestamp in nanoseconds at which this message started being processed. + pub processed_timestamp_unix_nanos: usize, +} + +pub enum WorkerToCoordinatorMsg { + /// Sends the metrics collected during task execution back to the coordinator. + /// This is sent after all partitions of a task have finished (or been dropped), + /// ensuring metrics are never lost due to early stream termination. + /// metrics[i] is the set of metrics for plan node i in pre-order traversal order. + TaskMetrics(TaskMetrics), + /// Load information reported by a task. This information is used for dynamically + /// sizing the number of workers involved in a query. + LoadInfo(LoadInfo), + LoadInfoEos, +} + +#[derive(Clone, Debug)] +pub struct TaskMetrics { + /// Metrics for a single task's plan nodes in pre-order traversal order. + /// The TaskKey is implicit — it is determined by the SetPlanRequest that + /// opened this coordinator channel connection. + pub pre_order_plan_metrics: Vec, + /// Metrics related to the execution of a task within a stage. This metrics, instead of being + /// associated to a specific node, they are global to the task, like the time at which the plan + /// was fed by the coordinator to the worker. + pub task_metrics: MetricsSet, +} + +#[derive(Default)] +pub struct LoadInfo { + /// The partition index to which this message belongs to. + pub partition: usize, + /// The amount of rows ready to be returned. + pub rows_ready: usize, + /// The estimated velocity at which rows will flow through the node. If all the rows were + /// already accumulated, they will be reported by `rows_ready`, and this field will be 0. + pub rows_per_second: usize, + /// The amount of bytes ready to be returned per column. + pub per_column_bytes_ready: Vec, + /// The estimated velocity at which data will flow through each column. If all the bytes were + /// already accumulated, they will be reported by `bytes_ready`, and this field will be 0. + pub per_column_bytes_per_second: Vec, + /// Approximate ratio of NDV for each column. + pub per_column_ndv_percentage: Vec, + /// Approximate ratio of null count for each column. + pub per_column_null_percentage: Vec, +} + +pub struct ExecuteTaskRequest { + /// The unique identifier of the task that is going to get executed. + pub task_key: TaskKey, + /// The start of the partition range of the specified task that is going to be executed. + pub target_partition_start: usize, + /// The end of the partition range of the specified task that is going to be executed. + pub target_partition_end: usize, + /// The head node the requested task should have. Depending on the network boundary executing + /// the task, the head node should be prepared differently, for example: + /// - A RepartitionExecHead implies a RepartitionExec at the head of the task. + /// - A BroadcastExecHead implies a BroadcastExec at the head of the task. + /// - A NoneHead does not need any specific head. + pub producer_head_spec: ProducerHeadSpec, +} + +#[derive(Clone)] +pub enum ProducerHeadSpec { + /// No specific head node is necessary. + None, + /// The head node should be a [BroadcastExec]. + BroadcastExec { output_partitions: usize }, + /// The head node should be a [RepartitionExec]. + RepartitionExec { partitioning: Vec }, +} + +pub struct GetWorkerInfoRequest {} + +pub struct GetWorkerInfoResponse { + pub version: String, +} diff --git a/src/stage.rs b/src/stage.rs index 94b3ef26..7903a48a 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -219,10 +219,9 @@ impl DistributedTaskContext { } } -use crate::common::serialize_uuid; -use crate::metrics::proto::metric_proto_to_df; -use crate::worker::generated::worker as pb; -use crate::{DistributedMetricsFormat, NetworkShuffleExec, rewrite_distributed_plan_with_metrics}; +use crate::{ + DistributedMetricsFormat, NetworkShuffleExec, TaskKey, rewrite_distributed_plan_with_metrics, +}; use crate::{NetworkBoundary, NetworkBoundaryExt}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::DataFusionError; @@ -437,21 +436,24 @@ fn display_inner_distributed_leaf( /// Gathers the metrics global to a stage. These metrics are not specific to any plan node, and /// are instead global to a whole stage. fn gather_stage_header_metrics(stage: &Stage, metrics_store: &MetricsStore) -> MetricsSet { - let mut task_key = pb::TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, + let mut task_key = TaskKey { + query_id: stage.query_id(), + stage_id: stage.num(), task_number: 0, }; let mut all_metrics = stage.metrics(); - while let Some(metrics_set) = metrics_store.get(&task_key).and_then(|v| v.task_metrics) { - for mut metric in metrics_set.metrics { - metric.labels.push(pb::Label { - name: DISTRIBUTED_DATAFUSION_TASK_ID_LABEL.to_string(), - value: task_key.task_number.to_string(), - }); - if let Ok(metric) = metric_proto_to_df(metric) { - all_metrics.push(metric) - }; + while let Some(metrics_set) = metrics_store.get(&task_key).map(|v| v.task_metrics) { + for metric in metrics_set.iter() { + let mut labels = metric.labels().to_vec(); + labels.push(Label::new( + DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, + task_key.task_number.to_string(), + )); + all_metrics.push(Arc::new(Metric::new_with_labels( + metric.value().clone(), + metric.partition(), + labels, + ))); } task_key.task_number += 1; } diff --git a/src/test_utils/in_memory_channel_resolver.rs b/src/test_utils/in_memory_channel_resolver.rs index 12ad84bb..92f710c7 100644 --- a/src/test_utils/in_memory_channel_resolver.rs +++ b/src/test_utils/in_memory_channel_resolver.rs @@ -1,8 +1,6 @@ -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; use crate::{ - BoxCloneSyncChannel, ChannelResolver, DefaultSessionBuilder, DistributedExt, - MappedWorkerSessionBuilderExt, SessionStateBuilderExt, Worker, WorkerResolver, - WorkerSessionBuilder, create_worker_client, + ChannelResolver, DefaultSessionBuilder, DistributedExt, MappedWorkerSessionBuilderExt, + SessionStateBuilderExt, Worker, WorkerChannel, WorkerResolver, WorkerSessionBuilder, grpc, }; use async_trait::async_trait; use datafusion::common::DataFusionError; @@ -19,7 +17,7 @@ const DUMMY_URL_PREFIX: &str = "http://url-"; /// tokio duplex rather than a TCP connection. #[derive(Clone)] pub struct InMemoryChannelResolver { - channel: WorkerServiceClient, + channel: grpc::BoxCloneSyncChannel, } impl InMemoryChannelResolver { @@ -50,7 +48,7 @@ impl InMemoryChannelResolver { })); let this = Self { - channel: create_worker_client(BoxCloneSyncChannel::new(channel)), + channel: grpc::BoxCloneSyncChannel::new(channel), }; let this_clone = this.clone(); @@ -83,8 +81,8 @@ impl ChannelResolver for InMemoryChannelResolver { async fn get_worker_client_for_url( &self, _: &url::Url, - ) -> Result, DataFusionError> { - Ok(self.channel.clone()) + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) } } diff --git a/src/test_utils/metrics.rs b/src/test_utils/metrics.rs index 0677ffca..e8e5fe46 100644 --- a/src/test_utils/metrics.rs +++ b/src/test_utils/metrics.rs @@ -1,6 +1,7 @@ use crate::coordinator::DistributedExec; -use crate::worker::generated::worker as pb; +use chrono::{DateTime, Utc}; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::metrics::{Count, Metric, MetricValue, MetricsSet, Time, Timestamp}; use std::sync::Arc; /// Waits until all worker tasks have reported their metrics back via the coordinator channel. @@ -11,43 +12,41 @@ pub async fn wait_for_all_metrics(plan: &Arc) { } /// creates a "distinct" set of metrics from the provided seed -pub fn make_test_metrics_set_proto_from_seed(seed: u64, num_metrics: usize) -> pb::MetricsSet { +pub fn make_test_metrics_set_from_seed(seed: u64, num_metrics: usize) -> MetricsSet { const TEST_TIMESTAMP: i64 = 1758200400000000000; // 2025-09-18 13:00:00 UTC - let mut result = pb::MetricsSet { metrics: vec![] }; + let mut result = MetricsSet::new(); for i in 0..num_metrics { - let value = seed + i as u64; - result.metrics.push(match i % 4 { - 0 => pb::Metric { - value: Some(pb::metric::Value::OutputRows(pb::OutputRows { value })), - labels: vec![], - partition: None, + let value = (seed + i as u64) as usize; + result.push(Arc::new(Metric::new( + match i % 4 { + 0 => { + let count = Count::new(); + count.add(value); + MetricValue::OutputRows(count) + } + 1 => { + let time = Time::new(); + time.add_duration(std::time::Duration::from_nanos(value as u64)); + MetricValue::ElapsedCompute(time) + } + 2 => MetricValue::StartTimestamp(timestamp_from_nanos( + TEST_TIMESTAMP + (value as i64 * 1_000_000_000), + )), + 3 => MetricValue::EndTimestamp(timestamp_from_nanos( + TEST_TIMESTAMP + (value as i64 * 1_000_000_000), + )), + _ => unreachable!(), }, - - 1 => pb::Metric { - value: Some(pb::metric::Value::ElapsedCompute(pb::ElapsedCompute { - value, - })), - labels: vec![], - partition: None, - }, - 2 => pb::Metric { - value: Some(pb::metric::Value::StartTimestamp(pb::StartTimestamp { - value: Some(TEST_TIMESTAMP + (value as i64 * 1_000_000_000)), - })), - labels: vec![], - partition: None, - }, - 3 => pb::Metric { - value: Some(pb::metric::Value::EndTimestamp(pb::EndTimestamp { - value: Some(TEST_TIMESTAMP + (value as i64 * 1_000_000_000)), - })), - labels: vec![], - partition: None, - }, - _ => unreachable!(), - }) + None, + ))) } result } + +fn timestamp_from_nanos(nanos: i64) -> Timestamp { + let timestamp = Timestamp::new(); + timestamp.set(DateTime::::from_timestamp_nanos(nanos)); + timestamp +} diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index 7d299489..aa56338e 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -1,15 +1,13 @@ #[cfg(test)] use super::parquet::register_parquet_tables; -use crate::NetworkBoundaryExt; -use crate::common::serialize_uuid; use crate::coordinator::DistributedExec; use crate::stage::Stage; -use crate::worker::generated::worker::TaskKey; #[cfg(test)] use crate::{ DistributedConfig, DistributedExt, SessionStateBuilderExt, TaskEstimation, TaskEstimator, display_plan_ascii, test_utils::in_memory_channel_resolver::InMemoryWorkerResolver, }; +use crate::{NetworkBoundaryExt, TaskKey}; #[cfg(test)] use datafusion::{ common::Result, @@ -67,9 +65,9 @@ pub fn get_stages_and_task_keys( // Add each task. for j in 0..stage.task_count() { task_keys.insert(TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, - task_number: j as u64, + query_id: stage.query_id(), + stage_id: stage.num(), + task_number: j, }); } diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 1526508c..18c51ffe 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -1,11 +1,11 @@ -use crate::common::{now_ns, serialize_uuid}; -use crate::worker::generated::worker as pb; -use crate::{BytesMetricExt, LatencyMetricExt, WorkUnit}; +use crate::common::now_ns; +use crate::{ + BytesMetricExt, CoordinatorToWorkerMsg, LatencyMetricExt, WorkUnit, WorkUnitBatch, WorkUnitMsg, +}; use datafusion::common::{HashMap, Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::MetricBuilder; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion_proto::protobuf::proto_error; use futures::StreamExt; use futures::stream::BoxStream; use std::sync::{Arc, Mutex}; @@ -13,8 +13,8 @@ use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; use uuid::Uuid; -pub(crate) type WorkUnitTx = UnboundedSender>; -pub(crate) type WorkUnitRx = UnboundedReceiver>; +pub(crate) type WorkUnitTx = UnboundedSender>; +pub(crate) type WorkUnitRx = UnboundedReceiver>; pub(crate) type RemoteWorkUnitFeedRxs = HashMap<(Uuid, usize), Mutex>>; pub(crate) type RemoteWorkUnitFeedTxs = HashMap<(Uuid, usize), WorkUnitTx>; @@ -36,7 +36,7 @@ pub(crate) struct RemoteWorkUnitFeedRegistry { } impl RemoteWorkUnitFeedRegistry { - /// Creates all the receivers and senders for a specific [WorkUnit] Feed id. One feed per + /// Creates all the receivers and senders for a specific [WorkUnitMsg] Feed id. One feed per /// partition is created. /// /// Calling this twice with the same `id` is a coordinator bug — duplicate declarations @@ -60,36 +60,27 @@ impl RemoteWorkUnitFeedRegistry { pub(crate) fn build_work_unit_batch_msg( id: &Uuid, work_unit_batch: Vec<(usize, Result>)>, -) -> Result { - Ok(pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch( - pb::WorkUnitBatch { - batch: work_unit_batch - .into_iter() - .map(|(partition, work_unit)| { - Ok(pb::WorkUnit { - id: serialize_uuid(id), - partition: partition as u64, - body: work_unit?.encode_to_bytes(), - created_timestamp_unix_nanos: now_ns(), - sent_timestamp_unix_nanos: 0, - received_timestamp_unix_nanos: 0, - processed_timestamp_unix_nanos: 0, - }) - }) - .collect::>()?, - }, - )), - }) +) -> Result { + Ok(CoordinatorToWorkerMsg::WorkUnitBatch(WorkUnitBatch { + batch: work_unit_batch + .into_iter() + .map(|(partition, work_unit)| { + Ok(WorkUnitMsg { + id: *id, + partition, + body: work_unit?.encode_to_bytes(), + created_timestamp_unix_nanos: now_ns(), + sent_timestamp_unix_nanos: 0, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + }) + }) + .collect::>()?, + })) } -pub(crate) fn set_work_unit_send_time( - mut msg: pb::CoordinatorToWorkerMsg, -) -> pb::CoordinatorToWorkerMsg { - if let pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(work_unit_batch)), - } = &mut msg - { +pub(crate) fn set_work_unit_send_time(mut msg: CoordinatorToWorkerMsg) -> CoordinatorToWorkerMsg { + if let CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) = &mut msg { for work_unit in &mut work_unit_batch.batch { work_unit.sent_timestamp_unix_nanos = now_ns(); } @@ -98,12 +89,9 @@ pub(crate) fn set_work_unit_send_time( } pub(crate) fn set_work_unit_received_time( - mut msg: pb::CoordinatorToWorkerMsg, -) -> pb::CoordinatorToWorkerMsg { - if let pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(work_unit_batch)), - } = &mut msg - { + mut msg: CoordinatorToWorkerMsg, +) -> CoordinatorToWorkerMsg { + if let CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) = &mut msg { for work_unit in &mut work_unit_batch.batch { work_unit.received_timestamp_unix_nanos = now_ns(); } @@ -111,7 +99,7 @@ pub(crate) fn set_work_unit_received_time( msg } -/// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnit]s coming over +/// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnitMsg]s coming over /// the wire from a [RemoteWorkUnitFeedRegistry]. /// /// Deserializing a [crate::WorkUnitFeed] with [crate::WorkUnitFeed::from_proto] always returns a @@ -126,7 +114,7 @@ pub(crate) struct RemoteFeedProvider { } impl RemoteFeedProvider { - pub(crate) fn feed( + pub(crate) fn feed( &self, partition: usize, ctx: Arc, @@ -168,36 +156,36 @@ impl RemoteFeedProvider { }; Ok(UnboundedReceiverStream::new(receiver) - .map(move |work_unit_or_err| { - let mut work_unit = work_unit_or_err?; + .map(move |work_unit_msg_or_err| { + let mut work_unit_msg = work_unit_msg_or_err?; let timer = elapsed_compute.timer(); - let result = T::decode(work_unit.body.as_slice()) - .map_err(|err| proto_error(format!("{err}"))); + let work_unit = T::decode(work_unit_msg.body.as_slice()) + .map_err(|err| datafusion_proto::protobuf::proto_error(format!("{err}"))); timer.done(); - work_unit.processed_timestamp_unix_nanos = now_ns(); + work_unit_msg.processed_timestamp_unix_nanos = now_ns(); + let body_len = work_unit_msg.body.len(); - let pb::WorkUnit { + let WorkUnitMsg { created_timestamp_unix_nanos: base, sent_timestamp_unix_nanos, received_timestamp_unix_nanos, processed_timestamp_unix_nanos, - body, .. - } = work_unit; + } = work_unit_msg; - bytes_transferred.add_bytes(body.len()); + bytes_transferred.add_bytes(body_len); msg_count.add(1); - send_latency_max.add_nanos((sent_timestamp_unix_nanos - base) as usize); - send_latency_p50.add_nanos((sent_timestamp_unix_nanos - base) as usize); + send_latency_max.add_nanos(sent_timestamp_unix_nanos - base); + send_latency_p50.add_nanos(sent_timestamp_unix_nanos - base); - received_latency_max.add_nanos((received_timestamp_unix_nanos - base) as usize); - received_latency_p50.add_nanos((received_timestamp_unix_nanos - base) as usize); + received_latency_max.add_nanos(received_timestamp_unix_nanos - base); + received_latency_p50.add_nanos(received_timestamp_unix_nanos - base); - processed_latency_max.add_nanos((processed_timestamp_unix_nanos - base) as usize); - processed_latency_p50.add_nanos((processed_timestamp_unix_nanos - base) as usize); + processed_latency_max.add_nanos(processed_timestamp_unix_nanos - base); + processed_latency_p50.add_nanos(processed_timestamp_unix_nanos - base); - result + work_unit }) .boxed()) } diff --git a/src/work_unit_feed/work_unit_feed_provider.rs b/src/work_unit_feed/work_unit_feed_provider.rs index aab04c88..7647b6a1 100644 --- a/src/work_unit_feed/work_unit_feed_provider.rs +++ b/src/work_unit_feed/work_unit_feed_provider.rs @@ -21,7 +21,7 @@ use std::sync::Arc; /// /// See [`WorkUnitFeedProvider::feed`] for the per-call contract. pub trait WorkUnitFeedProvider: Send + Sync + Debug { - type WorkUnit: WorkUnit + Default; + type WorkUnit: WorkUnit + Default + 'static; /// Builds a [`WorkUnit`] stream for the given `partition`. /// diff --git a/src/worker/gen/regen.sh b/src/worker/gen/regen.sh deleted file mode 100755 index 39e04fdc..00000000 --- a/src/worker/gen/regen.sh +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env bash - -set -e - -repo_root=$(git rev-parse --show-toplevel) -cd "$repo_root" && cargo run --manifest-path src/worker/gen/Cargo.toml diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index c3beb3e0..cf16e890 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,81 +1,67 @@ -use crate::common::{TreeNodeExt, deserialize_uuid}; +use crate::common::TreeNodeExt; use crate::execution_plans::SamplerExec; -use crate::metrics::proto::df_metrics_set_to_proto; -use crate::protobuf::datafusion_error_to_tonic_status; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; -use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; -use crate::worker::generated::worker::worker_service_server::WorkerService; -use crate::worker::generated::worker::{ - CoordinatorToWorkerMsg, TaskMetrics, WorkerToCoordinatorMsg, worker_to_coordinator_msg, -}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, Worker, - WorkerQueryContext, + CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, + DistributedTaskContext, TaskData, TaskMetrics, Worker, WorkerQueryContext, + WorkerToCoordinatorMsg, }; -use datafusion::common::DataFusionError; use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_proto::physical_plan::AsExecutionPlan; use datafusion_proto::protobuf::PhysicalPlanNode; -use futures::stream::FuturesUnordered; +use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; +use http::HeaderMap; use std::sync::{Arc, OnceLock}; use tokio::sync::oneshot; use tokio::sync::oneshot::Sender; -use tonic::{Request, Response, Status, Streaming}; -use url::Url; impl Worker { - pub(super) async fn impl_coordinator_channel( + pub async fn coordinator_channel( &self, - request: Request>, - ) -> Result::CoordinatorChannelStream>, Status> { - let (grpc_headers, _ext, mut body) = request.into_parts(); - + headers: HeaderMap, + mut stream: BoxStream<'static, Result>, + ) -> Result>> { // The first message must be a SetPlanRequest. - let Some(msg) = body.next().await else { - return Err(Status::internal("Empty Coordinator stream")); + let Some(msg) = stream.try_next().await? else { + return internal_err!("Empty Coordinator stream"); }; - let Some(Inner::SetPlanRequest(request)) = msg?.inner else { - return Err(Status::internal( - "First Coordinator message must be SetPlanRequest", - )); + + let CoordinatorToWorkerMsg::SetPlanRequest(request) = msg else { + return internal_err!("First Coordinator message must be SetPlanRequest"); }; - let key = request.task_key.ok_or_else(missing("task_key"))?; + + let key = request.task_key; let entry = self .task_data_entries - .get_with(key.clone(), async { Default::default() }) + .get_with(key, async { Default::default() }) .await; let mut remote_work_unit_feed_registry = RemoteWorkUnitFeedRegistry::default(); - for WorkUnitFeedDeclaration { id, partitions } in &request.work_unit_feed_declarations { - if let Ok(id) = deserialize_uuid(id) { - remote_work_unit_feed_registry.add(id, *partitions as usize); - } + for decl in request.work_unit_feed_declarations { + remote_work_unit_feed_registry.add(decl.id, decl.partitions); } let (metrics_tx, metrics_rx) = oneshot::channel(); let mut load_info_rxs = vec![]; let task_data = || async { - let headers = grpc_headers.into_headers(); - let mut cfg = SessionConfig::default() .with_extension(Arc::new(remote_work_unit_feed_registry.receivers)) .with_extension(Arc::new(DistributedTaskContext { - task_index: key.task_number as usize, - task_count: request.task_count as usize, + task_index: request.task_key.task_number, + task_count: request.task_count, })) .with_extension(Arc::new(LocalWorkerContext { task_data_entries: Arc::clone(&self.task_data_entries), - self_url: Url::parse(&request.target_worker_url) - .map_err(|e| DataFusionError::External(Box::new(e)))?, + self_url: request.target_worker_url, })) .with_distributed_option_extension_from_headers::(&headers)?; @@ -125,13 +111,9 @@ impl Worker { entry .write(task_data_result.clone()) - .map_err(|_| Status::internal(format!( - "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it." - )))?; + .map_err(|e| exec_datafusion_err!("{e}"))?; - let task_data = task_data_result - .map_err(DataFusionError::Shared) - .map_err(datafusion_error_to_tonic_status)?; + let task_data = task_data_result.map_err(DataFusionError::Shared)?; // Continue reading remaining messages (work unit feed data) in the background. let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); @@ -150,27 +132,22 @@ impl Worker { // is closed. #[allow(clippy::disallowed_methods)] tokio::spawn(async move { - let mut body = body.map_ok(set_work_unit_received_time); - while let Some(Ok(msg)) = body.next().await { - let Some(msg) = msg.inner else { - continue; - }; + let mut stream = stream.map_ok(set_work_unit_received_time); + while let Some(Ok(msg)) = stream.next().await { match msg { - Inner::SetPlanRequest(_) => { + CoordinatorToWorkerMsg::SetPlanRequest(_) => { // SetPlanRequest should be the first already polled message in the stream, // if some reached here it means that something is wrong. continue; } - Inner::WorkUnitBatch(msg) => { + CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) => { let Some(work_unit_senders) = work_unit_senders.as_mut() else { continue; }; - for wu in msg.batch { - let Ok(id) = deserialize_uuid(&wu.id) else { - continue; - }; - let partition = wu.partition as usize; - let Some(tx) = work_unit_senders.get(&(id, partition)) else { + for wu in work_unit_batch.batch { + let id = wu.id; + let partition = wu.partition; + let Some(tx) = work_unit_senders.get(&(wu.id, partition)) else { continue; }; if tx.send(Ok(wu)).is_err() { @@ -181,7 +158,7 @@ impl Worker { } } } - Inner::WorkUnitEos(_) => { + CoordinatorToWorkerMsg::WorkUnitEos => { // No further work unit message will be received here, so drop all the // sender sides so that receiver sides see an EOS upon draining the // remaining messages. @@ -197,8 +174,8 @@ impl Worker { let metrics_tx = task_data.metrics_tx.lock().unwrap().take(); if let Some(Ok(plan)) = task_data.final_plan.get() { let d_ctx = DistributedTaskContext { - task_index: key.task_number as usize, - task_count: request.task_count as usize, + task_index: key.task_number, + task_count: request.task_count, }; let task_data_metrics = &task_data.task_data_metrics; task_data_metrics.mark_execution_finished(); @@ -211,16 +188,12 @@ impl Worker { let load_info_stream = FuturesUnordered::from_iter(load_info_rxs) .filter_map(async |load_info_or_channel_dropped| { - // This error can only happen if the pb::LoadInfo sender was dropped, which is fine. + // This error can only happen if the LoadInfo sender was dropped, which is fine. let load_info = load_info_or_channel_dropped.ok()?; - Some(Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::LoadInfo(load_info)), - })) + Some(WorkerToCoordinatorMsg::LoadInfo(load_info)) }) .chain(futures::stream::once(async move { - Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::LoadInfoEos(true)), - }) + WorkerToCoordinatorMsg::LoadInfoEos })); // Stream back metrics when the coordinator channel reaches EOS. At that point the @@ -229,21 +202,15 @@ impl Worker { let metrics_stream = metrics_rx.into_stream(); let metrics_stream = metrics_stream.filter_map(async |task_metrics_or_channel_dropped| { let task_metrics = task_metrics_or_channel_dropped.ok()?; - Some(Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)), - })) + Some(WorkerToCoordinatorMsg::TaskMetrics(task_metrics)) }); - Ok(Response::new( - futures::stream::select(load_info_stream, metrics_stream).boxed(), - )) + Ok(futures::stream::select(load_info_stream, metrics_stream) + .map(Ok) + .boxed()) } } -fn missing(field: &'static str) -> impl FnOnce() -> Status { - move || Status::invalid_argument(format!("Missing field '{field}'")) -} - /// Collects metrics from the plan in pre-order traversal order and sends them via the /// coordinator channel oneshot. fn send_metrics_via_channel( @@ -254,17 +221,13 @@ fn send_metrics_via_channel( ) { 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(), - ); + pre_order_plan_metrics.push(node.metrics().unwrap_or_default()); Ok(TreeNodeRecursion::Continue) }); // Ignore send errors — the coordinator channel may have been dropped (e.g. query cancelled). let _ = metrics_tx.send(TaskMetrics { pre_order_plan_metrics, - task_metrics: Some(task_data_metrics.to_proto_metrics_set()), + task_metrics: task_data_metrics.to_metrics_set(), }); } diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index ef3b9e2f..eff0bd98 100644 --- a/src/worker/impl_execute_task.rs +++ b/src/worker/impl_execute_task.rs @@ -1,31 +1,13 @@ -use crate::DistributedConfig; -use crate::common::now_ns; -use crate::protobuf::datafusion_error_to_tonic_status; -use crate::worker::generated::worker::ExecuteTaskRequest; -use crate::worker::generated::worker::FlightAppMetadata; -use crate::worker::generated::worker::worker_service_server::WorkerService; -use crate::worker::spawn_select_all::spawn_select_all; +use crate::ExecuteTaskRequest; use crate::worker::worker_service::{TaskDataEntries, Worker}; -use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; -use arrow_flight::error::FlightError; -use arrow_select::dictionary::garbage_collect_any_dictionary; -use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; -use datafusion::arrow::ipc::CompressionType; -use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::common::exec_datafusion_err; -use datafusion::common::{Result, exec_err, internal_err}; +use datafusion::common::{Result, exec_err}; use datafusion::error::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::TryStreamExt; -use prost::Message; use std::sync::Arc; use std::time::Duration; -use tokio_stream::StreamExt; -use tonic::{Request, Response, Status}; -/// How many record batches to buffer from the plan execution. -const RECORD_BATCH_BUFFER_SIZE: usize = 2; const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; /// Builds several per-partition streams by retrieving the appropriate entry from [TaskDataEntries] @@ -33,175 +15,52 @@ const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; /// /// This method is async mainly for the key retrieval operation from [TaskDataEntries], but it does /// not start polling any stream, it just instantiates them. -pub(crate) async fn execute_local_task( - task_data_entries: &Arc, - body: ExecuteTaskRequest, -) -> Result<(Vec, Arc)> { - let Some(key) = body.task_key.as_ref().cloned() else { - return internal_err!("Missing task_key in LocalWorkerConnection"); - }; - let Some(producer_head) = body.producer_head.as_ref().cloned() else { - return internal_err!("Missing producer_head"); - }; - let entry = task_data_entries - .get_with(key.clone(), async { Default::default() }) - .await; - - // Other request is responsible for writing the plan that belongs to this TaskKey, so - // we'll resolve immediately if it was already there, or wait until it's ready. - let task_data = entry - .read(Duration::from_secs(WAIT_PLAN_TIMEOUT_SECS)) - .await - .map_err(|e| exec_datafusion_err!("Worker::execute_task timed-out while waiting for the plan to be set by the coordinator. ({e})"))? - .map_err(DataFusionError::Shared)?; - task_data.task_data_metrics.mark_execution_started_once(); - - let plan = task_data.plan(producer_head)?; - let task_ctx = task_data.task_ctx; - - let partition_count = plan.properties().partitioning.partition_count(); - let plan_name = plan.name(); - - // Execute all the requested partitions at once, and collect all the streams so that they - // can be merged into a single one at the end of this function. - let n_streams = body.target_partition_end - body.target_partition_start; - let mut streams = Vec::with_capacity(n_streams as usize); - for partition in body.target_partition_start..body.target_partition_end { - if partition >= partition_count as u64 { - return exec_err!( - "partition {partition} not available. The head plan {plan_name} of the stage just has {partition_count} partitions" - ); - } - - let stream = plan.execute(partition as usize, Arc::clone(&task_ctx))?; - let stream_schema = plan.schema(); - - streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _); +impl Worker { + pub async fn execute_task( + &self, + request: ExecuteTaskRequest, + ) -> Result<(Vec, Arc)> { + Self::execute_task_static(Arc::clone(&self.task_data_entries), request).await } - Ok((streams, task_ctx)) -} - -/// Builds several per-partition streams by retrieving the appropriate entry from [TaskDataEntries] -/// based on the task key extracted from the gRPC request. -/// -/// This method eagerly starts streaming data from the task, and communicates via channels the -/// produced [RecordBatch]s already encoded as Arrow Flight data. -pub(crate) async fn execute_remote_task( - task_data_entries: &Arc, - request: Request, -) -> Result::ExecuteTaskStream>, Status> { - let body = request.into_inner(); - let partition_range = body.target_partition_start..body.target_partition_end; - - let (arrow_streams, task_ctx) = execute_local_task(task_data_entries, body) - .await - .map_err(datafusion_error_to_tonic_status)?; - let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options()) - .map_err(datafusion_error_to_tonic_status)?; + pub(crate) async fn execute_task_static( + task_data_entries: Arc, + request: ExecuteTaskRequest, + ) -> Result<(Vec, Arc)> { + let entry = task_data_entries + .get_with(request.task_key, async { Default::default() }) + .await; + + // Other request is responsible for writing the plan that belongs to this TaskKey, so + // we'll resolve immediately if it was already there, or wait until it's ready. + let task_data = entry + .read(Duration::from_secs(WAIT_PLAN_TIMEOUT_SECS)) + .await + .map_err(|e| exec_datafusion_err!("Worker::execute_task timed-out while waiting for the plan to be set by the coordinator. ({e})"))? + .map_err(DataFusionError::Shared)?; + task_data.task_data_metrics.mark_execution_started_once(); + + let plan = task_data.plan(&request.producer_head_spec)?; + let task_ctx = task_data.task_ctx; + let partition_count = plan.properties().partitioning.partition_count(); + let plan_name = plan.name(); + + // Execute all the requested partitions at once, and collect all the streams so that they + // can be merged into a single one at the end of this function. + let n_streams = request.target_partition_end - request.target_partition_start; + let mut streams = Vec::with_capacity(n_streams); + for partition in request.target_partition_start..request.target_partition_end { + if partition >= partition_count { + return exec_err!( + "partition {partition} not available. The head plan {plan_name} of the stage just has {partition_count} partitions" + ); + } - let compression = match d_cfg.compression.as_str() { - "lz4" => Some(CompressionType::LZ4_FRAME), - "zstd" => Some(CompressionType::ZSTD), - "none" => None, - v => Err(Status::invalid_argument(format!( - "Unknown compression type {v}" - )))?, - }; - let mut flight_streams = Vec::with_capacity(arrow_streams.len()); - for (partition, arrow_stream) in partition_range.zip(arrow_streams) { - let flight_stream = build_flight_data_stream(arrow_stream, compression)?.map(move |msg| { - // For each FlightData produced by this stream, mark it with the appropriate - // partition. This stream will be merged with several others from other partitions, - // so marking it with the original partition allows it to be deconstructed into - // the original per-partition streams in later steps. - let flight_data = FlightAppMetadata { - partition, - created_timestamp_unix_nanos: now_ns(), - }; - msg.map(|v| v.with_app_metadata(flight_data.encode_to_vec())) - }); + let stream = plan.execute(partition, Arc::clone(&task_ctx))?; + let stream_schema = plan.schema(); - flight_streams.push(flight_stream); + streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _); + } + Ok((streams, task_ctx)) } - - // Merge all the per-partition streams into one. Each message in the stream is marked with - // the original partition, so they can be reconstructed at the other side of the boundary. - let memory_pool = Arc::clone(&task_ctx.runtime_env().memory_pool); - let stream = spawn_select_all(flight_streams, memory_pool, RECORD_BATCH_BUFFER_SIZE); - - Ok(Response::new(Box::pin(stream.map_err(|err| match err { - FlightError::Tonic(status) => *status, - _ => Status::internal(format!("Error during flight stream: {err}")), - })))) -} - -fn build_flight_data_stream( - stream: SendableRecordBatchStream, - compression_type: Option, -) -> Result { - let stream = FlightDataEncoderBuilder::new() - .with_options( - IpcWriteOptions::default() - .try_with_compression(compression_type) - .map_err(|err| Status::internal(err.to_string()))?, - ) - .with_schema(stream.schema()) - // This tells the encoder to send dictionaries across the wire as-is. - // The alternative (`DictionaryHandling::Hydrate`) would expand the dictionaries - // into their value types, which can potentially blow up the size of the data transfer. - // The main reason to use `DictionaryHandling::Hydrate` is for compatibility with clients - // that do not support dictionaries, but since we are using the same server/client on both - // sides, we can safely use `DictionaryHandling::Resend`. - // Note that we do garbage collection of unused dictionary values above, so we are not sending - // unused dictionary values over the wire. - .with_dictionary_handling(DictionaryHandling::Resend) - // Set max flight data size to unlimited. - // This requires servers and clients to also be configured to handle unlimited sizes. - // Using unlimited sizes avoids splitting RecordBatches into multiple FlightData messages, - // which could add significant overhead for large RecordBatches. - // The only reason to split them really is if the client/server are configured with a message size limit, - // which mainly makes sense in a public network scenario where you want to avoid DoS attacks. - // Since all of our Arrow Flight communication happens within trusted data plane networks, - // we can safely use unlimited sizes here. - .with_max_flight_data_size(usize::MAX) - .build( - stream - // Apply garbage collection of dictionary and view arrays before sending over the network - .and_then(|rb| std::future::ready(garbage_collect_arrays(rb))) - .map_err(|err| FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(err)))), - ); - Ok(stream) -} - -/// Garbage collects values sub-arrays. -/// -/// We apply this before sending RecordBatches over the network to avoid sending -/// values that are not referenced by any dictionary keys or buffers that are not used. -/// -/// Unused values can arise from operations such as filtering, where -/// some keys may no longer be referenced in the filtered result. -fn garbage_collect_arrays(batch: RecordBatch) -> Result { - let (schema, arrays, row_count) = batch.into_parts(); - - let arrays = arrays - .into_iter() - .map(|array| { - if let Some(array) = array.as_any_dictionary_opt() { - garbage_collect_any_dictionary(array) - } else if let Some(array) = array.as_string_view_opt() { - Ok(Arc::new(array.gc()) as Arc) - } else if let Some(array) = array.as_binary_view_opt() { - Ok(Arc::new(array.gc()) as Arc) - } else { - Ok(array) - } - }) - .collect::, _>>()?; - - Ok(RecordBatch::try_new_with_options( - schema, - arrays, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - )?) } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index e89921fc..6e8ce7b9 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -1,9 +1,7 @@ -pub(crate) mod generated; mod impl_coordinator_channel; mod impl_execute_task; mod session_builder; mod single_write_multi_read; -mod spawn_select_all; mod task_data; #[cfg(any(test, feature = "integration"))] pub(crate) mod test_utils; diff --git a/src/worker/single_write_multi_read.rs b/src/worker/single_write_multi_read.rs index b9632352..30c4ce08 100644 --- a/src/worker/single_write_multi_read.rs +++ b/src/worker/single_write_multi_read.rs @@ -63,6 +63,7 @@ impl SingleWriteMultiRead { } /// Reads the current value, if any, not waiting for it to be set by a writer. + #[cfg(feature = "grpc")] pub(crate) fn read_now(&self) -> Option { self.rx.borrow().clone() } diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index c29d60aa..c9085263 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,12 +1,13 @@ -use crate::MaxLatencyMetric; use crate::common::OnceLockResult; use crate::common::now_ns; use crate::distributed_planner::ProducerHead; -use crate::worker::generated::worker as pb; +use crate::protocol::ProducerHeadSpec; +use crate::{MaxLatencyMetric, TaskMetrics}; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; -use datafusion::physical_expr_common::metrics::CustomMetricValue; -use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet}; +use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; use tokio::sync::oneshot; @@ -16,13 +17,13 @@ use tokio::sync::oneshot; /// by concurrent requests for the same task which execute separate partitions. pub struct TaskData { /// Task context suitable for execute different partitions from the same task. - pub(super) task_ctx: Arc, + pub(crate) task_ctx: Arc, pub(crate) base_plan: Arc, pub(crate) final_plan: Arc>>, - /// Sender half of the metrics channel. `impl_execute_task` takes this (via `Option::take`) - /// once all partitions have finished or been dropped, sending the collected metrics back to - /// the coordinator through the `CoordinatorChannel` side channel. - pub(super) metrics_tx: Arc>>>, + /// Sender half of the metrics channel. `impl_coordinator_channel` takes this (via + /// `Option::take`) when the coordinator channel reaches EOS, sending the collected metrics + /// back to the coordinator through the `CoordinatorChannel` side channel. + pub(super) metrics_tx: Arc>>>, /// Metrics related to the execution of a task within a stage. This metrics, instead of being /// associated to a specific node, they are global to the task, like the time at which the plan /// was fed by the coordinator to the worker. @@ -35,7 +36,7 @@ pub(crate) const PLAN_FINISHED_AT_METRIC: &str = "plan_finished_at"; #[derive(Debug)] pub(super) struct TaskDataMetrics { - pub(super) query_start_time_ns: u64, + pub(super) query_start_time_ns: usize, /// When the plan was set by the coordinator. pub(super) plan_added_at: MaxLatencyMetric, /// When the plan execution was triggered by the parent worker. @@ -45,10 +46,10 @@ pub(super) struct TaskDataMetrics { } impl TaskDataMetrics { - pub(super) fn new(query_start_time_ns: u64) -> Self { + pub(super) fn new(query_start_time_ns: usize) -> Self { let plan_added_at = MaxLatencyMetric::default(); plan_added_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(query_start_time_ns), + now_ns::().saturating_sub(query_start_time_ns as u64), )); Self { query_start_time_ns, @@ -61,67 +62,57 @@ impl TaskDataMetrics { pub(super) fn mark_execution_started_once(&self) { if self.plan_executed_at.value() == 0 { self.plan_executed_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(self.query_start_time_ns), + now_ns::().saturating_sub(self.query_start_time_ns as u64), )) } } pub(super) fn mark_execution_finished(&self) { self.plan_finished_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(self.query_start_time_ns), + now_ns::().saturating_sub(self.query_start_time_ns as u64), )) } - pub(super) fn to_proto_metrics_set(&self) -> pb::MetricsSet { - let mut task_metrics_set = pb::MetricsSet { metrics: vec![] }; - - fn new_metric(name: &str, value: usize) -> pb::Metric { - pb::Metric { - partition: None, - labels: vec![], - value: Some(pb::metric::Value::CustomMaxLatency(pb::MaxLatency { - name: name.to_string(), - value: value as u64, - })), - } - } - task_metrics_set.metrics.push(new_metric( + pub(super) fn to_metrics_set(&self) -> MetricsSet { + let mut metrics_set = MetricsSet::new(); + metrics_set.push(max_latency_metric( PLAN_ADDED_AT_METRIC, - self.plan_added_at.as_usize(), + &self.plan_added_at, )); - task_metrics_set.metrics.push(new_metric( + metrics_set.push(max_latency_metric( PLAN_EXECUTED_AT_METRIC, - self.plan_executed_at.as_usize(), + &self.plan_executed_at, )); - task_metrics_set.metrics.push(new_metric( + metrics_set.push(max_latency_metric( PLAN_FINISHED_AT_METRIC, - self.plan_finished_at.as_usize(), + &self.plan_finished_at, )); - task_metrics_set + metrics_set } } -impl TaskData { - /// Returns the total number of partitions in this task. - pub(crate) fn total_partitions(&self) -> usize { - match self.final_plan.get() { - Some(Ok(plan)) => plan.output_partitioning().partition_count(), - _ => self - .base_plan - .properties() - .output_partitioning() - .partition_count(), - } - } +fn max_latency_metric(name: &'static str, value: &MaxLatencyMetric) -> Arc { + Arc::new(Metric::new( + MetricValue::Custom { + name: Cow::Borrowed(name), + value: Arc::new(MaxLatencyMetric::from_nanos(value.value())), + }, + None, + )) +} +impl TaskData { pub(crate) fn plan( &self, - producer_head: pb::execute_task_request::ProducerHead, + producer_head_spec: &ProducerHeadSpec, ) -> Result> { let result = self.final_plan.get_or_init(|| { - let producer_head = - ProducerHead::from_proto(producer_head, &self.base_plan.schema(), &self.task_ctx)?; + let producer_head = ProducerHead::from_spec( + producer_head_spec, + self.base_plan.schema(), + &self.task_ctx, + )?; Ok(producer_head.insert(Arc::clone(&self.base_plan))?) }); diff --git a/src/worker/test_utils/worker_handles.rs b/src/worker/test_utils/worker_handles.rs index f328145f..4fe19c5c 100644 --- a/src/worker/test_utils/worker_handles.rs +++ b/src/worker/test_utils/worker_handles.rs @@ -1,7 +1,7 @@ use crate::config_extension_ext::set_distributed_option_extension; -use crate::worker::generated::worker::TaskKey; +use crate::grpc::BoxCloneSyncChannel; use crate::worker::task_data::TaskDataMetrics; -use crate::{BoxCloneSyncChannel, DistributedConfig, DistributedExt, TaskData, Worker}; +use crate::{DistributedConfig, DistributedExt, TaskData, TaskKey, Worker}; use arrow_ipc::CompressionType; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -16,9 +16,9 @@ use tonic::transport::{Endpoint, Server}; use url::Url; use uuid::Uuid; -pub fn test_task_key_with_query(query_id: Uuid, task_number: u64) -> TaskKey { +pub fn test_task_key_with_query(query_id: Uuid, task_number: usize) -> TaskKey { TaskKey { - query_id: query_id.as_bytes().to_vec(), + query_id, stage_id: 0, task_number, } diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 7acd629c..4ccba411 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,48 +1,22 @@ -use crate::common::{OnceLockResult, on_drop_stream, serialize_uuid}; use crate::distributed_planner::ProducerHead; -use crate::metrics::LatencyMetricExt; -use crate::networking::get_distributed_channel_resolver; use crate::passthrough_headers::get_passthrough_headers; -use crate::protobuf::{datafusion_error_to_tonic_status, map_flight_to_datafusion_error}; use crate::stage::RemoteStage; -use crate::worker::generated::worker::FlightAppMetadata; -use crate::worker::generated::worker::{ExecuteTaskRequest, TaskKey}; -use crate::worker::impl_execute_task::execute_local_task; use crate::worker::worker_service::TaskDataEntries; -use crate::{BytesMetricExt, ChannelResolver, DistributedConfig}; -use arrow_flight::FlightData; -use arrow_flight::decode::FlightRecordBatchStream; -use arrow_flight::error::FlightError; -use dashmap::DashMap; +use crate::{ + ChannelResolver, ExecuteTaskRequest, TaskKey, Worker, get_distributed_channel_resolver, +}; use datafusion::arrow::array::RecordBatch; -use datafusion::common::instant::Instant; use datafusion::common::runtime::SpawnedTask; -use datafusion::common::{ - DataFusionError, Result, exec_err, internal_datafusion_err, internal_err, -}; +use datafusion::common::{DataFusionError, Result, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; -use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, MetricValue}; -use datafusion::physical_plan::metrics::{MetricBuilder, Time}; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::metrics::MetricBuilder; +use futures::future::{BoxFuture, Shared}; use futures::stream::BoxStream; -use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; -use http::Extensions; -use pin_project::{pin_project, pinned_drop}; -use prost::Message; -use std::borrow::Cow; +use futures::{FutureExt, StreamExt, TryFutureExt}; use std::fmt::{Debug, Formatter}; use std::ops::Range; -use std::pin::Pin; -use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; -use std::task::{Context, Poll}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::sync::Notify; -use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; -use tokio_stream::wrappers::UnboundedReceiverStream; -use tokio_util::sync::CancellationToken; -use tonic::metadata::MetadataMap; -use tonic::{Request, Status}; use url::Url; /// Context set by [crate::Worker::coordinator_channel] in DataFusion's @@ -59,441 +33,79 @@ pub(crate) struct LocalWorkerContext { pub(crate) self_url: Url, } -/// Holds a list of lazily initialized [WorkerConnection]s. Each position in the underlying -/// `connections` vector corresponds to the connection to one worker. It assumes a 1:1 mapping -/// between worker and tasks, and upon calling [WorkerConnectionPool::get_or_init_worker_connection] -/// it will initialize the corresponding position in the vector matching the provided `target_task` -/// index. +/// Manages connections to remote workers. +/// - Handles a range of partitions at a time in other to give the chance to [crate::WorkerChannel] +/// implementations to batch RecordBatch streams up, avoiding the overhead of multiple small +/// streams over an IO interface. +/// - Short circuits to a local in-memory connection if the remote worker that should be reached +/// happens to be the same one issuing the request. +/// - Lazy inits connections to a remote worker on first call to [WorkerConnectionPool::execute]. pub(crate) struct WorkerConnectionPool { - connections: Vec>>, + lazy_stream_groups: Vec>>, pub(crate) metrics: ExecutionPlanMetricsSet, } +/// A list of consumable RecordBatch streams, each one wrapped by `Mutex>` for +/// exactly-once consumption semantics. +type StreamGroup = Arc>>>>>; +/// Just some boilerplate for a shared future. +type SharedBoxFuture = Shared>>>; + impl WorkerConnectionPool { - /// Builds a new [WorkerConnectionPool] with as many empty slots for [WorkerConnection]s as + /// Builds a new [WorkerConnectionPool] with as many empty slots for worker stream entries as /// the provided `input_tasks`. pub(crate) fn new(input_tasks: usize) -> Self { - let mut connections = Vec::with_capacity(input_tasks); - for _ in 0..input_tasks { - connections.push(OnceLock::new()); - } Self { - connections, + lazy_stream_groups: (0..input_tasks).map(|_| OnceLock::new()).collect(), metrics: ExecutionPlanMetricsSet::default(), } } - /// Lazily initializes the [WorkerConnection] corresponding to the provided `target_task` - /// (therefore maintaining one independent [WorkerConnection] per `target_task`), and - /// returns it. - pub(crate) fn get_or_init_worker_connection( + pub(crate) fn execute( &self, input_stage: &RemoteStage, target_partitions: Range, target_task: usize, + target_partition: usize, producer_head: ProducerHead, ctx: &Arc, - ) -> Result<&(dyn WorkerConnection + Sync + Send)> { - let Some(worker_connection) = self.connections.get(target_task) else { - return internal_err!( - "WorkerConnections: Task index {target_task} not found, only have {} tasks", - self.connections.len() - ); - }; - - let conn = worker_connection.get_or_init(|| { - let Some(target_url) = input_stage.workers.get(target_task) else { - internal_err!("input_stage.workers[{target_task}] out of range.")? - }; - if let Some(lw_ctx) = ctx.session_config().get_extension::() - && &lw_ctx.self_url == target_url - { - // Instead of making a gRPC call to ourselves, better to just use local comms. - LocalWorkerConnection::init( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } else { - // We are trying to reach a URL different from ours, so use normal gRPC streams. - RemoteWorkerConnection::init( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } - }); - - match conn { - Ok(v) => Ok(v.as_ref()), - Err(err) => Err(DataFusionError::Shared(Arc::clone(err))), - } - } -} - -type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; - -/// Abstraction that allows treating remote and local comms as equal. Network boundaries do not -/// care if the stream comes over the wire or locally. -pub(crate) trait WorkerConnection { - /// Streams the specified partition. Consumers do not care if the implementation pulls data - /// from in-memory or from local comms. - fn execute(&self, partition: usize) -> Result>>; -} - -/// Represents a connection to one [Worker]. Network boundaries will use this for streaming -/// data from single partitions while the actual network communication is handling all the partitions -/// under the hood. -/// -/// This is done so that, rather than issuing one gRPC stream per partition, we issue one gRPC stream -/// per group of partitions, and we multiplex streamed record batches locally to in-memory channels. -/// -/// Even if Tonic can perfectly multiplex and interleave messages from different gRPC streams through -/// the same underlying TCP connection, there do is some overhead in having one gRPC stream per -/// partition VS a single gRPC stream interleaving multiple partitions. The whole serialized plan -/// needs to be sent over the wire on every gRPC call, so the less gRPC calls we do the better. -struct RemoteWorkerConnection { - task: Arc>, - not_consumed_streams: Arc, - cancel_token: CancellationToken, - per_partition_rx: DashMap>, - - first_poll_notify: Arc, - // Signals the demux task that buffered memory has been freed by a consumer. - mem_available_notify: Arc, - - // Metrics collection stuff. - memory_reservation: Arc, - elapsed_compute: Time, -} - -impl RemoteWorkerConnection { - fn init( - input_stage: &RemoteStage, - target_partition_range: Range, - target_task: usize, - producer_head: ProducerHead, - ctx: &Arc, - metrics: &ExecutionPlanMetricsSet, - ) -> Result { - let channel_resolver = get_distributed_channel_resolver(ctx.as_ref()); - let buffer_budget_bytes = - DistributedConfig::from_config_options(ctx.session_config().options())? - .worker_connection_buffer_budget_bytes; - // We are retaining record batches in memory until they are consumed, so we need to account - // for them in the memory pool. - let memory_reservation = - Arc::new(MemoryConsumer::new("WorkerConnection").register(ctx.memory_pool())); - let memory_reservation_clone = Arc::clone(&memory_reservation); - - // Track the maximum memory used to buffer recieved messages. - let mut curr_max_mem = 0; - let max_mem_used = MetricBuilder::new(metrics).global_gauge("max_mem_used"); - // Track the total encoded size of all recieved messages. - let bytes_transferred = MetricBuilder::new(metrics).bytes_counter("bytes_transferred"); - let msg_count = MetricBuilder::new(metrics).global_counter("msg_count"); - // Track end-to-end network latency distribution for all messages. - let min_latency = MetricBuilder::new(metrics).min_latency("network_latency_min"); - let max_latency = MetricBuilder::new(metrics).max_latency("network_latency_max"); - let p50_latency = MetricBuilder::new(metrics).p50_latency("network_latency_p50"); - let p95_latency = MetricBuilder::new(metrics).p95_latency("network_latency_p95"); - let first_latency = MetricBuilder::new(metrics).first_latency("network_latency_first"); - let sum_latency = Time::new(); - MetricBuilder::new(metrics).build(MetricValue::Time { - name: Cow::Borrowed("network_latency_sum"), - time: sum_latency.clone(), - }); - let latency_count = MetricBuilder::new(metrics).counter("network_latency_count", 0); - // Track the total CPU time spent in polling messages over the network + decoding them. - let elapsed_compute = Time::new(); - let elapsed_compute_clone = elapsed_compute.clone(); - MetricBuilder::new(metrics).build(MetricValue::ElapsedCompute(elapsed_compute.clone())); - - // Building the actual request that will be sent to the worker. - let headers = get_passthrough_headers(ctx.session_config()); - let request = Request::from_parts( - MetadataMap::from_headers(headers), - Extensions::default(), - ExecuteTaskRequest { - target_partition_start: target_partition_range.start as u64, - target_partition_end: target_partition_range.end as u64, - task_key: Some(TaskKey { - query_id: serialize_uuid(&input_stage.query_id), - stage_id: input_stage.num as u64, - task_number: target_task as u64, - }), - producer_head: Some(producer_head.to_proto(ctx)?), - }, - ); + ) -> Result>> { + let ch_resolver = get_distributed_channel_resolver(ctx.as_ref()); - let Some(url) = input_stage.workers.get(target_task).cloned() else { - return internal_err!("ProgrammingError: Task {target_task} not found"); + let Some(target_url) = input_stage.workers.get(target_task).cloned() else { + internal_err!("input_stage.workers[{target_task}] out of range.")? + }; + let task_key = TaskKey { + query_id: input_stage.query_id, + stage_id: input_stage.num, + task_number: target_task, }; - // The senders and receivers are unbounded queues used for multiplexing the record - // batches sent through the single gRPC stream into one stream per partition. They - // are unbounded to avoid head-of-line blocking: a single bounded queue could block - // the demux task and starve all sibling partitions even though they have capacity, - // which deadlocks queries with cross-partition dependencies. - // Total memory is bounded globally below via `mem_available_notify`. - let mut per_partition_tx = Vec::with_capacity(target_partition_range.len()); - let per_partition_rx = DashMap::with_capacity(target_partition_range.len()); - for partition in target_partition_range.clone() { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); - per_partition_tx.push(tx); - per_partition_rx.insert(partition, rx); - } - - let mem_available_notify = Arc::new(Notify::new()); - let mem_available_notify_for_task = Arc::clone(&mem_available_notify); - - let first_poll_notify = Arc::new(Notify::new()); - let first_poll_notify_for_task = Arc::clone(&first_poll_notify); - - // Cancellation token allows us to stop the background task promptly when all partition - // streams are dropped (e.g., when the query is cancelled). - let cancel_token = CancellationToken::new(); - let cancel = cancel_token.clone(); - - // This task will pull data from all the partitions in `target_partition_range`, and will - // fan them out to the appropriate `per_partition_rx` based on the "partition" declared - // in each individual record batch flight metadata. - let task = SpawnedTask::spawn(async move { - let mut client = match channel_resolver.get_worker_client_for_url(&url).await { - Ok(v) => v, - Err(err) => return fanout(&per_partition_tx, datafusion_error_to_tonic_status(&err)) - }; - - tokio::select! { - biased; - _ = cancel.cancelled() => { - // If all SendableRecordBatchStreams canceled before any poll, we need to - // anyway trigger the task execution and cancel it immediately so that the - // cancellation is propagated also in the remote worker. Otherwise, it might - // hang forever waiting for someone to execute it. - let _ = client.execute_task(request).await; - return - }, - _ = first_poll_notify_for_task.notified() => {} - } - - let mut interleaved_stream = match client.execute_task(request).await { - Ok(v) => v.into_inner(), - Err(err) => return fanout(&per_partition_tx, err), - }; - - loop { - // Backpressure gate. Per-partition channels are unbounded, so we cap - // total in-flight buffered bytes here by pausing the gRPC pull when - // consumers haven't drained enough. This propagates flow control all - // the way back to the worker without coupling sibling partitions. - // We always allow a message through when reservation == 0 to avoid - // livelock if a single message is larger than the budget. - while memory_reservation.size() >= buffer_budget_bytes { - tokio::select! { - biased; - _ = cancel.cancelled() => return, - _ = mem_available_notify_for_task.notified() => {} - } - } - - // Check for cancellation while waiting for the next message. - let flight_data = tokio::select! { - biased; - _ = cancel.cancelled() => return, - msg = interleaved_stream.next() => { - match msg { - Some(Ok(v)) => v, - Some(Err(err)) => return fanout(&per_partition_tx, err), - None => return, // Stream exhausted - } - } - }; - - // Earliest time at which the msg was received. - let msg_received_time = SystemTime::now(); - - let flight_metadata = match FlightAppMetadata::decode(flight_data.app_metadata.as_ref()) { - Ok(v) => v, - Err(err) => { - return fanout(&per_partition_tx, Status::internal(err.to_string())); - } - }; - - // Update the running latency tracker. - let sent_time = UNIX_EPOCH + Duration::from_nanos(flight_metadata.created_timestamp_unix_nanos); - if flight_metadata.created_timestamp_unix_nanos > 0 - && let Ok(delta) = msg_received_time.duration_since(sent_time) { - min_latency.add_duration(delta); - max_latency.add_duration(delta); - p50_latency.add_duration(delta); - p95_latency.add_duration(delta); - first_latency.add_duration(delta); - sum_latency.add_duration(delta); - latency_count.add(1); - } - - let partition = flight_metadata.partition as usize; - // the `per_partition_tx` variable is using a normal `Vec` for storing the - // channel transmitters, so we need to subtract the `target_partition_range.start` - // to the `partition` in order to offset it to the appropriate index. - let sender_i = partition - target_partition_range.start; - - let Some(o_tx) = per_partition_tx.get(sender_i) else { - let msg = format!( - "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" - ); - return fanout(&per_partition_tx, Status::internal(msg)); - }; - - // We need to send the memory reservation in the same tuple as the actual message - // so that it gets dropped as soon as the message leaves the queue. Dropping the - // memory reservation means releasing the memory from the pool for that specific - // message - let size = flight_data.encoded_len(); - memory_reservation.grow(size); - - // Update memory related metrics. - msg_count.add(1); - bytes_transferred.add_bytes(size); - let curr_mem = memory_reservation.size(); - if curr_mem > curr_max_mem { - curr_max_mem = curr_mem; - max_mem_used.set(curr_max_mem); - } - - if o_tx.send(Ok((flight_data, flight_metadata))).is_err() { - // The receiver for this partition was dropped (e.g. a hash join partition - // completed early without consuming its probe side). Don't exit: other - // partitions multiplexed over the same gRPC stream still need their data. - // Undo the memory reservation that was grown for this dropped batch. - memory_reservation.shrink(size); - continue; - }; - } - }.with_elapsed_compute(elapsed_compute)); - - Ok(Self { - task: Arc::new(task), - cancel_token, - not_consumed_streams: Arc::new(AtomicUsize::new(per_partition_rx.len())), - per_partition_rx, - mem_available_notify, - first_poll_notify, - - // metrics stuff - memory_reservation: memory_reservation_clone, - elapsed_compute: elapsed_compute_clone, - }) - } -} - -impl WorkerConnection for RemoteWorkerConnection { - /// Streams the provided `partition` from the remote worker. - /// - /// This method does not handle any network connection. Instead, the network comms are delegated - /// to the task spawned by [WorkerConnection::init], who is in charge of polling data not only - /// from the requested `partition`, but from any other partition in `target_partition_range`. - /// This method just streams all the record batches belonging to the provided `partition` from - /// an in-memory queue. - /// - /// The task that polls data over the network is held inactive until the first poll to the - /// stream returned by this method. - /// - /// When the returned stream is dropped (e.g., due to query cancellation), the background task - /// pulling from the Flight stream will be canceled promptly. - fn execute(&self, partition: usize) -> Result>> { - let Some((_, partition_receiver)) = self.per_partition_rx.remove(&partition) else { + // If we are physically in the same worker, short circuit into a local connection without + // going through the `WorkerChannel`. + if let Some(result) = self.local_stream_short_circuit( + task_key, + target_partition, + &target_url, + &producer_head, + ctx, + )? { + return Ok(result); + } + + // Otherwise, we need to reach the remote worker through the `WorkerChannel`. Unlike local + // connections, these remote connections span a range of partitions so that `WorkerChannel` + // implementations have the option of batch them. + let Some(worker_connection) = self.lazy_stream_groups.get(target_task) else { return internal_err!( - "WorkerConnection has no stream for target partition {partition}. Was it already consumed?" + "WorkerConnections: Task index {target_task} not found, only have {} tasks", + self.lazy_stream_groups.len() ); }; - let task = Arc::clone(&self.task); - let cancel_token = self.cancel_token.clone(); - - let first_poll_notify = Arc::clone(&self.first_poll_notify); - let stream = async move { - first_poll_notify.notify_one(); - UnboundedReceiverStream::new(partition_receiver) - } - .flatten_stream(); - - let stream = stream.map_err(|err| FlightError::Tonic(Box::new(err))); - let reservation = Arc::clone(&self.memory_reservation); - let mem_available_notify = Arc::clone(&self.mem_available_notify); - let stream = stream.map_ok(move |(data, _meta)| { - reservation.shrink(data.encoded_len()); - // Wake the demux task in case it is blocked on the byte budget. - mem_available_notify.notify_one(); - let _ = &task; // <- keep the task that polls data from the network alive. - data - }); - let stream = FlightRecordBatchStream::new_from_flight_data(stream); - let stream = stream.map_err(map_flight_to_datafusion_error); - let stream = stream.with_elapsed_compute(self.elapsed_compute.clone()); - // When the stream is dropped, cancel the background task to ensure prompt cleanup. - let not_consumed_streams = Arc::clone(&self.not_consumed_streams); - Ok(on_drop_stream(stream, move || { - let remaining_streams = not_consumed_streams.fetch_sub(1, Ordering::SeqCst) - 1; - if remaining_streams == 0 { - cancel_token.cancel(); - } - }) - .boxed()) - } -} - -/// Equivalent to [RemoteWorkerConnection], but that pulls data from the local registry of tasks -/// rather than doing it across a gRPC interface. -pub(crate) struct LocalWorkerConnection { - partition_start: usize, - local_streams: Vec>>>>, -} - -impl LocalWorkerConnection { - fn init( - input_stage: &RemoteStage, - target_partition_range: Range, - target_task: usize, - producer_head: ProducerHead, - ctx: &Arc, - metrics: &ExecutionPlanMetricsSet, - ) -> Result { - MetricBuilder::new(metrics) - .global_counter("local_connections_used") - .add(1); - let Some(lw_ctx) = ctx.session_config().get_extension::() else { - return exec_err!("Missing LocalWorkerContext extension"); - }; - - let task_key = TaskKey { - query_id: serialize_uuid(&input_stage.query_id), - stage_id: input_stage.num as u64, - task_number: target_task as u64, - }; - - let partition_start = target_partition_range.start; - let mut local_streams = Vec::with_capacity(target_partition_range.len()); - for partition_i in target_partition_range { - let request = ExecuteTaskRequest { - task_key: Some(task_key.clone()), - target_partition_start: partition_i as u64, - target_partition_end: (partition_i + 1) as u64, - producer_head: Some(producer_head.to_proto(ctx)?), - }; - - let task_data_entries = Arc::clone(&lw_ctx.task_data_entries); + let streams_shared_future = worker_connection.get_or_init(|| { + let metrics = self.metrics.clone(); + let ctx = Arc::clone(ctx); // The relevant entry from `task_data_entries` needs to be eagerly retrieved, it cannot be // left for until someone decides to start polling the returned `BoxStream`, otherwise, @@ -501,264 +113,119 @@ impl LocalWorkerConnection { // is polled, the entry might not be there. // // Note that this does not start polling the returned streams, it just instantiates them. - let streams_future = SpawnedTask::spawn(async move { - let (streams, _) = execute_local_task(&task_data_entries, request).await?; - Ok::<_, DataFusionError>(streams) + let streams_task = SpawnedTask::spawn(async move { + let request = ExecuteTaskRequest { + task_key, + target_partition_start: target_partitions.start, + target_partition_end: target_partitions.end, + producer_head_spec: producer_head.to_spec(ctx.session_config())?, + }; + let mut client = ch_resolver.get_worker_client_for_url(&target_url).await?; + let headers = get_passthrough_headers(ctx.session_config()); + let streams = client.execute_task(headers, request, metrics, &ctx).await?; + Ok(streams) }); - let stream = async move { - let mut streams = streams_future - .await - .map_err(|err| internal_datafusion_err!("{err}"))??; - if streams.len() != 1 { - return internal_err!("Expected exactly 1 local stream"); + async move { + match streams_task.await { + Ok(Ok(v)) => Ok(Arc::new( + v.into_iter().map(|v| Mutex::new(Some(v))).collect(), + )), + Ok(Err(e)) => Err(Arc::new(e)), + Err(e) => Err(Arc::new(internal_datafusion_err!( + "JoinError instantiating streams {e}" + ))), } - Ok(streams.swap_remove(0)) } - .try_flatten_stream() - .boxed(); + .boxed() + .shared() + }); - local_streams.push(Mutex::new(Some(stream))); + let streams_future = streams_shared_future.clone(); + Ok(async move { + let streams = streams_future.await.map_err(DataFusionError::Shared)?; + let Some(slot) = streams.get(target_partition - target_partitions.start) else { + return internal_err!( + "WorkerConnections has no stream for partition {target_partition}. Was it already consumed?" + ); + }; + slot.lock().unwrap().take().ok_or_else(|| { + internal_datafusion_err!( + "WorkerConnections stream for partition {target_partition} was already consumed" + ) + }) } - - Ok(Self { - partition_start, - local_streams, - }) + .try_flatten_stream() + .boxed()) } -} -impl WorkerConnection for LocalWorkerConnection { - fn execute(&self, partition: usize) -> Result>> { - let Some(relative_i) = partition.checked_sub(self.partition_start) else { - return internal_err!( - "LocalWorkerConnection received an invalid partition {partition}, the starting partition is {}", - self.partition_start - ); + fn local_stream_short_circuit( + &self, + task_key: TaskKey, + target_partition: usize, + target_url: &Url, + producer_head: &ProducerHead, + ctx: &Arc, + ) -> Result>>> { + let Some(task_data_entries) = ctx + .session_config() + .get_extension::() + .and_then(|lw_ctx| match &lw_ctx.self_url == target_url { + true => Some(Arc::clone(&lw_ctx.task_data_entries)), + false => None, + }) + else { + return Ok(None); }; - let Some(slot) = self.local_streams.get(relative_i) else { - return internal_err!( - "LocalWorkerConnection has no stream for partition {partition}. Was it already consumed?" - ); + + MetricBuilder::new(&self.metrics) + .global_counter("local_connections_used") + .add(1); + let request = ExecuteTaskRequest { + task_key, + target_partition_start: target_partition, + target_partition_end: target_partition + 1, + producer_head_spec: producer_head.to_spec(ctx.session_config())?, }; - slot.lock().unwrap().take().ok_or_else(|| { - internal_datafusion_err!( - "LocalWorkerConnection stream for partition {partition} was already consumed" - ) - }) - } -} + // The relevant entry from `task_data_entries` needs to be eagerly retrieved, it cannot be + // left for until someone decides to start polling the returned `BoxStream`, otherwise, + // there's risk that the entry is evicted by Moka's TTL, and by the time the returned stream + // is polled, the entry might not be there. + // + // Note that this does not start polling the returned streams, it just instantiates them. + let stream_task = SpawnedTask::spawn(async move { + let (mut streams, _) = Worker::execute_task_static(task_data_entries, request).await?; + if streams.len() != 1 { + return internal_err!( + "Expected exactly 1 stream out of Worker::execute_task_static, but got {}", + streams.len() + ); + } + Ok::<_, DataFusionError>(streams.swap_remove(0)) + }); -fn fanout(o_txs: &[UnboundedSender], err: Status) { - for o_tx in o_txs { - let _ = o_tx.send(Err(err.clone())); + Ok(Some( + async move { + stream_task + .await + .map_err(|err| internal_datafusion_err!("{err}"))? + } + .try_flatten_stream() + .boxed(), + )) } } impl Debug for WorkerConnectionPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("WorkerConnections") - .field("num_connections", &self.connections.len()) + .field("num_connections", &self.lazy_stream_groups.len()) .finish() } } impl Clone for WorkerConnectionPool { fn clone(&self) -> Self { - Self::new(self.connections.len()) - } -} - -impl Debug for RemoteWorkerConnection { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkerConnection").finish() - } -} - -trait ElapsedComputeFutureExt: Future + Sized { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture; -} - -trait ElapsedComputeStreamExt: Stream + Sized { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream; -} - -impl> ElapsedComputeFutureExt for F { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture { - ElapsedComputeFuture { - inner: self, - curr: Duration::default(), - elapsed_compute, - } - } -} - -impl> ElapsedComputeStreamExt for S { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { - ElapsedComputeStream { - inner: self, - curr: Duration::default(), - elapsed_compute, - } - } -} - -#[pin_project(PinnedDrop)] -struct ElapsedComputeStream { - #[pin] - inner: T, - curr: Duration, - elapsed_compute: Time, -} - -/// Drop implementation that ensures that any accumulated time is properly dumped to the metric -/// in case the stream gets dropped before completion. -#[pinned_drop] -impl PinnedDrop for ElapsedComputeStream { - fn drop(self: Pin<&mut Self>) { - if self.curr > Duration::default() { - let self_projected = self.project(); - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - } - } -} - -impl> Stream for ElapsedComputeStream { - type Item = O; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let self_projected = self.project(); - let start = Instant::now(); - let result = self_projected.inner.poll_next(cx); - *self_projected.curr += start.elapsed(); - if result.is_ready() { - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - *self_projected.curr = Duration::default(); - } - result - } -} - -#[pin_project(PinnedDrop)] -struct ElapsedComputeFuture { - #[pin] - inner: T, - curr: Duration, - elapsed_compute: Time, -} - -/// Drop implementation that ensures that any accumulated time is properly dumped to the metric -/// in case the future gets dropped before completion. -#[pinned_drop] -impl PinnedDrop for ElapsedComputeFuture { - fn drop(self: Pin<&mut Self>) { - if self.curr > Duration::default() { - let self_projected = self.project(); - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - } - } -} - -impl> Future for ElapsedComputeFuture { - type Output = O; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let self_projected = self.project(); - let start = Instant::now(); - let result = self_projected.inner.poll(cx); - *self_projected.curr += start.elapsed(); - if result.is_ready() { - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - *self_projected.curr = Duration::default(); - } - result - } -} - -#[cfg(test)] -mod tests { - use super::*; - use futures::StreamExt; - use futures::stream::unfold; - - #[tokio::test] - async fn elapsed_compute_future() { - async fn cheap() { - tokio::time::sleep(Duration::from_millis(1)).await; - } - - async fn expensive() { - let mut _count = 0f64; - for i in 0..100000 { - tokio::task::yield_now().await; - _count /= i as f64 - } - } - - let cheap_time = Time::new(); - cheap().with_elapsed_compute(cheap_time.clone()).await; - println!("cheap future: {}", cheap_time.value()); - - let expensive_time = Time::new(); - expensive() - .with_elapsed_compute(expensive_time.clone()) - .await; - println!("expensive future: {}", expensive_time.value()); - - assert!(expensive_time.value() > cheap_time.value()); - } - - #[tokio::test] - async fn elapsed_compute_stream() { - fn cheap() -> impl Stream { - unfold(0i64, |state| async move { - if state < 10 { - tokio::time::sleep(Duration::from_micros(10)).await; - Some((state, state + 1)) - } else { - None - } - }) - } - - fn expensive() -> impl Stream { - unfold(0i64, |state| async move { - if state < 10 { - // Simulate expensive computation - let mut _count = 0f64; - for i in 1..100000 { - _count += (i as f64).sqrt(); - } - tokio::task::yield_now().await; - Some((state, state + 1)) - } else { - None - } - }) - } - - let cheap_time = Time::new(); - cheap() - .with_elapsed_compute(cheap_time.clone()) - .collect::>() - .await; - println!("cheap future: {}", cheap_time.value()); - - let expensive_time = Time::new(); - expensive() - .with_elapsed_compute(expensive_time.clone()) - .collect::>() - .await; - println!("expensive future: {}", expensive_time.value()); - - assert!(expensive_time.value() > cheap_time.value()); + Self::new(self.lazy_stream_groups.len()) } } diff --git a/src/worker/worker_service.rs b/src/worker/worker_service.rs index 5d9c375f..3a285f11 100644 --- a/src/worker/worker_service.rs +++ b/src/worker/worker_service.rs @@ -1,17 +1,5 @@ -use crate::worker::WorkerSessionBuilder; -use crate::worker::generated::worker::worker_service_server::{WorkerService, WorkerServiceServer}; -use crate::worker::generated::worker::{ - CoordinatorToWorkerMsg, ExecuteTaskRequest, TaskKey, WorkerToCoordinatorMsg, -}; -use crate::worker::impl_execute_task::execute_remote_task; -use crate::worker::single_write_multi_read::SingleWriteMultiRead; -use crate::worker::task_data::TaskData; -use crate::{ - DefaultSessionBuilder, GetWorkerInfoRequest, GetWorkerInfoResponse, ObservabilityServiceImpl, - ObservabilityServiceServer, WorkerResolver, -}; -use arrow_flight::FlightData; -use async_trait::async_trait; +use crate::worker::{SingleWriteMultiRead, WorkerSessionBuilder}; +use crate::{DefaultSessionBuilder, TaskData, TaskKey}; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ExecutionPlan; @@ -20,8 +8,6 @@ use moka::future::Cache; use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; -use tonic::codegen::BoxStream; -use tonic::{Request, Response, Status, Streaming}; const TASK_CACHE_TTI: Duration = Duration::from_mins(10); @@ -44,10 +30,10 @@ pub struct Worker { /// TTL-based cache for task execution data. Entries are automatically evicted after /// TASK_CACHE_TTI seconds. This prevents memory leaks from abandoned or incomplete queries /// while allowing concurrent access to task results across multiple partition requests. - pub(super) task_data_entries: Arc, + pub(crate) task_data_entries: Arc, pub(super) session_builder: Arc, pub(super) hooks: WorkerHooks, - pub(super) max_message_size: Option, + pub(crate) max_message_size: Option, pub(super) version: Cow<'static, str>, } @@ -124,60 +110,17 @@ impl Worker { self } - /// Converts this [Worker] into a [`WorkerServiceServer`] with high default message size limits. - /// - /// This is a convenience method that wraps the endpoint in a [`WorkerServiceServer`] and - /// configures it with `max_decoding_message_size(usize::MAX)` and - /// `max_encoding_message_size(usize::MAX)` to avoid message size limitations for internal - /// communication. - /// - /// You can further customize the returned server by chaining additional tonic methods. - /// - /// # Example - /// - /// ``` - /// # use datafusion_distributed::Worker; - /// # use tonic::transport::Server; - /// # use std::net::{IpAddr, Ipv4Addr, SocketAddr}; - /// # async fn f() { - /// - /// let worker = Worker::default(); - /// let server = worker.into_worker_server(); - /// - /// Server::builder() - /// .add_service(Worker::default().into_worker_server()) - /// .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)) - /// .await; - /// - /// # } - /// ``` - pub fn into_worker_server(self) -> WorkerServiceServer { - WorkerServiceServer::new(self) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX) - } - - /// Creates an [`ObservabilityServiceServer`] that exposes task progress and cluster - /// worker discovery via the provided [`WorkerResolver`]. - /// - /// The returned server is meant to be added to the same [`tonic::transport::Server`] as the - /// Flight service — gRPC multiplexes both services on a single port. - pub fn with_observability_service( - &self, - worker_resolver: Arc, - ) -> ObservabilityServiceServer { - ObservabilityServiceServer::new(ObservabilityServiceImpl::new( - self.task_data_entries.clone(), - worker_resolver, - )) - } - /// Sets a version string reported by the `GetWorkerInfo` gRPC endpoint. pub fn with_version(mut self, version: impl Into>) -> Self { self.version = version.into(); self } + /// Returns the version set by [Self::with_version]. + pub fn version(&self) -> &str { + &self.version + } + /// Returns the number of cached task entries currently held by this worker. #[cfg(any(test, feature = "integration"))] pub async fn tasks_running(&self) -> usize { @@ -187,37 +130,3 @@ impl Worker { self.task_data_entries.entry_count() as usize } } - -/// Implementation of the `worker.proto` specification based on the generated Rust stubs. -/// -/// The methods are delegated to plan `impl Worker` implementations so that they can be implemented -/// in different files. -#[async_trait] -impl WorkerService for Worker { - type CoordinatorChannelStream = BoxStream; - - async fn coordinator_channel( - &self, - request: Request>, - ) -> Result, Status> { - self.impl_coordinator_channel(request).await - } - - type ExecuteTaskStream = BoxStream; - - async fn execute_task( - &self, - request: Request, - ) -> Result, Status> { - execute_remote_task(&self.task_data_entries, request).await - } - - async fn get_worker_info( - &self, - _request: Request, - ) -> Result, Status> { - Ok(Response::new(GetWorkerInfoResponse { - version: self.version.to_string(), - })) - } -} diff --git a/src/networking/worker_resolver.rs b/src/worker_resolver.rs similarity index 100% rename from src/networking/worker_resolver.rs rename to src/worker_resolver.rs