You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/08/04 23:35:09 UTC

(datafusion-ballista) branch main updated: fix(executor): bind the gRPC listener before registering with the scheduler (#2225)

This is an automated email from the ASF dual-hosted git repository.

andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new 984513ee7 fix(executor): bind the gRPC listener before registering with the scheduler (#2225)
984513ee7 is described below

commit 984513ee7eb4cbf346da5252703a8fe8485556b7
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Aug 4 17:35:05 2026 -0600

    fix(executor): bind the gRPC listener before registering with the scheduler (#2225)
    
    * fix(executor): bind the gRPC listener before registering with the scheduler
    
    The executor spawned its tonic server and registered with the scheduler
    immediately afterwards. Server::serve binds inside the future it returns,
    so the socket did not exist yet, while the scheduler's registration
    handler dials that port back via a single connect with no retry. Losing
    that race gave ECONNREFUSED, failed registration, and killed the executor
    at startup.
    
    Bind the listener on the current task before spawning, and serve from it
    with serve_with_incoming_shutdown, so registration cannot run before the
    port accepts connections. tonic ignores the builder's tcp_nodelay and
    tcp_keepalive when serving from a pre-bound listener, so the new
    create_grpc_server_incoming helper applies them itself.
    
    * refactor: tighten create_grpc_server_incoming docs and tests
    
    * test: restore port-in-use test for create_grpc_server_incoming
---
 ballista/core/src/utils.rs               | 59 ++++++++++++++++++++++++++++++++
 ballista/executor/src/executor_server.rs | 14 ++++++--
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs
index 56cef55f4..07f5ca394 100644
--- a/ballista/core/src/utils.rs
+++ b/ballista/core/src/utils.rs
@@ -33,11 +33,13 @@ use datafusion::physical_plan::{ExecutionPlan, RecordBatchStream, metrics};
 use futures::StreamExt;
 use log::error;
 use std::io::BufWriter;
+use std::net::SocketAddr;
 use std::path::Path;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use std::{fs::File, pin::Pin};
 use tonic::codegen::StdError;
+use tonic::transport::server::TcpIncoming;
 use tonic::transport::{Channel, Endpoint, Error, Server};
 
 /// Configuration for gRPC client connections.
@@ -356,6 +358,30 @@ pub fn create_grpc_server(config: &GrpcServerConfig) -> Server {
         )))
 }
 
+/// Binds a gRPC server's listening socket, for use with tonic's
+/// `serve_with_incoming` / `serve_with_incoming_shutdown`. Unlike tonic's
+/// `serve`, which binds lazily inside the future it returns, the socket is
+/// listening by the time this returns — so use this whenever a peer may be
+/// told to connect as soon as the server is started.
+///
+/// tonic ignores the builder's `tcp_nodelay` and `tcp_keepalive` when serving
+/// from a pre-bound listener, so this applies the same values that
+/// [`create_grpc_server`] sets, for the same reasons. The remaining settings
+/// still come from the builder.
+///
+/// # Panics
+///
+/// The listener is registered with the Tokio reactor, so this must be called
+/// from within a Tokio runtime.
+pub fn create_grpc_server_incoming(
+    addr: SocketAddr,
+    config: &GrpcServerConfig,
+) -> Result<TcpIncoming> {
+    Ok(TcpIncoming::bind(addr)?
+        .with_nodelay(Some(true))
+        .with_keepalive(Some(Duration::from_secs(config.tcp_keepalive_seconds))))
+}
+
 /// Recursively collects metrics from an execution plan and all its children.
 pub fn collect_plan_metrics(plan: &dyn ExecutionPlan) -> Vec<MetricsSet> {
     let mut metrics_array = Vec::<MetricsSet>::new();
@@ -440,4 +466,37 @@ mod tests {
         let result = create_grpc_client_endpoint("not a valid url", None);
         assert!(result.is_err());
     }
+
+    /// The point of binding up front is that the port is reachable before
+    /// anything is served on it, so a peer told to connect back cannot arrive
+    /// too early.
+    #[tokio::test]
+    async fn test_create_grpc_server_incoming_binds_eagerly() {
+        let incoming = create_grpc_server_incoming(
+            "127.0.0.1:0".parse().unwrap(),
+            &GrpcServerConfig::default(),
+        )
+        .expect("bind");
+        let addr = incoming.local_addr().expect("local addr");
+
+        // `incoming` is never handed to a server, and yet:
+        tokio::net::TcpStream::connect(addr)
+            .await
+            .expect("port is already listening");
+    }
+
+    /// Binding eagerly means a port conflict surfaces here, as an error, rather
+    /// than later inside the spawned server task.
+    #[tokio::test]
+    async fn test_create_grpc_server_incoming_port_in_use() {
+        let first = create_grpc_server_incoming(
+            "127.0.0.1:0".parse().unwrap(),
+            &GrpcServerConfig::default(),
+        )
+        .expect("bind");
+        let addr = first.local_addr().expect("local addr");
+
+        let result = create_grpc_server_incoming(addr, &GrpcServerConfig::default());
+        assert!(result.is_err());
+    }
 }
diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs
index aecb6644d..0d56417e8 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -53,7 +53,9 @@ use ballista_core::serde::scheduler::TaskKey;
 use ballista_core::serde::scheduler::from_proto::{
     get_task_definition, get_task_definition_vec,
 };
-use ballista_core::utils::{create_grpc_client_endpoint, create_grpc_server};
+use ballista_core::utils::{
+    create_grpc_client_endpoint, create_grpc_server, create_grpc_server_incoming,
+};
 
 use dashmap::DashMap;
 use datafusion::execution::TaskContext;
@@ -133,12 +135,19 @@ pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
     );
 
     // 1. Start executor grpc service
+    //
+    // The listening socket is bound here rather than inside the spawned task,
+    // because step 2 registers with the scheduler and the scheduler dials this
+    // port back to check connectivity. Binding lazily inside the server future
+    // let that callback lose the race and get ECONNREFUSED, which fails
+    // registration and takes the executor down with it.
     let server = {
         let executor_meta = executor.metadata.clone();
         let addr = format!("{}:{}", config.bind_host, executor_meta.grpc_port);
         let addr = addr.parse().unwrap();
         let grpc_server_config = config.grpc_server_config.clone();
 
+        let incoming = create_grpc_server_incoming(addr, &grpc_server_config)?;
         info!(
             "Ballista v{BALLISTA_VERSION} Rust Executor Grpc Server listening on {addr:?}"
         );
@@ -150,7 +159,7 @@ pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
             let shutdown_signal = grpc_shutdown.recv();
             let grpc_server_future = create_grpc_server(&grpc_server_config)
                 .add_service(server)
-                .serve_with_shutdown(addr, shutdown_signal);
+                .serve_with_incoming_shutdown(incoming, shutdown_signal);
             grpc_server_future.await.map_err(|e| {
                 error!("Tonic error, Could not start Executor Grpc Server.");
                 BallistaError::TonicError(e)
@@ -159,7 +168,6 @@ pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
     };
 
     // 2. Do executor registration
-    // TODO the executor registration should happen only after the executor grpc server started.
     let executor_server = Arc::new(executor_server);
     match register_executor(&mut scheduler, executor.clone()).await {
         Ok(_) => {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org