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.service.deploy.worker.storage.segment;
  
  import java.util.Deque;
  import java.util.HashMap;
  import java.util.LinkedList;
  import java.util.Map;
  
  import javax.annotation.concurrent.GuardedBy;
  
  import com.google.common.base.Preconditions;
  import io.netty.buffer.ByteBuf;
  import io.netty.channel.Channel;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.meta.DiskFileInfo;
  import org.apache.celeborn.common.network.protocol.RequestMessage;
  import org.apache.celeborn.common.network.protocol.SubPartitionReadData;
  import org.apache.celeborn.service.deploy.worker.memory.BufferRecycler;
  import org.apache.celeborn.service.deploy.worker.memory.RecyclableBuffer;
  import org.apache.celeborn.service.deploy.worker.memory.RecyclableSegmentIdBuffer;
  import org.apache.celeborn.service.deploy.worker.storage.MapPartitionDataReader;
  
  public class SegmentMapPartitionDataReader extends MapPartitionDataReader {
  
    public static final Logger logger = LoggerFactory.getLogger(SegmentMapPartitionDataReader.class);
  
    // The Flink EndOfSegment buffer data type, which has been written
    // to the buffer header, is in the 16th bit(starting from 0).
    private static final int END_OF_SEGMENT_BUFFER_DATA_TYPE = 7;
  
    private final int startPartitionIndex;
  
    private final int endPartitionIndex;
  
    @GuardedBy("lock")
    private final Deque<Integer> backlogs = new LinkedList<>();
  
    @GuardedBy("lock")
    // subpartitionId -> segmentId, current required segmentId by client per subpartition
    private Map<Integer, Integer> subPartitionRequiredSegmentIds = new HashMap<>();
  
    @GuardedBy("lock")
    private boolean hasUpdateSegmentId = false;
  
    @GuardedBy("lock")
    // subpartitionId -> segmentId, segmentId of current reading buffer per subpartition
    private Map<Integer, Integer> subPartitionLastSegmentIds = new HashMap<>();
  
    @GuardedBy("lock")
    // subpartitionId -> buffer index, current reading buffer index per subpartition
    private Map<Integer, Integer> subPartitionNextBufferIndex = new HashMap<>();
  
    public SegmentMapPartitionDataReader(
        int startPartitionIndex,
        int endPartitionIndex,
        DiskFileInfo fileInfo,
        long streamId,
        Channel associatedChannel,
        Runnable recycleStream) {
      super(
          startPartitionIndex,
          endPartitionIndex,
          fileInfo,
          streamId,
          associatedChannel,
          recycleStream);
      this.startPartitionIndex = startPartitionIndex;
      this.endPartitionIndex = endPartitionIndex;
      for (int i = startPartitionIndex; i <= endPartitionIndex; i++) {
        subPartitionLastSegmentIds.put(i, -1);
        subPartitionRequiredSegmentIds.put(i, -1);
        subPartitionNextBufferIndex.put(i, 0);
      }
    }
  
    @Override
    protected void tryNotifyBacklog(int numDataBuffers) {
      notifyBacklog(getBacklog());
    }
  
    @Override
    public synchronized void sendData() {
      while (true) {
        synchronized (lock) {
          if (!hasUpdateSegmentId) {
            logger.debug(
                "The required segment id is not updated for {}, skip sending data.", getStreamId());
            return;
          }
  
          RecyclableBuffer buffer;
  
          // Verify if the client requires the segmentId of the first buffer; if so, send it to the
          // worker; otherwise, wait until the client sends a new segmentId.
          boolean breakLoop = false;
          while ((buffer = buffersToSend.peek()) instanceof RecyclableSegmentIdBuffer) {
            RecyclableSegmentIdBuffer recyclableSegmentIdBuffer = (RecyclableSegmentIdBuffer) buffer;
            int subPartitionId = recyclableSegmentIdBuffer.getSubPartitionId();
            int segmentId = recyclableSegmentIdBuffer.getSegmentId();
            int requiredSegmentId = subPartitionRequiredSegmentIds.get(subPartitionId);
            if (segmentId != requiredSegmentId) {
              // If the queued head buffer is not the required segment id, we do not sent it.
              logger.info(
                  "The queued head buffer is not the required segment id, "
                      + "do not sent it. details: streamId {}, subPartitionId: {}, current segment id: {}, required segment id: {}, reader: {}",
                  streamId,
                  subPartitionId,
                  segmentId,
                  requiredSegmentId,
                  this);
              breakLoop = true;
              break;
            } else {
              buffersToSend.poll();
            }
          }
          if (breakLoop) {
            break;
          }
  
          // fetch first buffer and send
          buffer = fetchBufferToSend();
          if (buffer == null) {
            break;
          }
          if (buffer instanceof RecyclableSegmentIdBuffer) {
            logger.warn("Wrong type of buffer, the RecyclableSegmentIdBuffer is not expected.");
            return;
          }
          sendDataInternal(buffer);
        }
      }
  
      tryRecycleReader();
    }
  
    @Override
    protected void addBuffer(ByteBuf buffer, BufferRecycler bufferRecycler) {
      if (buffer == null) {
        return;
      }
  
      buffer.markReaderIndex();
      int subPartitionId = buffer.readInt();
      // check the buffer type
      boolean isEndOfSegment = isEndOfSegment(buffer, subPartitionId);
      buffer.resetReaderIndex();
  
      synchronized (lock) {
        if (!isReleased) {
          buffersToSend.add(new RecyclableBuffer(buffer, subPartitionId, bufferRecycler));
          increaseBacklog();
        } else {
          bufferRecycler.recycle(buffer);
          numInUseBuffers.decrementAndGet();
          throw new RuntimeException("Partition reader has been failed or finished.", errorCause);
        }
  
        subPartitionNextBufferIndex.compute(subPartitionId, (k, v) -> v + 1);
        if (isEndOfSegment) {
          updateSegmentId(subPartitionId);
        }
      }
    }
  
    @Override
    protected RecyclableBuffer fetchBufferToSend() {
      synchronized (lock) {
        if (isReleased || buffersToSend.isEmpty()) {
          return null;
        }
  
        RecyclableBuffer buffer = null;
        int numCredit = credits.get();
        if (numCredit > 0) {
          buffer = buffersToSend.poll();
        }
        if (numCredit <= 1) {
          notifyBacklog(getBacklog());
        }
        return buffer;
      }
    }
  
    @Override
    protected void notifyBacklog(int backlog) {
      if (backlog == 0) {
        return;
      }
      super.notifyBacklog(backlog);
    }
  
    @Override
    public String toString() {
      final StringBuilder sb = new StringBuilder("SegmentMapPartitionDataReader{");
      sb.append("startPartitionIndex=").append(startPartitionIndex);
      sb.append(", endPartitionIndex=").append(endPartitionIndex);
      sb.append(", streamId=").append(streamId);
      sb.append('}');
      return sb.toString();
    }
  
    private boolean isEndOfSegment(ByteBuf buffer, int subPartitionId) {
      boolean isEndOfSegment = false;
      Preconditions.checkState(
          subPartitionId >= startPartitionIndex && subPartitionId <= endPartitionIndex);
      if (mapFileMeta.hasPartitionSegmentIds()) {
        // skip another 3 int fields, the write details are in
        // FlinkShuffleClientImpl#pushDataToLocation
        buffer.skipBytes(3 * 4);
        int dataType = buffer.readByte();
        isEndOfSegment = END_OF_SEGMENT_BUFFER_DATA_TYPE == dataType;
      }
      return isEndOfSegment;
    }
  
    private int getBacklog() {
      synchronized (lock) {
        Integer backlog = backlogs.peekFirst();
        while (backlog != null && backlog == 0) {
          backlogs.pollFirst();
          backlog = backlogs.peekFirst();
        }
        if (backlog != null) {
          backlogs.pollFirst();
        }
        return backlog == null ? 0 : backlog;
      }
    }
  
    @GuardedBy("lock")
    private void addNewBacklog() {
      backlogs.addLast(0);
    }
  
    @GuardedBy("lock")
    private void increaseBacklog() {
      Integer backlog = backlogs.pollLast();
      if (backlog == null) {
        backlogs.addLast(1);
      } else {
        backlogs.addLast(backlog + 1);
      }
    }
  
    private void updateSegmentId(int subPartitionId) {
      synchronized (lock) {
        // Note that only when writing buffers, it has the segment info, when loading buffers from
        // disk,
        // we do not know when the segment is started, so we try to get the segment id from the buffer
        // index here.
        Integer segmentId =
            mapFileMeta.getSegmentIdByFirstBufferIndex(
                subPartitionId, subPartitionNextBufferIndex.get(subPartitionId));
        if (segmentId == null) {
          return;
        }
        if (segmentId != -1) {
          // For the continuous segments, we use the same backlog
          int subPartitionLastSegmentId = subPartitionLastSegmentIds.get(subPartitionId);
          if (segmentId == 0
              || (!segmentId.equals(subPartitionLastSegmentId)
                  && !segmentId.equals(subPartitionLastSegmentId + 1))) {
            addNewBacklog();
          }
          subPartitionLastSegmentIds.put(subPartitionId, segmentId);
        }
        logger.debug(
            "Insert a buffer to indicate the current segment id "
                + "subPartitionId={}, segmentId={} for {}.",
            subPartitionId,
            segmentId,
            this);
        // Before adding buffers in this segment, add a new buffer to indicate the segment id. So this
        // buffer
        // is in the head of this segment.
        buffersToSend.add(new RecyclableSegmentIdBuffer(subPartitionId, segmentId));
      }
    }
  
    @Override
    public RequestMessage generateReadDataMessage(
        long streamId, int subPartitionId, ByteBuf byteBuf) {
      return new SubPartitionReadData(streamId, subPartitionId, byteBuf);
    }
  
    public void notifyRequiredSegmentId(int segmentId, int subPartitionId) {
      synchronized (lock) {
        logger.debug(
            "Update the required segment id to {}, {}, subPartitionId: {}",
            segmentId,
            this,
            subPartitionId);
        this.subPartitionRequiredSegmentIds.put(subPartitionId, segmentId);
      }
    }
  
    protected void updateSegmentId() {
      synchronized (lock) {
        hasUpdateSegmentId = true;
        for (int i = startPartitionIndex; i <= endPartitionIndex; i++) {
          if (subPartitionLastSegmentIds.get(i) < 0) {
            updateSegmentId(i);
          }
        }
      }
    }
  }
  