venkata91 commented on code in PR #3842:
URL: https://github.com/apache/celeborn/pull/3842#discussion_r4067616051
##########
common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java:
##########
@@ -184,6 +194,82 @@ public void neverReturnInactiveClients() throws
IOException, InterruptedExceptio
factory.close();
}
+ @Test
+ public void replaceClientAfterPushTimeout() throws Exception {
+ AtomicBoolean respond = new AtomicBoolean(false);
+ BaseMessageHandler handler =
+ new BaseMessageHandler() {
+ @Override
+ public void receive(TransportClient client, RequestMessage message) {
+ PushData pushData = (PushData) message;
+ if (respond.get()) {
+ client
+ .getChannel()
+ .writeAndFlush(
+ new RpcResponse(
+ pushData.requestId, new
NioManagedBuffer(ByteBuffer.allocate(0))));
+ }
+ }
+
+ @Override
+ public boolean checkRegistered() {
+ return true;
+ }
+ };
+ TransportConf conf =
+ new TransportConf(TransportModuleConstants.REPLICATE_MODULE, new
CelebornConf());
+
+ try (TransportContext timeoutContext = new TransportContext(conf, handler);
+ TransportServer timeoutServer = timeoutContext.createServer();
+ TransportClientFactory factory = timeoutContext.createClientFactory())
{
+ TransportClient timedOutClient =
+ factory.createClient(getLocalHost(), timeoutServer.getPort());
+ RpcResponseCallback timedOutCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback outstandingCallback =
mock(RpcResponseCallback.class);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-1",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ 0,
+ timedOutCallback);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-2",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ TimeUnit.MINUTES.toMillis(1),
+ outstandingCallback);
+
+ timedOutClient.getHandler().failExpiredPushRequest();
Review Comment:
Addressed. The first request now uses a 1 ms timeout and the test waits 10
ms before calling failExpiredPushRequest(), ensuring dueTime is
deterministically in the past even with a strict comparison.
##########
common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java:
##########
@@ -350,6 +350,11 @@ public void timeOut() {
this.timedOut = true;
}
+ private void invalidate() {
+ this.timedOut = true;
+ channel.close();
+ }
Review Comment:
Addressed. invalidate() now delegates timedOut state management to the
existing timeOut() method before closing the channel, avoiding duplicated
state-setting logic.
##########
common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java:
##########
@@ -184,6 +194,83 @@ public void neverReturnInactiveClients() throws
IOException, InterruptedExceptio
factory.close();
}
+ @Test
+ public void replaceClientAfterPushTimeout() throws Exception {
+ AtomicBoolean respond = new AtomicBoolean(false);
+ BaseMessageHandler handler =
+ new BaseMessageHandler() {
+ @Override
+ public void receive(TransportClient client, RequestMessage message) {
+ PushData pushData = (PushData) message;
+ if (respond.get()) {
+ client
+ .getChannel()
+ .writeAndFlush(
+ new RpcResponse(
+ pushData.requestId, new
NioManagedBuffer(ByteBuffer.allocate(0))));
+ }
+ }
+
+ @Override
+ public boolean checkRegistered() {
+ return true;
+ }
+ };
+ TransportConf conf =
+ new TransportConf(TransportModuleConstants.REPLICATE_MODULE, new
CelebornConf());
+
+ try (TransportContext timeoutContext = new TransportContext(conf, handler);
+ TransportServer timeoutServer = timeoutContext.createServer();
+ TransportClientFactory factory = timeoutContext.createClientFactory())
{
+ TransportClient timedOutClient =
+ factory.createClient(getLocalHost(), timeoutServer.getPort());
+ RpcResponseCallback timedOutCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback outstandingCallback =
mock(RpcResponseCallback.class);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-1",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ 0,
+ timedOutCallback);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-2",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ TimeUnit.MINUTES.toMillis(1),
+ outstandingCallback);
+
+ timedOutClient.getHandler().failExpiredPushRequest();
+
+ Mockito.verify(timedOutCallback, Mockito.after(200).times(1))
+ .onFailure(
+ Mockito.argThat(
+ error ->
+
error.getMessage().startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())));
+ Mockito.verify(outstandingCallback,
Mockito.after(200).times(1)).onFailure(Mockito.any());
+ assertFalse(timedOutClient.isActive());
Review Comment:
Addressed. Replaced the 200 ms Mockito.after(...) checks with
Mockito.timeout(5000), avoiding flaky failures under CI load.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
