fix

Author: kaybhutani

docs/configuration/ha.md

@@ -20,7 +20,7 @@ license: |
 | Key | Default | isDynamic | Description | Since | Deprecated |
 | --- | ------- | --------- | ----------- | ----- | ---------- |
 | celeborn.master.ha.enabled | false | false | When true, master nodes run as Raft cluster mode. | 0.3.0 | celeborn.ha.enabled | 
-| celeborn.master.ha.gracefulShutdown.enabled | false | false | When true, the master will run a shutdown hook and transfer Raft leadership before shutting down. This reduces chances of client side failures by avoiding the Raft election window where no leader is available. | 0.7.0 |  | 
+| celeborn.master.ha.gracefulShutdown.enabled | false | false | When true, the master will transfer Raft leadership before shutting down gracefully. This reduces chances of client side failures by avoiding the Raft election window where no leader is available. | 0.7.0 |  | 
 | celeborn.master.ha.node.&lt;id&gt;.host | &lt;required&gt; | false | Host to bind of master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.&lt;id&gt;.host | 
 | celeborn.master.ha.node.&lt;id&gt;.internal.port | 8097 | false | Internal port for the workers and other masters to bind to a master node <id> in HA mode. | 0.5.0 |  | 
 | celeborn.master.ha.node.&lt;id&gt;.port | 9097 | false | Port to bind of master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.&lt;id&gt;.port | 

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HARaftServer.java

@@ -25,6 +25,7 @@
 import java.util.*;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import javax.net.ssl.KeyManager;
@@ -99,6 +100,7 @@ public class HARaftServer {
   private Optional<LeaderPeerEndpoints> cachedLeaderPeerRpcEndpoints = Optional.empty();
 
   private final CelebornConf conf;
+  private final AtomicBoolean stopped = new AtomicBoolean(false);
   private long workerTimeoutDeadline;
   private long appTimeoutDeadline;
 
@@ -274,8 +276,16 @@ public void start() throws IOException {
   }
 
   public void stop() {
+    stop(false);
+  }
+
+  public void stop(boolean transferLeadership) {
+    if (!stopped.compareAndSet(false, true)) {
+      LOG.info("Raft server {} already stopped.", server.getId());
+      return;
+    }
     try {
-      if (isLeader()) {
+      if (transferLeadership && isLeader()) {
         LOG.info(
             "This node {} is the Raft leader. Transferring leadership before shutdown.",
             server.getId());
@@ -647,16 +657,13 @@ boolean stepDown() {
               REQUEST_TIMEOUT_MS);
       RaftClientReply reply = server.transferLeadership(request);
       if (reply.isSuccess()) {
-        LOG.info("Successfully step down leader {}.", server.getId());
         return true;
-      } else {
-        LOG.warn("Step down leader failed!");
-        return false;
       }
+      LOG.warn("Step down leader {} failed.", server.getId());
     } catch (Exception e) {
-      LOG.warn("Step down leader failed!", e);
-      return false;
+      LOG.warn("Step down leader {} failed.", server.getId(), e);
     }
+    return false;
   }
 
   public void setDeadlineTime(long increaseWorkerTime, long increaseAppTime) {

master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala

@@ -378,26 +378,6 @@ private[celeborn] class Master(
     }
     logInfo("Stopping Celeborn Master.")
 
-    // Transfer Raft leadership before shutting down so other masters can
-    // immediately take over without waiting for heartbeat timeout
-    if (conf.haMasterGracefulShutdownEnabled) {
-      statusSystem match {
-        case ha: HAMasterMetaManager =>
-          val ratisServer = ha.getRatisServer
-          if (ratisServer != null) {
-            logInfo("HA graceful shutdown enabled. Stopping Raft server (will transfer leadership if leader).")
-            try {
-              ratisServer.stop()
-            } catch {
-              case e: Exception =>
-                logError("Failed to stop Raft server during Master shutdown.", e)
-            }
-          }
-        case _ =>
-          logInfo("Single-master mode, skipping Raft shutdown.")
-      }
-    }
-
     Option(checkForWorkerTimeoutTask).foreach(_.cancel(true))
     Option(checkForUnavailableWorkerTimeOutTask).foreach(_.cancel(true))
     Option(checkForApplicationTimeOutTask).foreach(_.cancel(true))
@@ -1607,6 +1587,22 @@ private[celeborn] class Master(
   override def stop(exitKind: Int): Unit = synchronized {
     if (!stopped) {
       logInfo("Stopping Master")
+      // Transfer Raft leadership before shutting down so other masters can
+      // immediately take over without waiting for heartbeat timeout.
+      val transferLeadership = conf.haMasterGracefulShutdownEnabled
+      statusSystem match {
+        case ha: HAMasterMetaManager =>
+          val ratisServer = ha.getRatisServer
+          if (ratisServer != null) {
+            try {
+              ratisServer.stop(transferLeadership)
+            } catch {
+              case e: Exception =>
+                logError("Failed to stop Raft server during Master shutdown.", e)
+            }
+          }
+        case _ => // single-master mode, no Raft server to stop
+      }
       rpcEnv.stop(self)
       if (conf.internalPortEnabled) {
         internalRpcEnvInUse.stop(internalRpcEndpointRef)

master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisMasterStatusSystemSuiteJ.java

@@ -1892,11 +1892,31 @@ public void testGracefulLeaderShutdownStepDown() {
         .getMasterStateMachine()
         .notifyLogFailed(new Exception("test leader graceful step down"), null);
 
-    Assert.assertFalse("Leader should step down without closing the shared server", leader.isLeader());
+    Assert.assertFalse(
+        "Leader should step down without closing the shared server", leader.isLeader());
     Assert.assertNotEquals(
         "Raft server should remain available to the rest of the suite",
         "CLOSED",
         leader.getServer().getLifeCycleState().name());
+
+    // Wait for a new leader to be elected so subsequent tests are not affected.
+    boolean newLeaderElected = false;
+    for (int i = 0; i < 30; i++) {
+      for (HARaftServer server : Arrays.asList(RATISSERVER1, RATISSERVER2, RATISSERVER3)) {
+        if (server != leader && server.isLeader()) {
+          newLeaderElected = true;
+          break;
+        }
+      }
+      if (newLeaderElected) break;
+      try {
+        Thread.sleep(1000);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+    }
+    Assert.assertTrue("A new leader should be elected after step down", newLeaderElected);
   }
 
   @AfterClass