fix(executor): bind the gRPC listener before registering with the scheduler

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());
+    }
 }

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(_) => {