Allow configuring the buffer in the worker connection

Author: gabotechs

src/distributed_ext.rs

@@ -541,6 +541,22 @@ pub trait DistributedExt: Sized {
     /// Same as [DistributedExt::with_distributed_partial_reduce] but with an in-place mutation.
     fn set_distributed_partial_reduce(&mut self, enabled: bool) -> Result<(), DataFusionError>;
 
+    /// Sets the soft byte budget that each per-worker connection will buffer in memory before
+    /// pausing the gRPC pull from that worker. Per-partition channels are unbounded (to avoid
+    /// head-of-line blocking between sibling partitions), so backpressure is enforced globally
+    /// per worker connection using this budget.
+    fn with_distributed_worker_connection_buffer_budget_bytes(
+        self,
+        budget_bytes: usize,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_worker_connection_buffer_budget_bytes] but with
+    /// an in-place mutation.
+    fn set_distributed_worker_connection_buffer_budget_bytes(
+        &mut self,
+        budget_bytes: usize,
+    ) -> Result<(), DataFusionError>;
+
     /// Registers a [WorkUnitFeed] so that Distributed DataFusion can discover it while traversing
     /// plans. For more info, refer to [WorkUnitFeed] docs.
     ///
@@ -693,6 +709,15 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_worker_connection_buffer_budget_bytes(
+        &mut self,
+        budget_bytes: usize,
+    ) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.worker_connection_buffer_budget_bytes = budget_bytes;
+        Ok(())
+    }
+
     fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
     where
         T: ExecutionPlan + 'static,
@@ -775,6 +800,10 @@ impl DistributedExt for SessionConfig {
             #[expr($?;Ok(self))]
             fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
+
             #[call(set_distributed_work_unit_feed)]
             #[expr($;self)]
             fn with_distributed_work_unit_feed<T, P, F>(mut self, getter: F) -> Self
@@ -875,6 +904,11 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($?;Ok(self))]
             fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
+
             fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
             where
                 T: ExecutionPlan + 'static,
@@ -981,6 +1015,11 @@ impl DistributedExt for SessionState {
             #[expr($?;Ok(self))]
             fn with_distributed_partial_reduce(mut self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_worker_connection_buffer_budget_bytes(mut self, budget_bytes: usize) -> Result<Self, DataFusionError>;
+
             fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
             where
                 T: ExecutionPlan + 'static,
@@ -1087,6 +1126,11 @@ impl DistributedExt for SessionContext {
             #[expr($?;Ok(self))]
             fn with_distributed_partial_reduce(self, enabled: bool) -> Result<Self, DataFusionError>;
 
+            fn set_distributed_worker_connection_buffer_budget_bytes(&mut self, budget_bytes: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_worker_connection_buffer_budget_bytes)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_worker_connection_buffer_budget_bytes(self, budget_bytes: usize) -> Result<Self, DataFusionError>;
+
             fn set_distributed_work_unit_feed<T, P, F>(&mut self, getter: F)
             where
                 T: ExecutionPlan + 'static,

src/distributed_planner/distributed_config.rs

@@ -58,6 +58,13 @@ extensions_options! {
         /// Disabled by default because its effectiveness is workload-dependent: it helps when
         /// aggregation significantly reduces cardinality, but adds overhead when it does not.
         pub partial_reduce: bool, default = false
+        /// Soft byte budget that each per-worker connection will buffer in memory before pausing
+        /// the gRPC pull from that worker. Per-partition channels are unbounded (to avoid
+        /// head-of-line blocking between sibling partitions), so backpressure is enforced
+        /// globally per [WorkerConnection] using this budget. A single message larger than this
+        /// budget will still be admitted (otherwise we would livelock), so the actual peak per
+        /// connection is `worker_connection_buffer_budget_bytes + max_message_size`.
+        pub worker_connection_buffer_budget_bytes: usize, default = 64 * 1024 * 1024
         /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to
         /// estimate how many tasks should be spawned for the [Stage] containing the leaf node.
         pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default()

src/worker/worker_connection_pool.rs

@@ -5,7 +5,7 @@ use crate::passthrough_headers::get_passthrough_headers;
 use crate::protobuf::{datafusion_error_to_tonic_status, map_flight_to_datafusion_error};
 use crate::worker::generated::worker::FlightAppMetadata;
 use crate::worker::generated::worker::{ExecuteTaskRequest, TaskKey};
-use crate::{BytesMetricExt, ChannelResolver, Stage};
+use crate::{BytesMetricExt, ChannelResolver, DistributedConfig, Stage};
 use arrow_flight::FlightData;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
@@ -99,11 +99,6 @@ impl WorkerConnectionPool {
 
 type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>;
 
-/// Soft byte budget the demux task will buffer in memory before pausing the gRPC
-/// pull. Per-partition channels are unbounded (to avoid head-of-line blocking
-/// between sibling partitions), so backpressure is enforced globally here instead.
-const PER_CONNECTION_BUFFER_BUDGET_BYTES: usize = 64 * 1024 * 1024;
-
 /// 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.
@@ -138,6 +133,9 @@ impl WorkerConnection {
         metrics: &ExecutionPlanMetricsSet,
     ) -> Result<Self> {
         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 =
@@ -235,7 +233,7 @@ impl WorkerConnection {
                 // 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() >= PER_CONNECTION_BUFFER_BUDGET_BYTES {
+                while memory_reservation.size() >= buffer_budget_bytes {
                     tokio::select! {
                         biased;
                         _ = cancel.cancelled() => return,