/**
 * 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.hadoop.mapred;

import static org.apache.hadoop.mapred.Task.Counter.MAP_OUTPUT_BYTES;
import static org.apache.hadoop.mapred.Task.Counter.MAP_OUTPUT_RECORDS;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.LocalFileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.RawComparator;
import org.apache.hadoop.io.WritableUtils;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.DefaultCodec;
import org.apache.hadoop.mapred.IFile.Writer;
import org.apache.hadoop.mapred.Merger.Segment;
import org.apache.hadoop.mapred.Task.TaskReporter;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.hadoop.util.ResourceCalculatorPlugin.ProcResourceValues;

public class BlockMapOutputBuffer<K extends BytesWritable, V extends BytesWritable>
    implements BlockMapOutputCollector<K, V> {

  private static final Log LOG = LogFactory.getLog(BlockMapOutputBuffer.class.getName());

  private final Partitioner<K, V> partitioner;
  private final int partitions;
  private final JobConf job;
  private final TaskReporter reporter;
  private final Class<K> keyClass;
  private final Class<V> valClass;
  private final int softBufferLimit;
  // Compression for map-outputs
  private CompressionCodec codec = null;
  // main output buffer
  private byte[] kvbuffer;
  private int kvBufferSize;
  // spill accounting
  private volatile int numSpills = 0;
  // number of spills for big records
  private volatile int numBigRecordsSpills = 0;
  private volatile int numBigRecordsWarnThreshold = 500;

  private final FileSystem localFs;
  private final FileSystem rfs;
  private final Counters.Counter mapOutputByteCounter;
  private final Counters.Counter mapOutputRecordCounter;
  private MapSpillSortCounters mapSpillSortCounter;

  private MapTask task;
  private ReducePartition<K, V>[] reducePartitions;
  private ArrayList<SpillRecord> indexCacheList;
  // an array of memory segments, one for each reduce partition.
  private Segment<K,V>[] inMemorySegments;
  private boolean hasInMemorySpill;
  private boolean lastSpillInMem;

  private int totalIndexCacheMemory;
  private static final int INDEX_CACHE_MEMORY_LIMIT = 2 * 1024 * 1024;
  private final MemoryBlockAllocator memoryBlockAllocator;

  @SuppressWarnings( { "unchecked", "deprecation" })
  public BlockMapOutputBuffer(TaskUmbilicalProtocol umbilical, JobConf job,
      TaskReporter reporter, MapTask task) throws IOException,
      ClassNotFoundException {
    this.task = task;
    this.job = job;
    this.reporter = reporter;
    localFs = FileSystem.getLocal(job);
    partitions = job.getNumReduceTasks();
    indexCacheList = new ArrayList<SpillRecord>();
    if (partitions > 0) {
      partitioner = (Partitioner<K, V>) ReflectionUtils.newInstance(job
          .getPartitionerClass(), job);
    } else {
      partitioner = new Partitioner() {
        @Override
        public int getPartition(Object key, Object value, int numPartitions) {
          return -1;
        }

        @Override
        public void configure(JobConf job) {
        }
      };
    }
    rfs = ((LocalFileSystem) localFs).getRaw();

    float spillper = job.getFloat("io.sort.spill.percent", (float) 0.9);
    if (spillper > (float) 1.0 || spillper < (float) 0.0) {
      LOG.error("Invalid \"io.sort.spill.percent\": " + spillper);
      spillper = 0.8f;
    }
    
    lastSpillInMem = job.getBoolean("mapred.map.lastspill.memory", true);
    numBigRecordsWarnThreshold =
        job.getInt("mapred.map.bigrecord.spill.warn.threshold", 500);

    int sortmb = job.getInt("io.sort.mb", 100);
    boolean localMode = job.get("mapred.job.tracker", "local").equals("local");
    if (localMode) {
      sortmb = job.getInt("io.sort.mb.localmode", 100);
    }
    if ((sortmb & 0x7FF) != sortmb) {
      throw new IOException("Invalid \"io.sort.mb\": " + sortmb);
    }
    LOG.info("io.sort.mb = " + sortmb);
    // buffers and accounting
    kvBufferSize = sortmb << 20;
    kvbuffer = new byte[kvBufferSize];
    softBufferLimit = (int) (kvbuffer.length * spillper);
    // k/v serialization
    keyClass = (Class<K>) job.getMapOutputKeyClass();
    valClass = (Class<V>) job.getMapOutputValueClass();
    if (!BytesWritable.class.isAssignableFrom(keyClass)
        || !BytesWritable.class.isAssignableFrom(valClass)) {
      throw new IOException(this.getClass().getName()
          + "  only support " + BytesWritable.class.getName()
          + " as key and value classes, MapOutputKeyClass is "
          + keyClass.getName() + ", MapOutputValueClass is "
          + valClass.getName());
    }

    int numMappers = job.getNumMapTasks();
    memoryBlockAllocator =
        new MemoryBlockAllocator(kvBufferSize, softBufferLimit, numMappers,
            partitions, this);

    // counters
    mapOutputByteCounter = re