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 17159eb34a [CELEBORN-2369] Retry stopped master Outbox failures
17159eb34a is described below
commit 17159eb34ac8132e53e8f19b2c8d38d63d7853cb
Author: Chao Sun <[email protected]>
AuthorDate: Tue Jun 30 11:23:05 2026 +0800
[CELEBORN-2369] Retry stopped master Outbox failures
## Why are the changes needed?
When a Celeborn client loses its cached master connection, the associated
Outbox can be stopped before queued or late master RPC messages are sent. Those
requests fail with `CelebornException: Message is dropped because Outbox is
stopped`.
`MasterClient` currently retries I/O and RPC timeout failures, but it does
not recognize this stopped-Outbox failure when it is wrapped by the normal
`awaitResult` path. In HA mode, a request can therefore fail instead of
clearing the stale `RpcEndpointRef` and reconnecting to another available
master.
The same legacy message is also used when the local `RpcEnv` shuts down
permanently. Treating every occurrence as retryable without distinguishing
shutdown would cause futile reconnect attempts through an already-stopped
environment.
JIRA: https://issues.apache.org/jira/browse/CELEBORN-2369
## What changes were proposed in this PR?
- Recognize the exact legacy stopped-Outbox failure through its cause chain
and reset the cached master endpoint so the existing HA retry path can
reconnect.
- Use `RpcEnvStoppedException` for terminal local shutdown so shutdown
failures remain non-retryable.
- Preserve each Outbox's original stop cause for queued and late messages.
- Fail messages submitted after `RpcEnv` shutdown immediately with the
terminal cause.
- Add focused coverage for HA reconnection, terminal no-reconnect behavior,
and transient versus terminal Outbox causes.
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes. Master RPCs can now fail over after a transient stopped-Outbox
failure instead of failing immediately. No public API or configuration changes
are introduced.
## How was this PR tested?
Using JDK 17:
- `build/mvn --no-transfer-progress -DskipTests spotless:apply`
- `build/mvn --no-transfer-progress -pl common -DskipTests test-compile`
- `build/mvn --no-transfer-progress -pl common -DargLine=
-Dtest=org.apache.celeborn.common.client.MasterClientSuiteJ surefire:test` — 17
tests passed.
- `build/mvn --no-transfer-progress -pl common -DargLine=
-DwildcardSuites=org.apache.celeborn.common.rpc.netty.OutboxSuite
scalatest:test` — 3 tests passed.
- `build/mvn --no-transfer-progress -pl common -DargLine=
-DwildcardSuites=org.apache.celeborn.common.rpc.netty.NettyRpcEnvSuite
scalatest:test` — 33 tests passed.
- `build/mvn --no-transfer-progress -DskipTests spotless:check` — all 10
reactor modules passed.
Closes #3744 from sunchao/codex/retry-stopped-master-outbox-oss.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../celeborn/common/client/MasterClient.java | 16 ++-
.../common/rpc/OutboxStoppedException.scala | 27 +++++
.../celeborn/common/rpc/netty/NettyRpcEnv.scala | 8 +-
.../apache/celeborn/common/rpc/netty/Outbox.scala | 36 ++++---
.../celeborn/common/client/MasterClientSuiteJ.java | 114 +++++++++++++++++++++
.../common/rpc/netty/NettyRpcEnvSuite.scala | 29 ++++++
.../celeborn/common/rpc/netty/OutboxSuite.scala | 81 +++++++++++++++
7 files changed, 294 insertions(+), 17 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 e57460c49c..252ad018a5 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
@@ -186,13 +186,27 @@ public class MasterClient {
resetRpcEndpointRef(oldRef);
}
return true;
- } else if (e.getCause() instanceof IOException || e instanceof
RpcTimeoutException) {
+ } else if (isRetryableRpcFailure(e)) {
resetRpcEndpointRef(oldRef);
return true;
}
return false;
}
+ private boolean isRetryableRpcFailure(Throwable throwable) {
+ if (throwable.getCause() instanceof IOException || throwable instanceof
RpcTimeoutException) {
+ return true;
+ }
+ Throwable current = throwable;
+ while (current != null) {
+ if (current instanceof OutboxStoppedException) {
+ return true;
+ }
+ current = current.getCause();
+ }
+ return false;
+ }
+
@Nullable
private MasterNotLeaderException findMasterNotLeaderException(Throwable
throwable) {
Throwable current = throwable;
diff --git
a/common/src/main/scala/org/apache/celeborn/common/rpc/OutboxStoppedException.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/OutboxStoppedException.scala
new file mode 100644
index 0000000000..14e97ad8a6
--- /dev/null
+++
b/common/src/main/scala/org/apache/celeborn/common/rpc/OutboxStoppedException.scala
@@ -0,0 +1,27 @@
+/*
+ * 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.celeborn.common.rpc
+
+import org.apache.celeborn.common.exception.CelebornException
+
+private[celeborn] class OutboxStoppedException()
+ extends CelebornException(OutboxStoppedException.MESSAGE)
+
+private[celeborn] object OutboxStoppedException {
+ val MESSAGE = "Message is dropped because Outbox is stopped"
+}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
index ec72194916..ed002d3b0b 100644
---
a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
@@ -105,6 +105,8 @@ class NettyRpcEnv(
private val stopped = new AtomicBoolean(false)
+ private[netty] def isStopped: Boolean = stopped.get()
+
/**
* A map for [[RpcAddress]] and [[Outbox]]. When we are connecting to a
remote [[RpcAddress]],
* we just put messages to its [[Outbox]] to implement a non-blocking `send`
method.
@@ -208,7 +210,9 @@ class NettyRpcEnv(
if (stopped.get) {
// It's possible that we put `targetOutbox` after stopping. So we need
to clean it.
outboxes.remove(receiver.address)
- targetOutbox.stop()
+ val cause = new RpcEnvStoppedException()
+ targetOutbox.stop(cause)
+ message.onFailure(cause)
} else {
targetOutbox.send(message)
}
@@ -336,7 +340,7 @@ class NettyRpcEnv(
while (iter.hasNext()) {
val outbox = iter.next()
outboxes.remove(outbox.address)
- outbox.stop()
+ outbox.stop(new RpcEnvStoppedException())
}
if (timeoutScheduler != null) {
timeoutScheduler.shutdownNow()
diff --git
a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/Outbox.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/Outbox.scala
index 948278b887..52b8e481f8 100644
--- a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/Outbox.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/Outbox.scala
@@ -23,10 +23,9 @@ import javax.annotation.concurrent.GuardedBy
import scala.util.control.NonFatal
-import org.apache.celeborn.common.exception.CelebornException
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.network.client.{RpcResponseCallback,
TransportClient}
-import org.apache.celeborn.common.rpc.{RpcAddress, RpcEnvStoppedException}
+import org.apache.celeborn.common.rpc.{OutboxStoppedException, RpcAddress,
RpcEnvStoppedException}
sealed private[celeborn] trait OutboxMessage {
@@ -104,6 +103,9 @@ private[celeborn] class Outbox(nettyEnv: NettyRpcEnv, val
address: RpcAddress) {
@GuardedBy("this")
private var stopped = false
+ @GuardedBy("this")
+ private var stopCause: Throwable = null
+
/**
* If there is any thread draining the message queue
*/
@@ -112,19 +114,20 @@ private[celeborn] class Outbox(nettyEnv: NettyRpcEnv, val
address: RpcAddress) {
/**
* Send a message. If there is no active connection, cache it and launch a
new connection. If
- * [[Outbox]] is stopped, the sender will be notified with a
[[CelebornException]].
+ * [[Outbox]] is stopped, the sender will be notified with the cause that
stopped it.
*/
def send(message: OutboxMessage): Unit = {
- val dropped = synchronized {
+ val failure = synchronized {
if (stopped) {
- true
+ assert(stopCause != null)
+ stopCause
} else {
messages.add(message)
- false
+ null
}
}
- if (dropped) {
- message.onFailure(new CelebornException("Message is dropped because
Outbox is stopped"))
+ if (failure != null) {
+ message.onFailure(failure)
} else {
drainOutbox()
}
@@ -225,6 +228,7 @@ private[celeborn] class Outbox(nettyEnv: NettyRpcEnv, val
address: RpcAddress) {
return
}
stopped = true
+ stopCause = e
closeClient()
}
// Remove this Outbox from nettyEnv so that the further messages will
create a new Outbox along
@@ -248,16 +252,20 @@ private[celeborn] class Outbox(nettyEnv: NettyRpcEnv, val
address: RpcAddress) {
client = null
}
- /**
- * Stop [[Outbox]]. The remaining messages in the [[Outbox]] will be
notified with a
- * [[CelebornException]].
- */
- def stop(): Unit = {
+ /** Stop [[Outbox]] using a terminal cause when its owning RPC environment
is shutting down. */
+ def stop(): Unit =
+ stop(
+ if (nettyEnv.isStopped) new RpcEnvStoppedException()
+ else new OutboxStoppedException())
+
+ /** Stop [[Outbox]] and notify the remaining messages with the supplied
cause. */
+ def stop(cause: Throwable): Unit = {
synchronized {
if (stopped) {
return
}
stopped = true
+ stopCause = cause
if (connectFuture != null) {
connectFuture.cancel(true)
}
@@ -268,7 +276,7 @@ private[celeborn] class Outbox(nettyEnv: NettyRpcEnv, val
address: RpcAddress) {
// update messages and it's safe to just drain the queue.
var message = messages.poll()
while (message != null) {
- message.onFailure(new CelebornException("Message is dropped because
Outbox is stopped"))
+ message.onFailure(cause)
message = messages.poll()
}
}
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 26ad11c432..e9781a9f4a 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
@@ -18,6 +18,7 @@
package org.apache.celeborn.common.client;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
@@ -44,9 +45,11 @@ import
org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFrom
import
org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFromWorker;
import
org.apache.celeborn.common.protocol.message.ControlMessages.HeartbeatFromWorkerResponse;
import
org.apache.celeborn.common.protocol.message.ControlMessages.OneWayMessageResponse$;
+import org.apache.celeborn.common.rpc.OutboxStoppedException;
import org.apache.celeborn.common.rpc.RpcAddress;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.rpc.RpcEnv;
+import org.apache.celeborn.common.rpc.RpcEnvStoppedException;
import org.apache.celeborn.common.rpc.RpcTimeoutException;
public class MasterClientSuiteJ {
@@ -245,6 +248,84 @@ public class MasterClientSuiteJ {
checkOneMasterAskFailedInHA(new RpcTimeoutException("test", new
TimeoutException("test")));
}
+ @Test
+ public void testStoppedOutboxFailureReconnectsToAnotherMasterInHA() {
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+ final RpcEndpointRef master2 = Mockito.mock(RpcEndpointRef.class);
+ final AtomicInteger master1Attempts = new AtomicInteger(0);
+
+ Mockito.doAnswer(
+ invocation -> {
+ assertEquals(0, master1Attempts.getAndIncrement());
+ return Future$.MODULE$.failed(new OutboxStoppedException());
+ })
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ 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":
+ return master1;
+ case "host2":
+ return master2;
+ default:
+ fail(
+ "Should reconnect from host1 to 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 reconnect after a stopped outbox failure.", t);
+ fail("It should reconnect after a stopped outbox failure.");
+ }
+
+ assertEquals(mockResponse, response);
+ assertEquals(1, master1Attempts.get());
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host1:9097")),
Mockito.anyString());
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host2:9097")),
Mockito.anyString());
+ }
+
+ @Test
+ public void testStoppedRpcEnvFailureDoesNotReconnectInHA() {
+ checkMasterAskFailureDoesNotReconnectInHA(new RpcEnvStoppedException());
+ }
+
+ @Test
+ public void testNestedIOExceptionDoesNotReconnectInHA() {
+ checkMasterAskFailureDoesNotReconnectInHA(
+ new CelebornException("Permanent failure", new IOException("test")));
+ }
+
+ @Test
+ public void testNestedRpcTimeoutDoesNotReconnectInHA() {
+ checkMasterAskFailureDoesNotReconnectInHA(
+ new CelebornException(
+ "Permanent failure", new RpcTimeoutException("test", new
TimeoutException("test"))));
+ }
+
@Test
public void testBootstrapMasterNotLeaderRedirectsToSuggestedLeaderInHA() {
final CelebornConf conf =
@@ -545,6 +626,39 @@ public class MasterClientSuiteJ {
assertEquals(mockResponse, response);
}
+ private void checkMasterAskFailureDoesNotReconnectInHA(Exception exception) {
+ final CelebornConf conf = prepareForCelebornConfWithHA();
+ final RpcEndpointRef master1 = Mockito.mock(RpcEndpointRef.class);
+
+ Mockito.doReturn(Future$.MODULE$.failed(exception))
+ .when(master1)
+ .ask(Mockito.any(), Mockito.any(), Mockito.any());
+ Mockito.doAnswer(
+ invocation -> {
+ RpcAddress address = invocation.getArgument(0, RpcAddress.class);
+ if ("host1".equals(address.host())) {
+ return master1;
+ }
+ fail("Should not reconnect after a non-retryable failure: " +
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);
+
+ CelebornException thrown =
+ assertThrows(
+ CelebornException.class,
+ () -> client.askSync(message, HeartbeatFromWorkerResponse.class));
+ assertSame(exception, thrown.getCause());
+ Mockito.verify(master1, Mockito.times(1)).ask(Mockito.any(),
Mockito.any(), Mockito.any());
+ Mockito.verify(rpcEnv, Mockito.times(1))
+ .setupEndpointRef(
+ Mockito.eq(RpcAddress.fromHostAndPort("host1:9097")),
Mockito.anyString());
+ }
+
private void checkOneMasterAskFailedInHA(Exception exception) {
final CelebornConf conf = prepareForCelebornConfWithHA();
diff --git
a/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
index f42ffab4cf..3cdb3fcfe5 100644
---
a/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
+++
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
@@ -65,6 +65,35 @@ class NettyRpcEnvSuite extends RpcEnvSuite with TimeLimits {
assert(e.getCause.getMessage.contains(uri))
}
+ test("ask through a stopped RPC environment fails immediately") {
+ val endpointName = "stopped-rpc-env"
+ env.setupEndpoint(
+ endpointName,
+ new RpcEndpoint {
+ override val rpcEnv: RpcEnv = env
+ override def receiveAndReply(context: RpcCallContext):
PartialFunction[Any, Unit] = {
+ case message => context.reply(message)
+ }
+ })
+ val clientEnv = createRpcEnv(createCelebornConf(), "stopped-client", 0,
clientMode = true)
+ try {
+ val endpointRef = clientEnv.setupEndpointRef(env.address, endpointName)
+
+ clientEnv.shutdown()
+ clientEnv.awaitTermination()
+
+ failAfter(5.seconds) {
+ val e = intercept[CelebornException] {
+ endpointRef.askSync[String]("hello")
+ }
+ assert(e.getCause.isInstanceOf[RpcEnvStoppedException])
+ }
+ } finally {
+ clientEnv.shutdown()
+ clientEnv.awaitTermination()
+ }
+ }
+
test("advertise address different from bind address") {
val celebornConf = createCelebornConf()
val config = RpcEnvConfig(
diff --git
a/common/src/test/scala/org/apache/celeborn/common/rpc/netty/OutboxSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/OutboxSuite.scala
new file mode 100644
index 0000000000..ff37f50cf0
--- /dev/null
+++
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/OutboxSuite.scala
@@ -0,0 +1,81 @@
+/*
+ * 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.celeborn.common.rpc.netty
+
+import java.nio.ByteBuffer
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicReference
+
+import org.mockito.Mockito.{mock, when}
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.rpc.{OutboxStoppedException, RpcAddress,
RpcEnvStoppedException}
+
+class OutboxSuite extends CelebornFunSuite {
+
+ private def failureMessage(
+ failure: AtomicReference[Throwable],
+ failed: CountDownLatch): RpcOutboxMessage =
+ RpcOutboxMessage(
+ ByteBuffer.allocate(0),
+ e => {
+ failure.set(e)
+ failed.countDown()
+ },
+ (_, _) => ())
+
+ test("send after terminal stop uses the original cause") {
+ val outbox = new Outbox(mock(classOf[NettyRpcEnv]),
RpcAddress("localhost", 12345))
+ val cause = new RpcEnvStoppedException()
+ val failure = new AtomicReference[Throwable]()
+ val failed = new CountDownLatch(1)
+
+ outbox.stop(cause)
+ outbox.send(failureMessage(failure, failed))
+
+ assert(failed.await(10, TimeUnit.SECONDS))
+ assert(failure.get() eq cause)
+ }
+
+ test("send after transient stop remains retryable") {
+ val outbox = new Outbox(mock(classOf[NettyRpcEnv]),
RpcAddress("localhost", 12345))
+ val failure = new AtomicReference[Throwable]()
+ val failed = new CountDownLatch(1)
+
+ outbox.stop()
+ outbox.send(failureMessage(failure, failed))
+
+ assert(failed.await(10, TimeUnit.SECONDS))
+ assert(failure.get().isInstanceOf[OutboxStoppedException])
+ assert(failure.get().getMessage === OutboxStoppedException.MESSAGE)
+ }
+
+ test("default stop after RPC environment shutdown uses the terminal cause") {
+ val nettyEnv = mock(classOf[NettyRpcEnv])
+ val outbox = new Outbox(nettyEnv, RpcAddress("localhost", 12345))
+ val failure = new AtomicReference[Throwable]()
+ val failed = new CountDownLatch(1)
+ when(nettyEnv.isStopped).thenReturn(true)
+
+ outbox.stop()
+ outbox.send(failureMessage(failure, failed))
+
+ assert(failed.await(10, TimeUnit.SECONDS))
+ assert(failure.get().isInstanceOf[RpcEnvStoppedException])
+ }
+}
