Skip to content

Commit 364ff1f

Browse files
JackieTien97JackieTien97
authored
Clean up query execution rejected by IoTDBLocal read-only check (#19240)
1 parent 9f9479c commit 364ff1f

2 files changed

Lines changed: 157 additions & 1 deletion

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -347,7 +347,14 @@ private ExecutionResult execution(
347347
queryContext.setDebug(debug);
348348
IQueryExecution execution = iQueryExecutionFactory.apply(queryContext, startTime);
349349
if (readOnlyInternalQuery) {
350-
InternalQueryExecutor.validateReadOnlyQuery(execution);
350+
try {
351+
InternalQueryExecutor.validateReadOnlyQuery(execution);
352+
} catch (RuntimeException e) {
353+
// The execution is neither registered nor started, so nothing else would release what
354+
// its analysis acquired, such as the UDF class loaders pinned in its query scope.
355+
execution.stopAndCleanup(e);
356+
throw e;
357+
}
351358
}
352359
if (execution.isQuery()) {
353360
queryExecutionMap.put(queryId, execution);

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/CoordinatorTest.java‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.db.queryengine.plan;
21+
22+
import org.apache.iotdb.commons.exception.SemanticException;
23+
import org.apache.iotdb.commons.queryengine.common.SessionInfo;
24+
import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
25+
import org.apache.iotdb.db.conf.IoTDBDescriptor;
26+
import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
27+
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
28+
import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis;
29+
import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
30+
import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
31+
import org.apache.iotdb.db.queryengine.plan.execution.QueryExecution;
32+
import org.apache.iotdb.db.queryengine.plan.planner.IPlanner;
33+
34+
import org.junit.AfterClass;
35+
import org.junit.BeforeClass;
36+
import org.junit.Test;
37+
38+
import java.lang.reflect.Field;
39+
import java.lang.reflect.InvocationTargetException;
40+
import java.lang.reflect.Method;
41+
import java.time.ZoneId;
42+
import java.util.concurrent.ExecutorService;
43+
import java.util.concurrent.Executors;
44+
import java.util.function.BiFunction;
45+
46+
import static org.junit.Assert.assertEquals;
47+
import static org.junit.Assert.assertNotNull;
48+
import static org.junit.Assert.assertNull;
49+
import static org.junit.Assert.assertThrows;
50+
import static org.mockito.ArgumentMatchers.any;
51+
import static org.mockito.Mockito.mock;
52+
import static org.mockito.Mockito.when;
53+
54+
public class CoordinatorTest {
55+
56+
private static final long QUERY_ID = 10_000L;
57+
58+
private static int previousDataNodeId;
59+
private static ExecutorService stateMachineExecutor;
60+
61+
@BeforeClass
62+
public static void setUp() throws Exception {
63+
// the coordinator builds its query id generator from the data node id of this node
64+
previousDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
65+
IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0);
66+
UDFClassLoaderManager.setupAndGetInstance();
67+
stateMachineExecutor = Executors.newSingleThreadExecutor();
68+
}
69+
70+
@AfterClass
71+
public static void tearDown() {
72+
stateMachineExecutor.shutdownNow();
73+
IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId);
74+
}
75+
76+
@Test
77+
public void testRejectedReadOnlyInternalQueryClosesUDFQueryScope() throws Exception {
78+
// Only the analysis tells that an IoTDBLocal statement is not a query. By then it may have
79+
// pinned the class loaders of the UDFs it uses in the query scope of the execution.
80+
IPlanner planner = mock(IPlanner.class);
81+
when(planner.analyze(any(MPPQueryContext.class)))
82+
.thenAnswer(
83+
invocation -> {
84+
MPPQueryContext queryContext = invocation.getArgument(0);
85+
queryContext.setQueryType(QueryType.WRITE);
86+
return mock(IAnalysis.class);
87+
});
88+
QueryExecution[] createdExecution = new QueryExecution[1];
89+
90+
SemanticException thrown =
91+
assertThrows(
92+
SemanticException.class,
93+
() ->
94+
executeReadOnlyInternalQuery(
95+
(queryContext, startTime) -> {
96+
createdExecution[0] =
97+
new QueryExecution(planner, queryContext, stateMachineExecutor);
98+
return createdExecution[0];
99+
}));
100+
101+
assertEquals(
102+
DataNodeQueryMessages
103+
.EXCEPTION_ONLY_QUERY_IS_SUPPORTED_FOR_IOTDBLOCAL_QUERY_INTERFACE_09195570,
104+
thrown.getMessage());
105+
// the rejected execution is never registered, so no later cleanup can find it
106+
assertNull(Coordinator.getInstance().getQueryExecution(QUERY_ID));
107+
assertUDFQueryScopeClosed(createdExecution[0]);
108+
}
109+
110+
private static void executeReadOnlyInternalQuery(
111+
BiFunction<MPPQueryContext, Long, IQueryExecution> queryExecutionFactory) throws Throwable {
112+
Method execution =
113+
Coordinator.class.getDeclaredMethod(
114+
"execution",
115+
long.class,
116+
SessionInfo.class,
117+
String.class,
118+
boolean.class,
119+
boolean.class,
120+
boolean.class,
121+
BiFunction.class);
122+
execution.setAccessible(true);
123+
try {
124+
execution.invoke(
125+
Coordinator.getInstance(),
126+
QUERY_ID,
127+
new SessionInfo(0, "test", ZoneId.systemDefault()),
128+
"INSERT INTO t2(time, s) SELECT time, f(s) FROM t1",
129+
false,
130+
false,
131+
true,
132+
queryExecutionFactory);
133+
} catch (InvocationTargetException e) {
134+
throw e.getCause();
135+
}
136+
}
137+
138+
private static void assertUDFQueryScopeClosed(QueryExecution execution) throws Exception {
139+
assertNotNull(execution);
140+
Field queryScopeField = QueryExecution.class.getDeclaredField("udfQueryScope");
141+
queryScopeField.setAccessible(true);
142+
UDFClassLoaderManager.QueryScope queryScope =
143+
(UDFClassLoaderManager.QueryScope) queryScopeField.get(execution);
144+
assertNotNull(queryScope);
145+
assertThrows(
146+
IllegalStateException.class,
147+
() -> UDFClassLoaderManager.getInstance().withQueryScope(queryScope, () -> null));
148+
}
149+
}

0 commit comments

Comments
 (0)