alluxio-master-2.x/stress/shell/src/main/java/alluxio/stress/cli/MultiOperationStressMasterBench.java
/*
 * 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.stress.cli;

import static alluxio.stress.master.MultiOperationMasterBenchParameters.BASES_OPTION_NAME;
import static alluxio.stress.master.MultiOperationMasterBenchParameters.OPERATIONS_OPTION_NAME;
import static alluxio.stress.master.MultiOperationMasterBenchParameters.OPERATIONS_RATIO_OPTION_NAME;
import static alluxio.stress.master.MultiOperationMasterBenchParameters.TARGET_THROUGHPUTS_OPTION_NAME;
import static alluxio.stress.master.MultiOperationMasterBenchParameters.THREADS_RATIO_OPTION_NAME;

import alluxio.AlluxioURI;
import alluxio.annotation.SuppressFBWarnings;
import alluxio.conf.InstancedConfiguration;
import alluxio.conf.PropertyKey;
import alluxio.conf.Source;
import alluxio.exception.AlluxioException;
import alluxio.exception.status.InvalidArgumentException;
import alluxio.grpc.CreateDirectoryPOptions;
import alluxio.grpc.DeletePOptions;
import alluxio.stress.StressConstants;
import alluxio.stress.common.FileSystemClientType;
import alluxio.stress.master.MasterBenchParameters;
import alluxio.stress.master.MasterBenchTaskResultStatistics;
import alluxio.stress.master.MultiOperationMasterBenchParameters;
import alluxio.stress.master.MultiOperationMasterBenchTaskResult;
import alluxio.stress.master.Operation;
import alluxio.util.CommonUtils;
import alluxio.util.FormatUtils;
import alluxio.util.io.PathUtils;

import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.RateLimiter;
import org.HdrHistogram.Histogram;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.Random;
import java.util.concurrent.Callable;
import java.util.function.Function;

/**
 * A stress master bench that benchmark a set of combined operations.
 */
// TODO(jiacheng): avoid the implicit casts and @SuppressFBWarnings
public class MultiOperationStressMasterBench
    extends StressMasterBenchBase<MultiOperationMasterBenchTaskResult,
    MultiOperationMasterBenchParameters> {
  private static final Logger LOG = LoggerFactory.getLogger(MultiOperationStressMasterBench.class);

  /**
   * Creates instance.
   */
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  public MultiOperationStressMasterBench() {
    super(new MultiOperationMasterBenchParameters());
  }

  /**
   * @param args command-line arguments
   */
  public static void main(String[] args) {
    mainInternal(args, new MultiOperationStressMasterBench());
  }

  @Override
  public String getBenchDescription() {
    return String.join("\n", ImmutableList.of(
        "A benchmarking tool to measure the master performance of Alluxio",
        "This benchmark is a variant of StressMasterBench, where a user can specify multiple ",
        "operations to run concurrently with some constraints to limit the throughput of ",
        "each operations.",
        "",
        "Example:",
        "# this would continuously run `GetFileStatus` and CreateFile "
        + "for 30s and record the throughput after 5s warmup. "
        + "The CreateFile throughput will be 4x more than GetFileStatus.",
        "$ bin/alluxio runClass alluxio.stress.cli.MultiOperationStressMasterBench \\",
        "--operations GetFileStatus,CreateFile --operations-ratio 1,4  --fixed-counts 100,1 \\",
        "--warmup 1s  --duration 10s -client-type AlluxioNative --threads 16 ",
        "",
        "Use the following command to prepare the test directory to list with before running test:",
        "$ bin/alluxio runClass alluxio.stress.cli.StressMasterBench \\",
        "--base alluxio:///stress-master-base-0 --fixed-count 100 --stop-count 100 --warmup 5s \\",
        "--operation CreateFile --duration 30s --client-type AlluxioNative"
    ));
  }

  @Override
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  public void validateParams() throws Exception {
    int numOperations = mParameters.mOperations.length;
    if (numOperations < 2) {
      throw new InvalidArgumentException(String.format(
          "%s operations must contain at least two operations.", OPERATIONS_OPTION_NAME));
    }

    if (mParameters.mBasePaths.length != mParameters.mOperations.length) {
      throw new InvalidArgumentException(
          String.format("%s must contain the same number of params as %s .",
              BASES_OPTION_NAME, OPERATIONS_OPTION_NAME));
    }

    InvalidArgumentException duplicatedModeSpecifiedException = new InvalidArgumentException(
        String.format("Exact one of the following params %s, %s, %s must be specified",
            TARGET_THROUGHPUTS_OPTION_NAME, THREADS_RATIO_OPTION_NAME,
            OPERATIONS_RATIO_OPTION_NAME));

    boolean modeSpecified = false;
    if (mParameters.mTargetThroughputs != null) {
      if (mParameters.mTargetThroughputs.length != numOperations) {
        throw new InvalidArgumentException(
            String.format("If specified, %s must contain the same number of args as %s.",
                TARGET_THROUGHPUTS_OPTION_NAME, OPERATIONS_OPTION_NAME));
      }
      modeSpecified = true;
    }

    if (mParameters.mThreadsRatio != null) {
      if (mParameters.mThreadsRatio.length != numOperations) {
        throw new InvalidArgumentException(
            String.format("If specified, %s must contain the same number of args as %s.",
                THREADS_RATIO_OPTION_NAME, OPERATIONS_OPTION_NAME));
      }
      if (modeSpecified) {
        throw duplicatedModeSpecifiedException;
      }
      modeSpecified = true;
    }

    if (mParameters.mOperationsRatio != null) {
      if (mParameters.mOperationsRatio.length != numOperations) {
        throw new InvalidArgumentException(
            String.format("If specified, %s must contain the same number of args as %s.",
                OPERATIONS_RATIO_OPTION_NAME, OPERATIONS_OPTION_NAME));
      }
      if (modeSpecified) {
        throw duplicatedModeSpecifiedException;
      }
      modeSpecified = true;
    }

    if (!modeSpecified) {
      throw duplicatedModeSpecifiedException;
    }
  }

  @Override
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  public void prepare() throws Exception {
    if (!mBaseParameters.mDistributed) {
      // set hdfs conf for preparation client
      Configuration hdfsConf = new Configuration();
      // force delete, create dirs through to UFS
      hdfsConf.set(PropertyKey.Name.USER_FILE_DELETE_UNCHECKED, "true");
      hdfsConf.set(PropertyKey.Name.USER_FILE_WRITE_TYPE_DEFAULT, mParameters.mWriteType);
      // more threads for parallel deletes for cleanup
      hdfsConf.set(PropertyKey.Name.USER_FILE_MASTER_CLIENT_POOL_SIZE_MAX, "256");
      FileSystem prepareFs = FileSystem.get(new URI(mParameters.mBasePath), hdfsConf);

      for (int i = 0; i < mParameters.mOperations.length; i++) {
        Operation op = mParameters.mOperations[i];
        // initialize the base, for only the non-distributed task (the cluster launching task)
        Path path = new Path(mParameters.mBasePaths[i]);
        // the base path depends on the operation
        Path basePath;
        if (op == Operation.CREATE_DIR) {
          basePath = new Path(path, "dirs");
        } else {
          basePath = new Path(path, "files");
        }
        if (!mParameters.mSkipPrepare) {
          if (op == Operation.CREATE_FILE
              || op == Operation.CREATE_DIR) {
            LOG.info("Cleaning base path: {}", basePath);
            long start = CommonUtils.getCurrentMs();
            deletePaths(prepareFs, basePath);
            long end = CommonUtils.getCurrentMs();
            LOG.info("Cleanup took: {} s", (end - start) / 1000.0);
            prepareFs.mkdirs(basePath);
          } else {
            // these are read operations. the directory must exist
            if (!prepareFs.exists(basePath)) {
              throw new IllegalStateException(String
                  .format("base path (%s) must exist for operation (%s)", basePath,
                      op));
            }
          }
        }
        if (!prepareFs.isDirectory(basePath)) {
          throw new IllegalStateException(String
              .format("base path (%s) must be a directory for operation (%s)", basePath,
                  op));
        }
      }
    }

    if (mParameters.mClientType == FileSystemClientType.ALLUXIO_HDFS) {
      throw new RuntimeException(
          "ALLUXIO HDFS Compatible API is not support for multi operations stress master bench");
    }

    // set hdfs conf for all test clients
    LOG.info("Using ALLUXIO Native API to perform the test.");
    InstancedConfiguration alluxioProperties = alluxio.conf.Configuration.copyGlobal();
    alluxioProperties.set(
        PropertyKey.USER_FILE_WRITE_TYPE_DEFAULT, mParameters.mWriteType, Source.RUNTIME);
    mCachedNativeFs = new alluxio.client.file.FileSystem[mParameters.mClients];
    for (int i = 0; i < mCachedNativeFs.length; i++) {
      mCachedNativeFs[i] = alluxio.client.file.FileSystem.Factory
          .create(alluxioProperties);
    }

    for (int i = 0; i < mParameters.mOperations.length; ++i) {
      if (mParameters.mOperations[i] == Operation.CRURD
          || mParameters.mOperations[i] == Operation.CREATE_DELETE_FILE) {
        AlluxioURI uri = new AlluxioURI(
            PathUtils.concatPath(mParameters.mBasePaths[i], mFilesDir, mBaseParameters.mId));
        if (mCachedNativeFs[0].exists(uri)) {
          mCachedNativeFs[0].delete(uri, DeletePOptions.newBuilder().setRecursive(true).build());
        }
        mCachedNativeFs[0].createDirectory(uri, CreateDirectoryPOptions.getDefaultInstance());
      }
    }
  }

  @Override
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  protected BenchThread getBenchThread(BenchContext context, int index) {
    if (mParameters.mClientType == FileSystemClientType.ALLUXIO_NATIVE) {
      return new AlluxioNativeBenchThread(context,
          mCachedNativeFs[index % mCachedNativeFs.length], index);
    }
    throw new RuntimeException(
        "ALLUXIO HDFS and POSIX API is not support for multi operations stress master bench");
  }

  private RateLimiter createUnlimitedRateLimiter() {
    return RateLimiter.create(Double.MAX_VALUE);
  }

  @Override
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  protected StressMasterBenchBase<
      MultiOperationMasterBenchTaskResult, MultiOperationMasterBenchParameters>.BenchContext
      getContext() throws IOException, AlluxioException {
    final BenchContext benchContext;
    if (mParameters.mTargetThroughputs == null) {
      RateLimiter[] rateLimiters = new RateLimiter[mParameters.mOperations.length];
      for (int i = 0; i < mParameters.mOperations.length; i++) {
        rateLimiters[i] = createUnlimitedRateLimiter();
      }
      benchContext = new BenchContext(
          createUnlimitedRateLimiter(), rateLimiters,
          mParameters.mOperations, mParameters.mBasePaths, mParameters.mDuration);
    } else {
      int sum = 0;
      RateLimiter[] rateLimiters = new RateLimiter[mParameters.mOperations.length];
      for (int i = 0; i < mParameters.mTargetThroughputs.length; i++) {
        rateLimiters[i] = RateLimiter.create(mParameters.mTargetThroughputs[i]);
        sum += mParameters.mTargetThroughputs[i];
      }
      benchContext = new BenchContext(
          RateLimiter.create(sum), rateLimiters, mParameters.mOperations, mParameters.mBasePaths,
          mParameters.mDuration);
    }
    return benchContext;
  }

  protected abstract class BenchThread implements Callable<Void> {
    private final BenchContext mContext;
    private final Histogram[] mResponseTimeNs;
    protected final Path[] mBasePaths;
    protected final Path[] mFixedBasePaths;
    private final MultiOperationMasterBenchTaskResult mResult;
    private final int mThreadIndex;
    private final Random mRandom = new Random();

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private BenchThread(
        BenchContext context, int threadIndex) {
      mContext = context;
      mResponseTimeNs = new Histogram[mParameters.mOperations.length];
      mBasePaths = mContext.getBasePaths();
      mFixedBasePaths = mContext.getFixedPaths();

      mResult = new MultiOperationMasterBenchTaskResult(mParameters.mOperations);
      for (int i = 0; i < mParameters.mOperations.length; i++) {
        mResponseTimeNs[i] = new Histogram(StressConstants.TIME_HISTOGRAM_MAX,
            StressConstants.TIME_HISTOGRAM_PRECISION);
      }
      mThreadIndex = threadIndex;
    }

    @Override
    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    public Void call() {
      try {
        if (mParameters.mOperationsRatio != null) {
          runOperationsRatioBased();
        } else if (mParameters.mThreadsRatio != null) {
          runThreadsRatioBased();
        } else {
          runInternalTargetThroughputBased();
        }
      } catch (Exception e) {
        e.printStackTrace();
        LOG.warn("Exception during bench thread runInternal", e);
        mResult.addErrorMessage(e.getMessage());
      }

      // Update local thread result
      mResult.setEndMs(CommonUtils.getCurrentMs());
      List<MasterBenchTaskResultStatistics> statistics = mResult.getAllStatistics();
      for (int i = 0; i < statistics.size(); i++) {
        statistics.get(i).encodeResponseTimeNsRaw(mResponseTimeNs[i]);
      }
      mResult.setParameters(mParameters);
      mResult.setBaseParameters(mBaseParameters);

      // merge local thread result with full result
      mContext.mergeThreadResult(mResult);

      return null;
    }

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private void runInternal(Function<Integer, Integer> operationIndexSelector) throws Exception {
      // When to start recording measurements
      long recordMs = mContext.getStartMs() + FormatUtils.parseTimeSize(mParameters.mWarmup);
      mResult.setRecordStartMs(recordMs);

      boolean useStopCount = mParameters.mStopCount != MasterBenchParameters.STOP_COUNT_INVALID;

      long bucketSize = (mContext.getEndMs() - recordMs) / StressConstants.MAX_TIME_COUNT;

      long waitMs = mContext.getStartMs() - CommonUtils.getCurrentMs();
      if (waitMs < 0) {
        throw new IllegalStateException(String.format(
            "Thread missed barrier. Increase the start delay. start: %d current: %d",
            mContext.getStartMs(), CommonUtils.getCurrentMs()));
      }
      CommonUtils.sleepMs(waitMs);

      long totalCounter;
      int operationIndex = 0;
      while (true) {
        if (Thread.currentThread().isInterrupted()) {
          break;
        }
        if (!useStopCount && CommonUtils.getCurrentMs() >= mContext.getEndMs()) {
          break;
        }
        totalCounter = mContext.getTotalCounter().incrementAndGet() - 1;
        if (useStopCount && totalCounter >= mParameters.mStopCount) {
          break;
        }

        long operationCount;
        operationIndex = operationIndexSelector.apply(operationIndex);
        operationCount = mContext.getOperationCounter(operationIndex).incrementAndGet() - 1;
        long startNs = System.nanoTime();
        applyOperation(operationIndex, operationCount);
        long endNs = System.nanoTime();

        long currentMs = CommonUtils.getCurrentMs();
        // Start recording after the warmup
        if (currentMs > recordMs) {
          mResult.incrementNumSuccess(operationIndex, 1);

          // record response times
          long responseTimeNs = endNs - startNs;
          mResponseTimeNs[operationIndex].recordValue(responseTimeNs);

          // track max response time
          long[] maxResponseTimeNs = mResult.getStatistics(operationIndex).mMaxResponseTimeNs;
          int bucket =
              Math.min(maxResponseTimeNs.length - 1, (int) ((currentMs - recordMs) / bucketSize));
          if (responseTimeNs > maxResponseTimeNs[bucket]) {
            maxResponseTimeNs[bucket] = responseTimeNs;
          }
        }
      }
    }

    /**
     * Gets the ratio array and normalizes it into probability distribution
     * e.g. [4,6,10] will be normalized into [0.2, 0.5, 1.0]
     * @return the probability distribution array
     */
    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private double[] getProbabilityDistribution(double[] ratios) {
      double ratioSum = 0;
      double last = 0;
      double[] probabilities = new double[ratios.length];
      for (int i = 0; i < mParameters.mOperations.length; i++) {
        ratioSum += ratios[i];
      }
      for (int i = 0; i < mParameters.mOperations.length; i++) {
        probabilities[i] = last + ratios[i] / ratioSum;
        last = probabilities[i];
      }
      return probabilities;
    }

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private void runOperationsRatioBased() throws Exception {
      double[] probabilities = getProbabilityDistribution(mParameters.mOperationsRatio);
      //Each run randomly picks an operation based on the probability array
      runInternal(lastOperationIndex -> {
        double r = mRandom.nextDouble();
        int operationIndex = 0;
        for (int i = 0; i < mParameters.mOperations.length; i++) {
          if (probabilities[i] >= r) {
            operationIndex = i;
            break;
          }
        }
        mContext.getRateLimiter(0).acquire();
        return operationIndex;
      });
    }

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private void runThreadsRatioBased() throws Exception {
      double[] probabilities = getProbabilityDistribution(mParameters.mThreadsRatio);
      int operationIndex = 0;
      //Each thread randomly picks an operation to do based on the probability array
      for (int i = 0; i < mParameters.mOperations.length; i++) {
        if (probabilities[i] > (double) mThreadIndex / mParameters.mThreads) {
          operationIndex = i;
          break;
        }
      }
      int finalOperationIndex = operationIndex;
      runInternal(lastOperationIndex -> {
        mContext.getRateLimiter(finalOperationIndex).acquire();
        return finalOperationIndex;
      });
    }

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    private void runInternalTargetThroughputBased() throws Exception {
      RateLimiter[] rls = mContext.getRateLimiters();
      // Each run picks the next available operation with the following 2 constraints:
      // 1. the overall throughput should not exceed the user specified limit
      // 2. the throughput of the operation to execute should not exceed the user specified limit,
      //    otherwise, try the next eligible one.
      runInternal((lastOperationIndex) -> {
        int operationIndex = (lastOperationIndex + 1) % mParameters.mOperations.length;
        mContext.getGrandRateLimiter().acquire();
        while (!rls[operationIndex].tryAcquire()) {
          operationIndex++;
          operationIndex %= mParameters.mOperations.length;
        }
        return operationIndex;
      });
    }

    protected abstract void applyOperation(int operationIndex, long operationCounter)
        throws IOException, AlluxioException;
  }

  private final class AlluxioNativeBenchThread extends BenchThread {
    private final alluxio.client.file.FileSystem mFs;

    private AlluxioNativeBenchThread(
        BenchContext context,
        alluxio.client.file.FileSystem fs, int threadIndex) {
      super(context, threadIndex);
      mFs = fs;
    }

    @Override
    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    protected void applyOperation(int operationIndex, long operationCounter)
        throws IOException, AlluxioException {
      MultiOperationStressMasterBench.this.applyNativeOperation(
          mFs, mParameters.mOperations[operationIndex], operationCounter,
          mBasePaths[operationIndex], mFixedBasePaths[operationIndex],
          mParameters.mFixedCounts[operationIndex]);
    }
  }
}
