This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new a70c8fddc [CELEBORN-2306] Master adds shutdown hook for ratis stepdown
a70c8fddc is described below
commit a70c8fddc456aeedcf7a5bb94d91fd9e23be278e
Author: Kartikay Bhutani <[email protected]>
AuthorDate: Tue May 12 10:19:34 2026 +0800
[CELEBORN-2306] Master adds shutdown hook for ratis stepdown
### What changes were proposed in this pull request?
- Adds master shutdown hook
- Updates RAFT stepdown to return bool
- Calls RAFT stepdown on manager shutdown to do graceful stepdown
### Why are the changes needed?
- Manager.stop() isnt being called from anywhere, it emmits some logs as
well for "Stopping manager" but since there is no shutdown hook defined, none
of them are logged or the function is called at all
- We faced a certain issue where the leader got removed from service mesh
before shutdown and followers redirected to it because it was still running
(able to send requests but not receive) for a brief period. This method adds a
graceful shutdown option to do a RATIS stepdown before shutting down.
### Does this PR resolve a correctness bug?
No.
### Does this PR introduce _any_ user-facing change?
Yes, adds an additional config.
### How was this patch tested?
Added a tests and validated that.
Closes #3659 from kaybhutani/kartikay/graceful-master-shutdown.
Lead-authored-by: Kartikay Bhutani <[email protected]>
Co-authored-by: Zaynt <[email protected]>
Co-authored-by: kartikay <[email protected]>
Signed-off-by: SteNicholas <[email protected]>
---
.../org/apache/celeborn/common/CelebornConf.scala | 23 +++++++++
docs/configuration/ha.md | 2 +
.../deploy/master/clustermeta/ha/HARaftServer.java | 43 +++++++++++++---
.../celeborn/service/deploy/master/Master.scala | 36 ++++++++++++-
.../ha/RatisMasterStatusSystemSuiteJ.java | 60 ++++++++++++++++++++++
5 files changed, 157 insertions(+), 7 deletions(-)
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 0b35f9ed2..796fdf7d4 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -728,6 +728,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def masterHttpIdleTimeout: Long = get(MASTER_HTTP_IDLE_TIMEOUT)
def haEnabled: Boolean = get(HA_ENABLED)
+ def haMasterGracefulShutdownEnabled: Boolean =
get(HA_MASTER_GRACEFUL_SHUTDOWN_ENABLED)
+ def haMasterGracefulShutdownTimeoutMs: Long =
get(HA_MASTER_GRACEFUL_SHUTDOWN_TIMEOUT)
def haMasterNodeId: Option[String] = get(HA_MASTER_NODE_ID)
@@ -2779,6 +2781,27 @@ object CelebornConf extends Logging {
.booleanConf
.createWithDefault(false)
+ val HA_MASTER_GRACEFUL_SHUTDOWN_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.master.ha.graceful.shutdown.enabled")
+ .categories("ha")
+ .version("0.7.0")
+ .doc("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.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val HA_MASTER_GRACEFUL_SHUTDOWN_TIMEOUT: ConfigEntry[Long] =
+ buildConf("celeborn.master.ha.graceful.shutdown.timeout")
+ .categories("ha")
+ .version("0.7.0")
+ .doc("Timeout for the master graceful shutdown process including " +
+ "Raft leadership transfer. Used as the shutdown hook timeout " +
+ "and the transfer-leadership request timeout.")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("30s")
+
val HA_MASTER_NODE_ID: OptionalConfigEntry[String] =
buildConf("celeborn.master.ha.node.id")
.withAlternative("celeborn.ha.master.node.id")
diff --git a/docs/configuration/ha.md b/docs/configuration/ha.md
index c88c1d2ad..ed8ec7dbe 100644
--- a/docs/configuration/ha.md
+++ b/docs/configuration/ha.md
@@ -20,6 +20,8 @@ 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.graceful.shutdown.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.graceful.shutdown.timeout | 30s | false | Timeout for the
master graceful shutdown process including Raft leadership transfer. Used as
the shutdown hook timeout and the transfer-leadership request timeout. | 0.7.0
| |
| celeborn.master.ha.node.<id>.host | <required> | false | Host to
bind of master node <id> in HA mode. | 0.3.0 |
celeborn.ha.master.node.<id>.host |
| celeborn.master.ha.node.<id>.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.<id>.port | 9097 | false | Port to bind of
master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.<id>.port
|
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HARaftServer.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HARaftServer.java
index ee02d8596..3054f8a03 100644
---
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HARaftServer.java
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HARaftServer.java
@@ -25,6 +25,7 @@ import java.nio.charset.StandardCharsets;
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,9 +276,37 @@ public class HARaftServer {
}
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 (transferLeadership && isLeader()) {
+ LOG.info(
+ "This node {} is the Raft leader. Transferring leadership before
shutdown.",
+ server.getId());
+ long startTime = System.currentTimeMillis();
+ boolean success = stepDown();
+ long elapsed = System.currentTimeMillis() - startTime;
+ if (success) {
+ LOG.info("Successfully transferred leadership from {} in {}ms.",
server.getId(), elapsed);
+ } else {
+ LOG.warn(
+ "Leadership transfer from {} failed after {}ms. "
+ + "Proceeding with shutdown anyway.",
+ server.getId(),
+ elapsed);
+ }
+ }
server.close();
+ LOG.info("Raft server {} closed.", server.getId());
} catch (IOException e) {
+ stopped.set(false);
+ LOG.error("Error while stopping Raft server {}.", server.getId(), e);
throw new RuntimeException(e);
}
}
@@ -616,7 +646,8 @@ public class HARaftServer {
return this.internalRpcEndpoint;
}
- void stepDown() {
+ boolean stepDown() {
+ long timeoutMs = conf.haMasterGracefulShutdownTimeoutMs();
try {
TransferLeadershipRequest request =
new TransferLeadershipRequest(
@@ -625,16 +656,16 @@ public class HARaftServer {
raftGroup.getGroupId(),
CallId.getAndIncrement(),
null,
- REQUEST_TIMEOUT_MS);
+ timeoutMs);
RaftClientReply reply = server.transferLeadership(request);
if (reply.isSuccess()) {
- LOG.info("Successfully step down leader {}.", server.getId());
- } else {
- LOG.warn("Step down leader failed!");
+ return true;
}
+ LOG.warn("Step down leader {} failed.", server.getId());
} catch (Exception e) {
- LOG.warn("Step down leader failed!", e);
+ LOG.warn("Step down leader {} failed.", server.getId(), e);
}
+ return false;
}
public void setDeadlineTime(long increaseWorkerTime, long increaseAppTime) {
diff --git
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
index 2ff7cfade..c87f53739 100644
---
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
+++
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
@@ -51,7 +51,7 @@ import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.quota.ResourceConsumption
import org.apache.celeborn.common.rpc._
import org.apache.celeborn.common.rpc.{RpcSecurityContextBuilder,
ServerSaslContextBuilder}
-import org.apache.celeborn.common.util.{CelebornHadoopUtils, JavaUtils,
PbSerDeUtils, SignalUtils, ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{CelebornExitKind, CelebornHadoopUtils,
JavaUtils, PbSerDeUtils, ShutdownHookManager, SignalUtils, ThreadUtils, Utils}
import org.apache.celeborn.server.common.{HttpService, Service}
import org.apache.celeborn.service.deploy.master.audit.ShuffleAuditLogger
import
org.apache.celeborn.service.deploy.master.clustermeta.SingleMasterMetaManager
@@ -377,6 +377,7 @@ private[celeborn] class Master(
return
}
logInfo("Stopping Celeborn Master.")
+
Option(checkForWorkerTimeoutTask).foreach(_.cancel(true))
Option(checkForUnavailableWorkerTimeOutTask).foreach(_.cancel(true))
Option(checkForApplicationTimeOutTask).foreach(_.cancel(true))
@@ -1564,6 +1565,23 @@ private[celeborn] class Master(
override def initialize(): Unit = {
super.initialize()
logInfo("Master started.")
+
+ // SIGTERM triggers leadership transfer (in stop()) then immediate
HTTP/RPC teardown.
+ // EXIT_IMMEDIATELY is intentional: once Raft leadership is transferred
and RPC is
+ // stopped, there is no need to drain HTTP connections.
+ ShutdownHookManager.get().addShutdownHook(
+ ThreadUtils.newThread(
+ new Runnable {
+ override def run(): Unit = {
+ logInfo("Shutdown hook called for Master.")
+ stop(CelebornExitKind.EXIT_IMMEDIATELY)
+ }
+ },
+ "master-shutdown-hook-thread"),
+ 100,
+ conf.haMasterGracefulShutdownTimeoutMs,
+ java.util.concurrent.TimeUnit.MILLISECONDS)
+
rpcEnv.awaitTermination()
if (conf.internalPortEnabled) {
internalRpcEnvInUse.awaitTermination()
@@ -1573,6 +1591,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)
diff --git
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisMasterStatusSystemSuiteJ.java
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisMasterStatusSystemSuiteJ.java
index 2b52e20d3..01a16a31e 100644
---
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisMasterStatusSystemSuiteJ.java
+++
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisMasterStatusSystemSuiteJ.java
@@ -1872,6 +1872,66 @@ public class RatisMasterStatusSystemSuiteJ {
Assert.assertEquals(STATUSSYSTEM3.registeredShuffleCount(), 8);
}
+ @Test
+ public void testLeaderStepDownOnLogFailed() {
+ // Identify the current leader
+ HARaftServer leader = null;
+ for (HARaftServer server : Arrays.asList(RATISSERVER1, RATISSERVER2,
RATISSERVER3)) {
+ if (server.isLeader()) {
+ leader = server;
+ break;
+ }
+ }
+ Assert.assertNotNull("A leader should exist before the test", leader);
+ Assert.assertTrue("Leader node should report isLeader=true",
leader.isLeader());
+
+ // Do not stop() the shared static server used by the suite, since that
permanently
+ // closes one of RATISSERVER1/2/3 and can break later tests depending on
execution order.
+ // Instead, trigger the leader to step down without closing the underlying
server.
+ leader
+ .getMasterStateMachine()
+ .notifyLogFailed(new Exception("test leader graceful step down"),
null);
+
+ 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",
+ org.apache.ratis.util.LifeCycle.State.CLOSED,
+ leader.getServer().getLifeCycleState());
+
+ // Wait until the cluster has a ready leader again, so subsequent tests
can proceed.
+ // Any of the three peers may win the ensuing election (including the one
that
+ // just stepped down, since TransferLeadershipRequest with a null target
just
+ // demotes the leader to follower and lets the normal election protocol
run).
+ // We require isLeaderReady(), not just isLeader(): a newly elected leader
rejects
+ // writes with LeaderNotReadyException until its no-op log entry for the
new term
+ // has been committed.
+ boolean leaderReady = false;
+ for (int i = 0; i < 60; i++) {
+ for (HARaftServer server : Arrays.asList(RATISSERVER1, RATISSERVER2,
RATISSERVER3)) {
+ if (server.isLeader()) {
+ try {
+ if
(server.getServer().getDivision(server.getGroupId()).getInfo().isLeaderReady())
{
+ leaderReady = true;
+ break;
+ }
+ } catch (IOException e) {
+ // Division not available yet; keep polling.
+ }
+ }
+ }
+ if (leaderReady) break;
+ try {
+ Thread.sleep(1000);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ break;
+ }
+ }
+ Assert.assertTrue(
+ "A ready leader should be elected after step down so later tests can
run", leaderReady);
+ }
+
@AfterClass
public static void testNotifyLogFailed() {
List<HARaftServer> list = Arrays.asList(RATISSERVER1, RATISSERVER2,
RATISSERVER3);
