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.network;
  
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.function.Supplier;
  
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.channel.ChannelInboundHandlerAdapter;
  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.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.network.protocol.Message;
  import org.apache.celeborn.common.network.util.FrameDecoder;
  import org.apache.celeborn.plugin.flink.protocol.ReadData;
  import org.apache.celeborn.plugin.flink.utils.BufferUtils;
  
  public class TransportFrameDecoderWithBufferSupplier extends ChannelInboundHandlerAdapter
      implements FrameDecoder {
    public static final Logger logger =
        LoggerFactory.getLogger(TransportFrameDecoderWithBufferSupplier.class);
    private int msgSize = -1;
    private int bodySize = -1;
    private Message.Type curType = Message.Type.UNKNOWN_TYPE;
    private ByteBuf headerBuf = Unpooled.buffer(HEADER_SIZE, HEADER_SIZE);
    private io.netty.buffer.CompositeByteBuf bodyBuf = null;
    private ByteBuf externalBuf = null;
    private final ByteBuf msgBuf = Unpooled.buffer(8);
    private Message curMsg = null;
    private int remainingSize = -1;
    private int totalReadBytes = 0;
    private int largeBufferHeaderRemainingBytes = -1;
    private boolean isReadingLargeBuffer = false;
    private ByteBuf largeBufferHeaderBuffer;
    public static final int DISABLE_LARGE_BUFFER_SPLIT_SIZE = -1;
  
    /**
     * The flink buffer size bytes. If the received buffer size large than this value, means that we
     * need to divide the received buffer into multiple smaller buffers, each small than {@link
     * #bufferSizeBytes}. And when this value set to {@link #DISABLE_LARGE_BUFFER_SPLIT_SIZE},
     * indicates that large buffer splitting will not be checked.
     */
    private final int bufferSizeBytes;
  
    private final ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers;
  
    public TransportFrameDecoderWithBufferSupplier(
        ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers) {
      this(bufferSuppliers, DISABLE_LARGE_BUFFER_SPLIT_SIZE);
    }
  
    public TransportFrameDecoderWithBufferSupplier(
        ConcurrentHashMap<Long, Supplier<ByteBuf>> bufferSuppliers, int bufferSizeBytes) {
      this.bufferSuppliers = bufferSuppliers;
      this.bufferSizeBytes = bufferSizeBytes;
    }
  
    private int copyByteBuf(io.netty.buffer.ByteBuf source, ByteBuf target, int targetSize) {
      int bytes = Math.min(source.readableBytes(), targetSize - target.readableBytes());
      target.writeBytes(source.readSlice(bytes).nioBuffer());
      return bytes;
    }
  
    private void decodeHeader(io.netty.buffer.ByteBuf buf, ChannelHandlerContext ctx) {
      copyByteBuf(buf, headerBuf, HEADER_SIZE);
      if (!headerBuf.isWritable()) {
        msgSize = headerBuf.readInt();
        if (msgBuf.capacity() < msgSize) {
          msgBuf.capacity(msgSize);
        }
        msgBuf.clear();
        curType = Message.Type.decode(headerBuf.nioBuffer());
        // type byte is read
        headerBuf.readByte();
        bodySize = headerBuf.readInt();
        if (bufferSizeBytes != DISABLE_LARGE_BUFFER_SPLIT_SIZE && bodySize > bufferSizeBytes) {
          // if the message body size is larger than bufferSizeBytes, we need to split it into two
          // parts: celeborn header and data buffer
          isReadingLargeBuffer = true;
          // create a temporary buffer to store the celeborn header
          largeBufferHeaderBuffer =
              Unpooled.buffer(BufferUtils.HEADER_LENGTH, BufferUtils.HEADER_LENGTH);
          largeBufferHeaderRemainingBytes = BufferUtils.HEADER_LENGTH;
        }
        decodeMsg(buf, ctx);
      }
    }
  
    private void decodeMsg(io.netty.buffer.ByteBuf buf, ChannelHandlerContext ctx) {
      if (msgBuf.readableBytes() < msgSize) {
        copyByteBuf(buf, msgBuf, msgSize);
      }
      if (msgBuf.readableBytes() == msgSize) {
        curMsg = MessageDecoderExt.decode(curType, msgBuf, false);
        if (bodySize <= 0) {
          ctx.fireChannelRead(curMsg);
          clear();
        }
      }
    }
  
    private io.netty.buffer.ByteBuf decodeBody(
        io.netty.buffer.ByteBuf buf, ChannelHandlerContext ctx) {
      if (bodyBuf == null) {
        if (buf.readableBytes() >= bodySize) {
          io.netty.buffer.ByteBuf body = buf.retain().readSlice(bodySize);
          curMsg.setBody(body);
          ctx.fireChannelRead(curMsg);
          clear();
          return buf;
        } else {
          bodyBuf = buf.alloc().compositeBuffer(Integer.MAX_VALUE);
        }
      }
      int remaining = bodySize - bodyBuf.readableBytes();
      io.netty.buffer.ByteBuf next;
      if (remaining >= buf.readableBytes()) {
        next = buf;
        buf = null;
      } else {
        next = buf.retain().readSlice(remaining);
      }
      bodyBuf.addComponent(next).writerIndex(bodyBuf.writerIndex() + next.readableBytes());
      if (bodyBuf.readableBytes() == bodySize) {
        curMsg.setBody(bodyBuf);
        ctx.fireChannelRead(curMsg);
        clear();
      }
      return buf;
    }
  
    private io.netty.buffer.ByteBuf decodeBodyCopyOut(
        io.netty.buffer.ByteBuf buf, ChannelHandlerContext ctx) {
      if (remainingSize > 0) {
        dropUnusedBytes(buf);
        return buf;
      }
  
      ReadData readData = (ReadData) curMsg;
      long streamId = readData.getStreamId();
      if (externalBuf == null) {
        Supplier<ByteBuf> supplier = bufferSuppliers.get(streamId);
        if (supplier == null) {
          return needDropUnusedBytes(streamId, buf);
        } else {
          try {
            externalBuf = supplier.get();
          } catch (Exception e) {
            return needDropUnusedBytes(streamId, buf);
          }
        }
      }
  
      if (largeBufferHeaderRemainingBytes > 0) {
        // if largeBufferHeaderRemainingBytes larger than zero, means that we are reading the celeborn
        // header
        int headerReadBytes = copyByteBuf(buf, largeBufferHeaderBuffer, BufferUtils.HEADER_LENGTH);
        largeBufferHeaderRemainingBytes -= headerReadBytes;
        totalReadBytes += headerReadBytes;
      } else {
        // if largeBufferHeaderRemainingBytes less or equal to zero, means that we are reading the
        // data buffer
        totalReadBytes += copyByteBuf(buf, externalBuf, getTargetDataBufferReadSize());
      }
  
      if (totalReadBytes == bodySize) {
        ByteBuf resultByteBuf;
        if (largeBufferHeaderBuffer == null) {
          resultByteBuf = externalBuf;
        } else {
          // composite the celeborn header and data buffer together
          CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer();
          compositeByteBuf.addComponent(true, largeBufferHeaderBuffer);
          compositeByteBuf.addComponent(true, externalBuf);
          resultByteBuf = compositeByteBuf;
        }
  
        ((ReadData) curMsg).setFlinkBuffer(resultByteBuf);
        ctx.fireChannelRead(curMsg);
        clear();
      }
      return buf;
    }
  
    private io.netty.buffer.ByteBuf needDropUnusedBytes(
        long streamId, io.netty.buffer.ByteBuf byteBuf) {
      logger.warn("Need drop unused bytes, streamId: {}, bodySize: {}", streamId, bodySize);
      remainingSize = bodySize;
      dropUnusedBytes(byteBuf);
      return byteBuf;
    }
  
    private void dropUnusedBytes(io.netty.buffer.ByteBuf source) {
      if (source.readableBytes() > 0) {
        if (remainingSize > source.readableBytes()) {
          remainingSize = remainingSize - source.readableBytes();
          source.skipBytes(source.readableBytes());
        } else {
          source.skipBytes(remainingSize);
          clear();
        }
      }
    }
  
    @Override
    public void channelRead(ChannelHandlerContext ctx, Object data) {
      io.netty.buffer.ByteBuf nettyBuf = (io.netty.buffer.ByteBuf) data;
      try {
        while (nettyBuf != null && nettyBuf.isReadable()) {
          if (headerBuf.isWritable()) {
            decodeHeader(nettyBuf, ctx);
          } else if (curMsg == null) {
            decodeMsg(nettyBuf, ctx);
          } else if (bodySize > 0) {
            if (curMsg.needCopyOut()) {
              // Only read data will enter this branch
              nettyBuf = decodeBodyCopyOut(nettyBuf, ctx);
            } else {
              nettyBuf = decodeBody(nettyBuf, ctx);
            }
          }
        }
      } finally {
        if (nettyBuf != null) {
          nettyBuf.release();
        }
      }
    }
  
    private int getTargetDataBufferReadSize() {
      if (isReadingLargeBuffer) {
        return bodySize - BufferUtils.HEADER_LENGTH;
      }
      return bodySize;
    }
  
    private void clear() {
      externalBuf = null;
      curMsg = null;
      curType = Message.Type.UNKNOWN_TYPE;
      headerBuf.clear();
      bodyBuf = null;
      bodySize = -1;
      remainingSize = -1;
      totalReadBytes = 0;
      largeBufferHeaderRemainingBytes = -1;
      largeBufferHeaderBuffer = null;
      isReadingLargeBuffer = false;
    }
  
    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
      super.channelInactive(ctx);
    }
  
    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
      clear();
  
      headerBuf.release();
      msgBuf.release();
      super.handlerRemoved(ctx);
    }
  
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
      super.exceptionCaught(ctx, cause);
    }
  }
  