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.tez.runtime.library.sort;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.util.*;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  
  import org.apache.hadoop.io.RawComparator;
  import org.apache.hadoop.io.WritableUtils;
  import org.apache.hadoop.io.serializer.Serializer;
  import org.apache.tez.common.counters.TezCounter;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.CelebornTezWriter;
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.unsafe.Platform;
  import org.apache.celeborn.common.util.Utils;
  
  public class CelebornSortBasedPusher<K, V> extends OutputStream {
    private final Logger logger = LoggerFactory.getLogger(CelebornSortBasedPusher.class);
    private final CelebornTezWriter celebornTezWriter;
    private final int maxIOBufferSize;
    private final int spillIOBufferSize;
    private final Serializer<K> kSer;
    private final Serializer<V> vSer;
    private final RawComparator<K> comparator;
    private final AtomicReference<Exception> exception = new AtomicReference<>();
    private final int numOutputs;
    private final TezCounter mapOutputByteCounter;
    private final TezCounter mapOutputRecordCounter;
    private final Map<Integer, List<SerializedKV>> partitionedKVs;
    private int writePos;
    private byte[] serializedKV;
    private final int maxPushDataSize;
    private Map<Integer, AtomicInteger> recordsPerPartition = new HashMap<>();
    private Map<Integer, AtomicLong> bytesPerPartition = new HashMap<>();
    private final boolean needSort;
  
    public CelebornSortBasedPusher(
        Serializer<K> kSer,
        Serializer<V> vSer,
        int maxIOBufferSize,
        int spillIOBufferSize,
        RawComparator<K> comparator,
        TezCounter mapOutputByteCounter,
        TezCounter mapOutputRecordCounter,
        CelebornTezWriter celebornTezWriter,
        CelebornConf celebornConf,
        boolean needSort) {
      this.kSer = kSer;
      this.vSer = vSer;
      this.maxIOBufferSize = maxIOBufferSize;
      this.spillIOBufferSize = spillIOBufferSize;
      this.mapOutputByteCounter = mapOutputByteCounter;
      this.mapOutputRecordCounter = mapOutputRecordCounter;
      this.comparator = comparator;
      this.celebornTezWriter = celebornTezWriter;
      this.needSort = needSort;
      this.numOutputs = celebornTezWriter.getNumPartitions();
      partitionedKVs = new HashMap<>();
      serializedKV = new byte[maxIOBufferSize];
      maxPushDataSize = (int) celebornConf.clientMrMaxPushData();
      try {
        kSer.open(this);
        vSer.open(this);
      } catch (IOException e) {
        exception.compareAndSet(null, e);
      }
    }
  
    public void insert(K key, V value, int partition) {
      try {
        if (writePos >= spillIOBufferSize) {
          // needs to sort and flush data
          if (logger.isDebugEnabled()) {
            logger.debug(
                "Data is large enough {}/{}/{}, trigger sort and flush",
                Utils.bytesToString(writePos),
                Utils.bytesToString(spillIOBufferSize),
                Utils.bytesToString(maxIOBufferSize));
          }
          if (needSort) {
            sortKVs();
          }
          sendKVAndUpdateWritePos();
        }
        int dataLen = insertRecordInternal(key, value, partition);
        if (numOutputs == 1 && !needSort) {
          recordsPerPartition.putIfAbsent(0, new AtomicInteger());
          bytesPerPartition.putIfAbsent(0, new AtomicLong());
          recordsPerPartition.get(0).incrementAndGet();
          bytesPerPartition.get(0).incrementAndGet();
        } else {
          recordsPerPartition.computeIfAbsent(partition, p -> new AtomicInteger());
          bytesPerPartition.computeIfAbsent(partition, p -> new AtomicLong());
          recordsPerPartition.get(partition).incrementAndGet();
          bytesPerPartition.get(partition).incrementAndGet();
        }
        if (logger.isDebugEnabled()) {
          logger.debug(
              "Sort based pusher insert into partition:{} with {} bytes", partition, dataLen);
        }
        mapOutputRecordCounter.increment(1);
        mapOutputByteCounter.increment(dataLen);
      } catch (IOException e) {
        exception.compareAndSet(null, e);
      }
    }
  
    private void sendKVAndUpdateWritePos() throws IOException {
      Iterator<Map.Entry<Integer, List<SerializedKV>>> entryIter =
          partitionedKVs.entrySet().iterator();
      while (entryIter.hasNext()) {
        Map.Entry<Integer, List<SerializedKV>> entry = entryIter.next();
        entryIter.remove();
        int partition = entry.getKey();
        List<SerializedKV> kvs = entry.getValue();
        List<SerializedKV> localKVs = new ArrayList<>();
        int partitionKVTotalLen = 0;
        // process buffers for specific partition
        for (SerializedKV kv : kvs) {
          partitionKVTotalLen += kv.kLen + kv.vLen;
          localKVs.add(kv);
          if (partitionKVTotalLen > maxPushDataSize) {
            // limit max size of pushdata to avoid possible memory issue in Celeborn worker
            // data layout
            // pushdata header (16) + pushDataLen(4) +
            // [varKeyLen+varValLen+serializedRecord(x)][...]
            sendSortedBuffersPartition(partition, localKVs, partitionKVTotalLen, false);
            localKVs.clear();
            partitionKVTotalLen = 0;
          }
        }
        if (!localKVs.isEmpty()) {
          sendSortedBuffersPartition(partition, localKVs, partitionKVTotalLen, true);
        }
        kvs.clear();
      }
      // all data sent
      partitionedKVs.clear();
      writePos = 0;
    }
  
    private void sendSortedBuffersPartition(
        int partition, List<SerializedKV> localKVs, int partitionKVTotalLen, boolean isMerge)
        throws IOException {
      int extraSize = 0;
      for (SerializedKV localKV : localKVs) {
        extraSize += WritableUtils.getVIntSize(localKV.kLen);
        extraSize += WritableUtils.getVIntSize(localKV.vLen);
      }
      // copied from hadoop logic
      extraSize += WritableUtils.getVIntSize(-1);
      extraSize += WritableUtils.getVIntSize(-1);
      // whole buffer's size + [(keyLen+valueLen)+(serializedKey+serializedValue)]
      int length = 4 + extraSize + partitionKVTotalLen;
      byte[] pkvs = new byte[length];
      int pkvsPos = 4;
      Platform.putInt(pkvs, Platform.BYTE_ARRAY_OFFSET, partitionKVTotalLen + extraSize);
      for (SerializedKV kv : localKVs) {
        int recordLen = kv.kLen + kv.vLen;
        // write key len
        pkvsPos = writeVLong(pkvs, pkvsPos, kv.kLen);
        // write value len
        pkvsPos = writeVLong(pkvs, pkvsPos, kv.vLen);
        // write serialized record
        System.arraycopy(serializedKV, kv.offset, pkvs, pkvsPos, recordLen);
        pkvsPos += recordLen;
      }
      // finally write -1 two times
      pkvsPos = writeVLong(pkvs, pkvsPos, -1);
      writeVLong(pkvs, pkvsPos, -1);
      if (isMerge) {
        celebornTezWriter.mergeData(partition, pkvs, length);
      } else {
        celebornTezWriter.pushData(partition, pkvs, length);
      }
    }
  
    /**
     * Write variable length int to array Modified from
     * org.apache.hadoop.io.WritableUtils#writeVLong(java.io.DataOutput, long)
     */
    private int writeVLong(byte[] data, int offset, long dataInt) {
      if (dataInt >= -112L && dataInt <= 127L) {
        data[offset++] = (byte) ((int) dataInt);
        return offset;
      }
  
      int len = -112;
      if (dataInt < 0L) {
        dataInt ^= -1L;
        len = -120;
      }
  
      long tmp = dataInt;
      while (tmp != 0) {
        tmp = tmp >> 8;
        len--;
      }
  
      data[offset++] = (byte) len;
  
      len = len < -120 ? -(len + 120) : -(len + 112);
  
      for (int idx = len; idx != 0; --idx) {
        int shiftBits = (idx - 1) * 8;
        long mask = 0xFFL << shiftBits;
        data[offset++] = ((byte) ((int) ((dataInt & mask) >> shiftBits)));
      }
      return offset;
    }
  
    private void sortKVs() {
      for (Map.Entry<Integer, List<SerializedKV>> partitionKVEntry : partitionedKVs.entrySet()) {
        partitionKVEntry
            .getValue()
            .sort(
                (o1, o2) ->
                    comparator.compare(
                        serializedKV, o1.offset, o1.kLen, serializedKV, o2.offset, o2.kLen));
      }
    }
  
    private int insertRecordInternal(K key, V value, int partition) throws IOException {
      int offset = writePos;
      int keyLen;
      int valLen;
      kSer.serialize(key);
      keyLen = writePos - offset;
      vSer.serialize(value);
      valLen = writePos - keyLen - offset;
      List<SerializedKV> serializedKVs =
          partitionedKVs.computeIfAbsent(partition, v -> new ArrayList<>());
      serializedKVs.add(new SerializedKV(offset, keyLen, valLen));
      if (logger.isDebugEnabled()) {
        logger.debug(
            "Pusher insert into buffer partition:{} offset:{} keyLen:{} valueLen:{} size:{}",
            partition,
            offset,
            keyLen,
            valLen,
            partitionedKVs.size());
      }
      return keyLen + valLen;
    }
  
    public void checkException() throws IOException {
      if (exception.get() != null) {
        throw new IOException("Write data to celeborn failed", exception.get());
      }
    }
  
    @Override
    public void write(int b) throws IOException {
      if (writePos < maxIOBufferSize) {
        serializedKV[writePos] = (byte) b;
        writePos++;
      } else {
        logger.warn("Sort push memory high, write pos {} max size {}", writePos, maxIOBufferSize);
        throw new IOException("Sort pusher memory exhausted.");
      }
    }
  
    @Override
    public void flush() {
      logger.info("Sort based pusher called flush");
      try {
        if (needSort) {
          sortKVs();
        }
        sendKVAndUpdateWritePos();
      } catch (IOException e) {
        exception.compareAndSet(null, e);
      }
    }
  
    @Override
    public void close() {
      flush();
      try {
        celebornTezWriter.close();
      } catch (IOException e) {
        exception.compareAndSet(null, e);
      }
      partitionedKVs.clear();
      serializedKV = null;
    }
  
    public int[] getRecordsPerPartition() {
      int[] values = new int[numOutputs];
      for (int i = 0; i < numOutputs; i++) {
        AtomicInteger records = recordsPerPartition.get(i);
        if (records != null) {
          values[i] = recordsPerPartition.get(i).get();
        } else {
          values[i] = 0;
        }
      }
      return values;
    }
  
    public long[] getBytesPerPartition() {
      long[] values = new long[numOutputs];
      for (int i = 0; i < numOutputs; i++) {
        AtomicLong bytes = bytesPerPartition.get(i);
        if (bytes != null) {
          values[i] = bytes.get();
        } else {
          values[i] = 0;
        }
      }
      return values;
    }
  
    static class SerializedKV {
      final int offset;
      final int kLen;
      final int vLen;
  
      public SerializedKV(int offset, int kLen, int vLen) {
        this.offset = offset;
        this.kLen = kLen;
        this.vLen = vLen;
      }
    }
  }
  