[#23771] YSQL, ASH: Fix prepared statements not being tracked in ASH

Summary:
pg_stat_statements doesn't track the execution of PREPARE, EXECUTE and
DEALLOCATE, so it doesn't set a query id for these statements, however
it sets the query id for the underlying query of EXECUTE, the one which
is actually executed by the executor hooks.

Previously, the execution of prepared statements using PREPARE and
EXECUTE were not being tracked by default after D34773 because the
execution happens after the utility node of EXECUTE, which made it a
nested statement, and unless pg_stat_statements.track='all' is set,
these are not tracked. As a result, the underlying query of EXECUTE
were not being tracked by default (with pg_stat_statements.track='top')
in ASH but they were in pg_stat_statements.

This diff fixes the issue by skipping the nested level when the node being
processed is PREPARE, EXECUTE or DEALLOCATE, which is the same
behaviour as pg_stat_statements, and this means that either the
underlying query of EXECUTE will always be tracked by ASH when it's
also being tracked by pg_stat_statements

This diff also enables the ASH java tests in TSAN mode after D37049
fixed the test failures with ASH in TSAN mode
Jira: DB-12675

Test Plan: ./yb_build.sh --java-test TestYbAsh#testPreparedStatements

Reviewers: jason

Reviewed By: jason

Subscribers: amitanand, hbhanawat, yql

Differential Revision: https://phorge.dev.yugabyte.com/D37737

java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbAsh.java

@@ -30,11 +30,11 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.yb.util.YBTestRunnerNonTsanOnly;
+import org.yb.YBTestRunner;
 
 import com.yugabyte.util.PSQLException;
 
-@RunWith(value = YBTestRunnerNonTsanOnly.class)
+@RunWith(value = YBTestRunner.class)
 public class TestYbAsh extends BasePgSQLTest {
   private static final int ASH_SAMPLING_INTERVAL = 1000;
 
@@ -340,4 +340,30 @@ public void testYsqlPids() throws Exception {
       assertGreaterThan(res, 0);
     }
   }
+
+  /**
+   * Test that we are tracking the correct query ids with prepared statements.
+   */
+  @Test
+  public void testPreparedStatements() throws Exception {
+    setAshConfigAndRestartCluster(10, ASH_SAMPLE_SIZE);
+
+    try (Statement statement = connection.createStatement()) {
+      String table = "test_table";
+      String pstmt = String.format("PREPARE f(INT) AS SELECT * FROM %s WHERE k = $1", table);
+      statement.execute(String.format("CREATE TABLE %s(k INT, v TEXT)", table));
+      statement.execute(pstmt);
+      for (int i = 0; i < 100; ++i) {
+        statement.execute(String.format("INSERT INTO %s VALUES(%d, 'v-%d')", table, i, i));
+      }
+      for (int i = 0; i < 100; ++i) {
+        statement.execute(String.format("EXECUTE f(%d)", i));
+      }
+      long pstmtQueryId = getSingleRow(statement, String.format("SELECT queryid FROM " +
+          "pg_stat_statements WHERE query = '%s'", pstmt)).getLong(0);
+      int res = getSingleRow(statement, String.format("SELECT COUNT(*) FROM %s WHERE " +
+          "query_id = %d", ASH_VIEW, pstmtQueryId)).getLong(0).intValue();
+      assertGreaterThan(res, 0);
+    }
+  }
 }

src/postgres/src/backend/utils/misc/yb_ash.c

@@ -428,19 +428,32 @@ yb_ash_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 					  QueryEnvironment *queryEnv, DestReceiver *dest,
 					  char *completionTag)
 {
-	uint64 query_id;
+	uint64		query_id;
+	bool		skip_nested_level;
+	Node	   *parsetree = pstmt->utilityStmt;
 
-	if (yb_enable_ash)
+	/*
+	 * We don't want to set query id if the node is PREPARE, EXECUTE or
+	 * DEALLOCATE because pg_stat_statements also doesn't do it. Check
+	 * comments in pgss_ProcessUtility for more info.
+	 */
+	skip_nested_level = IsA(parsetree, PrepareStmt) || IsA(parsetree, ExecuteStmt) ||
+						IsA(parsetree, DeallocateStmt);
+
+	if (!skip_nested_level)
 	{
-		query_id = pstmt->queryId != 0
-				   ? pstmt->queryId
-				   : yb_ash_utility_query_id(queryString,
-					   						 pstmt->stmt_len,
-											 pstmt->stmt_location);
-		YbAshSetQueryId(query_id);
+		if (yb_enable_ash)
+		{
+			query_id = pstmt->queryId != 0
+					   ? pstmt->queryId
+					   : yb_ash_utility_query_id(queryString,
+												 pstmt->stmt_len,
+												 pstmt->stmt_location);
+			YbAshSetQueryId(query_id);
+		}
+		++nested_level;
 	}
 
-	++nested_level;
 	PG_TRY();
 	{
 		if (prev_ProcessUtility)
@@ -451,18 +464,21 @@ yb_ash_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 			standard_ProcessUtility(pstmt, queryString,
 									context, params, queryEnv,
 									dest, completionTag);
-		--nested_level;
-
-		if (yb_enable_ash)
-			YbAshResetQueryId(query_id);
+		if (!skip_nested_level)
+		{
+			--nested_level;
+			if (yb_enable_ash)
+				YbAshResetQueryId(query_id);
+		}
 	}
 	PG_CATCH();
 	{
-		--nested_level;
-
-		if (yb_enable_ash)
-			YbAshResetQueryId(query_id);
-
+		if (!skip_nested_level)
+		{
+			--nested_level;
+			if (yb_enable_ash)
+				YbAshResetQueryId(query_id);
+		}
 		PG_RE_THROW();
 	}
 	PG_END_TRY();