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 1c807d99a [CELEBORN-2330] Fix HA master bootstrap redirect handling
1c807d99a is described below
commit 1c807d99a16e59fca0022b1524ec158bfa869514
Author: Chao Sun <[email protected]>
AuthorDate: Mon May 25 13:59:36 2026 +0800
[CELEBORN-2330] Fix HA master bootstrap redirect handling
## Why are the changes needed?
In HA mode, clients can hit a failover window where the master they contact
is no longer the leader and returns a `MasterNotLeaderException` with a
suggested leader address. The client-facing symptom is an RPC failure surfaced
as:
`CelebornException: Exception thrown in awaitResult`
The redirect signal is still present underneath that wrapper, but the
existing bootstrap and retry logic does not consistently preserve and follow
it. In particular:
- bootstrap-time redirects can be treated like generic connection failures
instead of explicit leader hints
- suggested leaders can themselves redirect again or fail setup
- after such failures, the client may not continue cleanly to the remaining
configured masters
The consequence is that a client can fail to establish a master connection
during HA leader transitions even when a reachable leader or another configured
master is available. That turns a recoverable redirect/failover event into an
avoidable client-visible failure and makes rolling upgrades or leader changes
noisier than necessary.
## What changes were proposed in this PR?
This port brings the HA redirect handling fix from `openai/celeborn#70`
onto upstream `main`.
When a master tells the client which leader to use, the client now keeps
that redirect information even if it is wrapped inside another exception, and
it actively follows the suggested leader during bootstrap and failover. If that
suggested leader points to another leader, the client can continue along that
redirect chain instead of giving up too early.
If the redirect path is no longer useful, for example because the suggested
leader cannot be reached, no leader is currently presented, or redirects start
looping, the client falls back to the remaining configured masters and keeps
searching for a usable endpoint. The retry scan also advances correctly after
endpoint setup fails, so one bad redirect does not prevent the client from
trying the next viable master.
The PR also adds focused HA tests that cover the bootstrap redirect cases,
chained redirects, redirect cycles, and fallback to configured masters.
## How was this PR tested?
- `build/mvn test -pl common -am -Dtest=MasterClientSuiteJ
-DwildcardSuites=org.apache.celeborn.common.client.__NoSuchSuite__`
Closes #3691 from sunchao/dev/chao/codex/port-pr70-to-oss-main.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: SteNicholas <[email protected]>
---
.../celeborn/common/client/MasterClient.java | 78 ++++++-
.../celeborn/common/client/MasterClientSuiteJ.java | 259 ++++++++++++++++++++-
2 files changed, 326 insertions(+), 11 deletions(-)
diff --git
a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
index 626c112a5..e57460c49 100644
--- a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
+++ b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
@@ -169,16 +169,18 @@ public class MasterClient {
}
private boolean shouldRetry(@Nullable RpcEndpointRef oldRef, Throwable e) {
- // It will always throw celeborn exception , so we need to get the cause
- // 'CelebornException: Exception thrown in awaitResult'
- if (e.getCause() instanceof MasterNotLeaderException) {
- MasterNotLeaderException exception = (MasterNotLeaderException)
e.getCause();
+ // A redirect can arrive wrapped by awaitResult or by endpoint setup
during bootstrap.
+ // Search the full cause chain so both paths can retry against the
suggested leader.
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception != null) {
String leaderAddr =
isWorker
? exception.getSuggestedInternalLeaderAddress()
: exception.getSuggestedLeaderAddress();
if (!leaderAddr.equals(MasterNotLeaderException.LEADER_NOT_PRESENTED)) {
- setRpcEndpointRef(leaderAddr);
+ if (!setRpcEndpointRef(leaderAddr)) {
+ resetRpcEndpointRef(oldRef);
+ }
} else {
LOG.warn("Master leader is not present currently, please check
masters' status!");
resetRpcEndpointRef(oldRef);
@@ -191,13 +193,55 @@ public class MasterClient {
return false;
}
- private void setRpcEndpointRef(String masterEndpoint) {
+ @Nullable
+ private MasterNotLeaderException findMasterNotLeaderException(Throwable
throwable) {
+ Throwable current = throwable;
+ while (current != null) {
+ if (current instanceof MasterNotLeaderException) {
+ return (MasterNotLeaderException) current;
+ }
+ current = current.getCause();
+ }
+ return null;
+ }
+
+ private boolean setRpcEndpointRef(String masterEndpoint) {
// This method should never care newer or old value, we just set the
suggested master endpoint.
// If an error occurs when setting the suggested Master, it means that the
Master may be down.
// At this time, we just set `rpcEndpointRef` to null. Then next time, we
will re-select the
// Master and get the correct leader.
- rpcEndpointRef.set(setupEndpointRef(masterEndpoint));
- LOG.info("Fail over to master {}.", masterEndpoint);
+ String nextMasterEndpoint = masterEndpoint;
+ Set<String> triedMasterEndpoints = new HashSet<>();
+ while (triedMasterEndpoints.add(nextMasterEndpoint)) {
+ try {
+ RpcEndpointRef endpointRef = setupEndpointRef(nextMasterEndpoint);
+ if (endpointRef != null) {
+ rpcEndpointRef.set(endpointRef);
+ LOG.info("Fail over to master {}.", nextMasterEndpoint);
+ return true;
+ }
+ break;
+ } catch (RuntimeException e) {
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception == null) {
+ break;
+ }
+
+ String leaderAddr =
+ isWorker
+ ? exception.getSuggestedInternalLeaderAddress()
+ : exception.getSuggestedLeaderAddress();
+ if (MasterNotLeaderException.LEADER_NOT_PRESENTED.equals(leaderAddr)) {
+ break;
+ }
+ nextMasterEndpoint = leaderAddr;
+ }
+ }
+ rpcEndpointRef.set(null);
+ LOG.info(
+ "Fail over to master {} failed during endpoint setup; will retry with
another master.",
+ masterEndpoint);
+ return false;
}
private void resetRpcEndpointRef(@Nullable RpcEndpointRef oldRef) {
@@ -225,6 +269,8 @@ public class MasterClient {
* @param currentIndex current attempt master address index.
* @throws IllegalStateException If after several attempts, the non-empty
RpcEndpointRef still
* cannot be obtained.
+ * @throws RuntimeException If endpoint setup receives a leader redirect
that should be retried by
+ * the outer send loop.
* @return non-empty RpcEndpointRef.
*/
private RpcEndpointRef getOrSetupRpcEndpointRef(AtomicInteger currentIndex) {
@@ -241,7 +287,13 @@ public class MasterClient {
if (endpointRef == null) {
int index = currentIndex.get();
do {
- RpcEndpointRef tempEndpointRef =
setupEndpointRef(activeMasterEndpoints.get(index));
+ RpcEndpointRef tempEndpointRef;
+ try {
+ tempEndpointRef = setupEndpointRef(activeMasterEndpoints.get(index));
+ } catch (RuntimeException e) {
+ currentIndex.set((index + 1) % activeMasterEndpoints.size());
+ throw e;
+ }
if (rpcEndpointRef.compareAndSet(null, tempEndpointRef)) {
index = (index + 1) % activeMasterEndpoints.size();
}
@@ -269,6 +321,14 @@ public class MasterClient {
rpcEnv.setupEndpointRef(
RpcAddress.fromHostAndPort(endpoint),
masterEndpointResolver.masterEndpointName());
} catch (Exception e) {
+ MasterNotLeaderException exception = findMasterNotLeaderException(e);
+ if (exception != null
+ && !MasterNotLeaderException.LEADER_NOT_PRESENTED.equals(
+ isWorker
+ ? exception.getSuggestedInternalLeaderAddress()
+ : exception.getSuggestedLeaderAddress())) {
+ throw new RuntimeException(e);
+ }
// Catch all exceptions. Because we don't care whether this exception is
IOException or
// TimeoutException or other exceptions, so we just try to connect to
host:port, if fail,
// we try next address.
diff --git
a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
index bb74d99b2..26ad11c43 100644
---
a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
@@ -245,6 +245,257 @@ public class MasterClientSuiteJ {
checkOneMasterAskFailedInHA(new RpcTimeoutException("test", new
TimeoutException("test")));
}
+ @Test
+ public void testBootstrapMasterNotLeaderRedirectsToSuggestedLeaderInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(),
"host1:9097,host3:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master2)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host1:9097", "host2:9097",
null));
+ case "host2":
+ return master2;
+ case "host3":
+ throw new AssertionError("Should follow the suggested leader
before host3.");
+ default:
+ fail("Should use master host1/host2:" + masterPort + ", but
use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should redirect to the suggested leader during
bootstrap.", t);
+ fail("It should redirect to the suggested leader during bootstrap.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void
testBootstrapMasterNotLeaderWithoutLeaderContinuesToOtherConfiguredMastersInHA()
{
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+
+ final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master2)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException(
+ "host1:9097",
MasterNotLeaderException.LEADER_NOT_PRESENTED, null));
+ case "host2":
+ return master2;
+ default:
+ fail("Should use master host1/host2:" + masterPort + ", but
use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should continue to another configured master during
bootstrap.", t);
+ fail("It should continue to another configured master during
bootstrap.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void testSuggestedLeaderSetupFailureRetriesAnotherMasterInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(),
"host1:9097,host4:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+ final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(
+ Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097",
"host2:9097", null)))
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master3)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ return master1;
+ case "host2":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host2:9097", "host3:9097",
null));
+ case "host3":
+ return master3;
+ case "host4":
+ throw new AssertionError("Should follow chained suggested
leaders before host4.");
+ default:
+ fail(
+ "Should use master host1/host2/host3:" + masterPort + ",
but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry with another master when suggested leader
setup races.", t);
+ fail("It should retry with another master when suggested leader setup
races.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
+ @Test
+ public void testSuggestedLeaderRedirectCycleRetriesConfiguredMasterInHA() {
+ final CelebornConf conf =
+ conf()
+ .set(CelebornConf.HA_ENABLED().key(), "true")
+ .set(CelebornConf.MASTER_ENDPOINTS().key(),
"host1:9097,host4:9097")
+ .set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
+
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+ final RpcEndpointRef master4 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(
+ Future$.MODULE$.failed(new MasterNotLeaderException("host1:9097",
"host2:9097", null)))
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master4)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ return master1;
+ case "host2":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host2:9097", "host3:9097",
null));
+ case "host3":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host3:9097", "host2:9097",
null));
+ case "host4":
+ return master4;
+ default:
+ fail(
+ "Should use master host1/host2/host3/host4:"
+ + masterPort
+ + ", but use "
+ + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry another configured master after a redirect
cycle.", t);
+ fail("It should retry another configured master after a redirect
cycle.");
+ }
+
+ assertEquals(mockResponse, response);
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host2:9097")),
Mockito.anyString());
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host3:9097")),
Mockito.anyString());
+ }
+
+ @Test
+ public void
testBootstrapSuggestedLeaderConnectionFailureRetriesConfiguredMasterInHA() {
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+
+ final RpcEndpointRef master3 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(Future$.MODULE$.successful(mockResponse))
+ .when(master3)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ switch (address.host()) {
+ case "host1":
+ throw new RuntimeException(
+ new MasterNotLeaderException("host1:9097", "host2:9097",
null));
+ case "host2":
+ throw new IOException("test");
+ case "host3":
+ return master3;
+ default:
+ fail(
+ "Should use master host1/host2/host3:" + masterPort + ",
but use " + address);
+ }
+ return null;
+ })
+ .when(rpcEnv)
+ .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
+
+ MasterClient client = new MasterClient(rpcEnv, conf, false);
+ HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
+
+ HeartbeatFromWorkerResponse response = null;
+ try {
+ response = client.askSync(message, HeartbeatFromWorkerResponse.class);
+ } catch (Throwable t) {
+ LOG.error("It should retry another configured master after bootstrap
redirect fails.", t);
+ fail("It should retry another configured master after bootstrap redirect
fails.");
+ }
+
+ assertEquals(mockResponse, response);
+ }
+
private void checkOneMasterDownInHA(Exception causedByException) {
final CelebornConf conf = prepareForCelebornConfWithHA();
@@ -440,15 +691,19 @@ public class MasterClientSuiteJ {
}
private CelebornConf prepareForCelebornConfWithoutHA() {
- return conf.clone()
+ return conf()
.set(CelebornConf.HA_ENABLED().key(), "false")
.set(CelebornConf.MASTER_ENDPOINTS().key(), masterHost + ":" +
masterPort);
}
private CelebornConf prepareForCelebornConfWithHA() {
- return conf.clone()
+ return conf()
.set(CelebornConf.HA_ENABLED().key(), "true")
.set(CelebornConf.MASTER_ENDPOINTS().key(),
"host1:9097,host2:9097,host3:9097")
.set(CelebornConf.MASTER_CLIENT_MAX_RETRIES().key(), "5");
}
+
+ private CelebornConf conf() {
+ return conf.clone();
+ }
}
