Skip to content

Commit 0398f68

Browse files
LantaoJinLantaoJin
authored
ResourceMonitor only checks the memory health by calculating the memory usage after GC of old gen in v3 (opensearch-project#3983)
* ResourceMonitor only checks the memory health by calculating the memory usage after GC of old gen Signed-off-by: Lantao Jin <ltjin@amazon.com> * address comments Signed-off-by: Lantao Jin <ltjin@amazon.com> * address comments Signed-off-by: Lantao Jin <ltjin@amazon.com> * Add memory usage fallback Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix IT Signed-off-by: Lantao Jin <ltjin@amazon.com> * Generate script code lazily to avoid unnecessary usage in Calcite Signed-off-by: Lantao Jin <ltjin@amazon.com> * typo Signed-off-by: Lantao Jin <ltjin@amazon.com> * Address the comments of 3929 followup Signed-off-by: Lantao Jin <ltjin@amazon.com> * Catch throwable instead of MemoryUsageException Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix unit test when the code in script query expression generated lazily Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 8923551 commit 0398f68

20 files changed

Lines changed: 371 additions & 81 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public class CalcitePlanContext {
3636
public final QueryType queryType;
3737
public final Integer querySizeLimit;
3838

39+
/** This thread local variable is only used to skip script encoding in script pushdown. */
3940
public static final ThreadLocal<Boolean> skipEncoding = ThreadLocal.withInitial(() -> false);
4041

4142
@Getter @Setter private boolean isResolvingJoinCondition = false;

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -111,11 +111,14 @@ public void executeWithCalcite(
111111
log.warn("Fallback to V2 query engine since got exception", t);
112112
executeWithLegacy(plan, queryType, listener, Optional.of(t));
113113
} else {
114-
if (t instanceof Error) {
114+
if (t instanceof Exception) {
115+
listener.onFailure((Exception) t);
116+
} else if (t instanceof VirtualMachineError) {
117+
// throw and fast fail the VM errors such as OOM (same with v2).
118+
throw t;
119+
} else {
115120
// Calcite may throw AssertError during query execution.
116121
listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
117-
} else {
118-
listener.onFailure((Exception) t);
119122
}
120123
}
121124
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ public static void reset() throws IOException {
5555
* addressed by: https://github.com/opensearch-project/sql/issues/3981
5656
*/
5757
protected Set<Integer> ignored() {
58-
return Set.of(29, 30);
58+
return Set.of(29);
5959
}
6060

6161
@Test

‎integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
99
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXPLAIN_API_ENDPOINT;
10-
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXTENDED_EXPLAIN_API_ENDPOINT;
1110
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.QUERY_API_ENDPOINT;
1211

1312
import com.google.common.io.Resources;
@@ -34,6 +33,8 @@
3433

3534
/** OpenSearch Rest integration test base for PPL testing. */
3635
public abstract class PPLIntegTestCase extends SQLIntegTestCase {
36+
private static final String EXTENDED_EXPLAIN_API_ENDPOINT =
37+
"/_plugins/_ppl/_explain?format=extended";
3738
private static final Logger LOG = LogManager.getLogger();
3839
@Rule public final RetryProcessor retryProcessor = new RetryProcessor();
3940

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.monitor;
7+
8+
import com.sun.management.GarbageCollectionNotificationInfo;
9+
import java.lang.management.GarbageCollectorMXBean;
10+
import java.lang.management.ManagementFactory;
11+
import java.util.List;
12+
import java.util.Locale;
13+
import java.util.Map;
14+
import java.util.concurrent.atomic.AtomicLong;
15+
import javax.management.Notification;
16+
import javax.management.NotificationEmitter;
17+
import javax.management.NotificationListener;
18+
import javax.management.openmbean.CompositeData;
19+
import org.apache.logging.log4j.LogManager;
20+
import org.apache.logging.log4j.Logger;
21+
22+
/** Get memory usage from GC notification listener, which is used in Calcite engine. */
23+
public class GCedMemoryUsage implements MemoryUsage {
24+
private static final Logger LOG = LogManager.getLogger();
25+
private static final List<String> OLD_GEN_GC_ACTION_KEYWORDS =
26+
List.of("major", "concurrent", "old", "full", "marksweep");
27+
28+
private GCedMemoryUsage() {
29+
registerGCListener();
30+
}
31+
32+
// Lazy initialize the instance to avoid register GCListener in v2.
33+
private static class Holder {
34+
static final MemoryUsage INSTANCE = new GCedMemoryUsage();
35+
}
36+
37+
/**
38+
* Get the singleton instance of GCedMemoryUsage.
39+
*
40+
* @return GCedMemoryUsage instance
41+
*/
42+
public static MemoryUsage getInstance() {
43+
return Holder.INSTANCE;
44+
}
45+
46+
private final AtomicLong usage = new AtomicLong(-1);
47+
48+
@Override
49+
public long usage() {
50+
return usage.get();
51+
}
52+
53+
@Override
54+
public void setUsage(long value) {
55+
usage.set(value);
56+
}
57+
58+
private void registerGCListener() {
59+
boolean registered = false;
60+
List<GarbageCollectorMXBean> gcBeans = ManagementFactory.getGarbageCollectorMXBeans();
61+
for (GarbageCollectorMXBean gcBean : gcBeans) {
62+
if (gcBean instanceof NotificationEmitter && isOldGenGc(gcBean.getName())) {
63+
LOG.info("{} listener registered for memory usage monitor.", gcBean.getName());
64+
registered = true;
65+
NotificationEmitter emitter = (NotificationEmitter) gcBean;
66+
emitter.addNotificationListener(
67+
new OldGenGCListener(),
68+
notification -> {
69+
if (!notification
70+
.getType()
71+
.equals(GarbageCollectionNotificationInfo.GARBAGE_COLLECTION_NOTIFICATION)) {
72+
return false;
73+
}
74+
CompositeData cd = (CompositeData) notification.getUserData();
75+
GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from(cd);
76+
return isOldGenGc(info.getGcAction());
77+
},
78+
null);
79+
}
80+
}
81+
if (!registered) {
82+
// fallback to RuntimeMemoryUsage
83+
LOG.info("No old gen GC listener registered, fallback to RuntimeMemoryUsage");
84+
throw new OpenSearchMemoryHealthy.MemoryUsageException();
85+
}
86+
}
87+
88+
private boolean isOldGenGc(String gcKeyword) {
89+
String keyword = gcKeyword.toLowerCase(Locale.ROOT);
90+
return OLD_GEN_GC_ACTION_KEYWORDS.stream().anyMatch(keyword::contains);
91+
}
92+
93+
private static class OldGenGCListener implements NotificationListener {
94+
@Override
95+
public void handleNotification(Notification notification, Object handback) {
96+
CompositeData cd = (CompositeData) notification.getUserData();
97+
GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from(cd);
98+
Map<String, java.lang.management.MemoryUsage> memoryUsageAfterGc =
99+
info.getGcInfo().getMemoryUsageAfterGc();
100+
// Skip Metaspace and CodeHeap spaces which the GC scope is out of stack GC.
101+
long totalStackUsed =
102+
memoryUsageAfterGc.entrySet().stream()
103+
.filter(
104+
entry ->
105+
!entry.getKey().equals("Metaspace")
106+
&& !entry.getKey().equals("Compressed Class Space")
107+
&& !entry.getKey().startsWith("CodeHeap"))
108+
.mapToLong(entry -> entry.getValue().getUsed())
109+
.sum();
110+
getInstance().setUsage(totalStackUsed);
111+
if (LOG.isDebugEnabled()) {
112+
LOG.debug("Old Gen GC detected, memory usage after GC is {} bytes.", totalStackUsed);
113+
}
114+
}
115+
}
116+
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.monitor;
7+
8+
/** Memory usage interface. It is used to get the memory usage of the VM. */
9+
public interface MemoryUsage {
10+
long usage();
11+
12+
void setUsage(long usage);
13+
}

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

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,17 @@
99
import java.util.concurrent.ThreadLocalRandom;
1010
import lombok.NoArgsConstructor;
1111
import lombok.extern.log4j.Log4j2;
12+
import org.opensearch.sql.common.setting.Settings;
1213

1314
/** OpenSearch Memory Monitor. */
1415
@Log4j2
1516
public class OpenSearchMemoryHealthy {
1617
private final RandomFail randomFail;
1718
private final MemoryUsage memoryUsage;
1819

19-
public OpenSearchMemoryHealthy() {
20+
public OpenSearchMemoryHealthy(Settings settings) {
2021
randomFail = new RandomFail();
21-
memoryUsage = new MemoryUsage();
22+
memoryUsage = buildMemoryUsage(settings);
2223
}
2324

2425
@VisibleForTesting
@@ -27,6 +28,24 @@ public OpenSearchMemoryHealthy(RandomFail randomFail, MemoryUsage memoryUsage) {
2728
this.memoryUsage = memoryUsage;
2829
}
2930

31+
private MemoryUsage buildMemoryUsage(Settings settings) {
32+
try {
33+
return isCalciteEnabled(settings)
34+
? GCedMemoryUsage.getInstance()
35+
: RuntimeMemoryUsage.getInstance();
36+
} catch (Throwable e) {
37+
return RuntimeMemoryUsage.getInstance();
38+
}
39+
}
40+
41+
private boolean isCalciteEnabled(Settings settings) {
42+
if (settings != null) {
43+
return settings.getSettingValue(Settings.Key.CALCITE_ENGINE_ENABLED);
44+
} else {
45+
return false;
46+
}
47+
}
48+
3049
/** Is Memory Healthy. Calculate based on the current heap memory usage. */
3150
public boolean isMemoryHealthy(long limitBytes) {
3251
final long memoryUsage = this.memoryUsage.usage();
@@ -50,17 +69,12 @@ public boolean shouldFail() {
5069
}
5170
}
5271

53-
static class MemoryUsage {
54-
public long usage() {
55-
final long freeMemory = Runtime.getRuntime().freeMemory();
56-
final long totalMemory = Runtime.getRuntime().totalMemory();
57-
return totalMemory - freeMemory;
58-
}
59-
}
72+
@NoArgsConstructor
73+
public static class MemoryUsageExceedFastFailureException extends MemoryUsageException {}
6074

6175
@NoArgsConstructor
62-
public static class MemoryUsageExceedFastFailureException extends RuntimeException {}
76+
public static class MemoryUsageExceedException extends MemoryUsageException {}
6377

6478
@NoArgsConstructor
65-
public static class MemoryUsageExceedException extends RuntimeException {}
79+
public static class MemoryUsageException extends RuntimeException {}
6680
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.monitor;
7+
8+
/** Get memory usage from runtime, which is used in v2. */
9+
public class RuntimeMemoryUsage implements MemoryUsage {
10+
private RuntimeMemoryUsage() {}
11+
12+
private static class Holder {
13+
static final MemoryUsage INSTANCE = new RuntimeMemoryUsage();
14+
}
15+
16+
public static MemoryUsage getInstance() {
17+
return Holder.INSTANCE;
18+
}
19+
20+
@Override
21+
public long usage() {
22+
final long freeMemory = Runtime.getRuntime().freeMemory();
23+
final long totalMemory = Runtime.getRuntime().totalMemory();
24+
return totalMemory - freeMemory;
25+
}
26+
27+
@Override
28+
public void setUsage(long usage) {
29+
throw new UnsupportedOperationException("Cannot set usage in RuntimeMemoryUsage");
30+
}
31+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java‎

Lines changed: 40 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import java.util.Locale;
5656
import java.util.Map;
5757
import java.util.Set;
58+
import java.util.function.Supplier;
5859
import lombok.Getter;
5960
import lombok.Setter;
6061
import org.apache.calcite.DataContext.Variable;
@@ -175,6 +176,15 @@ public static QueryBuilder analyze(
175176
return analyzeExpression(expression, schema, fieldTypes, rowType, cluster).builder();
176177
}
177178

179+
/**
180+
* Analyzes the expression and returns a {@link QueryExpression}.
181+
*
182+
* @param expression expression to analyze
183+
* @param schema current schema of scan operator
184+
* @param fieldTypes mapping of OpenSearch field name to ExprType, nested fields are flattened
185+
* @return search query which can be used to query OS cluster
186+
* @throws ExpressionNotAnalyzableException when expression can't processed by this analyzer
187+
*/
178188
public static QueryExpression analyzeExpression(
179189
RexNode expression,
180190
List<String> schema,
@@ -183,10 +193,28 @@ public static QueryExpression analyzeExpression(
183193
RelOptCluster cluster)
184194
throws ExpressionNotAnalyzableException {
185195
requireNonNull(expression, "expression");
196+
return analyzeExpression(
197+
expression,
198+
schema,
199+
fieldTypes,
200+
rowType,
201+
cluster,
202+
new Visitor(schema, fieldTypes, rowType, cluster));
203+
}
204+
205+
/** For test only, passing a customer Visitor */
206+
public static QueryExpression analyzeExpression(
207+
RexNode expression,
208+
List<String> schema,
209+
Map<String, ExprType> fieldTypes,
210+
RelDataType rowType,
211+
RelOptCluster cluster,
212+
Visitor visitor)
213+
throws ExpressionNotAnalyzableException {
214+
requireNonNull(expression, "expression");
186215
try {
187216
// visits expression tree
188-
QueryExpression queryExpression =
189-
(QueryExpression) expression.accept(new Visitor(schema, fieldTypes, rowType, cluster));
217+
QueryExpression queryExpression = (QueryExpression) expression.accept(visitor);
190218
return queryExpression;
191219
} catch (Throwable e) {
192220
if (e instanceof UnsupportedScriptException) {
@@ -201,14 +229,14 @@ public static QueryExpression analyzeExpression(
201229
}
202230

203231
/** Traverses {@link RexNode} tree and builds OpenSearch query. */
204-
private static class Visitor extends RexVisitorImpl<Expression> {
232+
static class Visitor extends RexVisitorImpl<Expression> {
205233

206234
List<String> schema;
207235
Map<String, ExprType> fieldTypes;
208236
RelDataType rowType;
209237
RelOptCluster cluster;
210238

211-
private Visitor(
239+
Visitor(
212240
List<String> schema,
213241
Map<String, ExprType> fieldTypes,
214242
RelDataType rowType,
@@ -736,7 +764,7 @@ private QueryExpression andOr(RexCall call) {
736764
}
737765
}
738766

739-
private Expression tryAnalyzeOperand(RexNode node) {
767+
public Expression tryAnalyzeOperand(RexNode node) {
740768
try {
741769
Expression expr = node.accept(this);
742770
if (expr instanceof NamedFieldExpression) {
@@ -1411,18 +1439,20 @@ private static String ipValueForPushDown(String value) {
14111439
}
14121440

14131441
public static class ScriptQueryExpression extends QueryExpression {
1414-
private final String code;
14151442
private RexNode analyzedNode;
1443+
// use lambda to generate code lazily to avoid store generated code
1444+
private final Supplier<String> codeGenerator;
14161445

14171446
public ScriptQueryExpression(
14181447
RexNode rexNode,
14191448
RelDataType rowType,
14201449
Map<String, ExprType> fieldTypes,
14211450
RelOptCluster cluster) {
14221451
RelJsonSerializer serializer = new RelJsonSerializer(cluster);
1423-
this.code =
1424-
SerializationWrapper.wrapWithLangType(
1425-
ScriptEngineType.CALCITE, serializer.serialize(rexNode, rowType, fieldTypes));
1452+
this.codeGenerator =
1453+
() ->
1454+
SerializationWrapper.wrapWithLangType(
1455+
ScriptEngineType.CALCITE, serializer.serialize(rexNode, rowType, fieldTypes));
14261456
}
14271457

14281458
@Override
@@ -1439,7 +1469,7 @@ public Script getScript() {
14391469
return new Script(
14401470
DEFAULT_SCRIPT_TYPE,
14411471
COMPOUNDED_LANG_NAME,
1442-
code,
1472+
codeGenerator.get(),
14431473
Collections.emptyMap(),
14441474
Map.of(Variable.UTC_TIMESTAMP.camelName, currentTime));
14451475
}

0 commit comments

Comments
 (0)