content: /*
   * 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.client;
  
  import alluxio.AlluxioURI;
  import alluxio.Constants;
  import alluxio.client.file.FileInStream;
  import alluxio.client.file.FileOutStream;
  import alluxio.client.file.FileSystem;
  import alluxio.client.file.URIStatus;
  import alluxio.conf.AlluxioProperties;
  import alluxio.conf.Configuration;
  import alluxio.conf.InstancedConfiguration;
  import alluxio.conf.PropertyKey;
  import alluxio.exception.AlluxioException;
  import alluxio.exception.FileAlreadyExistsException;
  import alluxio.exception.FileDoesNotExistException;
  import alluxio.grpc.CreateDirectoryPOptions;
  import alluxio.grpc.DeletePOptions;
  import alluxio.grpc.ListStatusPOptions;
  import alluxio.retry.CountingRetry;
  import alluxio.retry.RetryPolicy;
  import alluxio.stress.StressConstants;
  import alluxio.stress.cli.Benchmark;
  import alluxio.stress.client.CompactionParameters;
  import alluxio.stress.client.CompactionTaskResult;
  import alluxio.util.CommonUtils;
  import alluxio.util.FormatUtils;
  import alluxio.util.executor.ExecutorServiceFactories;
  
  import com.beust.jcommander.ParametersDelegate;
  import com.google.common.base.Preconditions;
  import com.google.common.base.Stopwatch;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.Lists;
  import org.HdrHistogram.Histogram;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.Callable;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.stream.Collectors;
  
  /**
   * Benchmark that simulates a workload that compacts many small files into a bigger file.
   */
  public class CompactionBench extends Benchmark<CompactionTaskResult> {
    private static final Logger LOG = LoggerFactory.getLogger(CompactionBench.class);
  
    protected ExecutorService mPool = null;
    @ParametersDelegate
    protected final CompactionParameters mParameters = new CompactionParameters();
    protected FileSystem[] mCachedFs;
    private AlluxioURI mRealSourceBase;
  
    /**
     * @param args command-line arguments
     */
    public static void main(String[] args) {
      mainInternal(args, new CompactionBench());
    }
  
    @Override
    public String getBenchDescription() {
      return String.join("\n", ImmutableList.of(
          "A benchmark that simulates the workload of compacting many small files into a bigger "
              + "file.",
          "",
          "Example:",
          "# This example creates 4 source directories each containing 100 source files that are "
              + "10KB each,",
          "# and compacts every 10 source files into 1 output file, resulting in 40 output files in",
          "# a single output directory.",
          "# The compaction is done on 1 job worker with 5 threads, meaning each thread will "
              + "process 80 source",
          "# files in 8 sequential batches.",
          "$ bin/alluxio runClass alluxio.stress.cli.client.CompactionBench "
              + "--cluster "
              + "--cluster-limit 1 "
              + "--base alluxio:///compaction-base "
              + "--source-files 1000 "
              + "--source-dirs 4 "
              + "--source-file-size 10kb "
              + "--threads 5 "
              + "--compact-ratio 10 "
      ));
    }
  
    @Override
    public CompactionTaskResult runLocal() throws Exception {
      mCachedFs = new FileSystem[mParameters.mThreads];
      AlluxioProperties properties = getCustomProperties(mParameters.mCompactProperties);
      for (int i = 0; i < mParameters.mThreads; i++) {
        mCachedFs[i] = FileSystem.Factory.create(new InstancedConfiguration(properties));
      }
      FileSystem fs = mCachedFs[0];
      AlluxioURI baseUri = new AlluxioURI(mParameters.mBase);
      AlluxioURI destBaseUri = baseUri.join(mParameters.mOutputBase);
      AlluxioURI stagingBaseUri = baseUri.join(mParameters.mStagingBase);
      // Scan base dir to get all subdirectories that contain files to compact
      List<AlluxioURI> subDirs =
          fs.listStatus(mRealSourceBase)
              .stream()
              .filter(URIStatus::isFolder)
              .map(uri -> new AlluxioURI(mRealSourceBase, uri.getPath(), false))
              .collect(Collectors.toList());
      // Partition them into batches and each thread will work on one batch
      List<List<AlluxioURI>> partitions = exactPartition(subDirs, mParameters.mThreads);
      List<CompletableFuture<CompactionTaskResult>> futures = new ArrayList<>(mParameters.mThreads);
      try {
        for (int i = 0; i < mParameters.mThreads; i++) {
          List<AlluxioURI> partition = partitions.get(i);
          Map<AlluxioURI, AlluxioURI> srcDestDirMap = partition
              .stream()
              .collect(Collectors.toMap(
                  src -> src,
                  src -> mParameters.mOutputInPlace ? src : destBaseUri));
          BenchThread thread = new BenchThread(
              mCachedFs[i],
              srcDestDirMap,
              stagingBaseUri,
              mParameters.mCompactRatio,
              FormatUtils.parseTimeSize(mParameters.mDelayMs),
              (int) FormatUtils.parseSpaceSize(mParameters.mBufSize),
              mParameters.mPreserveSource,
              mParameters.mDeleteByDir,
              mBaseParameters.mId);
          CompletableFuture<CompactionTaskResult> future = CompletableFuture.supplyAsync(() -> {
            CompactionTaskResult result;
            try {
              result = thread.call();
            } catch (Exception e) {
              LOG.error("Failed to run compaction thread", e);
              result = new CompactionTaskResult();
              result.addError(e.getMessage());
            }
            return result;
          }, getPool());
          futures.add(future);
        }
        LOG.info("{} jobs submitted", futures.size());
  
        // Collect the results
        CompactionTaskResult result = new CompactionTaskResult();
        result.setBaseParameters(mBaseParameters);
        result.setParameters(mParameters);
        for (CompletableFuture<CompactionTaskResult> future : futures) {
          CompactionTaskResult threadResult = future.join();
          result.merge(threadResult);
        }
        return result;
      } catch (Exception e) {
        LOG.error("Failed to execute RPC in pool", e);
        CompactionTaskResult result = new CompactionTaskResult();
        result.setBaseParameters(mBaseParameters);
        result.setParameters(mParameters);
        result.addError(e.getMessage());
        return result;
      }
    }
  
    @Override
    public void prepare() throws Exception {
      Preconditions.checkArgument(mParameters.mThreads > 0, "mThreads");
      AlluxioProperties properties = getCustomProperties(mParameters.mPrepareProperties);
      FileSystem prepareFs = FileSystem.Factory.create(new InstancedConfiguration(properties));
  
      // flags:
      // distributed    cluster   in-process    run type
      //  1              1         any             n/a
      //  1              0         1               at job worker
      //  1              0         0               n/a
      //  0              1         1               n/a
      //  0              1         0               at the local process invoking the benchmark
      //                                             and will continue to run at job worker
      //  0              0         1               at the local forked process
      //  0              0         0               at the local process invoking the benchmark
      //                                             and will continue to run in forked process
  
      int flags = Boolean.compare(mBaseParameters.mDistributed, false) << 2
          | Boolean.compare(mBaseParameters.mCluster, false) << 1
          | Boolean.compare(mBaseParameters.mInProcess, false);
      switch (flags) {
        case 0b000:
        case 0b010:
          if (!mParameters.mSkipPrepare) {
            prepareSourceBaseDir(prepareFs);
            prepareOutputBaseDir(prepareFs);
            prepareStagingBaseDir(prepareFs);
          }
          break;
        case 0b001:
          // set real base to "local"
          mRealSourceBase =
              new AlluxioURI(mParameters.mBase).join(mParameters.mSourceBase).join("local");
          if (!mParameters.mSkipPrepare) {
            prepareSourceFiles(prepareFs);
          }
          break;
        case 0b101:
          // set real base to the id of the job worker, to avoid sharing the same base
          // path with other job workers
          mRealSourceBase = new AlluxioURI(mParameters.mBase)
              .join(mParameters.mSourceBase).join(mBaseParameters.mId);
          if (!mParameters.mSkipPrepare) {
            prepareSourceFiles(prepareFs);
          }
          break;
        default:
          throw new IllegalStateException(String.format("Unknown combination of flags: %s",
              Integer.toBinaryString(flags)));
      }
    }
  
    private static AlluxioProperties getCustomProperties(List<String> propertyList) {
      AlluxioProperties properties = Configuration.copyProperties();
      for (String property : propertyList) {
        String[] parts = property.split("=", 2);
        Preconditions.checkArgument(parts.length == 2,
            "Property should be set as \"key=value\", got %s", property);
        properties.set(PropertyKey.fromString(parts[0]), parts[1]);
      }
      return properties;
    }
  
    private void prepareStagingBaseDir(FileSystem fs) throws IOException, AlluxioException {
      try {
        fs.createDirectory(new AlluxioURI(mParameters.mBase).join(mParameters.mStagingBase),
            CreateDirectoryPOptions.newBuilder().setRecursive(true).build());
      } catch (FileAlreadyExistsException ignored) { /* ignored */ }
    }
  
    private void prepareOutputBaseDir(FileSystem fs) throws IOException, AlluxioException {
      if (!mParameters.mOutputInPlace) {
        AlluxioURI path = new AlluxioURI(mParameters.mBase).join(mParameters.mOutputBase);
        try {
          fs.delete(path,
              DeletePOptions.newBuilder().setRecursive(true).build());
        } catch (FileDoesNotExistException ignored) { /* ignored */ }
        fs.createDirectory(path,
            CreateDirectoryPOptions.newBuilder().setRecursive(true).build());
      }
    }
  
    private void prepareSourceBaseDir(FileSystem fs) throws IOException, AlluxioException {
      try {
        fs.createDirectory(new AlluxioURI(mParameters.mBase).join(mParameters.mSourceBase),
            CreateDirectoryPOptions.newBuilder().setRecursive(true).build());
      } catch (FileAlreadyExistsException ignored) { /* ignored */ }
    }
  
    private void prepareSourceFiles(FileSystem fs) throws Exception {
      final int fileSize = (int) FormatUtils.parseSpaceSize(mParameters.mSourceFileSize);
      // cap to 1 MB as the assumption is the source files are small
      final byte[] fileData = new byte[Math.min(fileSize, Constants.MB)];
      Arrays.fill(fileData, (byte) 0x7A);
      try {
        fs.createDirectory(mRealSourceBase);
      } catch (FileAlreadyExistsException ignored) { /* ignored */ }
  
      final AtomicInteger numDirsCreated = new AtomicInteger();
      int createFilesParallelism = Runtime.getRuntime().availableProcessors() * 2;
      ExecutorService pool = ExecutorServiceFactories
          .fixedThreadPool("compact-bench-prepare-thread", createFilesParallelism)
          .create();
      List<CompletableFuture<Exception>> futures = new ArrayList<>(createFilesParallelism);
      for (int i = 0; i < createFilesParallelism; i++) {
        CompletableFuture<Exception> future = CompletableFuture.supplyAsync(() -> {
          try {
            int localNumDirsCreated;
            while ((localNumDirsCreated = numDirsCreated.getAndIncrement())
                < mParameters.mNumSourceDirs) {
              LOG.info("creating directory {}/{} ",
                  localNumDirsCreated, mParameters.mNumSourceDirs);
              AlluxioURI dir = mRealSourceBase.join(Integer.toString(localNumDirsCreated));
              try {
                fs.createDirectory(dir);
              } catch (FileAlreadyExistsException ignored) { /* ignored */ }
  
              for (int f = 0; f < mParameters.mNumSourceFiles; f++) {
                AlluxioURI path = dir.join(Integer.toString(f));
                try (FileOutStream stream = fs.createFile(path)) {
                  for (long offset = 0; offset < fileSize; offset += fileData.length) {
                    stream.write(fileData, 0, (int) Math.min(fileData.length, fileSize - offset));
                  }
                } catch (FileAlreadyExistsException e) {
                  fs.delete(path);
                  f--; // retry
                }
                // Print progress every 10% files have been created
                if (f % (mParameters.mNumSourceFiles / 10) == 0) {
                  LOG.info("{}/{} files created in dir {}",
                      f, mParameters.mNumSourceFiles, localNumDirsCreated);
                }
              }
              LOG.info("{}/{} directories created",
                  localNumDirsCreated, mParameters.mNumSourceDirs);
            }
          } catch (IOException | AlluxioException e) {
            return e;
          }
          return null;
        }, pool);
        futures.add(future);
      }
      try {
        for (CompletableFuture<Exception> future : futures) {
          Exception e = future.join();
          if (e != null) {
            LOG.error("Failed to prepare test directory and files", e);
            throw e;
          }
        }
      } finally {
        pool.shutdownNow();
        pool.awaitTermination(30, TimeUnit.SECONDS);
      }
    }
  
    @Override
    public void cleanup() throws Exception {
      super.cleanup();
      if (mPool != null) {
        LOG.debug("Terminating thread pool");
        mPool.shutdownNow();
        mPool.awaitTermination(30, TimeUnit.SECONDS);
      }
    }
  
    /**
     * If the thread pool is not yet initialized, creates the pool.
     *
     * @return the thread pool
     */
    public ExecutorService getPool() {
      if (mPool == null) {
        mPool = ExecutorServiceFactories
            .fixedThreadPool("compact-benchmark-thread", mParameters.mThreads).create();
      }
      return mPool;
    }
  
    /**
     * Splits a list into exactly {@code numPartitions} partitions. Let {@code S = L / N}
     * and {@code R = L % N}, so that {@code L = S * N + R},
     * where {@code L} is the number of items in the list, and {@code N} is the number of partitions.
     * The first {@code R} partitions have {@code S+1} items, and the remaining {@code N - R}
     * partitions have {@code S} items.
     * @param list the list to partition
     * @param numPartitions number of partitions
     * @param <T> item type
     * @return a partitioned list of sub-lists views
     */
    private static <T> List<List<T>> exactPartition(List<T> list, int numPartitions) {
      int length = list.size();
      int sizePerPartition = length / numPartitions;
      int leftover = length % numPartitions;
      int leftoverEndIndex = (sizePerPartition + 1) * leftover;
      List<List<T>> partitions = new ArrayList<>(numPartitions);
      for (int i = 0; i < leftover; i++) {
        partitions.add(list.subList(i * (sizePerPartition + 1), (i + 1) * (sizePerPartition + 1)));
      }
      for (int i = 0; i < numPartitions - leftover; i++) {
        partitions.add(list.subList(leftoverEndIndex + i * sizePerPartition,
            leftoverEndIndex + (i + 1) * sizePerPartition));
      }
      return partitions;
    }
  
    static class BenchThread implements Callable<CompactionTaskResult> {
      private final FileSystem mFs;
      /* input dir to output dir mapping */
      private final Map<AlluxioURI, AlluxioURI> mSrcDestMap;
      private final AlluxioURI mStagingDir;
      private final int mCompactRatio;
      private final long mDelayMs;
      private final int mBufSize;
      private final boolean mPreserveSource;
      private final boolean mDeleteByDir;
      private final String mWorkerId;
      private final CompactionTaskResult mResult;
      private final Histogram mRawRecords;
  
      public BenchThread(FileSystem fs, Map<AlluxioURI, AlluxioURI> dirMap, AlluxioURI stagingDir,
                         int compactRatio, long delayMs, int bufSize,
                         boolean preserveSource, boolean deleteByDir, String workerId) {
        Preconditions.checkArgument(compactRatio >= 1, "compactRatio should be 1 or greater");
        Preconditions.checkArgument(delayMs >= 0, "delayMs should be 0 or greater");
        Preconditions.checkArgument(bufSize > 0, "buffer size should be greater than 0");
        mFs = fs;
        mSrcDestMap = dirMap;
        mStagingDir = stagingDir;
        mCompactRatio = compactRatio;
        mDelayMs = delayMs;
        mBufSize = bufSize;
        mPreserveSource = preserveSource;
        mDeleteByDir = deleteByDir;
        mWorkerId = workerId;
        mResult = new CompactionTaskResult();
        mRawRecords = new Histogram(StressConstants.TIME_HISTOGRAM_MAX,
            StressConstants.TIME_HISTOGRAM_PRECISION);
      }
  
      @Override
      public CompactionTaskResult call() throws Exception {
        runInternal();
        return mResult;
      }
  
      private void runInternal() throws Exception {
        Stopwatch stopwatch = Stopwatch.createUnstarted();
        for (Map.Entry<AlluxioURI, AlluxioURI> entry : mSrcDestMap.entrySet()) {
          // Gather all input files by listing this subdirectory
          AlluxioURI srcDir = entry.getKey();
          AlluxioURI destDir = entry.getValue();
          List<AlluxioURI> files =
              mFs.listStatus(srcDir, ListStatusPOptions.newBuilder().setRecursive(false).build())
                  .stream()
                  .filter(uri -> !uri.isFolder() && uri.isCompleted())
                  .map(uri -> new AlluxioURI(srcDir, uri.getPath(), false))
                  .collect(Collectors.toList());
          // Partition files into batches
          List<List<AlluxioURI>> batches = Lists.partition(files, mCompactRatio);
          LOG.info("Partitioned {} files in dir {} into {} batches, each with {} files",
              files.size(), srcDir, batches.size(), mCompactRatio);
  
          for (int i = 0; i < batches.size(); i++) {
            // Process files from one batch
            List<AlluxioURI> batch = batches.get(i);
            String outputFileName = String.format("compact_output_part%d_dir%s_worker%s",
                i, srcDir.getName(), mWorkerId);
            Compactor compactor =
                new Compactor(mFs, batch.iterator(), destDir, outputFileName, mStagingDir, mBufSize);
  
            try {
              stopwatch.reset();
              stopwatch.start();
              compactor.run();
              stopwatch.stop();
              mRawRecords.recordValue(stopwatch.elapsed(TimeUnit.NANOSECONDS));
              mResult.incrementNumSuccess();
            } catch (Exception e) {
              LOG.warn("Batch {} in dir {} failed", i + 1, e);
              mResult.addError(e.getMessage());
            }
            LOG.info("Batch {}/{} in dir {} finished", i + 1, batches.size(), srcDir);
            // Sleep as needed
            Thread.sleep(mDelayMs);
          }
  
          // Delete input files
          if (!mPreserveSource) {
            if (!mDeleteByDir) {
              for (AlluxioURI file : files) {
                mFs.delete(file);
              }
              mFs.delete(srcDir, DeletePOptions.newBuilder().build());
            } else {
              mFs.delete(srcDir, DeletePOptions.newBuilder().setRecursive(true).build());
            }
          }
  
          mResult.getStatistics().encodeResponseTimeNsRaw(mRawRecords);
        }
      }
    }
  
    static class Compactor {
      private final FileSystem mFs;
      private final AlluxioURI mOutputBase;
      private final String mOutputFileName;
      private final Iterator<AlluxioURI> mInputs;
      private final AlluxioURI mStagingBase;
      private final byte[] mBuffer;
  
      public Compactor(FileSystem fs, Iterator<AlluxioURI> inputs, AlluxioURI outputBase,
                       String outputFileName, AlluxioURI stagingBase, int bufSize) {
        mFs = fs;
        mOutputBase = outputBase;
        mOutputFileName = outputFileName;
        mInputs = inputs;
        mStagingBase = stagingBase;
        mBuffer = new byte[bufSize];
      }
  
      public void run() throws IOException, AlluxioException {
        // First write an intermediate file
        AlluxioURI tempFile = mStagingBase.join(CommonUtils.randomAlphaNumString(8));
        try (FileOutStream out = mFs.createFile(tempFile)) {
          while (mInputs.hasNext()) {
            try (FileInStream input = mFs.openFile(mInputs.next())) {
              int bytesRead;
              while ((bytesRead = input.read(mBuffer)) >= 0) {
                out.write(mBuffer, 0, bytesRead);
              }
            }
          }
        }
        // Move intermediate file to output base dir
        RetryPolicy retry = new CountingRetry(5);
        String nameSuffix = "";
        boolean done = false;
        while (retry.attempt()) {
          try {
            mFs.rename(tempFile, mOutputBase.join(mOutputFileName + nameSuffix));
            done = true;
            break;
          } catch (FileAlreadyExistsException ignored) {
            nameSuffix = "_" + retry.getAttemptCount();
          }
        }
        if (!done) {
          throw new FileAlreadyExistsException(
              String.format("Output file %s already exists, renaming failed after %d attempts",
                  mOutputFileName, retry.getAttemptCount()));
        }
      }
    }
  }
  