Skip to content

Commit

Permalink
[#25865] YSQL, QueryDiagnostics: Utilize SPI to fetch ASH data
Browse files Browse the repository at this point in the history 
Summary:
Currently, we were fetching ASH data directly from ASH's circular by doing a binary search over it. This involves unnecessary complexity. This diff utilizes SPI framework to dump ASH data.
Jira: DB-15168

Test Plan:
yb_build.sh --java-test 'org.yb.pgsql.TestYbQueryDiagnostics#testLongQueryWith5000Constants'
yb_build.sh --java-test 'org.yb.pgsql.TestYbQueryDiagnostics#testComplexQuery'
yb_build.sh --java-test 'org.yb.pgsql.TestYbQueryDiagnostics#checkAshData'

Reviewers: asaha

Reviewed By: asaha

Subscribers: yql

Differential Revision: https://phorge.dev.yugabyte.com/D41747
  • Loading branch information
@IshanChhangani
IshanChhangani committed Mar 1, 2025
1 parent 79dc5da commit 0260740
Show file tree
Hide file tree
Showing 2 changed files with 216 additions and 105 deletions.
49 changes: 27 additions & 22 deletions java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,8 @@ public QueryDiagnosticsStatus(Path path, String status, String description,

private static final Logger LOG = LoggerFactory.getLogger(TestYbQueryDiagnostics.class);
private static final AtomicInteger queryIdGenerator = new AtomicInteger();
private static final int ASH_SAMPLING_INTERVAL_MS = 500;
// Smaller value gives higher chances of data capture
private static final int ASH_SAMPLING_INTERVAL_MS = 50;
private static final int YB_QD_MAX_EXPLAIN_PLAN_LEN = 16384;
private static final int YB_QD_MAX_BIND_VARS_LEN = 2048;
private static final int BG_WORKER_INTERVAL_MS = 1000;
Expand All @@ -122,7 +123,7 @@ public void setUp() throws Exception {

/* Required for some of the fields within schema details */
flagMap.put("ysql_beta_features", "true");
flagMap.put("ysql_yb_ash_sampling_interval_ms", "100");
flagMap.put("ysql_yb_ash_sampling_interval_ms", String.valueOf(ASH_SAMPLING_INTERVAL_MS));

restartClusterWithFlags(Collections.emptyMap(), flagMap);

Expand All @@ -140,7 +141,7 @@ public void setUp(int queryDiagnosticsCircularBufferSize) throws Exception {

/* Required for some of the fields within schema details */
flagMap.put("ysql_beta_features", "true");
flagMap.put("ysql_yb_ash_sampling_interval_ms", "50");
flagMap.put("ysql_yb_ash_sampling_interval_ms", String.valueOf(ASH_SAMPLING_INTERVAL_MS));

restartClusterWithFlags(Collections.emptyMap(), flagMap);

Expand Down Expand Up @@ -331,12 +332,16 @@ private Path runQueryDiagnostics(Statement statement, String queryId,
return Paths.get(resultSet.getString("yb_query_diagnostics"));
}

private void waitForBundleCompletion(String queryId, Statement statement) throws Exception {
waitForBundleCompletion(queryId, statement, 0);
}
/*
* Waits for the bundle to complete by checking the yb_query_diagnostics_status view.
*/
private void waitForBundleCompletion(String queryId, Statement statement,
int diagnosticsInterval) throws Exception {
Thread.sleep(diagnosticsInterval * 1000 + BG_WORKER_INTERVAL_MS);
statement.execute("SELECT pg_sleep(" +
(diagnosticsInterval + (BG_WORKER_INTERVAL_MS / 1000)) + ")");

long startTime = System.currentTimeMillis();

Expand Down Expand Up @@ -757,8 +762,6 @@ private String extractQueryTextFromPgssFile(Statement statement, Path pgssPath)

private void validateAgainstFile(String expectedFilePath, String actualData) throws Exception{

Files.write(Paths.get(expectedFilePath), actualData.getBytes());

Path expectedOutputPath = Paths.get(expectedFilePath);
String expectedOutput = new String(Files.readAllBytes(expectedOutputPath),
StandardCharsets.UTF_8);
Expand Down Expand Up @@ -1221,11 +1224,9 @@ public void testCircularBufferWrapAround() throws Exception {
}
}

/**
* Disable until #25865 is done.
*/
@Test
public void checkAshData() throws Exception {
int diagnosticsInterval = (5 * ASH_SAMPLING_INTERVAL_MS) / 1000; /* convert to seconds */
int diagnosticsInterval = 2;
QueryDiagnosticsParams params = new QueryDiagnosticsParams(
diagnosticsInterval /* diagnosticsInterval */,
100 /* explainSampleRate */,
Expand All @@ -1235,15 +1236,19 @@ public void checkAshData() throws Exception {
0 /* bindVarQueryMinDuration */);

try (Statement statement = connection.createStatement()) {
statement.execute("SELECT pg_sleep(0.5)");

String queryId = getQueryIdFromPgStatStatements(statement, "%pg_sleep%");
String queryId = getQueryIdFromPgStatStatements(statement, "%PREPARE%");
Path bundleDataPath = runQueryDiagnostics(statement, queryId, params);

/* Protects from "No query executed;" warning */
statement.execute("SELECT pg_sleep(0.1)");
/*
* Protects from "No query executed;" warning and ensures that we do have some data
* to dump in active_session_history.csv otherwise the file wouldn't be created
* and we would get a "File does not exist" error.
*/
for (int i = 0; i < 100; i++) {
statement.execute("EXECUTE stmt('var1', 1, 1.1)");
}

waitForBundleCompletion(queryId, statement, diagnosticsInterval);
waitForBundleCompletion(queryId, statement, 2 * diagnosticsInterval);

Path ashPath = getFilePathFromBaseDir(bundleDataPath,
"active_session_history.csv");
Expand All @@ -1254,7 +1259,7 @@ public void checkAshData() throws Exception {

@Test
public void checkPgssData() throws Exception {
int diagnosticsInterval = (5 * ASH_SAMPLING_INTERVAL_MS) / 1000; /* convert to seconds */
int diagnosticsInterval = 2;
QueryDiagnosticsParams queryDiagnosticsParams = new QueryDiagnosticsParams(
diagnosticsInterval,
100 /* explainSampleRate */,
Expand Down Expand Up @@ -1284,7 +1289,7 @@ public void checkPgssData() throws Exception {

@Test
public void testPgssResetBetweenDiagnostics() throws Exception {
int diagnosticsInterval = (5 * ASH_SAMPLING_INTERVAL_MS) / 1000; /* convert to seconds */
int diagnosticsInterval = 2;
QueryDiagnosticsParams queryDiagnosticsParams = new QueryDiagnosticsParams(
diagnosticsInterval,
100 /* explainSampleRate */,
Expand Down Expand Up @@ -1339,7 +1344,7 @@ public void testPgssResetBetweenDiagnostics() throws Exception {

@Test
public void emptyBundle() throws Exception {
int diagnosticsInterval = (5 * ASH_SAMPLING_INTERVAL_MS) / 1000; /* convert to seconds */
int diagnosticsInterval = 2;
QueryDiagnosticsParams params = new QueryDiagnosticsParams(
diagnosticsInterval,
100 /* explainSampleRate */,
Expand Down Expand Up @@ -1691,10 +1696,10 @@ public void testYbCancelQueryDiagnostics() throws Exception {
/*
* Test that yb_query_diagnostics works fine when diagnosing a query that has
* large number of joins, subqueries, and constants.
* Disable until #25865 is done
*/
@Test
public void testComplexQuery() throws Exception {
int diagnosticsInterval = 2;
int diagnosticsInterval = 5;
QueryDiagnosticsParams queryDiagnosticsParams = new QueryDiagnosticsParams(
diagnosticsInterval,
100 /* explainSampleRate */,
Expand Down Expand Up @@ -1747,8 +1752,8 @@ public void testComplexQuery() throws Exception {

/*
* Tests a Long query(around 100000 chars).
* Disable until #25865 is done
*/
@Test
public void testLongQueryWith5000Constants() throws Exception {
int diagnosticsInterval = 2;
QueryDiagnosticsParams queryDiagnosticsParams = new QueryDiagnosticsParams(
Expand Down
Loading

0 comments on commit 0260740

Please sign in to comment.