/*
 * 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.spark.shuffle.celeborn;

import java.io.IOException;
import java.util.LinkedList;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Consumer;

import org.apache.spark.TaskContext;
import org.apache.spark.memory.MemoryConsumer;
import org.apache.spark.memory.SparkOutOfMemoryError;
import org.apache.spark.memory.TaskMemoryManager;
import org.apache.spark.memory.TooLargePageException;
import org.apache.spark.unsafe.Platform;
import org.apache.spark.unsafe.UnsafeAlignedOffset;
import org.apache.spark.unsafe.array.LongArray;
import org.apache.spark.unsafe.memory.MemoryBlock;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.write.DataPusher;
import org.apache.celeborn.client.write.PushTask;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;

public class SortBasedPusher extends MemoryConsumer {

  class MemoryThresholdManager {

    private final long maxMemoryThresholdInBytes;
    private final double smallPushTolerateFactor;

    MemoryThresholdManager(double maxMemoryFactor, double smallPushTolerateFactor) {
      this.maxMemoryThresholdInBytes = (long) (Runtime.getRuntime().maxMemory() * maxMemoryFactor);
      this.smallPushTolerateFactor = smallPushTolerateFactor;
    }

    private boolean shouldGrow() {
      boolean enoughSpace = pushSortMemoryThreshold <= maxMemoryThresholdInBytes;
      double expectedPushSize = Long.MAX_VALUE;
      if (this.expectedPushedCount != 0) {
        expectedPushSize = this.expectedPushedBytes * 1.0 / this.expectedPushedCount;
      }
      boolean tooManyPushed =
          pushedMemorySizeInBytes * 1.0 / pushedCount * (1 + this.smallPushTolerateFactor)
              < expectedPushSize;
      return useAdaptiveThreshold && enoughSpace && tooManyPushed;
    }

    public void growThresholdIfNeeded() {
      if (shouldGrow()) {
        long oldThreshold = pushSortMemoryThreshold;
        pushSortMemoryThreshold = Math.min(pushSortMemoryThreshold * 2, maxMemoryThresholdInBytes);
        logger.info(
            "grow memory threshold from "
                + Utils.bytesToString(oldThreshold)
                + " to "
                + Utils.bytesToString(pushSortMemoryThreshold));
        pushedCount = 0;
        pushedMemorySizeInBytes = 0;
        expectedPushedBytes = 0;
        expectedPushedCount = 0;
      }
    }

    long pushedCount = 0;
    long pushedMemorySizeInBytes = 0;

    long expectedPushedCount = 0;
    long expectedPushedBytes = 0;

    public void updateStats(long pushedBytes, boolean updateExpected) {
      this.pushedMemorySizeInBytes += pushedBytes;
      this.pushedCount += 1;
      if (updateExpected) {
        this.expectedPushedBytes += pushedBytes;
        this.expectedPushedCount += 1;
      }
    }
  }

  private static final Logger logger = LoggerFactory.getLogger(SortBasedPusher.class);

  private static final int UAO_SIZE = UnsafeAlignedOffset.getUaoSize();

  /** Peak memory used by this sorter so far, in bytes. * */
  private long peakMemoryUsedBytes;

  private ShuffleInMemorySorter inMemSorter;
  private final LinkedList<MemoryBlock> allocatedPages = new LinkedList<>();
  private MemoryBlock currentPage = null;
  private long pageCursor = -1;

  private final ShuffleClient shuffleClient;
  private final TaskContext taskContext;
  private DataPusher dataPusher;
  private final int pushBufferMaxSize;
  private long pushSortMemoryThreshold;
  private final int shuffleId;
  private final int mapId;
  private final int attemptNumber;
  private final int numMappers;
  private final int numPartitions;
  private final Consumer<Integer> afterPush;
  private final LongAdder[] mapStatusLengths;
  private int[] shuffledPartitions = null;
  private int[] inversedShuffledPartitions = null;
  private final SendBufferPool sendBufferPool;

  final MemoryThresholdManager memoryThresholdManager;

  private final boolean useAdaptiveThreshold;

  private final double maxMemoryFactor;

  public SortBasedPusher(
      TaskMemoryManager memoryManager,
      ShuffleClient shuffleClient,
      TaskContext taskContext,
      int shuffleId,
      int mapId,
      int attemptNumber,
      long taskAttemptId,
      int numMappers,
      int numPartitions,
      CelebornConf conf,
      Consumer<Integer> afterPush,
      LongAdder[] mapStatusLengths,
      long pushSortMemoryThreshold,
      SendBufferPool sendBufferPool) {
    super(
        memoryManager,
        (int) Math.min(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, memoryManager.pageSizeBytes()),
        memoryManager.getTungstenMemoryMode());

    this.shuffleClient = shuffleClient;
    this.taskContext = taskContext;

    this.shuffleId = shuffleId;
    this.mapId = mapId;
    this.attemptNumber = attemptNumber;
    this.numMappers = numMappers;
    this.numPartitions = numPartitions;

    if (conf.clientPushSortRandomizePartitionIdEnabled()) {
      shuffledPartitions = new int[numPartitions];
      inversedShuffledPartitions = new int[numPartitions];
      JavaUtils.shuffleArray(shuffledPartitions, inversedShuffledPartitions);
    }

    this.afterPush = afterPush;
    this.mapStatusLengths = mapStatusLengths;

    this.sendBufferPool = sendBufferPool;

    try {
      LinkedBlockingQueue<PushTask> pushTaskQueue = sendBufferPool.acquirePushTaskQueue();
      dataPusher =
          new DataPusher(
              shuffleId,
              mapId,
              attemptNumber,
              taskAttemptId,
              numMappers,
              numPartitions,
              conf,
              shuffleClient,
              pushTaskQueue,
              afterPush,
              mapStatusLengths);
    } catch (InterruptedException e) {
      TaskInterruptedHelper.throwTaskKillException();
    }

    pushBufferMaxSize = conf.clientPushBufferMaxSize();
    useAdaptiveThreshold = conf.clientPushSortUseAdaptiveMemoryThreshold();
    maxMemoryFactor = conf.clientPushSortMaxMemoryFactor();
    this.pushSortMemoryThreshold = pushSortMemoryThreshold;

    this.memoryThresholdManager =
        new MemoryThresholdManager(maxMemoryFactor, conf.clientPushSortSmallPushTolerateFactor());

    int initialSize = Math.min((int) pushSortMemoryThreshold / 8, 1024 * 1024);
    this.inMemSorter = new ShuffleInMemorySorter(this, initialSize);
    this.peakMemoryUsedBytes = getMemoryUsage();
  }

  public long pushData(boolean growThreshold) throws IOException {
    final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
        inMemSorter.getSortedIterator();

    byte[] dataBuf = new byte[pushBufferMaxSize];
    int offSet = 0;
    int currentPartition = -1;
    while (sortedRecords.hasNext()) {
      sortedRecords.loadNext();
      final int partition =
          shuffledPartitions != null
              ? inversedShuffledPartitions[sortedRecords.packedRecordPointer.getPartitionId()]
              : sortedRecords.packedRecordPointer.getPartitionId();
      if (partition != currentPartition) {
        if (currentPartition == -1) {
          currentPartition = partition;
        } else {
          int bytesWritten =
              shuffleClient.mergeData(
                  shuffleId,
                  mapId,
                  attemptNumber,
                  currentPartition,
                  dataBuf,
                  0,
                  offSet,
                  numMappers,
                  numPartitions);
          mapStatusLengths[currentPartition].add(bytesWritten);
          afterPush.accept(bytesWritten);
          memoryThresholdManager.updateStats(offSet, offSet == pushBufferMaxSize);
          currentPartition = partition;
          offSet = 0;
        }
      }
      final long recordPointer = sortedRecords.packedRecordPointer.getRecordPointer();
      final Object recordPage = taskMemoryManager.getPage(recordPointer);
      final long recordOffsetInPage = taskMemoryManager.getOffsetInPage(recordPointer);
      int recordSize = UnsafeAlignedOffset.getSize(recordPage, recordOffsetInPage);

      if (offSet + recordSize > dataBuf.length) {
        try {
          dataPusher.addTask(partition, dataBuf, offSet);
          memoryThresholdManager.updateStats(offSet, true);
        } catch (InterruptedException e) {
          TaskInterruptedHelper.throwTaskKillException();
        }
        offSet = 0;
      }

      long recordReadPosition = recordOffsetInPage + UAO_SIZE;
      Platform.copyMemory(
          recordPage, recordReadPosition, dataBuf, Platform.BYTE_ARRAY_OFFSET + offSet, recordSize);
      offSet += recordSize;
    }
    if (offSet > 0) {
      try {
        dataPusher.addTask(currentPartition, dataBuf, offSet);
        memoryThresholdManager.updateStats(offSet, offSet == pushBufferMaxSize);
      } catch (InterruptedException e) {
        TaskInterruptedHelper.throwTaskKillException();
      }
    }

    if (growThreshold) {
      memoryThresholdManager.growThresholdIfNeeded();
    }
    long freedBytes = freeMemory();
    inMemSorter.freeMemory();
    taskContext.taskMetrics().incMemoryBytesSpilled(freedBytes);

    return freedBytes;
  }

  public boolean insertRecord(
      Object recordBase, long recordOffset, int recordSize, int partitionId, boolean copySize)
      throws IOException {
    int required;
    // Need 4 or 8 bytes to store the record recordSize.
    if (copySize) {
      required = recordSize + 4 + UAO_SIZE;
    } else {
      required = recordSize + UAO_SIZE;
    }

    long threshold = pushSortMemoryThreshold;
    if (getUsed() > threshold
        && pageCursor + required > currentPage.getBaseOffset() + currentPage.size()) {
      logger.debug(
          "Memory used {} exceeds threshold {}, need to trigger push. currentPage size: {}",
          Utils.bytesToString(getUsed()),
          Utils.bytesToString(pushSortMemoryThreshold),
          Utils.bytesToString(currentPage.size()));
      return false;
    }

    allocateMemoryForRecordIfNecessary(required);

    assert (currentPage != null);
    final Object base = currentPage.getBaseObject();
    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
    if (copySize) {
      UnsafeAlignedOffset.putSize(base, pageCursor, recordSize + 4);
      pageCursor += UAO_SIZE;
      Platform.putInt(base, pageCursor, Integer.reverseBytes(recordSize));
      pageCursor += 4;
      Platform.copyMemory(recordBase, recordOffset, base, pageCursor, recordSize);
      pageCursor += recordSize;
    } else {
      UnsafeAlignedOffset.putSize(base, pageCursor, recordSize);
      pageCursor += UAO_SIZE;
      Platform.copyMemory(recordBase, recordOffset, base, pageCursor, recordSize);
      pageCursor += recordSize;
    }
    if (shuffledPartitions != null) {
      inMemSorter.insertRecord(recordAddress, shuffledPartitions[partitionId]);
    } else {
      inMemSorter.insertRecord(recordAddress, partitionId);
    }

    return true;
  }

  private void growPointerArrayIfNecessary(long required) throws IOException {
    assert (inMemSorter != null);
    if (!inMemSorter.hasSpaceForAnotherRecord()) {
      if (inMemSorter.numRecords() <= 0) {
        // Spilling was triggered just before this method was called. The pointer array was freed
        // during the spill, so a new pointer array needs to be allocated here.
        LongArray array = allocateArray(inMemSorter.getInitialSize());
        inMemSorter.expandPointerArray(array);
        return;
      }

      long used = inMemSorter.getMemoryUsage();
      long requestedBytes = used / 8 * 2;
      int allocateMemoryRetryCount = 0;
      int maxMemoryAllocationRetry = 3;
      LongArray array = null;
      boolean continueRetry = true;
      while (allocateMemoryRetryCount < maxMemoryAllocationRetry && continueRetry) {
        try {
          // could trigger spilling
          logger.info("asking for " + requestedBytes + " more bytes to accommodate more records");
          array = allocateArray(requestedBytes);
          continueRetry = false;
        } catch (TooLargePageException e) {
          // The pointer array is too big to fix in a single page, spill.
          logger.info(
              "Pushdata in growPointerArrayIfNecessary, memory used {}",
              Utils.bytesToString(getUsed()));
          pushData(true);
          continueRetry = false;
        } catch (SparkOutOfMemoryError rethrow) {
          // should have trigger spilling
          allocateMemoryRetryCount += 1;
          if (inMemSorter.numRecords() > 0) {
            if (allocateMemoryRetryCount == maxMemoryAllocationRetry) {
              logger.error("OOM, unable to grow the pointer array");
              throw rethrow;
            } else {
              long oldReq = requestedBytes;
              requestedBytes = Math.max((long) (requestedBytes * 0.5), required);
              logger.warn(
                  "cannot allocate "
                      + oldReq
                      + " bytes, cut the request to "
                      + requestedBytes
                      + " bytes and retry",
                  rethrow);
              pushData(true);
            }
          } else {
            // The new array could not be allocated, but that is not an issue as it is longer
            // needed,
            // as all records were spilled.
            continueRetry = false;
          }
        }
      }

      if (inMemSorter.numRecords() <= 0) {
        // Spilling was triggered while trying to allocate the new array.
        if (array != null) {
          // We succeeded in allocating the new array, but, since all records were spilled, a
          // smaller array would also suffice.
          freeArray(array);
        }
        // The pointer array was freed during the spill, so a new pointer array needs to be
        // allocated here.
        array = allocateArray(inMemSorter.getInitialSize());
      }
      inMemSorter.expandPointerArray(array);
    }
  }

  /**
   * Allocates an additional page in order to insert an additional record. This will request
   * additional memory from the memory manager and spill if the requested memory can not be
   * obtained.
   *
   * @param required the required space in the data page, in bytes, including space for storing the
   *     record size.
   */
  private void acquireNewPageIfNecessary(int required) {
    if (currentPage == null
        || pageCursor + required > currentPage.getBaseOffset() + currentPage.size()) {
      currentPage = allocatePage(required);
      pageCursor = currentPage.getBaseOffset();
      allocatedPages.add(currentPage);
    }
  }

  /**
   * Allocates more memory in order to insert an additional record. This will request additional
   * memory from the memory manager and spill if the requested memory can not be obtained.
   *
   * @param required the required space in the data page, in bytes, including space for storing the
   *     record size.
   */
  private void allocateMemoryForRecordIfNecessary(int required) throws IOException {
    // Step 1:
    // Ensure that the pointer array has space for another record. This may cause a spill.
    growPointerArrayIfNecessary(required);
    // Step 2:
    // Ensure that the last page has space for another record. This may cause a spill.
    acquireNewPageIfNecessary(required);
    // Step 3:
    // The allocation in step 2 could have caused a spill, which would have freed the pointer
    // array allocated in step 1. Therefore we need to check again whether we have to allocate
    // a new pointer array.
    //
    // If the allocation in this step causes a spill event then it will not cause the page
    // allocated in the previous step to be freed. The function `spill` only frees memory if at
    // least one record has been inserted in the in-memory sorter. This will not be the case if
    // we have spilled in the previous step.
    //
    // If we did not spill in the previous step then `growPointerArrayIfNecessary` will be a
    // no-op that does not allocate any memory, and therefore can't cause a spill event.
    //
    // Thus there is no need to call `acquireNewPageIfNecessary` again after this step.
    growPointerArrayIfNecessary(required);
  }

  @Override
  public long spill(long l, MemoryConsumer memoryConsumer) throws IOException {
    logger.warn("SortBasedPusher not support spill yet");
    return 0;
  }

  private long getMemoryUsage() {
    long totalPageSize = 0;
    for (MemoryBlock page : allocatedPages) {
      totalPageSize += page.size();
    }
    return ((inMemSorter == null) ? 0 : inMemSorter.getMemoryUsage()) + totalPageSize;
  }

  private void updatePeakMemoryUsed() {
    long mem = getMemoryUsage();
    if (mem > peakMemoryUsedBytes) {
      peakMemoryUsedBytes = mem;
    }
  }

  /** Return the peak memory used so far, in bytes. */
  long getPeakMemoryUsedBytes() {
    updatePeakMemoryUsed();
    return peakMemoryUsedBytes;
  }

  private long freeMemory() {
    updatePeakMemoryUsed();
    long memoryFreed = 0;
    for (MemoryBlock block : allocatedPages) {
      memoryFreed += block.size();
      freePage(block);
    }
    allocatedPages.clear();
    currentPage = null;
    pageCursor = 0;
    return memoryFreed;
  }

  public void cleanupResources() {
    long freedBytes = freeMemory();
    if (inMemSorter != null) {
      inMemSorter.freeMemory();
      inMemSorter = null;
    }
    taskContext.taskMetrics().incMemoryBytesSpilled(freedBytes);
  }

  public long getPushSortMemoryThreshold() {
    return this.pushSortMemoryThreshold;
  }

  public void close(boolean throwTaskKilledOnInterruption) throws IOException {
    cleanupResources();
    try {
      dataPusher.waitOnTermination();
      sendBufferPool.returnPushTaskQueue(dataPusher.getAndResetIdleQueue());
    } catch (InterruptedException e) {
      if (throwTaskKilledOnInterruption) {
        TaskInterruptedHelper.throwTaskKillException();
      }
    }
  }

  // SPARK-29310 opens it to public in Spark 3.0, it's necessary to keep compatible with Spark 2
  @Override
  public long getUsed() {
    return super.getUsed();
  }
}
