/*
    * 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.giraph.ooc;
  
  import com.google.common.collect.Maps;
  import org.apache.giraph.conf.GiraphConstants;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.conf.IntConfOption;
  import org.apache.giraph.ooc.command.IOCommand.IOCommandType;
  import org.apache.log4j.Logger;
  
  import java.util.Map;
  import java.util.Queue;
  import java.util.concurrent.ArrayBlockingQueue;
  import java.util.concurrent.atomic.AtomicLong;
  
  /**
   * Class to collect statistics regarding IO operations done in out-of-core
   * mechanism.
   */
  public class OutOfCoreIOStatistics {
    /**
     * An estimate of disk bandwidth. This number is only used just at the start
     * of the computation, and will be updated/replaced later on once a few disk
     * operations happen.
     */
    public static final IntConfOption DISK_BANDWIDTH_ESTIMATE =
        new IntConfOption("giraph.diskBandwidthEstimate", 125,
            "An estimate of disk bandwidth (MB/s). This number is used just at " +
                "the beginning of the computation, and it will be " +
                "updated/replaced once a few disk operations happen.");
    /**
     * How many recent IO operations should we keep track of? Any report given out
     * of this statistics collector is only based on most recent IO operations.
     */
    public static final IntConfOption IO_COMMAND_HISTORY_SIZE =
        new IntConfOption("giraph.ioCommandHistorySize", 50,
            "Number of most recent IO operations to consider for reporting the" +
                "statistics.");
  
    /**
     * Use this option to control how frequently to print OOC statistics.
     */
    public static final IntConfOption STATS_PRINT_FREQUENCY =
        new IntConfOption("giraph.oocStatPrintFrequency", 200,
            "Number of updates before stats are printed.");
  
    /** Class logger */
    private static final Logger LOG =
        Logger.getLogger(OutOfCoreIOStatistics.class);
    /** Estimate of disk bandwidth (bytes/second) */
    private final AtomicLong diskBandwidthEstimate;
    /** Cached value for IO_COMMAND_HISTORY_SIZE */
    private final int maxHistorySize;
    /**
     * Coefficient/Weight of the most recent IO operation toward the disk
     * bandwidth estimate. Basically if the disk bandwidth estimate if d, and the
     * latest IO command happened at the rate of r, the new estimate of disk
     * bandwidth is calculated as:
     * d_new = updateCoefficient * r + (1 - updateCoefficient) * d
     */
    private final double updateCoefficient;
    /** Queue of all recent commands */
    private final Queue<StatisticsEntry> commandHistory;
    /**
     * Command statistics for each type of IO command. This is the statistics of
     * the recent commands in the history we keep track of (with 'maxHistorySize'
     * command in the history).
     */
    private final Map<IOCommandType, StatisticsEntry> aggregateStats;
    /** How many IO command completed? */
    private int numUpdates = 0;
    /** Cached value for {@link #STATS_PRINT_FREQUENCY} */
    private int statsPrintFrequency = 0;
  
    /**
     * Constructor
     *
     * @param conf configuration
     * @param numIOThreads number of disks/IO threads
     */
    public OutOfCoreIOStatistics(ImmutableClassesGiraphConfiguration conf,
                                 int numIOThreads) {
     this.diskBandwidthEstimate =
         new AtomicLong(DISK_BANDWIDTH_ESTIMATE.get(conf) *
             (long) GiraphConstants.ONE_MB);
     this.maxHistorySize = IO_COMMAND_HISTORY_SIZE.get(conf);
     this.updateCoefficient = 1.0 / maxHistorySize;
     // Adding more entry to the capacity of the queue to have a wiggle room
     // if all IO threads are adding/removing entries from the queue
     this.commandHistory =
         new ArrayBlockingQueue<>(maxHistorySize + numIOThreads);
     this.aggregateStats = Maps.newConcurrentMap();
     for (IOCommandType type : IOCommandType.values()) {
       aggregateStats.put(type, new StatisticsEntry(type, 0, 0, 0));
     }
     this.statsPrintFrequency = STATS_PRINT_FREQUENCY.get(conf);
   }
 
   /**
    * Update statistics with the last IO command that is executed.
    *
    * @param type type of the IO command that is executed
    * @param bytesTransferred number of bytes transferred in the last IO command
    * @param duration duration it took for the last IO command to complete
    *                 (milliseconds)
    */
   public void update(IOCommandType type, long bytesTransferred,
                      long duration) {
     StatisticsEntry entry = aggregateStats.get(type);
     synchronized (entry) {
       entry.setOccurrence(entry.getOccurrence() + 1);
       entry.setDuration(duration + entry.getDuration());
       entry.setBytesTransferred(bytesTransferred + entry.getBytesTransferred());
     }
     commandHistory.offer(
         new StatisticsEntry(type, bytesTransferred, duration, 0));
     if (type != IOCommandType.WAIT) {
       // If the current estimate is 'd', the new rate is 'r', and the size of
       // history is 'n', we can simply model all the past command's rate as:
       // d, d, d, ..., d, r
       // where 'd' happens for 'n-1' times. Hence the new estimate of the
       // bandwidth would be:
       // d_new = (d * (n-1) + r) / n = (1-1/n)*d + 1/n*r
       // where updateCoefficient = 1/n
       diskBandwidthEstimate.set((long)
           (updateCoefficient * (bytesTransferred / duration * 1000) +
               (1 - updateCoefficient) * diskBandwidthEstimate.get()));
     }
     if (commandHistory.size() > maxHistorySize) {
       StatisticsEntry removedEntry = commandHistory.poll();
       entry = aggregateStats.get(removedEntry.getType());
       synchronized (entry) {
         entry.setOccurrence(entry.getOccurrence() - 1);
         entry.setDuration(entry.getDuration() - removedEntry.getDuration());
         entry.setBytesTransferred(
             entry.getBytesTransferred() - removedEntry.getBytesTransferred());
       }
     }
     numUpdates++;
     // Outputting log every so many commands
     if (numUpdates % statsPrintFrequency == 0) {
       if (LOG.isInfoEnabled()) {
         LOG.info(this);
       }
     }
   }
 
   @Override
   public String toString() {
     StringBuffer sb = new StringBuffer();
     long waitTime = 0;
     long loadTime = 0;
     long bytesRead = 0;
     long storeTime = 0;
     long bytesWritten = 0;
     for (Map.Entry<IOCommandType, StatisticsEntry> entry :
         aggregateStats.entrySet()) {
       synchronized (entry.getValue()) {
         sb.append(entry.getKey() + ": " + entry.getValue() + ", ");
         if (entry.getKey() == IOCommandType.WAIT) {
           waitTime += entry.getValue().getDuration();
         } else if (entry.getKey() == IOCommandType.LOAD_PARTITION) {
           loadTime += entry.getValue().getDuration();
           bytesRead += entry.getValue().getBytesTransferred();
         } else {
           storeTime += entry.getValue().getDuration();
           bytesWritten += entry.getValue().getBytesTransferred();
         }
       }
     }
     sb.append(String.format("Average STORE: %.2f MB/s, ",
         (double) bytesWritten / storeTime * 1000 / 1024 / 1024));
     sb.append(String.format("DATA_INJECTION: %.2f MB/s, ",
         (double) (bytesRead - bytesWritten) /
             (waitTime + loadTime + storeTime) * 1000 / 1024 / 1024));
     sb.append(String.format("DISK_BANDWIDTH: %.2f MB/s",
         (double) diskBandwidthEstimate.get() / 1024 / 1024));
 
     return sb.toString();
   }
 
   /**
    * @return most recent estimate of the disk bandwidth
    */
   public long getDiskBandwidth() {
     return diskBandwidthEstimate.get();
   }
 
   /**
    * Get aggregate statistics of a given command type in the command history
    *
    * @param type type of the command to get the statistics for
    * @return aggregate statistics for the given command type
    */
   public BytesDuration getCommandTypeStats(IOCommandType type) {
     StatisticsEntry entry = aggregateStats.get(type);
     synchronized (entry) {
       return new BytesDuration(entry.getBytesTransferred(), entry.getDuration(),
           entry.getOccurrence());
     }
   }
 
   /**
    * Helper class to return results of statistics collector for a certain
    * command type
    */
   public static class BytesDuration {
     /** Number of bytes transferred in a few commands of the same type */
     private long bytes;
     /** Duration of it took to execute a few commands of the same type */
     private long duration;
     /** Number of commands executed of the same type */
     private int occurrence;
 
     /**
      * Constructor
      * @param bytes number of bytes transferred
      * @param duration duration it took to execute commands
      * @param occurrence number of commands
      */
     BytesDuration(long bytes, long duration, int occurrence) {
       this.bytes = bytes;
       this.duration = duration;
       this.occurrence = occurrence;
     }
 
     /**
      * @return number of bytes transferred for the same command type
      */
     public long getBytes() {
       return bytes;
     }
 
     /**
      * @return duration it took to execute a few commands of the same type
      */
     public long getDuration() {
       return duration;
     }
 
     /**
      * @return number of commands that are executed of the same type
      */
     public int getOccurrence() {
       return occurrence;
     }
   }
 
   /**
    * Helper class to keep statistics for a certain command type
    */
   private static class StatisticsEntry {
     /** Type of the command */
     private IOCommandType type;
     /**
      * Aggregate number of bytes transferred executing the particular command
      * type in the history we keep
      */
     private long bytesTransferred;
     /**
      * Aggregate duration it took executing the particular command type in the
      * history we keep
      */
     private long duration;
     /**
      * Number of occurrences of the particular command type in the history we
      * keep
      */
     private int occurrence;
 
     /**
      * Constructor
      *
      * @param type type of the command
      * @param bytesTransferred aggregate number of bytes transferred
      * @param duration aggregate execution time
      * @param occurrence number of occurrences of the particular command type
      */
     public StatisticsEntry(IOCommandType type, long bytesTransferred,
                            long duration, int occurrence) {
       this.type = type;
       this.bytesTransferred = bytesTransferred;
       this.duration = duration;
       this.occurrence = occurrence;
     }
 
     /**
      * @return type of the command
      */
     public IOCommandType getType() {
       return type;
     }
 
     /**
      * @return aggregate number of bytes transferred in the particular command
      *         type
      */
     public long getBytesTransferred() {
       return bytesTransferred;
     }
 
     /**
      * Update the aggregate number of bytes transferred
      *
      * @param bytesTransferred aggregate number of bytes
      */
     public void setBytesTransferred(long bytesTransferred) {
       this.bytesTransferred = bytesTransferred;
     }
 
     /**
      * @return aggregate duration it took to execute the particular command type
      */
     public long getDuration() {
       return duration;
     }
 
     /**
      * Update the aggregate duration
      *
      * @param duration aggregate duration
      */
     public void setDuration(long duration) {
       this.duration = duration;
     }
 
     /**
      * @return number of occurrences of the particular command type
      */
     public int getOccurrence() {
       return occurrence;
     }
 
     /**
      * Update the number of occurrences of the particular command type
      *
      * @param occurrence number of occurrences
      */
     public void setOccurrence(int occurrence) {
       this.occurrence = occurrence;
     }
 
     @Override
     public String toString() {
       if (type == IOCommandType.WAIT) {
         return String.format("%.2f sec", duration / 1000.0);
       } else {
         return String.format("%.2f MB/s",
             (double) bytesTransferred / duration * 1000 / 1024 / 2014);
       }
     }
   }
 }