This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 57f775a7 [CELEBORN-273] Move push data timeout checker into
TransportResponseHandler to keep callback status consistence (#1208)
57f775a7 is described below
commit 57f775a7e9db42d530fbe72e7ae4b0ac601725af
Author: Angerszhuuuu <[email protected]>
AuthorDate: Thu Feb 16 18:27:37 2023 +0800
[CELEBORN-273] Move push data timeout checker into TransportResponseHandler
to keep callback status consistence (#1208)
---
.../apache/celeborn/client/ShuffleClientImpl.java | 33 +++---
.../celeborn/client/ShuffleClientImplSuiteJ.java | 7 +-
.../celeborn/client/ShuffleClientSuiteJ.java | 5 +-
.../celeborn/common/network/TransportContext.java | 2 +-
.../common/network/client/TransportClient.java | 25 +++--
.../network/client/TransportResponseHandler.java | 76 +++++++++++---
.../common/network/util/TransportConf.java | 4 +
.../common/write/InFlightRequestTracker.java | 116 +++++----------------
.../celeborn/common/write/PushRequestInfo.java | 37 +++++++
.../apache/celeborn/common/write/PushState.java | 58 ++---------
.../org/apache/celeborn/common/CelebornConf.scala | 3 +-
.../network/TransportResponseHandlerSuiteJ.java | 48 +++++++--
.../celeborn/tests/spark/PushDataTimeoutTest.scala | 1 +
.../celeborn/tests/spark/SparkTestBase.scala | 2 +-
.../service/deploy/worker/PushDataHandler.scala | 34 +-----
15 files changed, 220 insertions(+), 231 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 78fd1aa9..d0dc3428 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -34,7 +34,6 @@ import com.google.common.util.concurrent.Uninterruptibles;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.CompositeByteBuf;
import io.netty.buffer.Unpooled;
-import io.netty.channel.ChannelFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -202,8 +201,7 @@ public class ShuffleClientImpl extends ShuffleClient {
PushData newPushData =
new PushData(MASTER_MODE, shuffleKey, newLoc.getUniqueId(),
newBuffer);
- ChannelFuture future = client.pushData(newPushData, callback);
- pushState.pushStarted(batchId, future, callback,
loc.hostAndPushPort(), pushDataTimeout);
+ client.pushData(newPushData, pushDataTimeout, callback);
} else {
throw new RuntimeException(
"Mock push data submit retry failed. remainReviveTimes = " +
remainReviveTimes + ".");
@@ -682,7 +680,7 @@ public class ShuffleClientImpl extends ShuffleClient {
attemptId,
nextBatchId);
splitPartition(shuffleId, partitionId, applicationId, loc);
- pushState.onSuccess(nextBatchId, loc.hostAndPushPort());
+ pushState.onSuccess(loc.hostAndPushPort());
callback.onSuccess(response);
} else if (reason == StatusCode.HARD_SPLIT.getValue()) {
logger.debug(
@@ -710,7 +708,7 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
nextBatchId);
- pushState.onCongestControl(nextBatchId,
loc.hostAndPushPort());
+ pushState.onCongestControl(loc.hostAndPushPort());
callback.onSuccess(response);
} else if (reason ==
StatusCode.PUSH_DATA_SUCCESS_SLAVE_CONGESTED.getValue()) {
logger.debug(
@@ -718,15 +716,15 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
nextBatchId);
- pushState.onCongestControl(nextBatchId,
loc.hostAndPushPort());
+ pushState.onCongestControl(loc.hostAndPushPort());
callback.onSuccess(response);
} else {
response.rewind();
- pushState.onSuccess(nextBatchId, loc.hostAndPushPort());
+ pushState.onSuccess(loc.hostAndPushPort());
callback.onSuccess(response);
}
} else {
- pushState.onSuccess(nextBatchId, loc.hostAndPushPort());
+ pushState.onSuccess(loc.hostAndPushPort());
callback.onSuccess(response);
}
}
@@ -794,9 +792,7 @@ public class ShuffleClientImpl extends ShuffleClient {
if (!testRetryRevive) {
TransportClient client =
dataClientFactory.createClient(loc.getHost(), loc.getPushPort(),
partitionId);
- ChannelFuture future = client.pushData(pushData, wrappedCallback);
- pushState.pushStarted(
- nextBatchId, future, wrappedCallback, loc.hostAndPushPort(),
pushDataTimeout);
+ client.pushData(pushData, pushDataTimeout, wrappedCallback);
} else {
wrappedCallback.onFailure(
new Exception(
@@ -1067,7 +1063,7 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
Arrays.toString(batchIds));
- pushState.onCongestControl(groupedBatchId, hostPort);
+ pushState.onCongestControl(hostPort);
callback.onSuccess(response);
} else if (reason ==
StatusCode.PUSH_DATA_SUCCESS_SLAVE_CONGESTED.getValue()) {
logger.debug(
@@ -1075,17 +1071,17 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
Arrays.toString(batchIds));
- pushState.onCongestControl(groupedBatchId, hostPort);
+ pushState.onCongestControl(hostPort);
callback.onSuccess(response);
} else {
// Should not happen in current architecture.
response.rewind();
logger.error("Push merged data should not receive this
response");
- pushState.onSuccess(groupedBatchId, hostPort);
+ pushState.onSuccess(hostPort);
callback.onSuccess(response);
}
} else {
- pushState.onSuccess(groupedBatchId, hostPort);
+ pushState.onSuccess(hostPort);
callback.onSuccess(response);
}
}
@@ -1141,8 +1137,7 @@ public class ShuffleClientImpl extends ShuffleClient {
try {
if (!testRetryRevive || remainReviveTimes < 1) {
TransportClient client = dataClientFactory.createClient(host, port);
- ChannelFuture future = client.pushMergedData(mergedData,
wrappedCallback);
- pushState.pushStarted(groupedBatchId, future, wrappedCallback,
hostPort, pushDataTimeout);
+ client.pushMergedData(mergedData, pushDataTimeout, wrappedCallback);
} else {
wrappedCallback.onFailure(
new Exception(
@@ -1514,9 +1509,7 @@ public class ShuffleClientImpl extends ShuffleClient {
// do push data
try {
TransportClient client = createClientWaitingInFlightRequest(location,
mapKey, pushState);
- ChannelFuture future = client.pushData(pushData, callback);
- pushState.pushStarted(
- nextBatchId, future, callback, location.hostAndPushPort(),
pushDataTimeout);
+ client.pushData(pushData, pushDataTimeout, callback);
} catch (Exception e) {
logger.warn("PushData byteBuf failed", e);
callback.onFailure(
diff --git
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
index ae368d2f..33b0c2af 100644
---
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
+++
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
@@ -18,6 +18,7 @@
package org.apache.celeborn.client;
import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
import static org.mockito.Mockito.when;
import java.io.IOException;
@@ -56,7 +57,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
@Test
public void testPushDataByteBufSuccess() throws IOException {
ByteBuf byteBuf = createByteBuf();
- when(client.pushData(any(), any()))
+ when(client.pushData(any(), anyLong(), any()))
.thenAnswer(
t -> {
RpcResponseCallback rpcResponseCallback =
@@ -82,7 +83,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
@Test
public void testPushDataByteBufHardSplit() throws IOException {
ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
- when(client.pushData(any(), any()))
+ when(client.pushData(any(), anyLong(), any()))
.thenAnswer(
t -> {
RpcResponseCallback rpcResponseCallback =
@@ -107,7 +108,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
@Test
public void testPushDataByteBufFail() throws IOException {
ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
- when(client.pushData(any(), any()))
+ when(client.pushData(any(), anyLong(), any()))
.thenAnswer(
t -> {
RpcResponseCallback rpcResponseCallback =
diff --git
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
index bb5f5aa4..24af088b 100644
--- a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
+++ b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
@@ -18,6 +18,7 @@
package org.apache.celeborn.client;
import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -312,12 +313,12 @@ public class ShuffleClientSuiteJ {
}
};
- when(client.pushData(any(), any())).thenAnswer(t -> mockedFuture);
+ when(client.pushData(any(), anyLong(), any())).thenAnswer(t ->
mockedFuture);
when(clientFactory.createClient(
masterLocation.getHost(), masterLocation.getPushPort(),
TEST_REDUCRE_ID))
.thenAnswer(t -> client);
- when(client.pushMergedData(any(), any())).thenAnswer(t -> mockedFuture);
+ when(client.pushMergedData(any(), anyLong(), any())).thenAnswer(t ->
mockedFuture);
when(clientFactory.createClient(masterLocation.getHost(),
masterLocation.getPushPort()))
.thenAnswer(t -> client);
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/TransportContext.java
b/common/src/main/java/org/apache/celeborn/common/network/TransportContext.java
index a4fd5035..69268de8 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/TransportContext.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/TransportContext.java
@@ -115,7 +115,7 @@ public class TransportContext {
private TransportChannelHandler createChannelHandler(
Channel channel, BaseMessageHandler msgHandler) {
- TransportResponseHandler responseHandler = new
TransportResponseHandler(channel);
+ TransportResponseHandler responseHandler = new
TransportResponseHandler(conf, channel);
TransportClient client = new TransportClient(channel, responseHandler);
TransportRequestHandler requestHandler =
new TransportRequestHandler(channel, client, msgHandler);
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
index 10304a5e..5631518b 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
@@ -38,6 +38,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
import org.apache.celeborn.common.network.protocol.*;
import org.apache.celeborn.common.network.util.NettyUtils;
+import org.apache.celeborn.common.write.PushRequestInfo;
/**
* Client for fetching consecutive chunks of a pre-negotiated stream. This API
is intended to allow
@@ -156,32 +157,40 @@ public class TransportClient implements Closeable {
return requestId;
}
- public ChannelFuture pushData(PushData pushData, RpcResponseCallback
callback) {
+ public ChannelFuture pushData(
+ PushData pushData, long pushDataTimeout, RpcResponseCallback callback) {
if (logger.isTraceEnabled()) {
logger.trace("Pushing data to {}", NettyUtils.getRemoteAddress(channel));
}
long requestId = requestId();
- handler.addPushRequest(requestId, callback);
-
+ long dueTime = System.currentTimeMillis() + pushDataTimeout;
+ PushRequestInfo info = new PushRequestInfo(dueTime, callback);
+ handler.addPushRequest(requestId, info);
pushData.requestId = requestId;
PushChannelListener listener = new PushChannelListener(requestId,
callback);
- return channel.writeAndFlush(pushData).addListener(listener);
+ ChannelFuture channelFuture =
channel.writeAndFlush(pushData).addListener(listener);
+ info.setChannelFuture(channelFuture);
+ return channelFuture;
}
- public ChannelFuture pushMergedData(PushMergedData pushMergedData,
RpcResponseCallback callback) {
+ public ChannelFuture pushMergedData(
+ PushMergedData pushMergedData, long pushDataTimeout, RpcResponseCallback
callback) {
if (logger.isTraceEnabled()) {
logger.trace("Pushing merged data to {}",
NettyUtils.getRemoteAddress(channel));
}
long requestId = requestId();
- handler.addPushRequest(requestId, callback);
-
+ long dueTime = System.currentTimeMillis() + pushDataTimeout;
+ PushRequestInfo info = new PushRequestInfo(dueTime, callback);
+ handler.addPushRequest(requestId, info);
pushMergedData.requestId = requestId;
PushChannelListener listener = new PushChannelListener(requestId,
callback);
- return channel.writeAndFlush(pushMergedData).addListener(listener);
+ ChannelFuture channelFuture =
channel.writeAndFlush(pushMergedData).addListener(listener);
+ info.setChannelFuture(channelFuture);
+ return channelFuture;
}
/**
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
index 07d56a49..337b8487 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
@@ -18,8 +18,11 @@
package org.apache.celeborn.common.network.client;
import java.io.IOException;
+import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import io.netty.channel.Channel;
@@ -29,6 +32,11 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.protocol.*;
import org.apache.celeborn.common.network.server.MessageHandler;
import org.apache.celeborn.common.network.util.NettyUtils;
+import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.protocol.TransportModuleConstants;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+import org.apache.celeborn.common.util.ThreadUtils;
+import org.apache.celeborn.common.write.PushRequestInfo;
/**
* Handler that processes server responses, in response to requests issued
from a
@@ -39,22 +47,62 @@ import org.apache.celeborn.common.network.util.NettyUtils;
public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
private static final Logger logger =
LoggerFactory.getLogger(TransportResponseHandler.class);
+ private final TransportConf conf;
private final Channel channel;
private final Map<StreamChunkSlice, ChunkReceivedCallback>
outstandingFetches;
private final Map<Long, RpcResponseCallback> outstandingRpcs;
- private final Map<Long, RpcResponseCallback> outstandingPushes;
+ private final ConcurrentHashMap<Long, PushRequestInfo> outstandingPushes;
/** Records the time (in system nanoseconds) that the last fetch or RPC
request was sent. */
private final AtomicLong timeOfLastRequestNs;
- public TransportResponseHandler(Channel channel) {
+ private final ScheduledExecutorService pushTimeoutChecker;
+ private final long pushTimeoutCheckerInterval;
+
+ public TransportResponseHandler(TransportConf conf, Channel channel) {
+ this.conf = conf;
this.channel = channel;
this.outstandingFetches = new ConcurrentHashMap<>();
this.outstandingRpcs = new ConcurrentHashMap<>();
this.outstandingPushes = new ConcurrentHashMap<>();
this.timeOfLastRequestNs = new AtomicLong(0);
+ pushTimeoutCheckerInterval = conf.pushDataTimeoutCheckIntervalMs();
+ pushTimeoutChecker =
+
ThreadUtils.newDaemonSingleThreadScheduledExecutor("push-timeout-checker-" +
this);
+ pushTimeoutChecker.scheduleAtFixedRate(
+ () -> failExpiredPushRequest(),
+ pushTimeoutCheckerInterval,
+ pushTimeoutCheckerInterval,
+ TimeUnit.MILLISECONDS);
+ }
+
+ public void failExpiredPushRequest() {
+ long currentTime = System.currentTimeMillis();
+ Iterator<Map.Entry<Long, PushRequestInfo>> iter =
outstandingPushes.entrySet().iterator();
+ while (iter.hasNext()) {
+ Map.Entry<Long, PushRequestInfo> entry = iter.next();
+ if (entry.getValue().dueTime <= currentTime) {
+ PushRequestInfo info = outstandingPushes.remove(entry.getKey());
+ if (info != null) {
+ if (info.channelFuture != null) {
+ info.channelFuture.cancel(true);
+ }
+ // When module name equals to DATA_MODULE, mean shuffle client push
data, else means
+ // do data replication.
+ if
(TransportModuleConstants.DATA_MODULE.equals(conf.getModuleName())) {
+ info.callback.onFailure(
+ new
IOException(StatusCode.PUSH_DATA_TIMEOUT_MASTER.getMessage()));
+ } else if
(TransportModuleConstants.PUSH_MODULE.equals(conf.getModuleName())) {
+ info.callback.onFailure(
+ new
IOException(StatusCode.PUSH_DATA_TIMEOUT_SLAVE.getMessage()));
+ }
+ info.channelFuture = null;
+ info.callback = null;
+ }
+ }
+ }
}
public void addFetchRequest(StreamChunkSlice streamChunkSlice,
ChunkReceivedCallback callback) {
@@ -81,12 +129,12 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
outstandingRpcs.remove(requestId);
}
- public void addPushRequest(long requestId, RpcResponseCallback callback) {
+ public void addPushRequest(long requestId, PushRequestInfo info) {
updateTimeOfLastRequest();
if (outstandingPushes.containsKey(requestId)) {
logger.warn("[addPushRequest] requestId {} already exists!", requestId);
}
- outstandingPushes.put(requestId, callback);
+ outstandingPushes.put(requestId, info);
}
public void removePushRequest(long requestId) {
@@ -112,9 +160,9 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
logger.warn("RpcResponseCallback.onFailure throws exception", e);
}
}
- for (Map.Entry<Long, RpcResponseCallback> entry :
outstandingPushes.entrySet()) {
+ for (Map.Entry<Long, PushRequestInfo> entry :
outstandingPushes.entrySet()) {
try {
- entry.getValue().onFailure(cause);
+ entry.getValue().callback.onFailure(cause);
} catch (Exception e) {
logger.warn("RpcResponseCallback.onFailure throws exception", e);
}
@@ -186,9 +234,9 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
}
} else if (message instanceof RpcResponse) {
RpcResponse resp = (RpcResponse) message;
- RpcResponseCallback listener = outstandingPushes.remove(resp.requestId);
- if (listener == null) {
- listener = outstandingRpcs.remove(resp.requestId);
+ PushRequestInfo info = outstandingPushes.remove(resp.requestId);
+ if (info == null) {
+ RpcResponseCallback listener = outstandingRpcs.remove(resp.requestId);
if (listener == null) {
logger.warn(
"Ignoring response for RPC {} from {} ({} bytes) since it is not
outstanding",
@@ -205,16 +253,16 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
}
} else {
try {
- listener.onSuccess(resp.body().nioByteBuffer());
+ info.callback.onSuccess(resp.body().nioByteBuffer());
} finally {
resp.body().release();
}
}
} else if (message instanceof RpcFailure) {
RpcFailure resp = (RpcFailure) message;
- RpcResponseCallback listener = outstandingPushes.remove(resp.requestId);
- if (listener == null) {
- listener = outstandingRpcs.remove(resp.requestId);
+ PushRequestInfo info = outstandingPushes.remove(resp.requestId);
+ if (info == null) {
+ RpcResponseCallback listener = outstandingRpcs.remove(resp.requestId);
if (listener == null) {
logger.warn(
"Ignoring response for RPC {} from {} ({}) since it is not
outstanding",
@@ -225,7 +273,7 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
listener.onFailure(new RuntimeException(resp.errorString));
}
} else {
- listener.onFailure(new RuntimeException(resp.errorString));
+ info.callback.onFailure(new RuntimeException(resp.errorString));
}
} else {
throw new IllegalStateException("Unknown response type: " +
message.type());
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
index 7590df81..29d40f2a 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
@@ -144,4 +144,8 @@ public class TransportConf {
public String decoderMode() {
return conf.networkIoDecoderMode(module);
}
+
+ public long pushDataTimeoutCheckIntervalMs() {
+ return conf.pushTimeoutCheckInterval();
+ }
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
index 336bba9c..c8c4d845 100644
---
a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
+++
b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
@@ -18,18 +18,15 @@
package org.apache.celeborn.common.write;
import java.io.IOException;
-import java.util.Map;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
-import io.netty.channel.ChannelFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.network.client.RpcResponseCallback;
-import org.apache.celeborn.common.protocol.message.StatusCode;
/*
* This class is for track in flight request and limit request.
@@ -43,8 +40,8 @@ public class InFlightRequestTracker {
private final PushStrategy pushStrategy;
private final AtomicInteger batchId = new AtomicInteger();
- private final ConcurrentHashMap<String, ConcurrentHashMap<Integer,
BatchInfo>>
- inflightBatchesPerAddress = new ConcurrentHashMap<>();
+ private final ConcurrentHashMap<String, Set<Integer>>
inflightBatchesPerAddress =
+ new ConcurrentHashMap<>();
public InFlightRequestTracker(CelebornConf conf, PushState pushState) {
this.waitInflightTimeoutMs = conf.pushLimitInFlightTimeoutMs();
@@ -54,34 +51,31 @@ public class InFlightRequestTracker {
}
public void addBatch(int batchId, String hostAndPushPort) {
- ConcurrentHashMap<Integer, BatchInfo> batchIdSetPerPair =
- inflightBatchesPerAddress.computeIfAbsent(hostAndPushPort, id -> new
ConcurrentHashMap<>());
- batchIdSetPerPair.computeIfAbsent(batchId, id -> new BatchInfo());
+ Set<Integer> batchIdSetPerPair =
+ inflightBatchesPerAddress.computeIfAbsent(
+ hostAndPushPort, id -> ConcurrentHashMap.newKeySet());
+ batchIdSetPerPair.add(batchId);
}
- public void onSuccess(int batchId, String hostAndPushPort) {
- pushStrategy.onSuccess(hostAndPushPort);
+ public void removeBatch(int batchId, String hostAndPushPort) {
+ Set<Integer> batchIdSet = inflightBatchesPerAddress.get(hostAndPushPort);
+ batchIdSet.remove(batchId);
+ if (batchIdSet.size() == 0) {
+ inflightBatchesPerAddress.remove(hostAndPushPort);
+ }
}
- public void onCongestControl(int batchId, String hostAndPushPort) {
- pushStrategy.onCongestControl(hostAndPushPort);
+ public void onSuccess(String hostAndPushPort) {
+ pushStrategy.onSuccess(hostAndPushPort);
}
- public void removeBatch(int batchId, String hostAndPushPort) {
- ConcurrentHashMap<Integer, BatchInfo> batchIdMap =
- inflightBatchesPerAddress.get(hostAndPushPort);
- BatchInfo info = batchIdMap.remove(batchId);
- if (info != null && info.channelFuture != null) {
- info.channelFuture.cancel(true);
- }
- if (batchIdMap.size() == 0) {
- inflightBatchesPerAddress.remove(hostAndPushPort);
- }
+ public void onCongestControl(String hostAndPushPort) {
+ pushStrategy.onCongestControl(hostAndPushPort);
}
- public ConcurrentHashMap<Integer, BatchInfo>
getBatchIdSetByAddressPair(String hostAndPort) {
+ public Set<Integer> getBatchIdSetByAddressPair(String hostAndPort) {
return inflightBatchesPerAddress.computeIfAbsent(
- hostAndPort, pair -> new ConcurrentHashMap<>());
+ hostAndPort, pair -> ConcurrentHashMap.newKeySet());
}
public boolean limitMaxInFlight(String hostAndPushPort) throws IOException {
@@ -92,17 +86,16 @@ public class InFlightRequestTracker {
pushStrategy.limitPushSpeed(pushState, hostAndPushPort);
int currentMaxReqsInFlight =
pushStrategy.getCurrentMaxReqsInFlight(hostAndPushPort);
- ConcurrentHashMap<Integer, BatchInfo> batchIdMap =
getBatchIdSetByAddressPair(hostAndPushPort);
+ Set batchIdSet = getBatchIdSetByAddressPair(hostAndPushPort);
long times = waitInflightTimeoutMs / delta;
try {
while (times > 0) {
- if (batchIdMap.size() <= currentMaxReqsInFlight) {
+ if (batchIdSet.size() <= currentMaxReqsInFlight) {
break;
}
if (pushState.exception.get() != null) {
throw pushState.exception.get();
}
- failExpiredBatch(true);
Thread.sleep(delta);
times--;
}
@@ -117,7 +110,7 @@ public class InFlightRequestTracker {
+ "for hostAndPushPort {}, "
+ "which exceeds the current limit {}.",
waitInflightTimeoutMs,
- batchIdMap.size(),
+ batchIdSet.size(),
hostAndPushPort,
currentMaxReqsInFlight);
}
@@ -139,7 +132,7 @@ public class InFlightRequestTracker {
try {
while (times > 0) {
Optional<Integer> inFlightSizeOptional =
-
inflightBatchesPerAddress.values().stream().map(Map::size).reduce(Integer::sum);
+
inflightBatchesPerAddress.values().stream().map(Set::size).reduce(Integer::sum);
inFlightSize = inFlightSizeOptional.orElse(0);
if (inFlightSize == 0) {
break;
@@ -147,7 +140,6 @@ public class InFlightRequestTracker {
if (pushState.exception.get() != null) {
throw pushState.exception.get();
}
- failExpiredBatch(true);
Thread.sleep(delta);
times--;
}
@@ -174,76 +166,18 @@ public class InFlightRequestTracker {
throw pushState.exception.get();
}
- ConcurrentHashMap<Integer, BatchInfo> batchIdMap =
getBatchIdSetByAddressPair(hostAndPushPort);
- return batchIdMap.size() > maxInFlight;
+ Set<Integer> batchIdSet = getBatchIdSetByAddressPair(hostAndPushPort);
+ return batchIdSet.size() > maxInFlight;
}
protected int nextBatchId() {
return batchId.incrementAndGet();
}
- /**
- * Trigger expired push request.
- *
- * @param isMaster When true means push data, when false means replicate
data.
- */
- public synchronized void failExpiredBatch(boolean isMaster) {
- long currentTime = System.currentTimeMillis();
- inflightBatchesPerAddress
- .values()
- .forEach(
- batchIdMap ->
- batchIdMap
- .values()
- .forEach(
- info -> {
- // Compatible with older versions of client
- if (info.pushDataTimeout > 0) {
- if (info.pushTime != -1
- && (currentTime - info.pushTime >
info.pushDataTimeout)) {
- if (info.callback != null) {
- info.channelFuture.cancel(true);
- if (isMaster) {
- info.callback.onFailure(
- new IOException(
-
StatusCode.PUSH_DATA_TIMEOUT_MASTER.getMessage()));
- } else {
- info.callback.onFailure(
- new IOException(
-
StatusCode.PUSH_DATA_TIMEOUT_SLAVE.getMessage()));
- }
- info.channelFuture = null;
- info.callback = null;
- }
- }
- }
- }));
- }
-
public void cleanup() {
if (!inflightBatchesPerAddress.isEmpty()) {
- logger.debug("Cancel all {} futures.", inflightBatchesPerAddress.size());
- inflightBatchesPerAddress
- .values()
- .forEach(
- batchIdMap ->
- batchIdMap
- .values()
- .forEach(
- info -> {
- if (info.channelFuture != null) {
- info.channelFuture.cancel(true);
- }
- }));
inflightBatchesPerAddress.clear();
}
pushStrategy.clear();
}
-
- static class BatchInfo {
- ChannelFuture channelFuture;
- long pushTime = -1;
- long pushDataTimeout;
- RpcResponseCallback callback;
- }
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
b/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
new file mode 100644
index 00000000..691b4335
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
@@ -0,0 +1,37 @@
+/*
+ * 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.write;
+
+import io.netty.channel.ChannelFuture;
+
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
+
+public class PushRequestInfo {
+ public ChannelFuture channelFuture;
+ public long dueTime;
+ public RpcResponseCallback callback;
+
+ public PushRequestInfo(long dueTime, RpcResponseCallback callback) {
+ this.dueTime = dueTime;
+ this.callback = callback;
+ }
+
+ public void setChannelFuture(ChannelFuture future) {
+ this.channelFuture = future;
+ }
+}
diff --git
a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
index dd342f2c..550989ca 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
@@ -19,52 +19,20 @@ package org.apache.celeborn.common.write;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
-import io.netty.channel.ChannelFuture;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.util.ThreadUtils;
public class PushState {
- private static final Logger logger =
LoggerFactory.getLogger(PushState.class);
private final int pushBufferMaxSize;
public AtomicReference<IOException> exception = new AtomicReference<>();
private final InFlightRequestTracker inFlightRequestTracker;
- private ScheduledExecutorService pushTimeoutChecker;
- private long pushTimeoutCheckerInterval;
public PushState(CelebornConf conf) {
pushBufferMaxSize = conf.pushBufferMaxSize();
inFlightRequestTracker = new InFlightRequestTracker(conf, this);
- pushTimeoutCheckerInterval = conf.pushTimeoutCheckInterval();
- }
-
- public void pushStarted(
- int batchId,
- ChannelFuture future,
- RpcResponseCallback callback,
- String hostAndPort,
- long pushDataTimeout) {
- InFlightRequestTracker.BatchInfo info =
-
inFlightRequestTracker.getBatchIdSetByAddressPair(hostAndPort).get(batchId);
- // In rare cases info could be null. For example, a speculative task has
one thread pushing,
- // and other thread retry-pushing. At time 1 thread 1 find StageEnded,
then it cleans up
- // PushState, at the same time thread 2 pushes data and calles pushStarted,
- // at this time info will be null
- if (info != null) {
- info.pushTime = System.currentTimeMillis();
- info.pushDataTimeout = pushDataTimeout;
- info.channelFuture = future;
- info.callback = callback;
- }
}
public void cleanup() {
@@ -101,16 +69,16 @@ public class PushState {
inFlightRequestTracker.addBatch(batchId, hostAndPushPort);
}
- public void onSuccess(int batchId, String hostAndPushPort) {
- inFlightRequestTracker.onSuccess(batchId, hostAndPushPort);
+ public void removeBatch(int batchId, String hostAndPushPort) {
+ inFlightRequestTracker.removeBatch(batchId, hostAndPushPort);
}
- public void onCongestControl(int batchId, String hostAndPushPort) {
- inFlightRequestTracker.onCongestControl(batchId, hostAndPushPort);
+ public void onSuccess(String hostAndPushPort) {
+ inFlightRequestTracker.onSuccess(hostAndPushPort);
}
- public void removeBatch(int batchId, String hostAndPushPort) {
- inFlightRequestTracker.removeBatch(batchId, hostAndPushPort);
+ public void onCongestControl(String hostAndPushPort) {
+ inFlightRequestTracker.onCongestControl(hostAndPushPort);
}
public boolean limitMaxInFlight(String hostAndPushPort) throws IOException {
@@ -124,18 +92,4 @@ public class PushState {
public boolean reachLimit(String hostAndPushPort, int maxInFlight) throws
IOException {
return inFlightRequestTracker.reachLimit(hostAndPushPort, maxInFlight);
}
-
- public void startChecker(boolean isMaster) {
- pushTimeoutChecker =
ThreadUtils.newDaemonSingleThreadScheduledExecutor("push-timeout-checker");
- pushTimeoutChecker.scheduleAtFixedRate(
- new Runnable() {
- @Override
- public void run() {
- inFlightRequestTracker.failExpiredBatch(isMaster);
- }
- },
- pushTimeoutCheckerInterval,
- pushTimeoutCheckerInterval,
- TimeUnit.MILLISECONDS);
- }
}
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 5a4c24de..64c4d102 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1325,7 +1325,7 @@ object CelebornConf extends Logging {
.doc("Interval for checking push data timeout.")
.version("0.3.0")
.timeConf(TimeUnit.MILLISECONDS)
- .createWithDefaultString("60s")
+ .createWithDefaultString("30s")
val FETCH_TIMEOUT: ConfigEntry[Long] =
buildConf("celeborn.fetch.timeout")
@@ -2196,6 +2196,7 @@ object CelebornConf extends Logging {
.version("0.2.0")
.doc(s"Timeout for a task to push data rpc message. This value should
better be more than twice of `${PUSH_TIMEOUT_CHECK_INTERVAL.key}`")
.timeConf(TimeUnit.MILLISECONDS)
+ .checkValue(_ > 0, "celeborn.push.data.timeout must be positive!")
.createWithDefaultString("120s")
val TEST_PUSH_MASTER_DATA_TIMEOUT: ConfigEntry[Boolean] =
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
index b4ae735d..c71411c3 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
@@ -22,21 +22,29 @@ import static org.mockito.Mockito.*;
import java.nio.ByteBuffer;
+import io.netty.channel.ChannelFuture;
import io.netty.channel.local.LocalChannel;
import org.junit.Test;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
import org.apache.celeborn.common.network.client.ChunkReceivedCallback;
import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportResponseHandler;
import org.apache.celeborn.common.network.protocol.*;
+import org.apache.celeborn.common.protocol.TransportModuleConstants;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.common.write.PushRequestInfo;
public class TransportResponseHandlerSuiteJ {
@Test
public void handleSuccessfulFetch() throws Exception {
StreamChunkSlice streamChunkSlice = new StreamChunkSlice(1, 0);
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.FETCH_MODULE, 8),
+ new LocalChannel());
ChunkReceivedCallback callback = mock(ChunkReceivedCallback.class);
handler.addFetchRequest(streamChunkSlice, callback);
assertEquals(1, handler.numOutstandingRequests());
@@ -49,7 +57,10 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void handleFailedFetch() throws Exception {
StreamChunkSlice streamChunkSlice = new StreamChunkSlice(1, 0);
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.FETCH_MODULE, 8),
+ new LocalChannel());
ChunkReceivedCallback callback = mock(ChunkReceivedCallback.class);
handler.addFetchRequest(streamChunkSlice, callback);
assertEquals(1, handler.numOutstandingRequests());
@@ -61,7 +72,10 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void clearAllOutstandingRequests() throws Exception {
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
ChunkReceivedCallback callback = mock(ChunkReceivedCallback.class);
handler.addFetchRequest(new StreamChunkSlice(1, 0), callback);
handler.addFetchRequest(new StreamChunkSlice(1, 1), callback);
@@ -80,7 +94,10 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void handleSuccessfulRPC() throws Exception {
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.RPC_MODULE, 8),
+ new LocalChannel());
RpcResponseCallback callback = mock(RpcResponseCallback.class);
handler.addRpcRequest(12345, callback);
assertEquals(1, handler.numOutstandingRequests());
@@ -97,7 +114,10 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void handleFailedRPC() throws Exception {
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.RPC_MODULE, 8),
+ new LocalChannel());
RpcResponseCallback callback = mock(RpcResponseCallback.class);
handler.addRpcRequest(12345, callback);
assertEquals(1, handler.numOutstandingRequests());
@@ -112,9 +132,14 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void handleSuccessfulPush() throws Exception {
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
RpcResponseCallback callback = mock(RpcResponseCallback.class);
- handler.addPushRequest(12345, callback);
+ PushRequestInfo info = new PushRequestInfo(System.currentTimeMillis() +
30000, callback);
+ info.setChannelFuture(mock(ChannelFuture.class));
+ handler.addPushRequest(12345, info);
assertEquals(1, handler.numOutstandingRequests());
// This response should be ignored.
@@ -129,9 +154,14 @@ public class TransportResponseHandlerSuiteJ {
@Test
public void handleFailedPush() throws Exception {
- TransportResponseHandler handler = new TransportResponseHandler(new
LocalChannel());
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
RpcResponseCallback callback = mock(RpcResponseCallback.class);
- handler.addPushRequest(12345, callback);
+ PushRequestInfo info = new PushRequestInfo(System.currentTimeMillis() +
30000L, callback);
+ info.setChannelFuture(mock(ChannelFuture.class));
+ handler.addPushRequest(12345, info);
assertEquals(1, handler.numOutstandingRequests());
handler.handle(new RpcFailure(54321, "uh-oh!")); // should be ignored
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
index 66700ff4..7238f501 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
@@ -49,6 +49,7 @@ class PushDataTimeoutTest extends AnyFunSuite
Seq("false", "true").foreach { enabled =>
val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
.set("spark.celeborn.push.data.timeout", "5s")
+ .set("spark.celeborn.push.timeoutCheck.interval", "2s")
.set("spark.celeborn.push.replicate.enabled", enabled)
.set("spark.celeborn.client.blacklistSlave.enabled", "false")
val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
index 481a2769..08dd18fa 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
@@ -36,7 +36,7 @@ trait SparkTestBase extends AnyFunSuite
override def beforeAll(): Unit = {
logInfo("test initialized , setup rss mini cluster")
- setUpMiniCluster()
+ setUpMiniCluster(workerNum = 5)
}
override def afterAll(): Unit = {
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
index 74937abd..24e19bd8 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -38,7 +38,6 @@ import
org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMod
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.unsafe.Platform
import org.apache.celeborn.common.util.PackedPartitionId
-import org.apache.celeborn.common.write.PushState
import
org.apache.celeborn.service.deploy.worker.congestcontrol.CongestionController
import org.apache.celeborn.service.deploy.worker.storage.{FileWriter,
HdfsFlusher, LocalFlusher, MapPartitionFileWriter, StorageManager}
@@ -62,7 +61,6 @@ class PushDataHandler extends BaseMessageHandler with Logging
{
var conf: CelebornConf = _
@volatile var pushMasterDataTimeoutTested = false
@volatile var pushSlaveDataTimeoutTested = false
- var pushState: PushState = _
def init(worker: Worker): Unit = {
workerSource = worker.workerSource
@@ -81,8 +79,6 @@ class PushDataHandler extends BaseMessageHandler with Logging
{
storageManager = worker.storageManager
shutdown = worker.shutdown
conf = worker.conf
- pushState = new PushState(conf)
- pushState.startChecker(false)
logInfo(s"diskReserveSize $diskReserveSize")
}
@@ -171,7 +167,6 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
// Fetch real batchId from body will add more cost and no meaning for
replicate.
val doReplicate = location != null && location.getPeer != null && isMaster
- val batchId = if (doReplicate) pushState.nextBatchId() else -1
val softSplit = new AtomicBoolean(false)
if (location == null) {
@@ -272,13 +267,10 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
s"${StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_SLAVE.getMessage}! Peer
$peerWorker unavailable for $location!"))
return
}
- pushState.addBatch(batchId, location.getPeer.hostAndPushPort())
// Handle the response from replica
val wrappedCallback = new RpcResponseCallback() {
override def onSuccess(response: ByteBuffer): Unit = {
- // Only master data will push data to slave
- pushState.removeBatch(batchId,
location.getPeer.hostAndPushPort())
if (response.remaining() > 0) {
val resp = ByteBuffer.allocate(response.remaining())
resp.put(response)
@@ -313,7 +305,6 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
override def onFailure(e: Throwable): Unit = {
logError(s"[handlePushData.onFailure] partitionLocation:
$location", e)
workerSource.incCounter(WorkerSource.PushDataFailCount)
- pushState.removeBatch(batchId,
location.getPeer.hostAndPushPort())
// Throw by slave peer worker
if
(e.getMessage.startsWith(StatusCode.PUSH_DATA_FAIL_SLAVE.getMessage)) {
callbackWithTimer.onFailure(e)
@@ -337,17 +328,10 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
shuffleKey,
pushData.partitionUniqueId,
pushData.body)
- val channelFuture = client.pushData(newPushData, wrappedCallback)
- pushState.pushStarted(
- batchId,
- channelFuture,
- wrappedCallback,
- location.getPeer.hostAndPushPort(),
- shufflePushDataTimeout.get(shuffleKey))
+ client.pushData(newPushData,
shufflePushDataTimeout.get(shuffleKey), wrappedCallback)
} catch {
case e: Exception =>
pushData.body().release()
- pushState.removeBatch(batchId,
location.getPeer.hostAndPushPort())
unavailablePeers.put(peerWorker, System.currentTimeMillis())
workerSource.incCounter(WorkerSource.PushDataFailCount)
callbackWithTimer.onFailure(
@@ -453,7 +437,6 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
// Fetch real batchId from body will add more cost and no meaning for
replicate.
val doReplicate = locations.head != null && locations.head.getPeer != null
&& isMaster
- val batchId = if (doReplicate) pushState.nextBatchId() else -1
// find FileWriters responsible for the data
locations.foreach { loc =>
@@ -544,13 +527,11 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
s"${StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_SLAVE.getMessage}! Peer
$peerWorker unavailable for $location!"))
return
}
- pushState.addBatch(batchId, location.getPeer.hostAndPushPort())
// Handle the response from replica
val wrappedCallback = new RpcResponseCallback() {
override def onSuccess(response: ByteBuffer): Unit = {
// Only master data enable replication will push data to slave
- pushState.removeBatch(batchId,
locations.head.getPeer.hostAndPushPort())
if (response.remaining() > 0) {
val resp = ByteBuffer.allocate(response.remaining())
resp.put(response)
@@ -577,7 +558,6 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
override def onFailure(e: Throwable): Unit = {
workerSource.incCounter(WorkerSource.PushDataFailCount)
- pushState.removeBatch(batchId,
locations.head.getPeer.hostAndPushPort())
// Throw by slave peer worker
if
(e.getMessage.startsWith(StatusCode.PUSH_DATA_FAIL_SLAVE.getMessage)) {
callbackWithTimer.onFailure(e)
@@ -605,19 +585,15 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
pushMergedData.partitionUniqueIds,
batchOffsets,
pushMergedData.body)
- val channelFuture = client.pushMergedData(newPushMergedData,
wrappedCallback)
- pushState.pushStarted(
- batchId,
- channelFuture,
- wrappedCallback,
- location.getPeer.hostAndPushPort(),
- shufflePushDataTimeout.get(shuffleKey))
+ client.pushMergedData(
+ newPushMergedData,
+ shufflePushDataTimeout.get(shuffleKey),
+ wrappedCallback)
} catch {
case e: Exception =>
pushMergedData.body().release()
unavailablePeers.put(peerWorker, System.currentTimeMillis())
workerSource.incCounter(WorkerSource.PushDataFailCount)
- pushState.removeBatch(batchId,
location.getPeer.hostAndPushPort())
callbackWithTimer.onFailure(
new Exception(
s"${StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_SLAVE.getMessage}! Create
connection to peer $peerWorker failed for $location",
