add single-transaction flag to examine rpc requests of a single rando…

…m transaction
yugabyte · Aug 18, 2021 · 512177d · 512177d
1 parent 125fff8
commit 512177d
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/src/com/oltpbenchmark/CommandLineOptions.java b/src/com/oltpbenchmark/CommandLineOptions.java
@@ -60,6 +60,7 @@ public class CommandLineOptions {
         COMMAND_LINE_OPTS.addOption(null, "merge-results", true, "Merge results from various output files");
         COMMAND_LINE_OPTS.addOption(null, "dir", true, "Directory containing the csv files");
         COMMAND_LINE_OPTS.addOption(null, "vv", false, "Output verbose execute results");
+        COMMAND_LINE_OPTS.addOption(null, "single-transaction", false, "Examine RPC requests of a single random transaction");
     }
 
     public CommandLineOptions() {}
@@ -173,4 +174,6 @@ public boolean getIsOutputMetricHistogramsSet() {
     }
 
     public boolean getShouldOutputVerboseExecuteResults() { return argsLine.hasOption("vv"); }
+
+    public boolean getSingleTransactionRPC() { return argsLine.hasOption("single-transaction"); }
 }
diff --git a/src/com/oltpbenchmark/DBWorkload.java b/src/com/oltpbenchmark/DBWorkload.java
@@ -214,7 +214,7 @@ public static void main(String[] args) throws Exception {
       // ----------------------------------------------------------------
       // CREATE BENCHMARK MODULE
       // ----------------------------------------------------------------
-      BenchmarkModule bench = new BenchmarkModule(wrkld);
+      BenchmarkModule bench = new BenchmarkModule(wrkld, options);
       Map<String, Object> initDebug = new ListOrderedMap<>();
       initDebug.put("Benchmark", String.format("%s {%s}", plugin.toUpperCase(), "BenchmarkModule"));
       initDebug.put("Configuration", configFile);

diff --git a/src/com/oltpbenchmark/api/BenchmarkModule.java b/src/com/oltpbenchmark/api/BenchmarkModule.java
@@ -32,6 +32,7 @@
 
 import org.apache.log4j.Logger;
 
+import com.oltpbenchmark.CommandLineOptions;
 import com.oltpbenchmark.WorkloadConfiguration;
 import com.oltpbenchmark.api.Loader.LoaderThread;
 import com.oltpbenchmark.util.ThreadUtil;
@@ -58,29 +59,34 @@ public class BenchmarkModule {
      */
     protected final WorkloadConfiguration workConf;
 
+    /**
+     * The command line options for this benchmark invocation
+     */
+    protected final CommandLineOptions options;
+
     /**
      * A single Random object that should be re-used by all a benchmark's components
      */
     private final Random rng = new Random();
 
-    public BenchmarkModule(WorkloadConfiguration workConf) {
+    public BenchmarkModule(WorkloadConfiguration workConf, CommandLineOptions options) {
         assert (workConf != null) : "The WorkloadConfiguration instance is null.";
 
         this.benchmarkName = "tpcc";
         this.workConf = workConf;
+        this.options = options;
         if (workConf.getNeedsExecution()) {
             try {
                 createDataSource();
             } catch (Exception e) {
                 LOG.error("Failed to create Data source", e);
-                throw e;
             }
         }
     }
 
     private final List<HikariDataSource> listDataSource = new ArrayList<>();
 
-    public void createDataSource() {
+    public void createDataSource() throws Exception {
         int numConnections =
             (workConf.getNumDBConnections() + workConf.getNodes().size() - 1) / workConf.getNodes().size();
         for (String ip : workConf.getNodes()) {
@@ -112,6 +118,16 @@ public void createDataSource() {
             config.setTransactionIsolation(workConf.getIsolationString());
             listDataSource.add(new HikariDataSource(config));
         }
+
+        // if the single transaction flag is set
+        // get a random datasource
+        if (options.getSingleTransactionRPC()) {
+          HikariDataSource dataSource = listDataSource.get(rng.nextInt(workConf.getNodes().size()));
+          try(Connection connection = dataSource.getConnection(); Statement statement = connection.createStatement()) {
+              statement.execute("SET yb_debug_log_docdb_requests=true");
+              statement.execute("SET log_statement='all'");
+          }
+        }
     }
 
     public final Connection makeConnection() throws SQLException {