[incubator-celeborn] branch main updated: [CELEBORN-314] Add ut for BufferPack and fix npe while logging (#1247)

[et...@apache.org]

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 b09b8552 [CELEBORN-314] Add ut for BufferPack and fix npe while logging (#1247)
b09b8552 is described below

commit b09b85521a3492eba7eac5b2a4c14bbecd3b8b46
Author: zhongqiangchen <zh...@alibaba-inc.com>
AuthorDate: Mon Feb 20 14:57:22 2023 +0800

    [CELEBORN-314] Add ut for BufferPack and fix npe while logging (#1247)
---
 .../celeborn/plugin/flink/buffer/BufferPacker.java |   3 +-
 .../celeborn/plugin/flink/BufferPackSuitJ.java     | 266 +++++++++++++++++++++
 2 files changed, 267 insertions(+), 2 deletions(-)

diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferPacker.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferPacker.java
index 0a2bf6ff..42678938 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferPacker.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/buffer/BufferPacker.java
@@ -109,9 +109,8 @@ public class BufferPacker {
   }
 
   public void drain() throws InterruptedException {
-    logBufferPack(true, cachedBuffer.getDataType(), cachedBuffer.readableBytes());
-
     if (cachedBuffer != null) {
+      logBufferPack(true, cachedBuffer.getDataType(), cachedBuffer.readableBytes());
       handleRipeBuffer(cachedBuffer, currentSubIdx);
     }
     cachedBuffer = null;
diff --git a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/BufferPackSuitJ.java b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/BufferPackSuitJ.java
new file mode 100644
index 00000000..b9b52175
--- /dev/null
+++ b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/BufferPackSuitJ.java
@@ -0,0 +1,266 @@
+/*
+ * 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.plugin.flink;
+
+import static org.apache.flink.runtime.io.network.buffer.Buffer.DataType.DATA_BUFFER;
+import static org.apache.flink.runtime.io.network.buffer.Buffer.DataType.EVENT_BUFFER;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+
+public class BufferPackSuitJ {
+  private static final int BUFFER_SIZE = 20 + 16;
+
+  private NetworkBufferPool networkBufferPool;
+
+  private BufferPool bufferPool;
+
+  @Before
+  public void setup() throws Exception {
+    networkBufferPool = new NetworkBufferPool(10, BUFFER_SIZE);
+    bufferPool = networkBufferPool.createBufferPool(10, 10);
+  }
+
+  @After
+  public void tearDown() {
+    bufferPool.lazyDestroy();
+    assertEquals(10, networkBufferPool.getNumberOfAvailableMemorySegments());
+    networkBufferPool.destroy();
+  }
+
+  @Test
+  public void testPackEmptyBuffers() throws Exception {
+    List<Buffer> buffers = requestBuffers(3);
+    setCompressed(buffers, true, true, false);
+    setDataType(buffers, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+
+    Integer subIdx = 2;
+
+    List<ByteBuf> output = new ArrayList<>();
+    BufferPacker.BiConsumerWithException<ByteBuf, Integer, InterruptedException> ripeBufferHandler =
+        (ripe, sub) -> {
+          assertEquals(subIdx, sub);
+          output.add(ripe);
+        };
+
+    BufferPacker packer = new BufferPacker(ripeBufferHandler);
+    packer.process(buffers.get(0), subIdx);
+    packer.process(buffers.get(1), subIdx);
+    packer.process(buffers.get(2), subIdx);
+    assertTrue(output.isEmpty());
+
+    packer.drain();
+    assertEquals(0, output.size());
+  }
+
+  @Test
+  public void testPartialBuffersForSameSubIdx() throws Exception {
+    List<Buffer> buffers = requestBuffers(3);
+    setCompressed(buffers, true, true, false);
+    setDataType(buffers, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+
+    List<Pair<ByteBuf, Integer>> output = new ArrayList<>();
+    BufferPacker.BiConsumerWithException<ByteBuf, Integer, InterruptedException> ripeBufferHandler =
+        (ripe, sub) -> output.add(Pair.of(ripe, sub));
+    BufferPacker packer = new BufferPacker(ripeBufferHandler);
+    fillBuffers(buffers, 0, 1, 2);
+
+    packer.process(buffers.get(0), 2);
+    packer.process(buffers.get(1), 2);
+    assertEquals(0, output.size());
+
+    packer.process(buffers.get(2), 2);
+    assertEquals(1, output.size());
+
+    packer.drain();
+    assertEquals(2, output.size());
+
+    List<Buffer> unpacked = new ArrayList<>();
+    output.forEach(
+        pair -> {
+          assertEquals(Integer.valueOf(2), pair.getRight());
+          unpacked.addAll(BufferPacker.unpack(pair.getLeft()));
+        });
+    checkIfCompressed(unpacked, true, true, false);
+    checkDataType(unpacked, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+    verifyBuffers(unpacked, 0, 1, 2);
+    unpacked.forEach(Buffer::recycleBuffer);
+  }
+
+  @Test
+  public void testPartialBuffersForMultipleSubIdx() throws Exception {
+    List<Buffer> buffers = requestBuffers(3);
+    setCompressed(buffers, true, true, false);
+    setDataType(buffers, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+
+    List<Pair<ByteBuf, Integer>> output = new ArrayList<>();
+    BufferPacker.BiConsumerWithException<ByteBuf, Integer, InterruptedException> ripeBufferHandler =
+        (ripe, sub) -> output.add(Pair.of(ripe, sub));
+    BufferPacker packer = new BufferPacker(ripeBufferHandler);
+    fillBuffers(buffers, 0, 1, 2);
+
+    packer.process(buffers.get(0), 0);
+    packer.process(buffers.get(1), 1);
+    assertEquals(1, output.size());
+
+    packer.process(buffers.get(2), 1);
+    assertEquals(1, output.size());
+
+    packer.drain();
+    assertEquals(2, output.size());
+
+    List<Buffer> unpacked = new ArrayList<>();
+    for (int i = 0; i < output.size(); i++) {
+      Pair<ByteBuf, Integer> pair = output.get(i);
+      assertEquals(Integer.valueOf(i), pair.getRight());
+      unpacked.addAll(BufferPacker.unpack(pair.getLeft()));
+    }
+
+    checkIfCompressed(unpacked, true, true, false);
+    checkDataType(unpacked, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+    verifyBuffers(unpacked, 0, 1, 2);
+    unpacked.forEach(Buffer::recycleBuffer);
+  }
+
+  @Test
+  public void testUnpackedBuffers() throws Exception {
+    List<Buffer> buffers = requestBuffers(3);
+    setCompressed(buffers, true, true, false);
+    setDataType(buffers, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+
+    List<Pair<ByteBuf, Integer>> output = new ArrayList<>();
+    BufferPacker.BiConsumerWithException<ByteBuf, Integer, InterruptedException> ripeBufferHandler =
+        (ripe, sub) -> output.add(Pair.of(ripe, sub));
+    BufferPacker packer = new BufferPacker(ripeBufferHandler);
+    fillBuffers(buffers, 0, 1, 2);
+
+    packer.process(buffers.get(0), 0);
+    packer.process(buffers.get(1), 1);
+    assertEquals(1, output.size());
+
+    packer.process(buffers.get(2), 2);
+    assertEquals(2, output.size());
+
+    packer.drain();
+    assertEquals(3, output.size());
+
+    List<Buffer> unpacked = new ArrayList<>();
+    for (int i = 0; i < output.size(); i++) {
+      Pair<ByteBuf, Integer> pair = output.get(i);
+      assertEquals(Integer.valueOf(i), pair.getRight());
+      unpacked.addAll(BufferPacker.unpack(pair.getLeft()));
+    }
+
+    checkIfCompressed(unpacked, true, true, false);
+    checkDataType(unpacked, EVENT_BUFFER, DATA_BUFFER, DATA_BUFFER);
+    verifyBuffers(unpacked, 0, 1, 2);
+    unpacked.forEach(Buffer::recycleBuffer);
+  }
+
+  @Test
+  public void testFailedToHandleRipeBufferAndClose() throws Exception {
+    List<Buffer> buffers = requestBuffers(1);
+    setCompressed(buffers, false);
+    setDataType(buffers, DATA_BUFFER);
+    fillBuffers(buffers, 0);
+
+    BufferPacker.BiConsumerWithException<ByteBuf, Integer, InterruptedException> ripeBufferHandler =
+        (ripe, sub) -> {
+          // ripe.release();
+          throw new RuntimeException("Test");
+        };
+    BufferPacker packer = new BufferPacker(ripeBufferHandler);
+    System.out.println(buffers.get(0).refCnt());
+    packer.process(buffers.get(0), 0);
+    try {
+      packer.drain();
+    } catch (RuntimeException e) {
+      e.printStackTrace();
+    } catch (Exception e) {
+      throw e;
+    }
+
+    // this should never throw any exception
+    packer.close();
+    assertEquals(0, bufferPool.bestEffortGetNumOfUsedBuffers());
+  }
+
+  private List<Buffer> requestBuffers(int n) {
+    List<Buffer> buffers = new ArrayList<>();
+    for (int i = 0; i < n; i++) {
+      Buffer buffer = bufferPool.requestBuffer();
+      buffers.add(buffer);
+    }
+    return buffers;
+  }
+
+  private void setCompressed(List<Buffer> buffers, boolean... values) {
+    for (int i = 0; i < buffers.size(); i++) {
+      buffers.get(i).setCompressed(values[i]);
+    }
+  }
+
+  private void setDataType(List<Buffer> buffers, Buffer.DataType... values) {
+    for (int i = 0; i < buffers.size(); i++) {
+      buffers.get(i).setDataType(values[i]);
+    }
+  }
+
+  private void checkIfCompressed(List<Buffer> buffers, boolean... values) {
+    for (int i = 0; i < buffers.size(); i++) {
+      assertEquals(values[i], buffers.get(i).isCompressed());
+    }
+  }
+
+  private void checkDataType(List<Buffer> buffers, Buffer.DataType... values) {
+    for (int i = 0; i < buffers.size(); i++) {
+      assertEquals(values[i], buffers.get(i).getDataType());
+    }
+  }
+
+  private void fillBuffers(List<Buffer> buffers, int... ints) {
+    for (int i = 0; i < buffers.size(); i++) {
+      Buffer buffer = buffers.get(i);
+      ByteBuf target = buffer.asByteBuf();
+      BufferUtils.setBufferHeader(target, buffer.getDataType(), buffer.isCompressed(), 4);
+      target.writerIndex(BufferUtils.HEADER_LENGTH);
+      target.writeInt(ints[i]);
+    }
+  }
+
+  private void verifyBuffers(List<Buffer> buffers, int... expects) {
+    for (int i = 0; i < buffers.size(); i++) {
+      ByteBuf actual = buffers.get(i).asByteBuf();
+      assertEquals(expects[i], actual.getInt(0));
+    }
+  }
+}