From 0ad9d4d3031c1c83f1e5d7b283ef72bc46cd0f56 Mon Sep 17 00:00:00 2001
From: Evgeniy Ivanov <i@eivanov.com>
Date: Fri, 15 Dec 2023 13:19:42 +0100
Subject: [PATCH] add the -vt/--virtual-threads option to use virtual threads

---
 .../java/com/oltpbenchmark/DBWorkload.java    | 14 +++++++++++---
 .../java/com/oltpbenchmark/ThreadBench.java   | 19 +++++++++++++++----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/oltpbenchmark/DBWorkload.java b/src/main/java/com/oltpbenchmark/DBWorkload.java
index f4108c27e..c56fa53d3 100644
--- a/src/main/java/com/oltpbenchmark/DBWorkload.java
+++ b/src/main/java/com/oltpbenchmark/DBWorkload.java
@@ -85,6 +85,11 @@ public static void main(String[] args) throws Exception {
       intervalMonitor = Integer.parseInt(argsLine.getOptionValue("im"));
     }
 
+    Boolean useVirtualThreads = false;
+    if (argsLine.hasOption("vt")) {
+      useVirtualThreads = Boolean.parseBoolean(argsLine.getOptionValue("vt"));
+    }
+
     // -------------------------------------------------------------------
     // GET PLUGIN LIST
     // -------------------------------------------------------------------
@@ -518,7 +523,7 @@ public static void main(String[] args) throws Exception {
     if (isBooleanOptionSet(argsLine, "execute")) {
       // Bombs away!
       try {
-        Results r = runWorkload(benchList, intervalMonitor);
+        Results r = runWorkload(benchList, intervalMonitor, useVirtualThreads);
         writeOutputs(r, activeTXTypes, argsLine, xmlConfig);
         writeHistograms(r);
 
@@ -567,6 +572,7 @@ private static Options buildOptions(XMLConfiguration pluginConfig) {
         "Base directory for the result files, default is current directory");
     options.addOption(null, "dialects-export", true, "Export benchmark SQL to a dialects file");
     options.addOption("jh", "json-histograms", true, "Export histograms to JSON file");
+    options.addOption("vt", "virtual-threads", true, "Use virtual threads instead of real threads");
     return options;
   }
 
@@ -733,7 +739,8 @@ private static void runLoader(BenchmarkModule bench)
     bench.loadDatabase();
   }
 
-  private static Results runWorkload(List<BenchmarkModule> benchList, int intervalMonitor)
+  private static Results runWorkload(
+      List<BenchmarkModule> benchList, int intervalMonitor, Boolean useVirtualThreads)
       throws IOException {
     List<Worker<?>> workers = new ArrayList<>();
     List<WorkloadConfiguration> workConfs = new ArrayList<>();
@@ -748,7 +755,8 @@ private static Results runWorkload(List<BenchmarkModule> benchList, int interval
               bench.getBenchmarkName().toUpperCase(), num_phases, (num_phases > 1 ? "s" : "")));
       workConfs.add(bench.getWorkloadConfiguration());
     }
-    Results r = ThreadBench.runRateLimitedBenchmark(workers, workConfs, intervalMonitor);
+    Results r =
+        ThreadBench.runRateLimitedBenchmark(workers, workConfs, intervalMonitor, useVirtualThreads);
     LOG.info(SINGLE_LINE);
     LOG.info("Rate limited reqs/s: {}", r);
     return r;
diff --git a/src/main/java/com/oltpbenchmark/ThreadBench.java b/src/main/java/com/oltpbenchmark/ThreadBench.java
index 9161ae904..88a91f061 100644
--- a/src/main/java/com/oltpbenchmark/ThreadBench.java
+++ b/src/main/java/com/oltpbenchmark/ThreadBench.java
@@ -37,23 +37,27 @@ public class ThreadBench implements Thread.UncaughtExceptionHandler {
   private final List<WorkloadConfiguration> workConfs;
   private final ArrayList<LatencyRecord.Sample> samples = new ArrayList<>();
   private final int intervalMonitor;
+  private final Boolean useVirtualThreads;
 
   private ThreadBench(
       List<? extends Worker<? extends BenchmarkModule>> workers,
       List<WorkloadConfiguration> workConfs,
-      int intervalMonitoring) {
+      int intervalMonitoring,
+      Boolean useVirtualThreads) {
     this.workers = workers;
     this.workConfs = workConfs;
     this.workerThreads = new ArrayList<>(workers.size());
     this.intervalMonitor = intervalMonitoring;
     this.testState = new BenchmarkState(workers.size() + 1);
+    this.useVirtualThreads = useVirtualThreads;
   }
 
   public static Results runRateLimitedBenchmark(
       List<Worker<? extends BenchmarkModule>> workers,
       List<WorkloadConfiguration> workConfs,
-      int intervalMonitoring) {
-    ThreadBench bench = new ThreadBench(workers, workConfs, intervalMonitoring);
+      int intervalMonitoring,
+      Boolean useVirtualThreads) {
+    ThreadBench bench = new ThreadBench(workers, workConfs, intervalMonitoring, useVirtualThreads);
     return bench.runRateLimitedMultiPhase();
   }
 
@@ -61,7 +65,14 @@ private void createWorkerThreads() {
 
     for (Worker<?> worker : workers) {
       worker.initializeState();
-      Thread thread = new Thread(worker);
+
+      Thread thread;
+      if (useVirtualThreads) {
+        thread = Thread.ofVirtual().unstarted(worker);
+      } else {
+        thread = new Thread(worker);
+      }
+
       thread.setUncaughtExceptionHandler(this);
       thread.start();
       this.workerThreads.add(thread);