Fixed FetchHandlerSuiteJ

apache · Nov 16, 2023 · 6534819 · 6534819
1 parent e7fc56b
commit 6534819
Showing 1 changed file with 28 additions and 4 deletions.
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
@@ -33,6 +33,8 @@
 import java.util.Random;
 import java.util.UUID;
 
+import com.google.common.base.Throwables;
+import io.netty.channel.Channel;
 import io.netty.channel.embedded.EmbeddedChannel;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -45,11 +47,13 @@
 import org.apache.celeborn.common.identity.UserIdentifier;
 import org.apache.celeborn.common.meta.FileInfo;
 import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
 import org.apache.celeborn.common.network.client.TransportClient;
 import org.apache.celeborn.common.network.client.TransportResponseHandler;
 import org.apache.celeborn.common.network.protocol.ChunkFetchSuccess;
 import org.apache.celeborn.common.network.protocol.Message;
 import org.apache.celeborn.common.network.protocol.OpenStream;
+import org.apache.celeborn.common.network.protocol.RpcFailure;
 import org.apache.celeborn.common.network.protocol.RpcRequest;
 import org.apache.celeborn.common.network.protocol.RpcResponse;
 import org.apache.celeborn.common.network.protocol.StreamHandle;
@@ -289,7 +293,9 @@ private void legacyOpenStreamAndCheck(
     ByteBuffer openStreamByteBuffer =
         new OpenStream(shuffleKey, fileName, startIndex, endIndex).toByteBuffer();
     fetchHandler.receive(
-        client, new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)));
+        client,
+        new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)),
+        createRpcResponseCallback(channel));
     RpcResponse result = channel.readOutbound();
     StreamHandle streamHandler = (StreamHandle) Message.decode(result.body().nioByteBuffer());
     if (endIndex == Integer.MAX_VALUE) {
@@ -318,7 +324,9 @@ private PbStreamHandler openStreamAndCheck(
                     .toByteArray())
             .toByteBuffer();
     fetchHandler.receive(
-        client, new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)));
+        client,
+        new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)),
+        createRpcResponseCallback(channel));
     RpcResponse result = channel.readOutbound();
     PbStreamHandler streamHandler =
         TransportMessage.fromByteBuffer(result.body().nioByteBuffer()).getParsedPayload();
@@ -352,7 +360,8 @@ private void fetchChunkAndCheck(
                                       .setLen(Integer.MAX_VALUE))
                               .build()
                               .toByteArray())
-                      .toByteBuffer())));
+                      .toByteBuffer())),
+          createRpcResponseCallback(channel));
       ChunkFetchSuccess chunkFetchSuccess = channel.readOutbound();
       chunkFetchSuccess.body().retain();
       // chunk size 8m
@@ -372,7 +381,8 @@ private void bufferStreamEnd(TransportClient client, FetchHandler fetchHandler,
                 .toByteArray());
     fetchHandler.receive(
         client,
-        new RpcRequest(dummyRequestId, new NioManagedBuffer(bufferStreamEnd.toByteBuffer())));
+        new RpcRequest(dummyRequestId, new NioManagedBuffer(bufferStreamEnd.toByteBuffer())),
+        createRpcResponseCallback(client.getChannel()));
   }
 
   private void checkOriginFileBeDeleted(FileInfo fileInfo) {
@@ -403,4 +413,18 @@ private ArrayList<Integer> generateMapIds(int batchCountPerMap) {
     Collections.shuffle(ids);
     return ids;
   }
+
+  private RpcResponseCallback createRpcResponseCallback(Channel channel) {
+    return new RpcResponseCallback() {
+      @Override
+      public void onSuccess(ByteBuffer response) {
+        channel.writeAndFlush(new RpcResponse(dummyRequestId, new NioManagedBuffer(response)));
+      }
+
+      @Override
+      public void onFailure(Throwable e) {
+        channel.writeAndFlush(new RpcFailure(dummyRequestId, Throwables.getStackTraceAsString(e)));
+      }
+    };
+  }
 }