From 79dc5da76b434ed51d910f8160b00f63374c7d50 Mon Sep 17 00:00:00 2001
From: IshanChhangani <f20200230@hyderabad.bits-pilani.ac.in>
Date: Fri, 21 Feb 2025 14:01:28 +0530
Subject: [PATCH] [#26116] YSQL, QueryDiagnostics: Fix
 TestYbQueryDiagnostics.testIntermediateFlushing

Summary:
TestYbQueryDiagnostics.testIntermediateFlushing has been failing in TSAN builds starting from #22612. This is a regression based on the stability history.
Jira: DB-15443

Test Plan: yb_build.sh --java-test 'org.yb.pgsql.TestYbQueryDiagnostics#testIntermediateFlushing'

Reviewers: asaha

Reviewed By: asaha

Subscribers: svc_phabricator

Differential Revision: https://phorge.dev.yugabyte.com/D42073
---
 .../org/yb/pgsql/TestYbQueryDiagnostics.java  | 56 +++++++++----------
 1 file changed, 26 insertions(+), 30 deletions(-)

diff --git a/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java b/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java
index 9e2333e08ff3..c2ce1bac391b 100644
--- a/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java
+++ b/java/yb-pgsql/src/test/java/org/yb/pgsql/TestYbQueryDiagnostics.java
@@ -338,25 +338,30 @@ private void waitForBundleCompletion(String queryId, Statement statement,
                                          int diagnosticsInterval) throws Exception {
         Thread.sleep(diagnosticsInterval * 1000 + BG_WORKER_INTERVAL_MS);
 
-        long start_time = System.currentTimeMillis();
-        while (true)
-        {
-            System.out.println("Waiting in forever loop");
+        long startTime = System.currentTimeMillis();
 
-            ResultSet resultSet = statement.executeQuery(
-                                  "SELECT * FROM yb_query_diagnostics_status where query_id = " +
-                                  queryId);
-
-            if (resultSet.next() && !resultSet.getString("status").equals("In Progress"))
-                break;
-
-            if (System.currentTimeMillis() - start_time > 60000) // 1 minute
-                fail("Bundle did not complete within the expected time");
-
-            Thread.sleep(BG_WORKER_INTERVAL_MS);
+        try {
+            TestUtils.waitFor(() -> {
+                ResultSet resultSet = statement.executeQuery(
+                        "SELECT * FROM yb_query_diagnostics_status where query_id = " +
+                                queryId);
+
+                if (resultSet.next() && !resultSet.getString("status").equals("In Progress"))
+                    return true;
+
+                return false;
+            },
+                    60000L, BG_WORKER_INTERVAL_MS);
+        } catch (Exception e) {
+            throw new AssertionError(
+                    "Bundle did not complete within the expected time");
         }
     }
 
+    private void waitForBundleCompletion(String queryId, Statement statement) throws Exception {
+        waitForBundleCompletion(queryId, statement, 0);
+    }
+
     /*
      * Counts the number of table sections in the schema details.
      */
@@ -1397,12 +1402,12 @@ public void testIntermediateFlushing() throws Exception {
              * we consider only half the length.
              */
             int varLen = minLen / 2;
-            String largeVariable = new String(new char[ varLen ]).replace('\0', 'a');
+            String largeVariable = new String(new char[varLen]).replace('\0', 'a');
 
             /* To ensure that the buffer overflows we do twice as many iterations */
             int noOfIterations = (maxLen / varLen) + 1;
 
-            int diagnosticsInterval = noOfIterations * (BG_WORKER_INTERVAL_MS / 1000);
+            int diagnosticsInterval = 2 * noOfIterations * (BG_WORKER_INTERVAL_MS / 1000);
             QueryDiagnosticsParams params = new QueryDiagnosticsParams(
                 diagnosticsInterval,
                 100 /* explainSampleRate */,
@@ -1420,19 +1425,7 @@ public void testIntermediateFlushing() throws Exception {
                 Thread.sleep(BG_WORKER_INTERVAL_MS);
             }
 
-            long start_time = System.currentTimeMillis();
-            while (true)
-            {
-                ResultSet resultSet = statement.executeQuery(
-                                      "SELECT * FROM yb_query_diagnostics_status");
-                if (resultSet.next() && resultSet.getString("status").equals("Success"))
-                    break;
-
-                if (System.currentTimeMillis() - start_time > 60000) // 1 minute
-                    fail("Bundle did not complete within the expected time");
-
-                Thread.sleep(BG_WORKER_INTERVAL_MS);
-            }
+            waitForBundleCompletion(queryId, statement);
 
             /* Bundle has expired */
             Path bindVariablesPath = bundleDataPath.resolve("constants_and_bind_variables.csv");
@@ -1444,6 +1437,9 @@ public void testIntermediateFlushing() throws Exception {
             String bindVariablesContent = new String(Files.readAllBytes(bindVariablesPath));
             String explainPlanContent = new String(Files.readAllBytes(explainPlanPath));
 
+            LOG.info("Explain plan content: \n" + explainPlanContent);
+            LOG.info("Bind variables content: \n" + bindVariablesContent);
+
             int bindVariablesLength = bindVariablesContent.length();
             int explainPlanLength = explainPlanContent.length();