Fix IT

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

integ-test/src/test/java/org/opensearch/sql/ppl/ResourceMonitorIT.java

@@ -13,6 +13,7 @@
 import org.junit.Test;
 import org.opensearch.client.ResponseException;
 import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.opensearch.monitor.GCedMemoryUsage;
 
 public class ResourceMonitorIT extends PPLIntegTestCase {
 
@@ -27,11 +28,7 @@ public void queryExceedResourceLimitShouldFail() throws IOException {
     // update plugins.ppl.query.memory_limit to 1%
     updateClusterSettings(
         new ClusterSetting("persistent", Settings.Key.QUERY_MEMORY_LIMIT.getKeyValue(), "1%"));
-    // ClickBench Q30 is high memory consumption query.
-    String query = sanitize(loadFromFile("clickbench/queries/q30.ppl"));
-
-    ResponseException exception =
-        expectThrows(ResponseException.class, () -> executeQueryMultipleTimes(query));
+    ResponseException exception = expectThrows(ResponseException.class, this::executeQuery);
     assertEquals(500, exception.getResponse().getStatusLine().getStatusCode());
     assertThat(
         exception.getMessage(),
@@ -40,12 +37,21 @@ public void queryExceedResourceLimitShouldFail() throws IOException {
     // update plugins.ppl.query.memory_limit to default value 85%
     updateClusterSettings(
         new ClusterSetting("persistent", Settings.Key.QUERY_MEMORY_LIMIT.getKeyValue(), "85%"));
-    executeQueryMultipleTimes(query);
+    executeQuery();
   }
 
-  /** Run the same query multiple times in case GC can be triggered in integration test */
-  private void executeQueryMultipleTimes(String query) throws IOException {
-    for (int i = 0; i < 10; i++) {
+  private void executeQuery() throws IOException {
+    if (GCedMemoryUsage.initialized()) {
+      // ClickBench Q30 is a high memory consumption query. Run 5 times to ensure GC triggered.
+      String query = sanitize(loadFromFile("clickbench/queries/q30.ppl"));
+      for (int i = 0; i < 5; i++) {
+        JSONObject result = executeQuery(query);
+        verifyNumOfRows(result, 1);
+      }
+    } else {
+      // Q2 is not a high memory consumption query.
+      // It cannot run in 1% resource but passed in 85%.
+      String query = sanitize(loadFromFile("clickbench/queries/q2.ppl"));
       JSONObject result = executeQuery(query);
       verifyNumOfRows(result, 1);
     }

opensearch/src/main/java/org/opensearch/sql/opensearch/monitor/GCedMemoryUsage.java

@@ -24,6 +24,7 @@ public class GCedMemoryUsage implements MemoryUsage {
   private static final Logger LOG = LogManager.getLogger();
   private static final List<String> OLD_GEN_GC_ACTION_KEYWORDS =
       List.of("major", "concurrent", "old", "full", "marksweep");
+  private static boolean initialized = false;
 
   private GCedMemoryUsage() {
     registerGCListener();
@@ -55,6 +56,10 @@ public void setUsage(long value) {
     usage.set(value);
   }
 
+  public static boolean initialized() {
+    return initialized;
+  }
+
   private void registerGCListener() {
     List<GarbageCollectorMXBean> gcBeans = ManagementFactory.getGarbageCollectorMXBeans();
     if (gcBeans.stream()
@@ -68,6 +73,7 @@ private void registerGCListener() {
     for (GarbageCollectorMXBean gcBean : gcBeans) {
       if (gcBean instanceof NotificationEmitter && isOldGenGc(gcBean.getName())) {
         LOG.info("{} listener registered for memory usage monitor.", gcBean.getName());
+        initialized = true;
         NotificationEmitter emitter = (NotificationEmitter) gcBean;
         emitter.addNotificationListener(
             new OldGenGCListener(),