Skip to content

Protect UDF classloader lifetime with query scopes #19107

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shuwenwei wants to merge 10 commits into
base: master
Choose a base branch
from
Open

Protect UDF classloader lifetime with query scopes #19107

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
c2b9625
Fix table model UDF classloader lifecycle
shuwenwei Nov 3, 2026
aa62c2f
Fix table UDF lifecycle cleanup guards
shuwenwei Nov 3, 2026
104cde7
Fix UDF classloader lifecycle with query scopes
shuwenwei Nov 5, 2026
145c528
Add UDF query scope lifecycle regression coverage
shuwenwei Nov 5, 2026
358408c
spotless
shuwenwei Nov 5, 2026
437d6e9
Fix checked exceptions in UDF scope tests
shuwenwei Nov 5, 2026
a4b1455
Fix UDF metadata function type classification
shuwenwei Nov 5, 2026
b27cace
spotless
shuwenwei Nov 6, 2026
3ccfa98
Fix UDF lifecycle test validation
shuwenwei Nov 6, 2026
3f6dd87
Fix generated UDF test compilation
shuwenwei Nov 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions .../src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
import static org.apache.iotdb.commons.conf.IoTDBConstant.FUNCTION_TYPE_USER_DEFINED_AGG_FUNC;
import static org.apache.iotdb.commons.conf.IoTDBConstant.FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

Expand Down Expand Up @@ -135,6 +136,24 @@ public void testCreateShowDropScalarFunction() {
}
}

@Test
public void testInformationSchemaFunctionsShowsExternalScalarType() {
try (Connection connection = EnvFactory.getEnv().getTableConnection();
Statement statement = connection.createStatement()) {
statement.execute(
"create function udsf as 'org.apache.iotdb.db.query.udf.example.relational.ContainNull'");
try (ResultSet resultSet =
statement.executeQuery(
"select function_type from information_schema.functions where function_table = 'UDSF'")) {
assertTrue(resultSet.next());
assertEquals(FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC, resultSet.getString(1));
assertFalse(resultSet.next());
}
} catch (SQLException throwable) {
fail(throwable.getMessage());
}
}

@Test
public void testCreateShowDropAggregateFunction() {
try (Connection connection = EnvFactory.getEnv().getTableConnection();
Expand Down
26 changes: 25 additions & 1 deletion ...e/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,22 +22,46 @@
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
import org.apache.iotdb.commons.request.IConsensusRequest;
import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
import org.apache.iotdb.consensus.IStateMachine;
import org.apache.iotdb.consensus.common.request.ByteBufferConsensusRequest;
import org.apache.iotdb.db.i18n.DataNodeMiscMessages;
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.ByteBuffer;

public abstract class BaseStateMachine implements IStateMachine, IStateMachine.EventApi {

private static final Logger logger = LoggerFactory.getLogger(BaseStateMachine.class);

protected FragmentInstance getFragmentInstance(IConsensusRequest request) {
FragmentInstance instance;
if (request instanceof ByteBufferConsensusRequest) {
instance = FragmentInstance.deserializeFrom(request.serializeToByteBuffer());
ByteBuffer payload = request.serializeToByteBuffer();
FragmentInstanceId fragmentInstanceId = FragmentInstanceId.deserialize(payload);
UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
UDFClassLoaderManager.QueryScope queryScope =
classLoaderManager == null
? null
: classLoaderManager.openQueryScope(fragmentInstanceId.getFullId());
try {
instance =
queryScope == null
? FragmentInstance.deserializeFrom(fragmentInstanceId, payload)
: classLoaderManager.withQueryScope(
queryScope,
() -> FragmentInstance.deserializeFrom(fragmentInstanceId, payload));
instance.attachUDFQueryScope(queryScope);
} catch (RuntimeException | Error e) {

Check warning on line 59 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Catch Exception instead of Error. 


See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGqwPkLjlQ3I3p-fvhS&open=AaGqwPkLjlQ3I3p-fvhS&pullRequest=19107
if (classLoaderManager != null && queryScope != null) {
classLoaderManager.finalizeUDFQuery(queryScope);
}
throw e;
}
} else if (request instanceof FragmentInstance) {
instance = (FragmentInstance) request;
} else {
Expand Down
6 changes: 5 additions & 1 deletion ...in/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,11 @@ public DataSet read(IConsensusRequest request) {
logger.error(DataNodeMiscMessages.GET_FRAGMENT_INSTANCE_FAILED, e);
return null;
}
return QUERY_INSTANCE_MANAGER.execDataQueryFragmentInstance(fragmentInstance, region);
try {
return QUERY_INSTANCE_MANAGER.execDataQueryFragmentInstance(fragmentInstance, region);
} finally {
fragmentInstance.closeUDFQueryScopeIfPresent();
}
}
}

Expand Down
6 changes: 5 additions & 1 deletion ...ava/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,10 @@ public DataSet read(final IConsensusRequest request) {
.PIPE_LOG_SCHEMAREGIONSTATEMACHINE_EXECUTE_READ_PLAN_FRAGMENTINSTANCE_F85A001F,
schemaRegion.getSchemaRegionId(),
fragmentInstance.getId());
return QUERY_INSTANCE_MANAGER.execSchemaQueryFragmentInstance(fragmentInstance, schemaRegion);
try {
return QUERY_INSTANCE_MANAGER.execSchemaQueryFragmentInstance(fragmentInstance, schemaRegion);
} finally {
fragmentInstance.closeUDFQueryScopeIfPresent();
}
}
}
73 changes: 52 additions & 21 deletions ...rc/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@
import org.apache.iotdb.commons.subscription.meta.topic.TopicMeta;
import org.apache.iotdb.commons.trigger.TriggerInformation;
import org.apache.iotdb.commons.udf.UDFInformation;
import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
import org.apache.iotdb.commons.utils.KillPoint.DataNodeKillPoints;
import org.apache.iotdb.commons.utils.KillPoint.KillPoint;
import org.apache.iotdb.commons.utils.PathUtils;
Expand Down Expand Up @@ -528,37 +529,67 @@ public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanc

// We deserialize here instead of the underlying state machine because parallelism is possible
// here but not at the underlying state machine
final FragmentInstance fragmentInstance;
final FragmentInstanceId fragmentInstanceId;
try {
fragmentInstance = FragmentInstance.deserializeFrom(req.fragmentInstance.body);
fragmentInstanceId = FragmentInstanceId.deserialize(req.fragmentInstance.body);
} catch (final Exception e) {
LOGGER.warn(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED, e);
TSendFragmentInstanceResp resp = new TSendFragmentInstanceResp(false);
resp.setMessage(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED + e.getMessage());
return resp;
}

RegionExecutionResult executionResult = null;
if (Objects.nonNull(groupId)) {
// For schema query, there may be a "write" before "read"
// the result is not null iff the "write" has failed
executionResult =
new TableSchemaQueryWriteVisitor().processFragment(fragmentInstance, groupId);
}
if (Objects.isNull(executionResult)) {
final RegionReadExecutor executor = new RegionReadExecutor();
executionResult =
groupId == null
? executor.execute(fragmentInstance)
: executor.execute(groupId, fragmentInstance);
final UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
final UDFClassLoaderManager.QueryScope queryScope =
classLoaderManager == null
? null
: classLoaderManager.openQueryScope(fragmentInstanceId.getFullId());
final FragmentInstance fragmentInstance;
try {
fragmentInstance =
queryScope == null
? FragmentInstance.deserializeFrom(fragmentInstanceId, req.fragmentInstance.body)
: classLoaderManager.withQueryScope(
queryScope,
() ->
FragmentInstance.deserializeFrom(
fragmentInstanceId, req.fragmentInstance.body));
fragmentInstance.attachUDFQueryScope(queryScope);
} catch (final Exception e) {
if (classLoaderManager != null && queryScope != null) {
classLoaderManager.finalizeUDFQuery(queryScope);
}
LOGGER.warn(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED, e);
TSendFragmentInstanceResp resp = new TSendFragmentInstanceResp(false);
resp.setMessage(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED + e.getMessage());
return resp;
}

final TSendFragmentInstanceResp resp = new TSendFragmentInstanceResp();
resp.setAccepted(executionResult.isAccepted());
resp.setMessage(executionResult.getMessage());
resp.setNeedRetry(executionResult.isReadNeedRetry());
resp.setStatus(executionResult.getStatus());
return resp;
try {
RegionExecutionResult executionResult = null;
if (Objects.nonNull(groupId)) {
// For schema query, there may be a "write" before "read"
// the result is not null iff the "write" has failed
executionResult =
new TableSchemaQueryWriteVisitor().processFragment(fragmentInstance, groupId);
}
if (Objects.isNull(executionResult)) {
final RegionReadExecutor executor = new RegionReadExecutor();
executionResult =
groupId == null
? executor.execute(fragmentInstance)
: executor.execute(groupId, fragmentInstance);
}

final TSendFragmentInstanceResp resp = new TSendFragmentInstanceResp();
resp.setAccepted(executionResult.isAccepted());
resp.setMessage(executionResult.getMessage());
resp.setNeedRetry(executionResult.isReadNeedRetry());
resp.setStatus(executionResult.getStatus());
return resp;
} finally {
fragmentInstance.closeUDFQueryScopeIfPresent();
}
}

@Override
Expand Down
26 changes: 26 additions & 0 deletions ...main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import org.apache.iotdb.commons.path.IFullPath;
import org.apache.iotdb.commons.queryengine.common.SessionInfo;
import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils;
import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
import org.apache.iotdb.commons.utils.TestOnly;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
Expand Down Expand Up @@ -116,6 +117,7 @@ public class FragmentInstanceContext extends QueryContext {
private static final long LONG_WAIT_DURATION = 5_000_000_000L;
private static final int MODS_MEMORY_ESTIMATE_READ_INTERVAL = 10_000;
private final FragmentInstanceId id;
private UDFClassLoaderManager.QueryScope udfQueryScope;

private final FragmentInstanceStateMachine stateMachine;

Expand Down Expand Up @@ -1348,6 +1350,7 @@ private void releaseTVListOwnedByQuery() {
* be decreased.
*/
public synchronized void releaseResource() {
finalizeUDFQuery();
try {
getQueryChunkReaders().close();
} catch (IOException e) {
Expand Down Expand Up @@ -1493,6 +1496,29 @@ private void releaseDataNodeQueryContext() {
}
}

synchronized void adoptUDFQueryScope(UDFClassLoaderManager.QueryScope transferredScope) {
if (udfQueryScope == null) {
udfQueryScope = transferredScope;
}
}

public synchronized UDFClassLoaderManager.QueryScope getOrCreateUDFQueryScope() {
UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
if (udfQueryScope == null && classLoaderManager != null) {
udfQueryScope = classLoaderManager.openQueryScope(id.getFullId());
}
return udfQueryScope;
}

synchronized void finalizeUDFQuery() {
UDFClassLoaderManager.QueryScope queryScope = udfQueryScope;
udfQueryScope = null;
UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
if (queryScope != null && classLoaderManager != null) {
classLoaderManager.finalizeUDFQuery(queryScope);
}
}

public void setMayHaveTmpFile(boolean mayHaveTmpFile) {
this.mayHaveTmpFile = mayHaveTmpFile;
}
Expand Down
42 changes: 24 additions & 18 deletions ...main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,16 +164,19 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
dataNodeQueryContexts[0] =
getOrCreateDataNodeQueryContext(
instanceId.getQueryId(), instance.getDataNodeFINum());
return createFragmentInstanceContext(
fragmentInstanceId,
stateMachine,
instance.getSessionInfo(),
dataRegion,
instance.getGlobalTimePredicate(),
dataNodeQueryContextMap,
DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()),
instance.isDebug(),
instance.isVerbose());
FragmentInstanceContext createdContext =
createFragmentInstanceContext(
fragmentInstanceId,
stateMachine,
instance.getSessionInfo(),
dataRegion,
instance.getGlobalTimePredicate(),
dataNodeQueryContextMap,
DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()),
instance.isDebug(),
instance.isVerbose());
createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope());
return createdContext;
});
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
context.setHighestPriority(instance.isHighestPriority());
Expand All @@ -185,7 +188,6 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
instance.getFragment().getTypeProvider(),
context,
dataNodeQueryContexts[0]);

List<IDriver> drivers = new ArrayList<>();
driverFactories.forEach(factory -> drivers.add(factory.createDriver()));
// For ShowQueries related instances, isHighestPriority == true
Expand Down Expand Up @@ -235,6 +237,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance(
// still being created. Capture the failure result from the local context before
// returning from computeIfAbsent instead of looking it up from the map later.
failedInstanceInfo.set(context.getInstanceInfo());
context.finalizeUDFQuery();
clearFIRelatedResources(instanceId);
return null;
}
Expand Down Expand Up @@ -318,12 +321,15 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance(
instanceId,
fragmentInstanceId -> {
contextCreated[0] = true;
return createFragmentInstanceContext(
fragmentInstanceId,
stateMachine,
instance.getSessionInfo(),
instance.isDebug(),
instance.isVerbose());
FragmentInstanceContext createdContext =
createFragmentInstanceContext(
fragmentInstanceId,
stateMachine,
instance.getSessionInfo(),
instance.isDebug(),
instance.isVerbose());
createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope());
return createdContext;
});
rejectIfRepeatedDispatch(contextCreated[0], instanceId);
context.setHighestPriority(instance.isHighestPriority());
Expand All @@ -335,7 +341,6 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance(
instance.getFragment().getTypeProvider(),
context,
schemaRegion);

List<IDriver> drivers = new ArrayList<>();
driverFactories.forEach(factory -> drivers.add(factory.createDriver()));
context.initializeNumOfDrivers(drivers.size());
Expand Down Expand Up @@ -370,6 +375,7 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance(
// See execDataQueryFragmentInstance for why this result must not be fetched from
// instanceContext after computeIfAbsent returns.
failedInstanceInfo.set(context.getInstanceInfo());
context.finalizeUDFQuery();
clearFIRelatedResources(instanceId);
return null;
}
Expand Down
37 changes: 14 additions & 23 deletions ...in/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@
import org.apache.iotdb.commons.queryengine.common.NodeRef;
import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.parameter.InputLocation;
import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService;
import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
import org.apache.iotdb.commons.utils.TestOnly;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
Expand Down Expand Up @@ -136,27 +134,20 @@ protected void initTransformers(
Map<String, List<InputLocation>> inputLocations,
Expression[] outputExpressions,
Map<NodeRef<Expression>, TSDataType> expressionTypes) {
UDFManagementService.getInstance().acquireLock();
try {
// This statement must be surrounded by the registration lock.
UDFClassLoaderManager.getInstance().initializeUDFQuery(udtfQueryId);
// UDF executors will be initialized at the same time
transformers =
new EvaluationDAGBuilder(
udtfQueryId,
inputLayer,
inputLocations,
outputExpressions,
expressionTypes,
udtfContext,
udfTransformerMemoryBudgetInMB + udfCollectorMemoryBudgetInMB)
.buildLayerMemoryAssigner()
.bindInputLayerColumnIndexWithExpression()
.buildResultColumnPointReaders()
.getOutputReaders();
} finally {
UDFManagementService.getInstance().releaseLock();
}
// UDF executors will be initialized at the same time
transformers =
new EvaluationDAGBuilder(
udtfQueryId,
inputLayer,
inputLocations,
outputExpressions,
expressionTypes,
udtfContext,
udfTransformerMemoryBudgetInMB + udfCollectorMemoryBudgetInMB)
.buildLayerMemoryAssigner()
.bindInputLayerColumnIndexWithExpression()
.buildResultColumnPointReaders()
.getOutputReaders();
}

protected YieldableState iterateAllColumnsToNextValid() throws Exception {
Expand Down
Loading
Loading