[SPARK-53296][CORE] ESS exit main thread in case boss thread exits

common/network-common/src/main/java/org/apache/spark/network/TransportContext.java

@@ -22,6 +22,7 @@
 
 import java.io.Closeable;
 import java.util.ArrayList;
+import java.util.function.Function;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -164,7 +165,12 @@ public TransportClientFactory createClientFactory() {
 
   /** Create a server which will attempt to bind to a specific port. */
   public TransportServer createServer(int port, List<TransportServerBootstrap> bootstraps) {
-    return new TransportServer(this, null, port, rpcHandler, bootstraps);
+    return createServer(port, bootstraps, null);
+  }
+
+  /** Create a server which will attempt to bind to a specific port */
+  public TransportServer createServer(int port, List<TransportServerBootstrap> bootstraps, Function<Throwable, Void> onUncaughtException) {
+    return new TransportServer(this, null, port, rpcHandler, bootstraps, onUncaughtException);
   }
 
   /** Create a server which will attempt to bind to a specific host and port. */

common/network-common/src/main/java/org/apache/spark/network/server/TransportServer.java

@@ -19,6 +19,7 @@
 
 import java.io.Closeable;
 import java.net.InetSocketAddress;
+import java.util.function.Function;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -67,6 +68,16 @@ public TransportServer(
       int portToBind,
       RpcHandler appRpcHandler,
       List<TransportServerBootstrap> bootstraps) {
+    this(context, hostToBind, portToBind, appRpcHandler, bootstraps, null);
+  }
+
+  public TransportServer(
+          TransportContext context,
+          String hostToBind,
+          int portToBind,
+          RpcHandler appRpcHandler,
+          List<TransportServerBootstrap> bootstraps,
+          Function<Throwable, Void> onUncaughtException) {
     this.context = context;
     this.conf = context.getConf();
     this.appRpcHandler = appRpcHandler;
@@ -81,7 +92,7 @@ public TransportServer(
 
     boolean shouldClose = true;
     try {
-      init(hostToBind, portToBind);
+      init(hostToBind, portToBind, onUncaughtException);
       shouldClose = false;
     } finally {
       if (shouldClose) {
@@ -97,11 +108,11 @@ public int getPort() {
     return port;
   }
 
-  private void init(String hostToBind, int portToBind) {
+  private void init(String hostToBind, int portToBind, Function<Throwable, Void> onUncaughtException) {
 
     IOMode ioMode = IOMode.valueOf(conf.ioMode());
-    EventLoopGroup bossGroup = NettyUtils.createEventLoop(ioMode, 1,
-      conf.getModuleName() + "-boss");
+    EventLoopGroup bossGroup = NettyUtils.createBossEventLoop(ioMode, 1,
+            conf.getModuleName() + "-boss", onUncaughtException);
     EventLoopGroup workerGroup =  NettyUtils.createEventLoop(ioMode, conf.serverThreads(),
       conf.getModuleName() + "-server");
 

common/network-common/src/main/java/org/apache/spark/network/util/BossThreadFactory.java

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.util;
+
+import io.netty.util.concurrent.DefaultThreadFactory;
+
+import java.util.function.Function;
+
+public class BossThreadFactory extends DefaultThreadFactory {
+    private final Function<Throwable, Void> onUncaughtException;
+
+    public BossThreadFactory(String threadPoolPrefix, boolean daemon, Function<Throwable, Void> onUncaughtException) {
+        super(threadPoolPrefix, daemon);
+        this.onUncaughtException = onUncaughtException;
+    }
+
+    @Override
+    public Thread newThread(Runnable r) {
+        Thread t = super.newThread(r);
+        t.setUncaughtExceptionHandler((thread, throwable) -> {
+            if (onUncaughtException != null) {
+                try {
+                    onUncaughtException.apply(throwable);
+                } catch (Exception e) {
+                    System.err.println("Exception in onUncaughtException handler: " + e);
+                }
+            }
+            System.err.println("Uncaught exception in thread " + thread.getName() + ": " + throwable);
+        });
+        return t;
+    }
+}

common/network-common/src/main/java/org/apache/spark/network/util/NettyUtils.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.util;
 
 import java.util.concurrent.ThreadFactory;
+import java.util.function.Function;
 
 import io.netty.buffer.PooledByteBufAllocator;
 import io.netty.channel.Channel;
@@ -61,6 +62,20 @@ public static ThreadFactory createThreadFactory(String threadPoolPrefix) {
     return new DefaultThreadFactory(threadPoolPrefix, true);
   }
 
+  public static EventLoopGroup createBossEventLoop(IOMode mode, int numThreads, String threadPrefix, Function<Throwable, Void> onUncaughtException) {
+    ThreadFactory threadFactory;
+    if (onUncaughtException == null) {
+      threadFactory = createThreadFactory(threadPrefix);
+    } else {
+      threadFactory = new BossThreadFactory(threadPrefix, true, onUncaughtException);
+    }
+
+    return switch (mode) {
+      case NIO -> new NioEventLoopGroup(numThreads, threadFactory);
+      case EPOLL -> new EpollEventLoopGroup(numThreads, threadFactory);
+    };
+  }
+
   /** Creates a Netty EventLoopGroup based on the IOMode. */
   public static EventLoopGroup createEventLoop(IOMode mode, int numThreads, String threadPrefix) {
     ThreadFactory threadFactory = createThreadFactory(threadPrefix);

common/network-common/src/test/java/org/apache/spark/network/util/BossThreadFactorySuite.java

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.util;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.jupiter.api.Test;
+import java.util.function.Function;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class BossThreadFactorySuite {
+
+    @Test
+    public void testBossThreadFactory() throws InterruptedException {
+        AtomicBoolean wasCalled = new AtomicBoolean(false);
+
+        Function<Throwable, Void> onUncaughtException = (throwable) -> {
+            System.out.println("Uncaught exception handler invoked!");
+            wasCalled.set(true);
+            return null;
+        };
+
+        BossThreadFactory factory = new BossThreadFactory("test", false, onUncaughtException);
+
+        Runnable faultyTask = () -> {
+            throw new RuntimeException("Test exception from thread!");
+        };
+
+        Thread thread = factory.newThread(faultyTask);
+        thread.start();
+        thread.join();
+
+        assertTrue(wasCalled.get(), "Test failed: onUncaughtException was not called.");
+    }
+}

core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala

@@ -19,6 +19,7 @@ package org.apache.spark.deploy
 
 import java.io.File
 import java.util.concurrent.CountDownLatch
+import java.util.function.Function
 
 import scala.jdk.CollectionConverters._
 
@@ -59,8 +60,9 @@ class ExternalShuffleService(sparkConf: SparkConf, securityManager: SecurityMana
       numUsableCores = 0,
       sslOptions = Some(securityManager.getRpcSSLOptions()))
   private val blockHandler = newShuffleBlockHandler(transportConf)
+  @volatile
   private var transportContext: TransportContext = _
-
+  @volatile
   private var server: TransportServer = _
 
   private val shuffleServiceSource = new ExternalShuffleServiceSource
@@ -103,7 +105,7 @@ class ExternalShuffleService(sparkConf: SparkConf, securityManager: SecurityMana
   }
 
   /** Start the external shuffle service */
-  def start(): Unit = {
+  def start(essServerStopFunc: Function[Throwable, Void] = null): Unit = {
     require(server == null, "Shuffle server already started")
     val authEnabled = securityManager.isAuthenticationEnabled()
     logInfo(log"Starting shuffle service on port ${MDC(PORT, port)}" +
@@ -115,7 +117,7 @@ class ExternalShuffleService(sparkConf: SparkConf, securityManager: SecurityMana
         Nil
       }
     transportContext = new TransportContext(transportConf, blockHandler, true)
-    server = transportContext.createServer(port, bootstraps.asJava)
+    server = transportContext.createServer(port, bootstraps.asJava, essServerStopFunc)
 
     shuffleServiceSource.registerMetricSet(server.getAllMetrics)
     blockHandler.getAllMetrics.getMetrics.put("numRegisteredConnections",
@@ -177,7 +179,16 @@ object ExternalShuffleService extends Logging {
     // and we assume the user really wants it to be running
     sparkConf.set(config.SHUFFLE_SERVICE_ENABLED.key, "true")
     server = newShuffleService(sparkConf, securityManager)
-    server.start()
+    val essServerStopFunc: Function[Throwable, Void] = (t: Throwable) => {
+      logError("Boss thread exited with uncaught exception, shutting down shuffle service.", t)
+      server.stop()
+      barrier.countDown()
+      // Exit with a non-zero code to indicate boss thread failure.
+      System.exit(1)
+      null
+    }
+
+    server.start(essServerStopFunc)
 
     logDebug("Adding shutdown hook") // force eager creation of logger
     ShutdownHookManager.addShutdownHook { () =>