alluxio-master-2.x/stress/shell/src/main/java/alluxio/stress/cli/GetPinnedFileIdsBench.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.ClientContext;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.FileSystemMasterClient;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.status.AlluxioStatusException;
import alluxio.grpc.CreateDirectoryPOptions;
import alluxio.grpc.CreateFilePOptions;
import alluxio.grpc.DeletePOptions;
import alluxio.grpc.FileSystemMasterWorkerServiceGrpc;
import alluxio.grpc.GetPinnedFileIdsPRequest;
import alluxio.grpc.SetAttributePOptions;
import alluxio.master.MasterClientContext;
import alluxio.resource.CloseableResource;
import alluxio.stress.rpc.GetPinnedFileIdsParameters;
import alluxio.stress.rpc.RpcTaskResult;
import alluxio.util.CommonUtils;
import alluxio.util.FormatUtils;

import com.beust.jcommander.ParametersDelegate;
import com.google.common.base.Stopwatch;
import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * A benchmarking tool for the GetPinnedFileIds RPC.
 */
public class GetPinnedFileIdsBench extends RpcBench<GetPinnedFileIdsParameters> {
  private static final Logger LOG = LoggerFactory.getLogger(GetPinnedFileIdsBench.class);

  @ParametersDelegate
  private final GetPinnedFileIdsParameters mParameters = new GetPinnedFileIdsParameters();

  private final FileSystemContext mFileSystemContext = FileSystemContext.create();
  private final ThreadLocal<Stopwatch> mDurationStopwatch =
      ThreadLocal.withInitial(Stopwatch::createUnstarted);
  private final ThreadLocal<Stopwatch> mPointStopwatch =
      ThreadLocal.withInitial(Stopwatch::createUnstarted);
  private final PinListFileSystemMasterClient mWorkerClient =
      new PinListFileSystemMasterClient(
          MasterClientContext.newBuilder(ClientContext.create()).build());

  @Override
  public String getBenchDescription() {
    return String.join("\n", ImmutableList.of(
        "A benchmarking tool for the GetPinnedFileIds RPC.",
        "The test will generate a specified number of test files and pin them. "
            + "Then it will keep calling the GetPinnedFileIds RPC by the specified load until the "
            + "specified duration has elapsed. The test files will be cleaned up in the end.",
        "",
        "Example:",
        "# 2 job workers will be chosen to run the benchmark",
        "# Each job worker runs 3 simulated clients",
        "# Each client keeps requesting a total number of 10k pinned files for "
            + "a total of 100 milliseconds",
        "$ bin/alluxio runClass alluxio.stress.cli.GetPinnedFileIdsBench --concurrency 3 \\",
        "--cluster --cluster-limit 2 --num-files 10000 --duration 100ms",
        ""
    ));
  }

  @Override
  public void prepare() throws Exception {
    // The task ID is different for local and cluster executions
    // So including that in the log can help associate the log to the run
    LOG.info("Task ID is {}", mBaseParameters.mId);

    // The preparation is done by the invoking shell process to ensure the preparation is only
    // done once, so skip preparation when running in job worker
    if (mBaseParameters.mDistributed) {
      LOG.info("Skipping preparation in distributed execution");
      return;
    }

    AlluxioURI baseUri = new AlluxioURI(mParameters.mBasePath);
    try (CloseableResource<alluxio.client.file.FileSystemMasterClient> client =
             mFileSystemContext.acquireMasterClientResource()) {
      LOG.info("Creating temporary directory {} for benchmark", baseUri);
      client.get().createDirectory(baseUri,
              CreateDirectoryPOptions.newBuilder().setAllowExists(true).build());
    }
    int numFiles = mParameters.mNumFiles;
    int fileNameLength = (int) Math.ceil(Math.max(8, Math.log10(numFiles)));

    CompletableFuture<Void>[] futures = new CompletableFuture[numFiles];
    LOG.info("Generating {} pinned test files at the master", numFiles);
    for (int i = 0; i < numFiles; i++) {
      AlluxioURI fileUri = baseUri.join(CommonUtils.randomAlphaNumString(fileNameLength));
      CompletableFuture<Void> future = CompletableFuture.supplyAsync((Supplier<Void>) () -> {
        try (CloseableResource<FileSystemMasterClient> client =
                 mFileSystemContext.acquireMasterClientResource()) {
          client.get().createFile(fileUri,
              CreateFilePOptions
                  .newBuilder()
                  .setBlockSizeBytes(Configuration.getBytes(
                      PropertyKey.USER_BLOCK_SIZE_BYTES_DEFAULT))
                  .build());
          client.get().setAttribute(fileUri,
              SetAttributePOptions
                  .newBuilder()
                  .setPinned(true)
                  .build());
        } catch (AlluxioStatusException e) {
          LOG.warn("Exception during file creation of {}", fileUri, e);
        }
        return null;
      }, getPool());
      futures[i] = (future);
    }
    CompletableFuture.allOf(futures).join();
    LOG.info("Test files generated");
  }

  @Override
  public void cleanup() throws Exception {
    // skip cleanup in job worker as the test files are to be cleaned up by the client
    if (mBaseParameters.mDistributed) {
      LOG.info("Skipping cleanup in distributed execution");
    } else {
      AlluxioURI baseUri = new AlluxioURI(mParameters.mBasePath);
      try (CloseableResource<FileSystemMasterClient> client =
               mFileSystemContext.acquireMasterClientResource()) {
        LOG.info("Deleting test directory {}", baseUri);
        client.get().delete(baseUri, DeletePOptions.newBuilder().setRecursive(true).build());
      } catch (AlluxioStatusException e) {
        LOG.warn("Failed to delete test directory {}, manual cleanup needed", baseUri, e);
      }
    }
    super.cleanup();
  }

  @Override
  public RpcTaskResult runRPC() throws Exception {
    RpcTaskResult result = new RpcTaskResult();

    mDurationStopwatch.get().reset().start();
    long durationMs = FormatUtils.parseTimeSize(mParameters.mDuration);
    LOG.info("Beginning benchmark, running for {} ms", durationMs);
    while (mDurationStopwatch.get().elapsed(TimeUnit.MILLISECONDS) < durationMs) {
      try {
        mPointStopwatch.get().reset().start();

        int numPinnedFiles = mWorkerClient.getPinListLength();

        mPointStopwatch.get().stop();

        if (numPinnedFiles != mParameters.mNumFiles) {
          result.addError(String.format("Unexpected number of files: %d, expected %d",
              numPinnedFiles, mParameters.mNumFiles));
          continue;
        }
        result.addPoint(
            new RpcTaskResult.Point(mPointStopwatch.get().elapsed(TimeUnit.MILLISECONDS)));
      } catch (Exception e) {
        LOG.error("Failed when running", e);
        result.addError(e.getMessage());
      }
    }
    return result;
  }

  @Override
  public GetPinnedFileIdsParameters getParameters() {
    return mParameters;
  }

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

  /**
   * A helper client that is used to minimize the deserialization overhead by calling
   * getPinnedFileIdsCount instead of getPinnedFileIdsList.
   *
   * By calling getPinnedFileIdsList, GRPC deserializes the IDs from the response, which in this
   * benchmark are irrelevant, since we only care about the length of the list. By getting only
   * the count of the IDs, we avoid the deserialization of the IDs.
   * We want to minimize the client side cost while measuring the RPC performance, so that the
   * collected results can get as close as possible to what the server side
   * (in this benchmark's case the master) can achieve.
   */
  private static class PinListFileSystemMasterClient
      extends alluxio.worker.file.FileSystemMasterClient {
    private static final Logger LOG =
        LoggerFactory.getLogger(PinListFileSystemMasterClient.class);
    private FileSystemMasterWorkerServiceGrpc.FileSystemMasterWorkerServiceBlockingStub mClient =
        null;

    public PinListFileSystemMasterClient(MasterClientContext conf) {
      super(conf);
    }

    @Override
    protected void afterConnect() throws IOException {
      mClient = FileSystemMasterWorkerServiceGrpc.newBlockingStub(mChannel);
    }

    public int getPinListLength() throws IOException {
      return retryRPC(() -> mClient
          .withDeadlineAfter(
              mContext.getClusterConf().getMs(PropertyKey.WORKER_MASTER_PERIODICAL_RPC_TIMEOUT),
              TimeUnit.MILLISECONDS)
          .getPinnedFileIds(GetPinnedFileIdsPRequest.newBuilder().build()).getPinnedFileIdsCount(),
          LOG, "GetPinList", "");
    }
  }
}
