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.celeborn.plugin.flink.utils.Utils.checkNotNull;
  import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.util.concurrent.CompletableFuture;
  
  import javax.annotation.Nullable;
  
  import org.apache.flink.annotation.VisibleForTesting;
  import org.apache.flink.runtime.checkpoint.CheckpointException;
  import org.apache.flink.runtime.event.AbstractEvent;
  import org.apache.flink.runtime.io.network.api.EndOfData;
  import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent;
  import org.apache.flink.runtime.io.network.api.StopMode;
  import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
  import org.apache.flink.runtime.io.network.buffer.Buffer;
  import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
  import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
  import org.apache.flink.runtime.io.network.buffer.BufferPool;
  import org.apache.flink.runtime.io.network.partition.BufferAvailabilityListener;
  import org.apache.flink.runtime.io.network.partition.ResultPartition;
  import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
  import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
  import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
  import org.apache.flink.runtime.io.network.partition.ResultSubpartitionView;
  import org.apache.flink.util.function.SupplierWithException;
  
  import org.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
  import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
  import org.apache.celeborn.plugin.flink.utils.BufferUtils;
  import org.apache.celeborn.plugin.flink.utils.Utils;
  
  /**
   * An abstract {@link ResultPartition} which appends records and events to {@link DataBuffer} and
   * after the {@link DataBuffer} is full, all data in the {@link DataBuffer} will be copied and
   * spilled to the remote shuffle service in subpartition index order sequentially. Large records
   * that can not be appended to an empty {@link DataBuffer} will be spilled directly.
   */
  public abstract class AbstractRemoteShuffleResultPartition extends ResultPartition {
  
    private final RemoteShuffleResultPartitionDelegation delegation;
  
    private final SupplierWithException<BufferPool, IOException> bufferPoolFactory;
  
    public AbstractRemoteShuffleResultPartition(
        String owningTaskName,
        int partitionIndex,
        ResultPartitionID partitionId,
        ResultPartitionType partitionType,
        int numSubpartitions,
        int numTargetKeyGroups,
        int networkBufferSize,
        ResultPartitionManager partitionManager,
        @Nullable BufferCompressor bufferCompressor,
        SupplierWithException<BufferPool, IOException> bufferPoolFactory,
        RemoteShuffleOutputGate outputGate) {
      super(
          owningTaskName,
          partitionIndex,
          partitionId,
          partitionType,
          numSubpartitions,
          numTargetKeyGroups,
          partitionManager,
          bufferCompressor,
          bufferPoolFactory);
      delegation =
          new RemoteShuffleResultPartitionDelegation(
              networkBufferSize, outputGate, this::updateStatistics, numSubpartitions);
      this.bufferPoolFactory = bufferPoolFactory;
    }
  
    @Override
    public void setup() throws IOException {
      // We can't call the `setup` method of the base class, otherwise it will cause a partition leak.
      // The reason is that this partition will be registered to the partition manager during
      // `super.setup()`.
      // Since this is a cluster/remote partition(i.e. resources are not stored on the Flink TM),
      // Flink does not trigger the resource releasing over TM. Therefore, the partition object is
      // leaked.
      // So we copy the logic of `setup` but don't register partition to partition manager.
      checkState(
          this.bufferPool == null,
          "Bug in result partition setup logic: Already registered buffer pool.");
      this.bufferPool = checkNotNull(bufferPoolFactory.get());
      // this is an empty method, but still call it in case of we implement it in the future.
      setupInternal();
      BufferUtils.reserveNumRequiredBuffers(bufferPool, 1);
      delegation.setup(
          bufferPool, bufferCompressor, this::canBeCompressed, this::checkInProduceState);
    }
  
    @Override
    protected void setupInternal() {
      // do not need to implement
    }
  
    @Override
    public void emitRecord(ByteBuffer record, int targetSubpartition) throws IOException {
      delegation.emit(record, targetSubpartition, DataType.DATA_BUFFER, false);
    }
  
    @Override
    public void broadcastRecord(ByteBuffer record) throws IOException {
      delegation.broadcast(record, DataType.DATA_BUFFER);
    }
  
    @Override
    public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
      Buffer buffer = EventSerializer.toBuffer(event, isPriorityEvent);
      try {
        ByteBuffer serializedEvent = buffer.getNioBufferReadable();
        delegation.broadcast(serializedEvent, buffer.getDataType());
      } finally {
        buffer.recycleBuffer();
      }
    }
  
    @Override
    public void alignedBarrierTimeout(long l) {}
  
    @Override
    public void abortCheckpoint(long l, CheckpointException e) {}
  
    @Override
    public void finish() throws IOException {
      Utils.checkState(!isReleased(), "Result partition is already released.");
      broadcastEvent(EndOfPartitionEvent.INSTANCE, false);
      delegation.finish();
      super.finish();
    }
  
    @Override
    public synchronized void close() {
      delegation.close(super::close);
    }
  
    @Override
    protected void releaseInternal() {
      // no-op
    }
  
    @Override
    public void flushAll() {
      delegation.flushAll();
    }
  
    @Override
    public void flush(int subpartitionIndex) {
      flushAll();
    }
  
    @Override
    public CompletableFuture<?> getAvailableFuture() {
      return AVAILABLE;
    }
  
    @Override
    public int getNumberOfQueuedBuffers() {
      return 0;
    }
  
    @Override
    public long getSizeOfQueuedBuffersUnsafe() {
      return 0;
    }
  
    @Override
    public int getNumberOfQueuedBuffers(int targetSubpartition) {
      return 0;
    }
  
    @Override
    public ResultSubpartitionView createSubpartitionView(
        int index, BufferAvailabilityListener availabilityListener) {
      throw new UnsupportedOperationException("Not supported.");
    }
  
    @Override
    public void notifyEndOfData(StopMode mode) throws IOException {
      if (!delegation.isEndOfDataNotified()) {
        broadcastEvent(new EndOfData(mode), false);
        delegation.setEndOfDataNotified(true);
      }
    }
  
    @Override
    public CompletableFuture<Void> getAllDataProcessedFuture() {
      return CompletableFuture.completedFuture(null);
    }
  
    @Override
    public String toString() {
      return "ResultPartition "
          + partitionId.toString()
          + " ["
          + partitionType
          + ", "
          + numSubpartitions
          + " subpartitions, shuffle-descriptor: "
          + delegation.getOutputGate().getShuffleDesc()
          + "]";
    }
  
    @VisibleForTesting
    public RemoteShuffleResultPartitionDelegation getDelegation() {
      return delegation;
    }
  
    public void updateStatistics(BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast) {
      numBuffersOut.inc(isBroadcast ? numSubpartitions : 1);
      long readableBytes =
          (long) bufferWithSubpartition.getBuffer().readableBytes() - BufferUtils.HEADER_LENGTH;
      updateReadableBytes(bufferWithSubpartition, isBroadcast, readableBytes);
      numBytesOut.inc(isBroadcast ? readableBytes * numSubpartitions : readableBytes);
    }
  
    public abstract void updateReadableBytes(
        BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes);
  }
  