Refactor impl_set_plan.rs into impl_coordinator_channel.rs

Author: gabotechs

src/work_unit_feed/mod.rs

@@ -5,7 +5,7 @@ mod work_unit_feed;
 mod work_unit_feed_provider;
 mod work_unit_feed_registry;
 
-pub(crate) use remote_work_unit_feed::{RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedTxs};
+pub(crate) use remote_work_unit_feed::RemoteWorkUnitFeedRegistry;
 pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed};
 
 pub use work_unit::WorkUnit;

src/worker/impl_set_plan.rs src/worker/impl_coordinator_channel.rssrc/worker/impl_set_plan.rs renamed to src/worker/impl_coordinator_channel.rs

@@ -1,71 +1,44 @@
 use crate::common::deserialize_uuid;
-use crate::config_extension_ext::set_distributed_option_extension_from_headers;
-use crate::protobuf::DistributedCodec;
-use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedTxs};
+use crate::work_unit_feed::RemoteWorkUnitFeedRegistry;
+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::{PreOrderTaskMetrics, SetPlanRequest};
-use crate::worker::worker_connection_pool::LocalWorkerContext;
-use crate::{DistributedConfig, DistributedTaskContext, Worker, WorkerQueryContext};
-use datafusion::error::DataFusionError;
-use datafusion::execution::{SessionStateBuilder, TaskContext};
-use datafusion::physical_plan::ExecutionPlan;
+use crate::worker::generated::worker::worker_service_server::WorkerService;
+use crate::worker::generated::worker::{
+    CoordinatorToWorkerMsg, WorkerToCoordinatorMsg, worker_to_coordinator_msg,
+};
+use crate::{
+    DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, Worker,
+    WorkerQueryContext,
+};
+use datafusion::common::DataFusionError;
+use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 use datafusion_proto::protobuf::PhysicalPlanNode;
+use futures::{FutureExt, StreamExt};
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
 use tokio::sync::oneshot;
-use tonic::Status;
-use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
 use url::Url;
 
-#[derive(Clone, Debug)]
-/// TaskData stores state for a single task being executed by this Endpoint. It may be shared
-/// 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<TaskContext>,
-    /// Plan to be executed.
-    pub(crate) plan: Arc<dyn ExecutionPlan>,
-    /// `num_partitions_remaining` is initialized to the total number of partitions in the task (not
-    /// only tasks in the partition group). This is decremented for each request to the endpoint
-    /// for this task. Once this count is zero, the task is likely complete. The task may not be
-    /// complete because it's possible that the same partition was retried and this count was
-    /// decremented more than once for the same partition.
-    pub(super) num_partitions_remaining: Arc<AtomicUsize>,
-    /// 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<std::sync::Mutex<Option<oneshot::Sender<PreOrderTaskMetrics>>>>,
-}
-
-impl TaskData {
-    /// Returns the number of partitions remaining to be processed.
-    pub(crate) fn num_partitions_remaining(&self) -> usize {
-        self.num_partitions_remaining
-            .load(std::sync::atomic::Ordering::Relaxed)
-    }
-
-    /// Returns the total number of partitions in this task.
-    pub(crate) fn total_partitions(&self) -> usize {
-        self.plan.properties().partitioning.partition_count()
-    }
-}
-
 impl Worker {
-    /// Sets the plan for a task and returns a receiver that will yield the collected metrics
-    /// once all partitions of that task have finished executing (or been dropped early).
-    pub(crate) async fn impl_set_plan(
+    pub(super) async fn impl_coordinator_channel(
         &self,
-        request: SetPlanRequest,
-        grpc_headers: MetadataMap,
-    ) -> Result<
-        (
-            RemoteWorkUnitFeedTxs,
-            oneshot::Receiver<PreOrderTaskMetrics>,
-        ),
-        Status,
-    > {
+        request: Request<Streaming<CoordinatorToWorkerMsg>>,
+    ) -> Result<Response<<Worker as WorkerService>::CoordinatorChannelStream>, Status> {
+        let (grpc_headers, _ext, mut body) = request.into_parts();
+
+        // The first message must be a SetPlanRequest.
+        let Some(msg) = body.next().await else {
+            return Err(Status::internal("Empty Coordinator stream"));
+        };
+        let Some(Inner::SetPlanRequest(request)) = msg?.inner else {
+            return Err(Status::internal(
+                "First Coordinator message must be SetPlanRequest",
+            ));
+        };
         let key = request.task_key.ok_or_else(missing("task_key"))?;
 
         let entry = self
@@ -94,11 +67,12 @@ impl Worker {
                 .with_extension(Arc::new(LocalWorkerContext {
                     self_url: Url::parse(&request.target_worker_url)
                         .map_err(|e| DataFusionError::External(Box::new(e)))?,
-                }));
-            set_distributed_option_extension_from_headers::<DistributedConfig>(&mut cfg, &headers)?;
+                }))
+                .with_distributed_option_extension_from_headers::<DistributedConfig>(&headers)?;
 
-            let shuffle_batch_size =
-                DistributedConfig::from_config_options(cfg.options())?.shuffle_batch_size;
+            let d_cfg = DistributedConfig::from_config_options(cfg.options())?;
+            let shuffle_batch_size = d_cfg.shuffle_batch_size;
+            let collect_metrics = d_cfg.collect_metrics;
             if shuffle_batch_size != 0 {
                 cfg = cfg.with_batch_size(shuffle_batch_size);
             }
@@ -129,7 +103,10 @@ impl Worker {
                 plan,
                 task_ctx,
                 num_partitions_remaining: Arc::new(AtomicUsize::new(total_partitions)),
-                metrics_tx: Arc::new(std::sync::Mutex::new(Some(metrics_tx))),
+                metrics_tx: match collect_metrics {
+                    true => Arc::new(std::sync::Mutex::new(Some(metrics_tx))),
+                    false => Arc::new(std::sync::Mutex::new(None)),
+                },
             })
         };
 
@@ -138,7 +115,39 @@ impl Worker {
                 "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it."
             ))
         })?;
-        Ok((remote_work_unit_feed_registry.senders, metrics_rx))
+
+        // Continue reading remaining messages (work unit feed data) in the background.
+        let work_unit_senders = remote_work_unit_feed_registry.senders;
+        tokio::spawn(async move {
+            while let Some(Ok(msg)) = body.next().await {
+                let Some(Inner::WorkUnit(msg)) = msg.inner else {
+                    continue;
+                };
+                let Ok(id) = deserialize_uuid(&msg.id) else {
+                    continue;
+                };
+                let Some(tx) = work_unit_senders.get(&(id, msg.partition as usize)) else {
+                    continue;
+                };
+                if tx.send(Ok(msg.body)).is_err() {
+                    break; // channel closed
+                }
+            }
+        });
+
+        // Stream back the metrics once the task finishes executing.
+        // The oneshot receiver resolves when impl_execute_task sends the collected
+        // metrics after all partitions have finished or been dropped.
+        let metrics_stream = metrics_rx.into_stream();
+        let metrics_stream = metrics_stream.filter_map(|task_metrics| async move {
+            match task_metrics {
+                Ok(task_metrics) => Some(WorkerToCoordinatorMsg {
+                    inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)),
+                }),
+                Err(_) => None, // channel dropped without sending any message
+            }
+        });
+        Ok(Response::new(metrics_stream.map(Ok).boxed()))
     }
 }
 

src/worker/mod.rs

@@ -1,9 +1,10 @@
 pub(crate) mod generated;
+mod impl_coordinator_channel;
 mod impl_execute_task;
-mod impl_set_plan;
 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;
 mod worker_connection_pool;
@@ -18,4 +19,4 @@ pub use session_builder::{
 };
 pub use worker_service::Worker;
 
-pub use impl_set_plan::TaskData;
+pub use task_data::TaskData;

src/worker/task_data.rs

@@ -0,0 +1,39 @@
+use crate::worker::generated::worker::PreOrderTaskMetrics;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use tokio::sync::oneshot;
+
+#[derive(Clone, Debug)]
+/// TaskData stores state for a single task being executed by this Endpoint. It may be shared
+/// 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<TaskContext>,
+    /// Plan to be executed.
+    pub(crate) plan: Arc<dyn ExecutionPlan>,
+    /// `num_partitions_remaining` is initialized to the total number of partitions in the task (not
+    /// only tasks in the partition group). This is decremented for each request to the endpoint
+    /// for this task. Once this count is zero, the task is likely complete. The task may not be
+    /// complete because it's possible that the same partition was retried and this count was
+    /// decremented more than once for the same partition.
+    pub(super) num_partitions_remaining: Arc<AtomicUsize>,
+    /// 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<std::sync::Mutex<Option<oneshot::Sender<PreOrderTaskMetrics>>>>,
+}
+
+impl TaskData {
+    /// Returns the number of partitions remaining to be processed.
+    pub(crate) fn num_partitions_remaining(&self) -> usize {
+        self.num_partitions_remaining
+            .load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    /// Returns the total number of partitions in this task.
+    pub(crate) fn total_partitions(&self) -> usize {
+        self.plan.properties().partitioning.partition_count()
+    }
+}

src/worker/worker_service.rs

@@ -1,13 +1,10 @@
-use crate::common::deserialize_uuid;
 use crate::worker::WorkerSessionBuilder;
-use crate::worker::generated::worker::coordinator_to_worker_msg::Inner;
 use crate::worker::generated::worker::worker_service_server::{WorkerService, WorkerServiceServer};
-use crate::worker::generated::worker::worker_to_coordinator_msg;
 use crate::worker::generated::worker::{
     CoordinatorToWorkerMsg, ExecuteTaskRequest, TaskKey, WorkerToCoordinatorMsg,
 };
-use crate::worker::impl_set_plan::TaskData;
 use crate::worker::single_write_multi_read::SingleWriteMultiRead;
+use crate::worker::task_data::TaskData;
 use crate::{
     DefaultSessionBuilder, ObservabilityServiceImpl, ObservabilityServiceServer, WorkerResolver,
 };
@@ -16,7 +13,6 @@ use async_trait::async_trait;
 use datafusion::common::DataFusionError;
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::physical_plan::ExecutionPlan;
-use futures::StreamExt;
 use moka::future::Cache;
 use std::borrow::Cow;
 use std::sync::Arc;
@@ -186,51 +182,7 @@ impl WorkerService for Worker {
         &self,
         request: Request<Streaming<CoordinatorToWorkerMsg>>,
     ) -> Result<Response<Self::CoordinatorChannelStream>, Status> {
-        let (grpc_headers, _ext, mut body) = request.into_parts();
-
-        // The first message must be a SetPlanRequest.
-        let Some(msg) = body.next().await else {
-            return Err(Status::internal("Empty Coordinator stream"));
-        };
-        let Some(Inner::SetPlanRequest(request)) = msg?.inner else {
-            return Err(Status::internal(
-                "First Coordinator message must be SetPlanRequest",
-            ));
-        };
-        let (work_unit_senders, metrics_rx) = self.impl_set_plan(request, grpc_headers).await?;
-
-        // Continue reading remaining messages (work unit feed data) in the background.
-        tokio::spawn(async move {
-            while let Some(Ok(msg)) = body.next().await {
-                let Some(Inner::WorkUnit(msg)) = msg.inner else {
-                    continue;
-                };
-                let Ok(id) = deserialize_uuid(&msg.id) else {
-                    continue;
-                };
-                let Some(tx) = work_unit_senders.get(&(id, msg.partition as usize)) else {
-                    continue;
-                };
-                if tx.send(Ok(msg.body)).is_err() {
-                    break; // channel closed
-                }
-            }
-        });
-
-        // Stream back the metrics once the task finishes executing.
-        // The oneshot receiver resolves when impl_execute_task sends the collected
-        // metrics after all partitions have finished or been dropped.
-        let stream = futures::stream::once(async move {
-            match metrics_rx.await {
-                Ok(task_metrics) => Ok(WorkerToCoordinatorMsg {
-                    inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)),
-                }),
-                Err(_) => Err(Status::internal(
-                    "Metrics channel closed before metrics were sent",
-                )),
-            }
-        });
-        Ok(Response::new(stream.boxed()))
+        self.impl_coordinator_channel(request).await
     }
 
     type ExecuteTaskStream = BoxStream<FlightData>;