/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.cli;

import alluxio.AlluxioURI;
import alluxio.Configuration;
import alluxio.Constants;
import alluxio.PropertyKey;
import alluxio.RuntimeConstants;
import alluxio.client.ClientContext;
import alluxio.client.file.FileInStream;
import alluxio.client.file.FileOutStream;
import alluxio.client.file.FileSystem;
import alluxio.util.FormatUtils;

import java.util.Arrays;

import javax.annotation.concurrent.ThreadSafe;

/**
 * Mini benchmark that writes/reads one file with a given size and operation type.
 */
@ThreadSafe
public final class MiniBenchmark {
  /**
   * The operation types to test.
   */
  enum OperationType {
    /**
     * Read from a file.
     */
    READ,
    /**
     * Write to a file.
     */
    WRITE,
  }

  private MiniBenchmark() {} // prevent instantiation

  private static void usage() {
    System.out.println("Usage:");
    System.out.println("To run a mini benchmark to read or write a file.");
    System.out.println(String
        .format("java -cp %s %s <master address> <[READ, WRITE]> <fileSize> <iterations>",
        RuntimeConstants.ALLUXIO_JAR, TestRunner.class.getCanonicalName()));
  }

  /** Directory for the test generated files. */
  public static final String TEST_PATH = "/default_mini_benchmark";

  /**
   * @param args there are no arguments needed
   * @throws Exception if error occurs during tests
   */
  public static void main(String[] args) throws Exception {
    if (args.length != 4) {
      usage();
      System.exit(-1);
    }
    AlluxioURI masterLocation = new AlluxioURI(args[0]);
    Configuration.set(PropertyKey.MASTER_HOSTNAME, masterLocation.getHost());
    Configuration.set(PropertyKey.MASTER_RPC_PORT, Integer.toString(masterLocation.getPort()));
    ClientContext.init();

    OperationType operationType = Enum.valueOf(OperationType.class, args[1]);
    long fileSize = FormatUtils.parseSpaceSize(args[2]);
    int iterations = Integer.parseInt(args[3]);

    switch (operationType) {
      case READ:
        readFile(fileSize, iterations);
        break;
      case WRITE:
        writeFile(fileSize, iterations);
        break;
    }
  }

  public static void readFile(long fileSize, int iterations) throws Exception {
    FileSystem fileSystem = FileSystem.Factory.get();
    byte[] buffer = new byte[(int) Math.min(fileSize, 4 * Constants.MB)];

    long start = System.nanoTime();
    for (int i = 0; i < iterations; i++) {
      try (FileInStream inStream = fileSystem.openFile(new AlluxioURI(TEST_PATH))) {
        while (inStream.read(buffer) != -1) {}
      }
    }
    System.out.printf("Runtime: %f seconds.\n",
        (System.nanoTime() - start) * 1.0 / Constants.SECOND_NANO);
  }

  public static void writeFile(long fileSize, int iterations) throws Exception {
    FileSystem fileSystem = FileSystem.Factory.get();
    byte[] buffer = new byte[(int) Math.min(fileSize, 4 * Constants.MB)];
    Arrays.fill(buffer, (byte) 'a');
    AlluxioURI path = new AlluxioURI(TEST_PATH);

    long runTime = 0;
    for (int i = 0; i < iterations; i++) {
      if (fileSystem.exists(path)) {
        fileSystem.delete(path);
      }

      long bytesWritten = 0;
      long start = System.nanoTime();
      try (FileOutStream outStream = fileSystem.createFile(new AlluxioURI(TEST_PATH))) {
        while (bytesWritten < fileSize) {
          outStream.write(buffer, 0, (int) Math.min(buffer.length, fileSize - bytesWritten));
          bytesWritten += buffer.length;
        }
      }
      runTime += System.nanoTime() - start;
    }

    System.out.printf("Runtime: %f seconds.\n", runTime * 1.0 / Constants.SECOND_NANO);
  }
}

@@ -852,26 +852,105 @@
 io.util.
-FormatUtil
+CommonUtils;%0Aimport alluxio.util.FormatUtils;%0A%0Aimport com.google.common.base.Precondition
 s;%0A%0Aimpo
@@ -1693,18 +1693,21 @@
 AR, 
-TestRunner
+MiniBenchmark
 .cla
@@ -2287,25 +2287,8 @@
 RT, 
-Integer.toString(
 mast
@@ -2308,17 +2308,16 @@
 tPort())
-)
 ;%0A    Cl
@@ -2517,24 +2517,55 @@
 (args%5B3%5D);%0A%0A
+    CommonUtils.warmUpLoop();%0A%0A
     switch (
@@ -2735,23 +2735,265 @@
     
-%7D%0A  %7D%0A%0A  public
+  default:%0A        throw new RuntimeException(%22Unsupported type.%22);%0A    %7D%0A  %7D%0A%0A  /**%0A   * Reads a file.%0A   *%0A   * @param fileSize the file size%0A   * @param iterations the number of iterations to run%0A   * @throws Exception if it fails to read%0A   */%0A  private
  sta
@@ -3537,22 +3537,183 @@
 %0A  %7D%0A%0A  
-public
+/**%0A   * Writes a file.%0A   *%0A   * @param fileSize the file size%0A   * @param iterations number of iterations%0A   * @throws Exception it if fails to write%0A   */%0A  private
  static 

Address Bin's comments