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.buffer;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayDeque;
  import java.util.Queue;
  
  import org.apache.flink.core.memory.MemorySegment;
  import org.apache.flink.runtime.io.network.buffer.Buffer;
  import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
  import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
  import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBufAllocator;
  import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.plugin.flink.utils.BufferUtils;
  import org.apache.celeborn.plugin.flink.utils.Utils;
  import org.apache.celeborn.reflect.DynMethods;
  
  /**
   * Harness used to pack multiple partial buffers together as a full one. There are two Flink
   * integration strategies: Remote Shuffle Service and Hybrid Shuffle. In Remote Shuffle Service
   * integration strategy, the {@link BufferPacker} will receive buffers containing both shuffle data
   * and the Celeborn header. In Hybrid Shuffle integration strategy employs the subclass {@link
   * ReceivedNoHeaderBufferPacker}, which receives buffers containing only shuffle data. In these two
   * integration strategies, the BufferPacker must utilize different methods to pack buffers, and the
   * result of the packed buffer should be same.
   */
  public class BufferPacker {
    private static Logger logger = LoggerFactory.getLogger(BufferPacker.class);
  
    public interface BiConsumerWithException<T, U, E extends Throwable> {
      void accept(T var1, U var2) throws E;
    }
  
    protected final BiConsumerWithException<ByteBuf, BufferHeader, InterruptedException>
        ripeBufferHandler;
  
    protected Buffer cachedBuffer;
  
    protected int currentSubIdx = -1;
  
    public BufferPacker(
        BiConsumerWithException<ByteBuf, BufferHeader, InterruptedException> ripeBufferHandler) {
      this.ripeBufferHandler = ripeBufferHandler;
    }
  
    public void process(Buffer buffer, int subIdx) throws InterruptedException {
      if (buffer == null) {
        return;
      }
  
      if (buffer.readableBytes() == 0) {
        buffer.recycleBuffer();
        return;
      }
  
      if (cachedBuffer == null) {
        cachedBuffer = buffer;
        currentSubIdx = subIdx;
      } else if (currentSubIdx != subIdx) {
        Buffer dumpedBuffer = cachedBuffer;
        cachedBuffer = buffer;
        int targetSubIdx = currentSubIdx;
        currentSubIdx = subIdx;
        logBufferPack(false, dumpedBuffer.getDataType(), dumpedBuffer.readableBytes());
        handleRipeBuffer(
            dumpedBuffer, targetSubIdx, dumpedBuffer.getDataType(), dumpedBuffer.isCompressed());
      } else {
        /**
         * this is an optimization. if cachedBuffer can contain other buffer, then other buffer can
         * reuse the same HEADER_LENGTH_PREFIX of the cachedBuffer, so cachedbuffer just reads data
         * whose length is buffer.readableBytes() - BufferUtils.HEADER_LENGTH_PREFIX
         */
        if (cachedBuffer.readableBytes() + buffer.readableBytes() - BufferUtils.HEADER_LENGTH_PREFIX
            <= cachedBuffer.getMaxCapacity()) {
          cachedBuffer
              .asByteBuf()
              .writeBytes(
                  buffer.asByteBuf(),
                  BufferUtils.HEADER_LENGTH_PREFIX,
                  buffer.readableBytes() - BufferUtils.HEADER_LENGTH_PREFIX);
          logBufferPack(
              false, buffer.getDataType(), buffer.readableBytes() - BufferUtils.HEADER_LENGTH_PREFIX);
  
          buffer.recycleBuffer();
        } else {
          Buffer dumpedBuffer = cachedBuffer;
          cachedBuffer = buffer;
          logBufferPack(false, dumpedBuffer.getDataType(), dumpedBuffer.readableBytes());
  
          handleRipeBuffer(
              dumpedBuffer, currentSubIdx, dumpedBuffer.getDataType(), dumpedBuffer.isCompressed());
        }
      }
    }
  
    protected void logBufferPack(boolean isDrain, Buffer.DataType dataType, int length) {
      logger.debug(
          "isDrain:{}, cachedBuffer pack partition:{} type:{}, length:{}",
          isDrain,
          currentSubIdx,
          dataType,
          length);
    }
  
    public void drain() throws InterruptedException {
      if (cachedBuffer != null) {
        logBufferPack(true, cachedBuffer.getDataType(), cachedBuffer.readableBytes());
        handleRipeBuffer(
            cachedBuffer, currentSubIdx, cachedBuffer.getDataType(), cachedBuffer.isCompressed());
      }
      cachedBuffer = null;
      currentSubIdx = -1;
    }
  
    protected void handleRipeBuffer(
        Buffer buffer, int subIdx, Buffer.DataType dataType, boolean isCompressed)
        throws InterruptedException {
      // Always set the compress flag to false, because the result buffer generated by {@link
      // BufferPacker} needs to be split into multiple buffers in unpack process,
      // If the compress flag is set to true for this result buffer, it will throw an exception during
      // the unpack process, as compressed buffer cannot be sliced.
      buffer.setCompressed(false);
      ripeBufferHandler.accept(
          buffer.asByteBuf(), new BufferHeader(subIdx, dataType, isCompressed, buffer.getSize()));
    }
  
    public boolean isEmpty() {
      return cachedBuffer == null;
    }
  
    public void close() {
      if (cachedBuffer != null) {
        cachedBuffer.recycleBuffer();
        cachedBuffer = null;
      }
      currentSubIdx = -1;
    }
  
    public static Queue<Buffer> unpack(ByteBuf byteBuf) {
      Queue<Buffer> buffers = new ArrayDeque<>();
      try {
        if (byteBuf instanceof CompositeByteBuf) {
          // If the received byteBuf is a CompositeByteBuf, it indicates that the byteBuf originates
          // from the Flink hybrid shuffle integration strategy. This byteBuf consists of two parts: a
          // celeborn header and a data buffer.
          CompositeByteBuf compositeByteBuf = (CompositeByteBuf) byteBuf;
          ByteBuf headerBuffer = compositeByteBuf.component(0).unwrap();
          ByteBuf dataBuffer = compositeByteBuf.component(1).unwrap();
          dataBuffer.retain();
          Utils.checkState(
              dataBuffer instanceof Buffer, "Illegal data buffer type for CompositeByteBuf.");
          BufferHeader bufferHeader = BufferUtils.getBufferHeaderFromByteBuf(headerBuffer, 0);
          Buffer slice = ((Buffer) dataBuffer).readOnlySlice(0, bufferHeader.getSize());
          buffers.add(
              new UnpackSlicedBuffer(
                  slice,
                  bufferHeader.getDataType(),
                  bufferHeader.isCompressed(),
                  bufferHeader.getSize()));
  
          return buffers;
        }
  
        Utils.checkState(byteBuf instanceof Buffer, "Illegal buffer type.");
  
        Buffer buffer = (Buffer) byteBuf;
        int position = 0;
        int totalBytes = buffer.readableBytes();
        boolean isFirst = true;
        while (position < totalBytes) {
          BufferHeader bufferHeader;
          if (isFirst) {
            bufferHeader = BufferUtils.getBufferHeader(buffer, position, isFirst);
            position += BufferUtils.HEADER_LENGTH;
          } else {
            // in the remaining data, the headlength is BufferUtils.HEADER_LENGTH -
            // BufferUtils.HEADER_LENGTH_PREFIX
            logger.debug("readbuffer: total: {}, position: {}", totalBytes, position);
            bufferHeader = BufferUtils.getBufferHeader(buffer, position);
            position += BufferUtils.HEADER_LENGTH - BufferUtils.HEADER_LENGTH_PREFIX;
          }
  
          Buffer slice = buffer.readOnlySlice(position, bufferHeader.getSize());
          position += bufferHeader.getSize();
  
          buffers.add(
              new UnpackSlicedBuffer(
                  slice,
                  bufferHeader.getDataType(),
                  bufferHeader.isCompressed(),
                  bufferHeader.getSize()));
          slice.retainBuffer();
          isFirst = false;
        }
        logger.debug(
            "Unpack buffer size {} get sliced buffers {} detail {}",
            buffer.getSize(),
            buffers.size(),
            buffers);
  
        return buffers;
      } catch (Throwable throwable) {
        buffers.forEach(Buffer::recycleBuffer);
        throw throwable;
      } finally {
        byteBuf.release();
      }
    }
  
    private static class UnpackSlicedBuffer implements Buffer {
  
      private final Buffer buffer;
  
      private DataType dataType;
  
      private boolean isCompressed;
  
      private final int size;
  
      UnpackSlicedBuffer(Buffer buffer, DataType dataType, boolean isCompressed, int size) {
        this.buffer = buffer;
        this.dataType = dataType;
        this.isCompressed = isCompressed;
        this.size = size;
      }
  
      @Override
      public boolean isBuffer() {
        return dataType.isBuffer();
      }
  
      @Override
      public MemorySegment getMemorySegment() {
        return buffer.getMemorySegment();
      }
  
      @Override
      public int getMemorySegmentOffset() {
        return buffer.getMemorySegmentOffset();
      }
  
      @Override
      public BufferRecycler getRecycler() {
        return buffer.getRecycler();
      }
  
      // Flink 1.18.0
      // [FLINK-32549][network] Tiered storage memory manager supports ownership transfer for buffers
      @SuppressWarnings("MissingOverride")
      public void setRecycler(BufferRecycler bufferRecycler) {
        DynMethods.builder("setRecycler")
            .impl(buffer.getClass().getName(), BufferRecycler.class)
            .build()
            .bind(buffer)
            .invoke(bufferRecycler);
      }
  
      @Override
      public void recycleBuffer() {
        buffer.recycleBuffer();
      }
  
      @Override
      public boolean isRecycled() {
        return buffer.isRecycled();
      }
  
      @Override
      public Buffer retainBuffer() {
        return buffer.retainBuffer();
      }
  
      @Override
      public Buffer readOnlySlice() {
        return buffer.readOnlySlice();
      }
  
      @Override
      public Buffer readOnlySlice(int i, int i1) {
        return buffer.readOnlySlice(i, i1);
      }
  
      @Override
      public int getMaxCapacity() {
        return buffer.getMaxCapacity();
      }
  
      @Override
      public int getReaderIndex() {
        return buffer.getReaderIndex();
      }
  
      @Override
      public void setReaderIndex(int i) throws IndexOutOfBoundsException {
        buffer.setReaderIndex(i);
      }
  
      @Override
      public int getSize() {
        return size;
      }
  
      @Override
      public void setSize(int i) {
        buffer.setSize(i);
      }
  
      @Override
      public int readableBytes() {
        return buffer.readableBytes();
      }
  
      @Override
      public ByteBuffer getNioBufferReadable() {
        return buffer.getNioBufferReadable();
      }
  
      @Override
      public ByteBuffer getNioBuffer(int i, int i1) throws IndexOutOfBoundsException {
        return buffer.getNioBuffer(i, i1);
      }
  
      @Override
      public void setAllocator(ByteBufAllocator byteBufAllocator) {
        buffer.setAllocator(byteBufAllocator);
      }
  
      @Override
      public ByteBuf asByteBuf() {
        return buffer.asByteBuf();
      }
  
      @Override
      public boolean isCompressed() {
        return isCompressed;
      }
  
      @Override
      public void setCompressed(boolean b) {
        isCompressed = b;
      }
  
      @Override
      public DataType getDataType() {
        return dataType;
      }
  
      @Override
      public void setDataType(DataType dataType) {
        this.dataType = dataType;
      }
  
      @Override
      public int refCnt() {
        return buffer.refCnt();
      }
  
      @Override
      public String toString() {
        return "UnpackSlicedBuffer{"
            + "dataType="
            + dataType
            + ", isCompressed="
            + isCompressed
            + ", size="
            + size
            + ", hash="
            + System.identityHashCode(this)
            + '}';
      }
    }
  }
  