ResourceMonitor only checks the memory health by calculating the memory usage after GC of old gen

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

Author: LantaoJin

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -112,11 +112,14 @@ public void executeWithCalcite(
         log.warn("Fallback to V2 query engine since got exception", t);
         executeWithLegacy(plan, queryType, listener, Optional.of(t));
       } else {
-        if (t instanceof Error) {
+        if (t instanceof Exception) {
+          listener.onFailure((Exception) t);
+        } else if (t instanceof VirtualMachineError) {
+          // throw and fast fail the VM errors such as OOM (same with v2).
+          throw t;
+        } else {
           // Calcite may throw AssertError during query execution.
           listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
-        } else {
-          listener.onFailure((Exception) t);
         }
       }
     }

integ-test/src/test/java/org/opensearch/sql/calcite/clickbench/PPLClickBenchIT.java

@@ -53,7 +53,7 @@ public static void reset() throws IOException {
    * down, which will cause ResourceMonitor restriction.
    */
   protected Set<Integer> ignored() {
-    return Set.of(29, 30);
+    return Set.of(29);
   }
 
   @Test

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

@@ -0,0 +1,107 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.monitor;
+
+import com.sun.management.GarbageCollectionNotificationInfo;
+import java.lang.management.GarbageCollectorMXBean;
+import java.lang.management.ManagementFactory;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
+import javax.management.Notification;
+import javax.management.NotificationEmitter;
+import javax.management.NotificationListener;
+import javax.management.openmbean.CompositeData;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+/** Get memory usage from GC notification listener, which is used in Calcite engine. */
+public class GCedMemoryUsage implements MemoryUsage {
+  private static final Logger LOG = LogManager.getLogger();
+
+  private GCedMemoryUsage() {
+    registerGCListener();
+  }
+
+  // Lazy initialize the instance to avoid register GCListener in v2.
+  private static class Holder {
+    static final MemoryUsage INSTANCE = new GCedMemoryUsage();
+  }
+
+  /**
+   * Get the singleton instance of GCedMemoryUsage.
+   *
+   * @return GCedMemoryUsage instance
+   */
+  public static MemoryUsage getInstance() {
+    return Holder.INSTANCE;
+  }
+
+  private final AtomicLong usage = new AtomicLong(-1);
+
+  @Override
+  public long usage() {
+    return usage.get();
+  }
+
+  @Override
+  public void setUsage(long value) {
+    usage.set(value);
+  }
+
+  private void registerGCListener() {
+    List<GarbageCollectorMXBean> gcBeans = ManagementFactory.getGarbageCollectorMXBeans();
+    for (GarbageCollectorMXBean gcBean : gcBeans) {
+      if (gcBean instanceof NotificationEmitter) {
+        NotificationEmitter emitter = (NotificationEmitter) gcBean;
+        emitter.addNotificationListener(new FullGCListener(), null, null);
+      }
+    }
+  }
+
+  private static class FullGCListener implements NotificationListener {
+    @Override
+    public void handleNotification(Notification notification, Object handback) {
+      if (notification
+          .getType()
+          .equals(GarbageCollectionNotificationInfo.GARBAGE_COLLECTION_NOTIFICATION)) {
+        CompositeData cd = (CompositeData) notification.getUserData();
+        GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from(cd);
+
+        String gcAction = info.getGcAction();
+        LOG.info("gcAction: {}", gcAction);
+        if (gcAction.contains("major")
+            || gcAction.contains("concurrent")
+            || gcAction.contains("old")
+            || gcAction.contains("full")) {
+          Map<String, java.lang.management.MemoryUsage> memoryUsageAfterGc =
+              info.getGcInfo().getMemoryUsageAfterGc();
+          String possibleOldGenKey = findOldGenKey(memoryUsageAfterGc.keySet());
+          if (possibleOldGenKey != null) {
+            java.lang.management.MemoryUsage oldGenUsage =
+                memoryUsageAfterGc.get(possibleOldGenKey);
+            getInstance().setUsage(oldGenUsage.getUsed());
+            LOG.info("Full GC detected, memory usage after GC is {} bytes.", getInstance().usage());
+          }
+        }
+      }
+    }
+
+    private String findOldGenKey(Set<String> memoryPoolNames) {
+      LOG.info("Memory pool names: {}", memoryPoolNames);
+      List<String> possibleOldGenKeys =
+          Arrays.asList("G1 Old Gen", "PS Old Gen", "CMS Old Gen", "Tenured Gen", "Old Gen");
+      for (String key : possibleOldGenKeys) {
+        if (memoryPoolNames.contains(key)) {
+          return key;
+        }
+      }
+      return null;
+    }
+  }
+}

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

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.monitor;
+
+public interface MemoryUsage {
+  long usage();
+
+  void setUsage(long usage);
+}

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

@@ -9,16 +9,20 @@
 import java.util.concurrent.ThreadLocalRandom;
 import lombok.NoArgsConstructor;
 import lombok.extern.log4j.Log4j2;
+import org.opensearch.sql.common.setting.Settings;
 
 /** OpenSearch Memory Monitor. */
 @Log4j2
 public class OpenSearchMemoryHealthy {
   private final RandomFail randomFail;
   private final MemoryUsage memoryUsage;
 
-  public OpenSearchMemoryHealthy() {
+  public OpenSearchMemoryHealthy(Settings settings) {
     randomFail = new RandomFail();
-    memoryUsage = new MemoryUsage();
+    memoryUsage =
+        isCalciteEnabled(settings)
+            ? GCedMemoryUsage.getInstance()
+            : RuntimeMemoryUsage.getInstance();
   }
 
   @VisibleForTesting
@@ -27,6 +31,14 @@ public OpenSearchMemoryHealthy(RandomFail randomFail, MemoryUsage memoryUsage) {
     this.memoryUsage = memoryUsage;
   }
 
+  private boolean isCalciteEnabled(Settings settings) {
+    if (settings != null) {
+      return settings.getSettingValue(Settings.Key.CALCITE_ENGINE_ENABLED);
+    } else {
+      return false;
+    }
+  }
+
   /** Is Memory Healthy. Calculate based on the current heap memory usage. */
   public boolean isMemoryHealthy(long limitBytes) {
     final long memoryUsage = this.memoryUsage.usage();
@@ -50,14 +62,6 @@ public boolean shouldFail() {
     }
   }
 
-  static class MemoryUsage {
-    public long usage() {
-      final long freeMemory = Runtime.getRuntime().freeMemory();
-      final long totalMemory = Runtime.getRuntime().totalMemory();
-      return totalMemory - freeMemory;
-    }
-  }
-
   @NoArgsConstructor
   public static class MemoryUsageExceedFastFailureException extends RuntimeException {}
 

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.monitor;
+
+/** Get memory usage from runtime, which is used in v2. */
+public class RuntimeMemoryUsage implements MemoryUsage {
+  private RuntimeMemoryUsage() {}
+
+  private static class Holder {
+    static final MemoryUsage INSTANCE = new RuntimeMemoryUsage();
+  }
+
+  public static MemoryUsage getInstance() {
+    return Holder.INSTANCE;
+  }
+
+  @Override
+  public long usage() {
+    final long freeMemory = Runtime.getRuntime().freeMemory();
+    final long totalMemory = Runtime.getRuntime().totalMemory();
+    return totalMemory - freeMemory;
+  }
+
+  @Override
+  public void setUsage(long usage) {
+    throw new UnsupportedOperationException("Cannot set usage in RuntimeMemoryUsage");
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java

@@ -260,7 +260,7 @@ public OpenSearchRequestBuilder createRequestBuilder() {
   }
 
   public OpenSearchResourceMonitor createOpenSearchResourceMonitor() {
-    return new OpenSearchResourceMonitor(getSettings(), new OpenSearchMemoryHealthy());
+    return new OpenSearchResourceMonitor(getSettings(), new OpenSearchMemoryHealthy(settings));
   }
 
   public OpenSearchRequest buildRequest(OpenSearchRequestBuilder requestBuilder) {

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java

@@ -28,7 +28,7 @@ public class OpenSearchStorageEngine implements StorageEngine {
   @Override
   public Table getTable(DataSourceSchemaName dataSourceSchemaName, String name) {
     if (isSystemIndex(name)) {
-      return new OpenSearchSystemIndex(client, name);
+      return new OpenSearchSystemIndex(client, settings, name);
     } else {
       return new OpenSearchIndex(client, settings, name);
     }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java

@@ -120,13 +120,19 @@ public boolean moveNext() {
 
   @Override
   public void reset() {
-    iterator = Collections.emptyIterator();
+    OpenSearchResponse response = client.search(request);
+    if (!response.isEmpty()) {
+      iterator = response.iterator();
+    } else {
+      iterator = Collections.emptyIterator();
+    }
     queryCount = 0;
   }
 
   @Override
   public void close() {
-    reset();
+    iterator = Collections.emptyIterator();
+    queryCount = 0;
     client.cleanup(request);
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/system/CalciteEnumerableSystemIndexScan.java

@@ -69,7 +69,9 @@ public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
       @Override
       public Enumerator<Object> enumerator() {
         return new OpenSearchSystemIndexEnumerator(
-            getFieldPath(), sysIndex.getSystemIndexBundle().getRight());
+            getFieldPath(),
+            sysIndex.getSystemIndexBundle().getRight(),
+            sysIndex.createOpenSearchResourceMonitor());
       }
     };
   }