From 55ec75fff50c46ff496d62350737a5daa84d8aa7 Mon Sep 17 00:00:00 2001
From: Sydney Munro <97561403+sydney-munro@users.noreply.github.com>
Date: Tue, 26 Mar 2024 10:14:20 -0700
Subject: [PATCH] chore(benchmarking): Change bidi workload to be write only
 and pass bidi flag (#2468)

---
 .../StorageSharedBenchmarkingCli.java         | 51 +++++++++----------
 1 file changed, 23 insertions(+), 28 deletions(-)

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
index 32d85b0792..3effab8c24 100644
--- 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
@@ -87,6 +87,9 @@ public final class StorageSharedBenchmarkingCli implements Runnable {
       defaultValue = "0")
   int warmup;
 
+  @Option(names = "-bidi", description = "If bidi should be enabled")
+  boolean bidiEnabled;
+
   Path tempDir;
 
   PrintWriter printWriter;
@@ -107,11 +110,8 @@ public void run() {
       case "w1r3":
         runWorkload1();
         break;
-      case "bidi":
-        runWorkloadBidi();
-        break;
-      case "default-nobidi":
-        runWorkloadNoBidi();
+      case "write-only":
+        runWorkloadWriteOnly();
         break;
       default:
         throw new IllegalStateException("Specify a workload to run");
@@ -158,36 +158,31 @@ private void runWorkload1DirectPath() {
     }
   }
 
-  private void runWorkloadBidi() {
-    StorageOptions options =
-        StorageOptions.grpc()
-            .setProjectId(project)
-            .setBlobWriteSessionConfig(BlobWriteSessionConfigs.bidiWrite())
-            .build();
-    Storage storageClient = options.getService();
+  private void runWorkloadWriteOnly() {
     try {
-      runBidi(storageClient);
+      if (bidiEnabled) {
+        StorageOptions options =
+            StorageOptions.grpc()
+                .setProjectId(project)
+                .setBlobWriteSessionConfig(BlobWriteSessionConfigs.bidiWrite())
+                .build();
+        Storage storageClient = options.getService();
+        runBidi(storageClient);
+      } else {
+        StorageOptions options =
+            StorageOptions.grpc()
+                .setProjectId(project)
+                .setBlobWriteSessionConfig(BlobWriteSessionConfigs.getDefault())
+                .build();
+        Storage storageClient = options.getService();
+        runBidi(storageClient);
+      }
     } catch (Exception e) {
       System.err.println("Failed to run workload bidi" + e.getMessage());
       System.exit(1);
     }
   }
 
-  private void runWorkloadNoBidi() {
-    StorageOptions options =
-        StorageOptions.grpc()
-            .setProjectId(project)
-            .setBlobWriteSessionConfig(BlobWriteSessionConfigs.getDefault())
-            .build();
-    Storage storageClient = options.getService();
-    try {
-      runBidi(storageClient);
-    } catch (Exception e) {
-      System.err.println("Failed to run workload no bidi" + e.getMessage());
-      System.exit(1);
-    }
-  }
-
   private void runW1R3(Storage storageClient) throws ExecutionException, InterruptedException {
     ListeningExecutorService executorService =
         MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(workers));