From 2ef1b7be576a17472327efd0cd3d42672a8c4665 Mon Sep 17 00:00:00 2001
From: Brad Baker <bbaker@atlassian.com>
Date: Fri, 1 Mar 2024 10:39:47 +1100
Subject: [PATCH] Now allows the benchmark to go into profiling mode

---
 .../java/benchmark/ComplexQueryBenchmark.java | 40 ++++++++++++++++---
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/test/java/benchmark/ComplexQueryBenchmark.java b/src/test/java/benchmark/ComplexQueryBenchmark.java
index 1dda48e196..7c35d6473e 100644
--- a/src/test/java/benchmark/ComplexQueryBenchmark.java
+++ b/src/test/java/benchmark/ComplexQueryBenchmark.java
@@ -29,6 +29,9 @@
 import org.openjdk.jmh.runner.options.Options;
 import org.openjdk.jmh.runner.options.OptionsBuilder;
 
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
@@ -102,21 +105,47 @@ public static void main(String[] args) throws Exception {
     @SuppressWarnings({"ConstantValue", "LoopConditionNotUpdatedInsideLoop"})
     private static void runAtStartup() {
         // set this to true if you want to hook in profiler say to a forever running JVM
-        boolean forever = Boolean.getBoolean("forever");
+        int runForMillis = getRunForMillis();
 
-        long then = System.currentTimeMillis();
-        System.out.printf("Running initial code before starting the benchmark in forever=%b mode \n", forever);
+        if (runForMillis <= 0) {
+            return;
+        }
+        long now, then = System.currentTimeMillis();
+        System.out.printf("Running initial code before starting the benchmark - runForMillis=%d  \n", runForMillis);
+        System.out.print("Get your profiler in order and press enter...  \n");
+        readLine();
+        System.out.print("Lets go...\n");
         ComplexQueryBenchmark complexQueryBenchmark = new ComplexQueryBenchmark();
         complexQueryBenchmark.setUp();
         do {
-            System.out.print("Running queries....\n");
+            System.out.printf("Running queries for %d millis....\n", System.currentTimeMillis() - then);
             complexQueryBenchmark.howManyItems = 100;
             complexQueryBenchmark.runManyQueriesToCompletion();
-        } while (forever);
+            now = System.currentTimeMillis();
+        } while ((now - then) < runForMillis);
         complexQueryBenchmark.tearDown();
 
         System.out.printf("This took %d millis\n", System.currentTimeMillis() - then);
+        System.exit(0);
+
+    }
+
+    private static void readLine() {
+        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
+        try {
+            br.readLine();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 
+    private static int getRunForMillis() {
+        String runFor = System.getenv("runForMillis");
+        try {
+            return Integer.parseInt(runFor);
+        } catch (NumberFormatException e) {
+            return -1;
+        }
     }
 
     @SuppressWarnings("UnnecessaryLocalVariable")
@@ -200,6 +229,7 @@ private void sleep(Integer howLong) {
     }
 
     AtomicInteger logCount = new AtomicInteger();
+
     private void logEvery(int every, String s) {
         int count = logCount.getAndIncrement();
         if (count == 0 || count % every == 0) {