Add newConnectionPerTxn config option and implementation (#208)

Co-authored-by: Andy Pavlo <[email protected]>

Author: mbutrovich

src/main/java/com/oltpbenchmark/DBWorkload.java

@@ -125,6 +125,7 @@ public static void main(String[] args) throws Exception {
             wrkld.setRandomSeed(xmlConfig.getInt("randomSeed", -1));
             wrkld.setBatchSize(xmlConfig.getInt("batchsize", 128));
             wrkld.setMaxRetries(xmlConfig.getInt("retries", 3));
+            wrkld.setNewConnectionPerTxn(xmlConfig.getBoolean("newConnectionPerTxn", false));
 
             int terminals = xmlConfig.getInt("terminals[not(@bench)]", 0);
             terminals = xmlConfig.getInt("terminals" + pluginTest, terminals);
@@ -170,6 +171,7 @@ public static void main(String[] args) throws Exception {
             initDebug.put("Batch Size", wrkld.getBatchSize());
             initDebug.put("Scale Factor", wrkld.getScaleFactor());
             initDebug.put("Terminals", wrkld.getTerminals());
+            initDebug.put("New Connection Per Txn", wrkld.getNewConnectionPerTxn());
 
             if (selectivity != -1) {
                 initDebug.put("Selectivity", selectivity);

src/main/java/com/oltpbenchmark/WorkloadConfiguration.java

@@ -50,6 +50,12 @@ public class WorkloadConfiguration {
     private String dataDir = null;
     private String ddlPath = null;
 
+    /**
+     * If true, establish a new connection for each transaction, otherwise use one persistent connection per client
+     * session. This is useful to measure the connection overhead.
+     */
+    private boolean newConnectionPerTxn = false;
+
     public String getBenchmarkName() {
         return benchmarkName;
     }
@@ -118,6 +124,23 @@ public void setMaxRetries(int maxRetries) {
         this.maxRetries = maxRetries;
     }
 
+    /**
+     * @return @see newConnectionPerTxn member docs for behavior.
+     */
+    public boolean getNewConnectionPerTxn() {
+        return newConnectionPerTxn;
+    }
+
+    /**
+     * Used by the configuration loader at startup. Changing it any other time is probably dangeroues. @see
+     * newConnectionPerTxn member docs for behavior.
+     *
+     * @param newConnectionPerTxn
+     */
+    public void setNewConnectionPerTxn(boolean newConnectionPerTxn) {
+        this.newConnectionPerTxn = newConnectionPerTxn;
+    }
+
     /**
      * Initiate a new benchmark and workload state
      */

src/main/java/com/oltpbenchmark/api/Worker.java

@@ -50,7 +50,7 @@ public abstract class Worker<T extends BenchmarkModule> implements Runnable {
 
     private final int id;
     private final T benchmark;
-    protected final Connection conn;
+    protected Connection conn = null;
     protected final WorkloadConfiguration configuration;
     protected final TransactionTypes transactionTypes;
     protected final Map<TransactionType, Procedure> procedures = new HashMap<>();
@@ -74,12 +74,14 @@ public Worker(T benchmark, int id) {
         this.currStatement = null;
         this.transactionTypes = this.configuration.getTransTypes();
 
-        try {
-            this.conn = this.benchmark.makeConnection();
-            this.conn.setAutoCommit(false);
-            this.conn.setTransactionIsolation(this.configuration.getIsolationMode());
-        } catch (SQLException ex) {
-            throw new RuntimeException("Failed to connect to database", ex);
+        if (!this.configuration.getNewConnectionPerTxn()) {
+            try {
+                this.conn = this.benchmark.makeConnection();
+                this.conn.setAutoCommit(false);
+                this.conn.setTransactionIsolation(this.configuration.getIsolationMode());
+            } catch (SQLException ex) {
+                throw new RuntimeException("Failed to connect to database", ex);
+            }
         }
 
         // Generate all the Procedures that we're going to need
@@ -391,6 +393,20 @@ protected final void doWork(DatabaseType databaseType, TransactionType transacti
 
                 TransactionStatus status = TransactionStatus.UNKNOWN;
 
+                if (this.conn == null) {
+                    try {
+                        this.conn = this.benchmark.makeConnection();
+                        this.conn.setAutoCommit(false);
+                        this.conn.setTransactionIsolation(this.configuration.getIsolationMode());
+                    } catch (SQLException ex) {
+                        if (LOG.isDebugEnabled()) {
+                            LOG.debug(String.format("%s failed to open a connection...", this));
+                        }
+                        retryCount++;
+                        continue;
+                    }
+                }
+
                 try {
 
                     if (LOG.isDebugEnabled()) {
@@ -438,6 +454,14 @@ protected final void doWork(DatabaseType databaseType, TransactionType transacti
                     }
 
                 } finally {
+                    if (this.configuration.getNewConnectionPerTxn() && this.conn != null) {
+                        try {
+                            this.conn.close();
+                            this.conn = null;
+                        } catch (SQLException e) {
+                            LOG.error("Connection couldn't be closed.", e);
+                        }
+                    }
 
                     switch (status) {
                         case UNKNOWN -> this.txnUnknown.put(transactionType);
@@ -511,10 +535,12 @@ protected void initialize() {
      * Called at the end of the test to do any clean up that may be required.
      */
     public void tearDown() {
-        try {
-            conn.close();
-        } catch (SQLException e) {
-            LOG.error("Connection couldn't be closed.", e);
+        if (!this.configuration.getNewConnectionPerTxn() && this.conn != null) {
+            try {
+                conn.close();
+            } catch (SQLException e) {
+                LOG.error("Connection couldn't be closed.", e);
+            }
         }
     }