pg stats with workload name(#31)

yugabyte · Nov 17, 2022 · 9286a5e · 9286a5e
1 parent bfc2402
commit 9286a5e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 39 deletions.
diff --git a/src/main/java/com/oltpbenchmark/api/Worker.java b/src/main/java/com/oltpbenchmark/api/Worker.java
@@ -342,52 +342,14 @@ public final void run() {
             workloadState.finishedWork();
         }
 
-        if (this.getWorkloadConfiguration().getXmlConfig().containsKey("collect_pg_stat_statements") &&
-            this.getWorkloadConfiguration().getXmlConfig().getBoolean("collect_pg_stat_statements")) {
-            LOG.info("Collecting pg_stat_statements");
-            try {
-                excutePgStatStatements(conn);
-            } catch (SQLException e) {
-                throw new RuntimeException(e);
-            }
-        }
+
 
         LOG.debug("worker calling teardown");
 
         tearDown();
     }
 
-    private void excutePgStatStatements(Connection conn) throws SQLException {
-        String pgStatDDL = "select * from pg_stat_statements;";
-        String PgStatsDir = "ResultsForPgStats";
-        FileUtil.makeDirIfNotExists("results" + "/" + PgStatsDir);
-        String fileForPgStats = PgStatsDir + "/" + TimeUtil.getCurrentTimeString() + ".json";
-        PrintStream ps;
-        try {
-            ps = new PrintStream(FileUtil.joinPath("results", fileForPgStats));
-        } catch (FileNotFoundException exc) {
-            throw new RuntimeException(exc);
-        }
 
-        Map<String, JSONObject> summaryMap = new TreeMap<>();
-        Statement stmt = conn.createStatement();
-        JSONObject outer = new JSONObject();
-        int count = 0;
-        ResultSet resultSet = stmt.executeQuery(pgStatDDL);
-        ResultSetMetaData rsmd = resultSet.getMetaData();
-        int columnsNumber = rsmd.getColumnCount();
-        while (resultSet.next()) {
-            JSONObject inner = new JSONObject();
-            for (int i = 1; i <= columnsNumber; i++) {
-                String columnValue = resultSet.getString(i);
-                inner.put(rsmd.getColumnName(i), columnValue);
-            }
-            outer.put("Record_" + count, inner);
-            count++;
-        }
-        summaryMap.put("PgStats", outer);
-        ps.println(JSONUtil.format(JSONUtil.toJSONString(summaryMap)));
-    }
 
 
     private TransactionType getTransactionType(SubmittedProcedure pieceOfWork, Phase phase, State state, WorkloadState workloadState) {

diff --git a/src/main/java/com/oltpbenchmark/benchmarks/featurebench/FeatureBenchWorker.java b/src/main/java/com/oltpbenchmark/benchmarks/featurebench/FeatureBenchWorker.java
@@ -63,6 +63,7 @@ public FeatureBenchWorker(FeatureBenchBenchmark benchmarkModule, int id) {
 
     protected void initialize() {
 
+
         if (this.getWorkloadConfiguration().getXmlConfig().containsKey("collect_pg_stat_statements") &&
             this.getWorkloadConfiguration().getXmlConfig().getBoolean("collect_pg_stat_statements")) {
             LOG.info("Resetting pg_stat_statements");
@@ -213,6 +214,16 @@ protected TransactionStatus executeWork(Connection conn, TransactionType txnType
     @Override
     public void tearDown() {
 
+        if (this.getWorkloadConfiguration().getXmlConfig().containsKey("collect_pg_stat_statements") &&
+            this.getWorkloadConfiguration().getXmlConfig().getBoolean("collect_pg_stat_statements")) {
+            LOG.info("Collecting pg_stat_statements");
+            try {
+                excutePgStatStatements(conn);
+            } catch (SQLException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
         if (!this.configuration.getNewConnectionPerTxn() && this.conn != null && ybm != null) {
             try {
                 if (this.configuration.getWorkloadState().getGlobalState() == State.EXIT && !isCleanUpDone.get()) {
@@ -240,4 +251,36 @@ public void tearDown() {
             }
         }
     }
+
+    private void excutePgStatStatements(Connection conn) throws SQLException {
+        String pgStatDDL = "select * from pg_stat_statements;";
+        String PgStatsDir = "ResultsForPgStats";
+        FileUtil.makeDirIfNotExists("results" + "/" + PgStatsDir);
+        String fileForPgStats = PgStatsDir + "/" + workloadName + "_"+TimeUtil.getCurrentTimeString()+".json";
+        PrintStream ps;
+        try {
+            ps = new PrintStream(FileUtil.joinPath("results", fileForPgStats));
+        } catch (FileNotFoundException exc) {
+            throw new RuntimeException(exc);
+        }
+
+        Map<String, JSONObject> summaryMap = new TreeMap<>();
+        Statement stmt = conn.createStatement();
+        JSONObject outer = new JSONObject();
+        int count = 0;
+        ResultSet resultSet = stmt.executeQuery(pgStatDDL);
+        ResultSetMetaData rsmd = resultSet.getMetaData();
+        int columnsNumber = rsmd.getColumnCount();
+        while (resultSet.next()) {
+            JSONObject inner = new JSONObject();
+            for (int i = 1; i <= columnsNumber; i++) {
+                String columnValue = resultSet.getString(i);
+                inner.put(rsmd.getColumnName(i), columnValue);
+            }
+            outer.put("Record_" + count, inner);
+            count++;
+        }
+        summaryMap.put("PgStats", outer);
+        ps.println(JSONUtil.format(JSONUtil.toJSONString(summaryMap)));
+    }
 }