[incubator-celeborn] 11/42: [CELEBORN-397][FLINK] Flink plugin support UnpooledByteBufAllocator. (#1324)

[re...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

rexxiong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git

commit e3e2a4f7918f1027fdc1bea32b7e9e4d982e11b3
Author: Ethan Feng <et...@apache.org>
AuthorDate: Mon Mar 13 11:36:13 2023 +0800

    [CELEBORN-397][FLINK] Flink plugin support UnpooledByteBufAllocator. (#1324)
---
 .../plugin/flink/network/FlinkTransportClientFactory.java         | 6 ++++++
 .../celeborn/common/network/client/TransportClientFactory.java    | 8 ++++++--
 2 files changed, 12 insertions(+), 2 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 61c44747b..58ec46562 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
@@ -21,6 +21,7 @@ import java.io.IOException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
+import io.netty.buffer.UnpooledByteBufAllocator;
 import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
 
 import org.apache.celeborn.common.network.TransportContext;
@@ -42,6 +43,11 @@ public class FlinkTransportClientFactory extends TransportClientFactory {
         remoteHost, remotePort, -1, new TransportFrameDecoderWithBufferSupplier(bufferSuppliers));
   }
 
+  @Override
+  protected void initializeMemoryAllocator() {
+    this.pooledAllocator = new UnpooledByteBufAllocator(true);
+  }
+
   public void registerSupplier(long streamId, Supplier<ByteBuf> supplier) {
     bufferSuppliers.put(streamId, supplier);
   }
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 0043d729c..79711111d 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
@@ -27,7 +27,7 @@ import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.base.Preconditions;
 import io.netty.bootstrap.Bootstrap;
-import io.netty.buffer.PooledByteBufAllocator;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.*;
 import io.netty.channel.socket.SocketChannel;
 import org.slf4j.Logger;
@@ -76,7 +76,7 @@ public class TransportClientFactory implements Closeable {
 
   private final Class<? extends Channel> socketChannelClass;
   private EventLoopGroup workerGroup;
-  private PooledByteBufAllocator pooledAllocator;
+  protected ByteBufAllocator pooledAllocator;
 
   public TransportClientFactory(TransportContext context) {
     this.context = Preconditions.checkNotNull(context);
@@ -90,6 +90,10 @@ public class TransportClientFactory implements Closeable {
     logger.info("mode " + ioMode + " threads " + conf.clientThreads());
     this.workerGroup =
         NettyUtils.createEventLoop(ioMode, conf.clientThreads(), conf.getModuleName() + "-client");
+    initializeMemoryAllocator();
+  }
+
+  protected void initializeMemoryAllocator() {
     this.pooledAllocator =
         NettyUtils.createPooledByteBufAllocator(
             conf.preferDirectBufs(), false /* allowCache */, conf.clientThreads());