content: /*
   * 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.Arrays;
  import java.util.Collection;
  import java.util.List;
  import java.util.Queue;
  import java.util.Random;
  
  import org.apache.commons.lang3.tuple.Pair;
  import org.apache.flink.core.memory.MemorySegment;
  import org.apache.flink.core.memory.MemorySegmentFactory;
  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.NetworkBuffer;
  import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
  import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
  import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf;
  import org.apache.flink.shaded.netty4.io.netty.buffer.Unpooled;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.Parameterized;
  
  import org.apache.celeborn.plugin.flink.buffer.BufferHeader;
  import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
  import org.apache.celeborn.plugin.flink.buffer.ReceivedNoHeaderBufferPacker;
  import org.apache.celeborn.plugin.flink.utils.BufferUtils;
  
  @RunWith(Parameterized.class)
  public class BufferPackSuiteJ {
    private static final int BUFFER_SIZE = 20 + 16;
  
    private NetworkBufferPool networkBufferPool;
  
    private BufferPool bufferPool;
  
    private boolean bufferPackerReceivedBufferHasHeader;
  
    public BufferPackSuiteJ(boolean bufferPackerReceivedBufferHasHeader) {
      this.bufferPackerReceivedBufferHasHeader = bufferPackerReceivedBufferHasHeader;
    }
  
    @Parameterized.Parameters
    public static Collection prepareData() {
      Object[][] object = {{true}, {false}};
      return Arrays.asList(object);
    }
  
    @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, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) -> {
                assertEquals(subIdx, Integer.valueOf(header.getSubPartitionId()));
                output.add(ripe);
              };
  
      BufferPacker packer = createBufferPakcer(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, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) ->
                  output.add(
                      Pair.of(addBufferHeaderPossible(ripe, header), header.getSubPartitionId()));
      BufferPacker packer = createBufferPakcer(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, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) ->
                  output.add(
                      Pair.of(addBufferHeaderPossible(ripe, header), header.getSubPartitionId()));
      BufferPacker packer = createBufferPakcer(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, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) ->
                  output.add(
                      Pair.of(addBufferHeaderPossible(ripe, header), header.getSubPartitionId()));
      BufferPacker packer = createBufferPakcer(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 testUnpackCompositeBuffer() throws Exception {
      Buffer dataBuffer = bufferPool.requestBuffer();
      fillBufferWithRandomByte(dataBuffer);
      ByteBuf bufferHeaderByteBuf = createBufferHeaderByteBuf(BUFFER_SIZE);
      bufferHeaderByteBuf.retain();
      CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer();
      compositeByteBuf.addComponent(true, bufferHeaderByteBuf);
      compositeByteBuf.addComponent(true, dataBuffer.asByteBuf());
  
      Queue<Buffer> unpackedBuffers = BufferPacker.unpack(compositeByteBuf);
      Assert.assertEquals(1, unpackedBuffers.size());
      Assert.assertEquals(dataBuffer.readableBytes(), unpackedBuffers.peek().readableBytes());
      Assert.assertEquals(BUFFER_SIZE, unpackedBuffers.peek().readableBytes());
      for (int i = 0; i < BUFFER_SIZE; ++i) {
        Assert.assertEquals(
            dataBuffer.getMemorySegment().get(i), unpackedBuffers.peek().getMemorySegment().get(i));
      }
      dataBuffer.recycleBuffer();
    }
  
    @Test
    public void testPackMultipleBuffers() throws Exception {
      int numBuffers = 7;
      List<Buffer> buffers = new ArrayList<>();
      buffers.add(buildSomeBuffer(100));
      buffers.addAll(requestBuffers(numBuffers - 1));
      setCompressed(buffers, true, true, true, false, false, false, true);
      setDataType(
          buffers,
          EVENT_BUFFER,
          DATA_BUFFER,
          DATA_BUFFER,
          EVENT_BUFFER,
          DATA_BUFFER,
          DATA_BUFFER,
          EVENT_BUFFER);
  
      List<Pair<ByteBuf, Integer>> output = new ArrayList<>();
      BufferPacker.BiConsumerWithException<ByteBuf, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) ->
                  output.add(
                      Pair.of(addBufferHeaderPossible(ripe, header), header.getSubPartitionId()));
      BufferPacker packer = createBufferPakcer(ripeBufferHandler);
      fillBuffers(buffers, 0, 1, 2, 3, 4, 5, 6, 7);
  
      for (int i = 0; i < buffers.size(); i++) {
        packer.process(buffers.get(i), 0);
      }
      packer.drain();
  
      List<Buffer> unpacked = new ArrayList<>();
      for (int i = 0; i < output.size(); i++) {
        Pair<ByteBuf, Integer> pair = output.get(i);
        assertEquals(Integer.valueOf(0), pair.getRight());
        unpacked.addAll(BufferPacker.unpack(pair.getLeft()));
      }
      assertEquals(7, unpacked.size());
  
      checkIfCompressed(unpacked, true, true, true, false, false, false, true);
      checkDataType(
          unpacked,
          EVENT_BUFFER,
          DATA_BUFFER,
          DATA_BUFFER,
          EVENT_BUFFER,
          DATA_BUFFER,
          DATA_BUFFER,
          EVENT_BUFFER);
      verifyBuffers(unpacked, 0, 1, 2, 3, 4, 5, 6, 7);
      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, BufferHeader, InterruptedException>
          ripeBufferHandler =
              (ripe, header) -> {
                // ripe.release();
                throw new RuntimeException("Test");
              };
      BufferPacker packer = createBufferPakcer(ripeBufferHandler);
      System.out.println(buffers.get(0).refCnt());
      packer.process(buffers.get(0), 0);
      try {
        packer.drain();
      } catch (RuntimeException ignored) {
      }
  
      // 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();
  
        if (bufferPackerReceivedBufferHasHeader) {
          // If the buffer includes a header, we need to leave space for the header, so we should
          // update the writer index to BufferUtils.HEADER_LENGTH.
          BufferUtils.setBufferHeader(target, buffer.getDataType(), buffer.isCompressed(), 4);
          target.writerIndex(BufferUtils.HEADER_LENGTH);
        } else {
          // if the buffer does not have a header, we can directly write data starting from the
          // beginning of the buffer.
          target.writerIndex(0);
        }
        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));
      }
    }
  
    public static Buffer buildSomeBuffer(int size) {
      final MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(size);
      return new NetworkBuffer(seg, MemorySegment::free, Buffer.DataType.DATA_BUFFER, size);
    }
  
    public ByteBuf addBufferHeaderPossible(ByteBuf byteBuf, BufferHeader bufferHeader) {
      // Try to add buffer header if bufferPackerReceivedBufferHasHeader set to false in BufferPacker
      // drain process
      if (bufferPackerReceivedBufferHasHeader) {
        return byteBuf;
      }
  
      CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer();
      // create a small buffer headerBuf to write the buffer header
      ByteBuf headerBuf = Unpooled.buffer(BufferUtils.HEADER_LENGTH);
  
      // write celeborn buffer header (subpartitionid(4) + attemptId(4) + nextBatchId(4) +
      // compressedsize)
      headerBuf.writeInt(bufferHeader.getSubPartitionId());
      headerBuf.writeInt(0);
      headerBuf.writeInt(0);
      headerBuf.writeInt(
          byteBuf.readableBytes() + (BufferUtils.HEADER_LENGTH - BufferUtils.HEADER_LENGTH_PREFIX));
  
      // write flink buffer header (dataType(1) + isCompress(1) + size(4))
      headerBuf.writeByte(bufferHeader.getDataType().ordinal());
      headerBuf.writeBoolean(bufferHeader.isCompressed());
      headerBuf.writeInt(bufferHeader.getSize());
  
      // composite the headerBuf and data buffer together
      compositeByteBuf.addComponents(true, headerBuf, byteBuf);
      ByteBuf packedByteBuf = Unpooled.wrappedBuffer(compositeByteBuf.nioBuffer());
      byteBuf.writerIndex(0);
      byteBuf.writeBytes(packedByteBuf, 0, packedByteBuf.readableBytes());
      return byteBuf;
    }
  
    public BufferPacker createBufferPakcer(
        BufferPacker.BiConsumerWithException<ByteBuf, BufferHeader, InterruptedException>
            ripeBufferHandler) {
      if (bufferPackerReceivedBufferHasHeader) {
        return new BufferPacker(ripeBufferHandler);
      } else {
        return new ReceivedNoHeaderBufferPacker(ripeBufferHandler);
      }
    }
  
    public ByteBuf createBufferHeaderByteBuf(int dataBufferSize) {
      ByteBuf headerBuf = Unpooled.directBuffer(BufferUtils.HEADER_LENGTH, BufferUtils.HEADER_LENGTH);
      // write celeborn buffer header (subpartitionid(4) + attemptId(4) + nextBatchId(4) +
      // compressedsize)
      headerBuf.writeInt(0);
      headerBuf.writeInt(0);
      headerBuf.writeInt(0);
      headerBuf.writeInt(
          dataBufferSize + (BufferUtils.HEADER_LENGTH - BufferUtils.HEADER_LENGTH_PREFIX));
  
      // write flink buffer header (dataType(1) + isCompress(1) + size(4))
      headerBuf.writeByte(DATA_BUFFER.ordinal());
      headerBuf.writeBoolean(false);
      headerBuf.writeInt(dataBufferSize);
      return headerBuf;
    }
  
    public void fillBufferWithRandomByte(Buffer buffer) {
      Random random = new Random();
      for (int i = 0; i < buffer.getMaxCapacity(); i++) {
        buffer.asByteBuf().writeByte(random.nextInt(255));
      }
    }
  }
  