alluxio-master-2.x/stress/shell/src/main/java/alluxio/stress/cli/StressMasterBenchBase.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 alluxio.AlluxioURI;
import alluxio.annotation.SuppressFBWarnings;
import alluxio.client.file.FileOutStream;
import alluxio.client.file.URIStatus;
import alluxio.exception.AlluxioException;
import alluxio.exception.FileAlreadyExistsException;
import alluxio.exception.UnexpectedAlluxioException;
import alluxio.grpc.Bits;
import alluxio.grpc.CreateDirectoryPOptions;
import alluxio.grpc.CreateFilePOptions;
import alluxio.grpc.DeletePOptions;
import alluxio.grpc.ListStatusPOptions;
import alluxio.grpc.LoadMetadataPType;
import alluxio.grpc.PMode;
import alluxio.grpc.SetAttributePOptions;
import alluxio.stress.BaseParameters;
import alluxio.stress.StressConstants;
import alluxio.stress.common.FileSystemClientType;
import alluxio.stress.master.MasterBenchBaseParameters;
import alluxio.stress.master.MasterBenchTaskResultBase;
import alluxio.stress.master.MasterBenchTaskResultStatistics;
import alluxio.stress.master.Operation;
import alluxio.util.CommonUtils;
import alluxio.util.FormatUtils;
import alluxio.util.executor.ExecutorServiceFactories;
import alluxio.util.io.PathUtils;

import com.google.common.util.concurrent.RateLimiter;
import org.apache.hadoop.fs.FileStatus;
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.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * The base class for master stress bench tests.
 *
 * @param <T> the type of task result
 * @param <P> the type of task parameter
 */
// TODO(jiacheng): avoid the implicit casts and @SuppressFBWarnings
public abstract class StressMasterBenchBase
    <T extends MasterBenchTaskResultBase<P>, P extends MasterBenchBaseParameters>
    extends AbstractStressBench<T, P> {
  private static final Logger LOG = LoggerFactory.getLogger(StressMasterBenchBase.class);

  protected byte[] mFiledata;

  /** Cached FS instances. */
  protected FileSystem[] mCachedFs;

  /** In case the Alluxio Native API is used,  use the following instead. */
  protected alluxio.client.file.FileSystem[] mCachedNativeFs;
  /* Directories where the stress bench creates files depending on the --operation chosen. */
  protected final String mDirsDir = "dirs";
  protected final String mFilesDir = "files";
  protected final String mFixedDir = "fixed";

  // vars for createTestTree
  protected int[] mPathNodeIds;
  protected int[] mTreeLevelNodeCount;
  protected int mTreeTotalNodeCount;

  /**
   * Creates instance.
   */
  protected StressMasterBenchBase(P parameters) {
    mParameters = parameters;
  }

  protected abstract BenchContext getContext() throws IOException, AlluxioException;

  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  protected void deletePaths(FileSystem fs, Path basePath) throws Exception {
    // the base dir has sub directories per task id
    if (!fs.exists(basePath)) {
      return;
    }
    FileStatus[] subDirs = fs.listStatus(basePath);
    if (subDirs.length == 0) {
      return;
    }

    if (mParameters.mClientType == FileSystemClientType.ALLUXIO_HDFS) {
      fs.delete(basePath, true);
      if (fs.exists(basePath)) {
        throw new UnexpectedAlluxioException(String.format("Unable to delete the files"
            + " in path %s.Please confirm whether it is HDFS file system."
            + " You may need to modify `--client-type` parameter", basePath));
      }
      return;
    }

    // Determine the fixed portion size. Each sub directory has a fixed portion.
    Path fixedPath = new Path(subDirs[0].getPath(), "fixed");
    int fixedSize = fs.exists(fixedPath) ? fs.listStatus(fixedPath).length : 0;

    long batchSize = 50_000;
    int deleteThreads = 256;
    ExecutorService service =
        ExecutorServiceFactories.fixedThreadPool("bench-delete-thread", deleteThreads).create();

    for (FileStatus subDir : subDirs) {
      LOG.info("Cleaning up all files in: {}", subDir.getPath());
      AtomicLong globalCounter = new AtomicLong();
      Path fixedBase = new Path(subDir.getPath(), "fixed");
      long runningLimit = 0;

      // delete individual files in batches, to avoid the recursive-delete problem
      while (!Thread.currentThread().isInterrupted()) {
        AtomicLong success = new AtomicLong();
        runningLimit += batchSize;
        long limit = runningLimit;

        List<Callable<Void>> callables = new ArrayList<>(deleteThreads);
        for (int i = 0; i < deleteThreads; i++) {
          callables.add(() -> {
            while (!Thread.currentThread().isInterrupted()) {
              long counter = globalCounter.getAndIncrement();
              if (counter >= limit) {
                globalCounter.getAndDecrement();
                return null;
              }
              Path deletePath;
              if (counter < fixedSize) {
                deletePath = new Path(fixedBase, Long.toString(counter));
              } else {
                deletePath = new Path(subDir.getPath(), Long.toString(counter));
              }
              if (fs.delete(deletePath, true)) {
                success.getAndIncrement();
              }
            }
            return null;
          });
        }
        // This may cancel some remaining threads, but that is fine, because any remaining paths
        // will be taken care of during the final recursive delete.
        service.invokeAll(callables, 1, TimeUnit.MINUTES);

        if (success.get() == 0) {
          // stop deleting one-by-one if none of the batch succeeded.
          break;
        }
        LOG.info("Removed {} files", success.get());
      }
    }

    service.shutdownNow();
    service.awaitTermination(10, TimeUnit.SECONDS);

    // Cleanup the rest recursively, which should be empty or much smaller than the full tree.
    LOG.info("Deleting base directory: {}", basePath);
    fs.delete(basePath, true);
  }

  @Override
  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  public T runLocal() throws Exception {
    ExecutorService service =
        ExecutorServiceFactories.fixedThreadPool("bench-thread", mParameters.mThreads).create();

    long fileSize = FormatUtils.parseSpaceSize(mParameters.mCreateFileSize);
    mFiledata = new byte[(int) Math.min(fileSize, StressConstants.WRITE_FILE_ONCE_MAX_BYTES)];
    Arrays.fill(mFiledata, (byte) 0x7A);

    BenchContext context = getContext();

    List<Callable<Void>> callables = new ArrayList<>(mParameters.mThreads);
    for (int i = 0; i < mParameters.mThreads; i++) {
      callables.add(getBenchThread(context, i));
    }
    LOG.info("Starting {} bench threads", callables.size());
    service.invokeAll(callables, FormatUtils.parseTimeSize(mBaseParameters.mBenchTimeout),
        TimeUnit.MILLISECONDS);
    LOG.info("Bench threads finished");

    service.shutdownNow();
    service.awaitTermination(30, TimeUnit.SECONDS);

    if (!mBaseParameters.mProfileAgent.isEmpty()) {
      context.addAdditionalResult();
    }

    return context.getResult();
  }

  protected abstract Callable<Void> getBenchThread(BenchContext context, int index);

  protected final class BenchContext {
    private final RateLimiter mGrandRateLimiter;
    private final RateLimiter[] mOperationRateLimiters;
    private final long mStartMs;
    private final long mEndMs;
    private final AtomicLong[] mOperationCounters;
    private final AtomicLong mTotalCounter;
    private final Path[] mBasePaths;
    private final Path[] mFixedBasePaths;
    /** The results. Access must be synchronized for thread safety. */
    private T mResult;

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    BenchContext(
        RateLimiter grandRateLimiter, RateLimiter[] rateLimiters,
        Operation[] operations, String[] basePaths, String duration) {
      mGrandRateLimiter = grandRateLimiter;
      mOperationRateLimiters = rateLimiters;
      long durationMs = FormatUtils.parseTimeSize(duration);
      long warmupMs = FormatUtils.parseTimeSize(mParameters.mWarmup);
      long startMs = mBaseParameters.mStartMs;
      if (mBaseParameters.mStartMs == BaseParameters.UNDEFINED_START_MS) {
        startMs = CommonUtils.getCurrentMs() + 1000;
      }
      mStartMs = startMs;
      mEndMs = startMs + warmupMs + durationMs;

      mOperationCounters = new AtomicLong[operations.length];
      mTotalCounter = new AtomicLong();

      mBasePaths = new Path[operations.length];
      mFixedBasePaths = new Path[operations.length];

      mPathNodeIds = new int[mParameters.mTreeDepth];
      mTreeLevelNodeCount = new int[mParameters.mTreeDepth];
      mTreeLevelNodeCount[mParameters.mTreeDepth - 1] = mParameters.mTreeWidth;
      for (int levelCount = mTreeLevelNodeCount.length - 2; levelCount >= 0; levelCount--) {
        mTreeLevelNodeCount[levelCount] =
            mTreeLevelNodeCount[levelCount + 1] * mParameters.mTreeWidth;
      }
      mTreeTotalNodeCount = mTreeLevelNodeCount[0] * mParameters.mTreeThreads;

      for (int i = 0; i < operations.length; i++) {
        mOperationCounters[i] = new AtomicLong();
        if (operations[i] == Operation.CREATE_DIR) {
          mBasePaths[i] =
              new Path(PathUtils.concatPath(basePaths[i], mDirsDir, mBaseParameters.mId));
        } else if (operations[i] == Operation.CREATE_TREE
            || operations[i] == Operation.LOAD_METADATA) {
          mBasePaths[i] = new Path(PathUtils.concatPath(basePaths[i],
              extractHostName(mBaseParameters.mId)));
        } else {
          mBasePaths[i] =
              new Path(PathUtils.concatPath(basePaths[i], mFilesDir, mBaseParameters.mId));
        }
        mFixedBasePaths[i] = new Path(mBasePaths[i], mFixedDir);
        LOG.info(
            "BenchContext: basePath: {}, fixedBasePath: {}", mBasePaths[i], mFixedBasePaths[i]);
      }
    }

    @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
    public BenchContext(RateLimiter rateLimiter, Operation operation, String duration) {
      this(rateLimiter, new RateLimiter[]{rateLimiter},
          new Operation[]{operation}, new String[]{mParameters.mBasePath}, duration);
    }

    public RateLimiter[] getRateLimiters() {
      return mOperationRateLimiters;
    }

    public RateLimiter getGrandRateLimiter() {
      return mGrandRateLimiter;
    }

    public RateLimiter getRateLimiter(int index) {
      return mOperationRateLimiters[index];
    }

    public long getStartMs() {
      return mStartMs;
    }

    public long getEndMs() {
      return mEndMs;
    }

    public AtomicLong getOperationCounter(int index) {
      return mOperationCounters[index];
    }

    AtomicLong getTotalCounter() {
      return mTotalCounter;
    }

    public Path getBasePath(int index) {
      return mBasePaths[index];
    }

    Path[] getBasePaths() {
      return mBasePaths;
    }

    Path[] getFixedPaths() {
      return mFixedBasePaths;
    }

    public Path getFixedBasePath(int index) {
      return mFixedBasePaths[index];
    }

    public synchronized void mergeThreadResult(T threadResult) {
      if (mResult == null) {
        mResult = threadResult;
        return;
      }
      try {
        mResult.merge(threadResult);
      } catch (Exception e) {
        LOG.warn("Exception during result merge", e);
        mResult.addErrorMessage(e.getMessage());
      }
    }

    @SuppressFBWarnings(value = "DMI_HARDCODED_ABSOLUTE_FILENAME")
    public synchronized void addAdditionalResult() throws IOException {
      if (mResult == null) {
        return;
      }
      Map<String, MethodStatistics> nameStatistics =
          processMethodProfiles(mResult.getRecordStartMs(), mResult.getEndMs(),
              profileInput -> {
              String method = profileInput.getMethod();
              if (profileInput.getType().contains("RPC")) {
                final int classNameDivider = profileInput.getMethod().lastIndexOf(".");
                method = profileInput.getMethod().substring(classNameDivider + 1);
              }
              return profileInput.getType() + ":" + method;
            });

      for (Map.Entry<String, MethodStatistics> entry : nameStatistics.entrySet()) {
        final MasterBenchTaskResultStatistics stats = new MasterBenchTaskResultStatistics();
        stats.encodeResponseTimeNsRaw(entry.getValue().getTimeNs());
        stats.mNumSuccesses = entry.getValue().getNumSuccess();
        stats.mMaxResponseTimeNs = entry.getValue().getMaxTimeNs();
        mResult.putStatisticsForMethod(entry.getKey(), stats);
      }
    }

    public synchronized T getResult() {
      return mResult;
    }
  }

  @SuppressFBWarnings("BC_UNCONFIRMED_CAST")
  protected void applyNativeOperation(
      alluxio.client.file.FileSystem fs, Operation operation, long counter,
      Path basePath, Path fixedBasePath, int fixedCount)
      throws IOException, AlluxioException {
    Path path;
    AlluxioURI uri;
    switch (operation) {
      case CREATE_DIR:
        if (counter < fixedCount) {
          path = new Path(fixedBasePath, Long.toString(counter));
        } else {
          path = new Path(basePath, Long.toString(counter));
        }

        fs.createDirectory(new AlluxioURI(path.toString()),
            CreateDirectoryPOptions.newBuilder().setRecursive(true).build());
        break;
      case CREATE_FILE:
        if (counter < fixedCount) {
          path = new Path(fixedBasePath, Long.toString(counter));
        } else {
          path = new Path(basePath, Long.toString(counter));
        }
        long fileSize = FormatUtils.parseSpaceSize(mParameters.mCreateFileSize);
        try (FileOutStream stream = fs.createFile(new AlluxioURI(path.toString()),
            CreateFilePOptions.newBuilder().setRecursive(true).build())) {
          for (long i = 0; i < fileSize; i += StressConstants.WRITE_FILE_ONCE_MAX_BYTES) {
            stream.write(mFiledata, 0,
                (int) Math.min(StressConstants.WRITE_FILE_ONCE_MAX_BYTES, fileSize - i));
          }
        }
        break;
      case GET_BLOCK_LOCATIONS:
        counter = counter % fixedCount;
        path = new Path(fixedBasePath, Long.toString(counter));
        fs.getBlockLocations(new AlluxioURI(path.toString()));
        break;
      case GET_FILE_STATUS:
        counter = counter % fixedCount;
        path = new Path(fixedBasePath, Long.toString(counter));
        fs.getStatus(new AlluxioURI(path.toString()));
        break;
      case LIST_DIR:
        List<alluxio.client.file.URIStatus> files
            = fs.listStatus(new AlluxioURI(fixedBasePath.toString()));
        if (files.size() != fixedCount) {
          throw new IOException(String
              .format("listing `%s` expected %d files but got %d files", fixedBasePath,
                  fixedCount, files.size()));
        }
        break;
      case LIST_DIR_LOCATED:
        throw new UnsupportedOperationException("LIST_DIR_LOCATED is not supported!");
      case OPEN_FILE:
        counter = counter % fixedCount;
        path = new Path(fixedBasePath, Long.toString(counter));
        fs.openFile(new AlluxioURI(path.toString())).close();
        break;
      case RENAME_FILE:
        if (counter < fixedCount) {
          path = new Path(fixedBasePath, Long.toString(counter));
        } else {
          path = new Path(basePath, Long.toString(counter));
        }
        Path dst = new Path(path + "-renamed");
        fs.rename(new AlluxioURI(path.toString()), new AlluxioURI(dst.toString()));
        break;
      case DELETE_FILE:
        if (counter < fixedCount) {
          path = new Path(fixedBasePath, Long.toString(counter));
        } else {
          path = new Path(basePath, Long.toString(counter));
        }

        fs.delete(new AlluxioURI(path.toString()),
            DeletePOptions.newBuilder().setRecursive(false).build());
        break;
      case CREATE_DELETE_FILE:
        path = new Path(basePath, Long.toString(counter));
        uri = new AlluxioURI(path.toString());
        fs.createFile(uri, CreateFilePOptions.newBuilder().setRecursive(true).build()).close();
        fs.delete(uri);
        break;
      case SET_ATTRIBUTE:
        if (counter < fixedCount) {
          path = new Path(fixedBasePath, Long.toString(counter));
        } else {
          path = new Path(basePath, Long.toString(counter));
        }
        fs.setAttribute(new AlluxioURI(path.toString()), SetAttributePOptions.newBuilder()
                .setMode(PMode.newBuilder()
                    .setOwnerBits(Bits.ALL)
                    .setGroupBits(Bits.ALL)
                    .setOtherBits(Bits.ALL).build()).build());
        break;
      case CRURD:
        path = new Path(basePath, Long.toString(counter));
        uri = new AlluxioURI(path.toString());
        fs.createFile(uri, CreateFilePOptions.newBuilder().setRecursive(true).build()).close();
        if (!fs.exists(uri)) {
          throw new IOException("[INCONSISTENCY] file doesn't exist after creation");
        }
        fs.setAttribute(new AlluxioURI(path.toString()), SetAttributePOptions.newBuilder()
            .setMode(PMode.newBuilder()
                .setOwnerBits(Bits.ALL)
                .setGroupBits(Bits.ALL)
                .setOtherBits(Bits.ALL).build()).build());
        URIStatus us = fs.getStatus(uri);
        if (us.getMode() != 0777) {
          throw new IOException("[INCONSISTENCY] file update doesn't reflect");
        }
        fs.delete(uri);
        if (fs.exists(uri)) {
          throw new IOException("[INCONSISTENCY] file still exists after deletion");
        }
        break;
      case LOAD_METADATA:
        fs.loadMetadata(new AlluxioURI(basePath + "/" + counter), ListStatusPOptions.newBuilder()
            .setLoadMetadataType(LoadMetadataPType.ALWAYS)
            .setRecursive(true)
            .setLoadMetadataOnly(true).build());
        break;
      case CREATE_TREE:
        String nodePath = "";
        int nodeNumber = (int) counter;
        for (int levelCount = 0; levelCount < mParameters.mTreeDepth; levelCount++) {
          mPathNodeIds[levelCount] = nodeNumber / mTreeLevelNodeCount[levelCount];
          nodeNumber = nodeNumber % mTreeLevelNodeCount[levelCount];
          nodePath += "/";
          nodePath += mPathNodeIds[levelCount];
        }
        for (int fileNumber = 0; fileNumber < mParameters.mTreeFiles; fileNumber++) {
          try {
            fs.createFile(new AlluxioURI((basePath + nodePath + "/"
                    + nodeNumber + "/" + fileNumber + ".txt")),
                CreateFilePOptions.newBuilder().setRecursive(true).build()).close();
          } catch (FileAlreadyExistsException e) {
            break;
          }
        }
        break;
      default:
        throw new IllegalStateException("Unknown operation: " + operation);
    }
  }

  protected String extractHostName(String mId) {
    String hostName = "";
    String[] splitedMid = mId.split("-");
    hostName += splitedMid[0];
    for (int i = 1; i < splitedMid.length - 1; i++) {
      hostName += "-";
      hostName += splitedMid[i];
    }
    return hostName;
  }
}
