feat: Follow-up CLI Improvements (#2184)

sydney-munro · web-flow · commit d9859768081e · 2023-08-30T21:08:14.000Z
1. Rename Workload1.java to W1R3 to prevent future confusion when it is used for other workloads
2. Take in temporary directory as a CLI argument
3. Utilize PrintWrite instead of System.out directly
diff --git a/storage-shared-benchmarking/src/main/java/com/google/cloud/storage/benchmarking/StorageSharedBenchmarkingCli.java b/storage-shared-benchmarking/src/main/java/com/google/cloud/storage/benchmarking/StorageSharedBenchmarkingCli.java
@@ -21,15 +21,12 @@
 import com.google.api.core.ListenableFutureToApiFuture;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.ApiExceptions;
-import com.google.cloud.storage.BlobInfo;
-import com.google.cloud.storage.DataGenerator;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.StorageOptions;
-import com.google.cloud.storage.TmpFile;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
-import java.io.IOException;
+import java.io.PrintWriter;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
@@ -81,6 +78,12 @@ public final class StorageSharedBenchmarkingCli implements Runnable {
       required = true)
   String testType;
 
+  @Option(
+      names = "-temp_dir_location",
+      defaultValue = "/tmp",
+      description = "Specify the path where the temporary directory should be located")
+  String tempDirLocation;
+
   public static void main(String[] args) {
     CommandLine cmd = new CommandLine(StorageSharedBenchmarkingCli.class);
     System.exit(cmd.execute(args));
@@ -103,23 +106,18 @@ private void runWorkload1() {
     StorageOptions retryStorageOptions =
         StorageOptions.newBuilder().setProjectId(project).setRetrySettings(retrySettings).build();
     Storage storageClient = retryStorageOptions.getService();
-    Path tempDir = Paths.get(System.getProperty("java.io.tmpdir"));
+    Path tempDir = Paths.get(tempDirLocation);
     ListeningExecutorService executorService =
         MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(workers));
     List<ApiFuture<String>> workloadRuns = new ArrayList<>();
     Range objectSizeRange = Range.of(objectSize);
     for (int i = 0; i < samples; i++) {
-      try {
-        TmpFile file =
-            DataGenerator.base64Characters()
-                .tempFile(tempDir, getRandomInt(objectSizeRange.min, objectSizeRange.max));
-        BlobInfo blob = BlobInfo.newBuilder(bucket, file.toString()).build();
-        workloadRuns.add(
-            convert(
-                executorService.submit(new Workload1(file, blob, storageClient, workers, api))));
-      } catch (IOException e) {
-        throw new RuntimeException(e);
-      }
+      int objectSize = getRandomInt(objectSizeRange.min, objectSizeRange.max);
+      PrintWriter pw = new PrintWriter(System.out, true);
+      workloadRuns.add(
+          convert(
+              executorService.submit(
+                  new W1R3(storageClient, workers, api, pw, objectSize, tempDir, bucket))));
     }
     ApiExceptions.callAndTranslateApiException(ApiFutures.allAsList(workloadRuns));
   }
diff --git a/storage-shared-benchmarking/src/main/java/com/google/cloud/storage/benchmarking/W1R3.java b/storage-shared-benchmarking/src/main/java/com/google/cloud/storage/benchmarking/W1R3.java
@@ -18,54 +18,69 @@
 
 import com.google.cloud.storage.Blob;
 import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.DataGenerator;
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.TmpFile;
+import java.io.PrintWriter;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.Callable;
 
-final class Workload1 implements Callable<String> {
-  private final TmpFile file;
-  private final BlobInfo blob;
+final class W1R3 implements Callable<String> {
+
   private final Storage storage;
   private final int workers;
   private final String api;
+  private final PrintWriter printWriter;
+  private final int objectSize;
+  private final Path tempDirectory;
+  private final String bucketName;
 
-  Workload1(TmpFile file, BlobInfo blob, Storage storage, int workers, String api) {
-    this.file = file;
-    this.blob = blob;
+  W1R3(
+      Storage storage,
+      int workers,
+      String api,
+      PrintWriter printWriter,
+      int objectSize,
+      Path tempDirectory,
+      String bucketName) {
     this.storage = storage;
     this.workers = workers;
     this.api = api;
+    this.printWriter = printWriter;
+    this.objectSize = objectSize;
+    this.tempDirectory = tempDirectory;
+    this.bucketName = bucketName;
   }
 
   @Override
   public String call() throws Exception {
-    Clock clock = Clock.systemDefaultZone();
+    // Create the file to be uploaded and fill it with data
+    TmpFile file = DataGenerator.base64Characters().tempFile(tempDirectory, objectSize);
+    BlobInfo blob = BlobInfo.newBuilder(bucketName, file.toString()).build();
 
     // Get the start time
+    Clock clock = Clock.systemDefaultZone();
     Instant startTime = clock.instant();
     Blob created = storage.createFrom(blob, file.getPath());
     Instant endTime = clock.instant();
     Duration elapsedTimeUpload = Duration.between(startTime, endTime);
-    System.out.println(
+    printWriter.println(
         generateCloudMonitoringResult(
                 "WRITE",
                 StorageSharedBenchmarkingUtils.calculateThroughput(
                     created.getSize().longValue(), elapsedTimeUpload),
                 created)
             .toString());
-    Path tempDir = Paths.get(System.getProperty("java.io.tmpdir"));
     for (int i = 0; i <= StorageSharedBenchmarkingUtils.DEFAULT_NUMBER_OF_READS; i++) {
-      TmpFile dest = TmpFile.of(tempDir, "prefix", "bin");
+      TmpFile dest = TmpFile.of(tempDirectory, "prefix", "bin");
       startTime = clock.instant();
       storage.downloadTo(created.getBlobId(), dest.getPath());
       endTime = clock.instant();
       Duration elapsedTimeDownload = Duration.between(startTime, endTime);
-      System.out.println(
+      printWriter.println(
           generateCloudMonitoringResult(
                   "READ[" + i + "]",
                   StorageSharedBenchmarkingUtils.calculateThroughput(
