This is an automated email from the ASF dual-hosted git repository.
ethanfeng 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 cb8df62e [CELEBORN-324][FLINK] Flink plugin needs reuse connections.
(#1257)
cb8df62e is described below
commit cb8df62ec580297f88d78fa57355a68f818684d7
Author: Ethan Feng <[email protected]>
AuthorDate: Tue Feb 21 18:32:00 2023 +0800
[CELEBORN-324][FLINK] Flink plugin needs reuse connections. (#1257)
---
.../flink/network/FlinkTransportClientFactory.java | 21 ++++++++++++++++-----
.../TransportFrameDecoderWithBufferSupplier.java | 18 +++++++++---------
.../plugin/flink/readclient/RssBufferStream.java | 8 +++-----
.../network/client/TransportClientFactory.java | 2 +-
.../common/network/server/BufferStreamManager.java | 2 +-
5 files changed, 30 insertions(+), 21 deletions(-)
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
index 0472fb8e..61c44747 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
@@ -18,7 +18,7 @@
package org.apache.celeborn.plugin.flink.network;
import java.io.IOException;
-import java.net.InetSocketAddress;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;
import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
@@ -28,14 +28,25 @@ import
org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.client.TransportClientFactory;
public class FlinkTransportClientFactory extends TransportClientFactory {
+
+ private ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers;
+
public FlinkTransportClientFactory(TransportContext context) {
super(context);
+ bufferSuppliers = new ConcurrentHashMap<>();
}
- public TransportClient createClient(String remoteHost, int remotePort,
Supplier<ByteBuf> supplier)
+ public TransportClient createClient(String remoteHost, int remotePort)
throws IOException, InterruptedException {
- final InetSocketAddress resolvedAddress = new
InetSocketAddress(remoteHost, remotePort);
- return internalCreateClient(
- resolvedAddress, new
TransportFrameDecoderWithBufferSupplier(supplier));
+ return createClient(
+ remoteHost, remotePort, -1, new
TransportFrameDecoderWithBufferSupplier(bufferSuppliers));
+ }
+
+ public void registerSupplier(long streamId, Supplier<ByteBuf> supplier) {
+ bufferSuppliers.put(streamId, supplier);
+ }
+
+ public void unregisterSupplier(long streamId) {
+ bufferSuppliers.remove(streamId);
}
}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplier.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplier.java
index 604366eb..2dc54162 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplier.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplier.java
@@ -17,6 +17,7 @@
package org.apache.celeborn.plugin.flink.network;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;
import io.netty.channel.ChannelHandlerContext;
@@ -30,7 +31,6 @@ import org.apache.celeborn.plugin.flink.protocol.ReadData;
public class TransportFrameDecoderWithBufferSupplier extends
ChannelInboundHandlerAdapter
implements FrameDecoder {
- private final Supplier<ByteBuf> bufferSupplier;
private int msgSize = -1;
private int bodySize = -1;
private Message.Type curType = Message.Type.UNKNOWN_TYPE;
@@ -39,9 +39,11 @@ public class TransportFrameDecoderWithBufferSupplier extends
ChannelInboundHandl
private ByteBuf externalBuf = null;
private final ByteBuf msgBuf = Unpooled.buffer(8);
private Message curMsg = null;
+ private final ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers;
- public TransportFrameDecoderWithBufferSupplier(Supplier<ByteBuf>
bufferSupplier) {
- this.bufferSupplier = bufferSupplier;
+ public TransportFrameDecoderWithBufferSupplier(
+ ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers) {
+ this.bufferSuppliers = bufferSuppliers;
}
private void copyByteBuf(io.netty.buffer.ByteBuf source, ByteBuf target, int
targetSize) {
@@ -112,16 +114,13 @@ public class TransportFrameDecoderWithBufferSupplier
extends ChannelInboundHandl
private io.netty.buffer.ByteBuf decodeBodyCopyOut(
io.netty.buffer.ByteBuf buf, ChannelHandlerContext ctx) {
+ ReadData readData = (ReadData) curMsg;
if (externalBuf == null) {
- externalBuf = bufferSupplier.get();
+ externalBuf = bufferSuppliers.get(readData.getStreamId()).get();
}
copyByteBuf(buf, externalBuf, bodySize);
if (externalBuf.readableBytes() == bodySize) {
- if (curMsg instanceof ReadData) {
- ((ReadData) curMsg).setFlinkBuffer(externalBuf);
- } else {
- curMsg.setBody(externalBuf.nioBuffer());
- }
+ ((ReadData) curMsg).setFlinkBuffer(externalBuf);
ctx.fireChannelRead(curMsg);
clear();
}
@@ -138,6 +137,7 @@ public class TransportFrameDecoderWithBufferSupplier
extends ChannelInboundHandl
decodeMsg(nettyBuf, ctx);
} else if (bodySize > 0) {
if (curMsg.needCopyOut()) {
+ // Only readdata will enter this branch
nettyBuf = decodeBodyCopyOut(nettyBuf, ctx);
} else {
nettyBuf = decodeBody(nettyBuf, ctx);
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
index db6559ef..b853d43d 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java
@@ -74,8 +74,7 @@ public class RssBufferStream {
this.client =
clientFactory.createClient(
locations[currentLocationIndex].getHost(),
- locations[currentLocationIndex].getFetchPort(),
- supplier);
+ locations[currentLocationIndex].getFetchPort());
OpenStreamWithCredit openBufferStream =
new OpenStreamWithCredit(
shuffleKey,
@@ -93,6 +92,7 @@ public class RssBufferStream {
public void onSuccess(ByteBuffer response) {
StreamHandle streamHandle = (StreamHandle)
Message.decode(response);
RssBufferStream.this.streamId = streamHandle.streamId;
+ clientFactory.registerSupplier(RssBufferStream.this.streamId,
supplier);
mapShuffleClient
.getReadClientHandler()
.registerHandler(streamId, messageConsumer, client);
@@ -157,8 +157,6 @@ public class RssBufferStream {
}
public void close() {
- if (client != null) {
- client.close();
- }
+ clientFactory.unregisterSupplier(this.getStreamId());
}
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
index 4d5e7ffb..de8e61de 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
@@ -188,7 +188,7 @@ public class TransportClientFactory implements Closeable {
*
* <p>As with {@link #createClient(String, int)}, this method is blocking.
*/
- protected TransportClient internalCreateClient(
+ private TransportClient internalCreateClient(
InetSocketAddress address, ChannelInboundHandlerAdapter decoder)
throws IOException, InterruptedException {
Bootstrap bootstrap = new Bootstrap();
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
index 2ede302b..45c83cd3 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
@@ -217,7 +217,7 @@ public class BufferStreamManager {
public MapDataPartition(FileInfo fileInfo) {
this.fileInfo = fileInfo;
- readExecutor =
storageFetcherPool.getExecutorPool(fileInfo.getFilePath());
+ readExecutor =
storageFetcherPool.getExecutorPool(fileInfo.getMountPoint());
}
public synchronized void setupDataPartitionReader(
