From c2b96252f871c0ecaac3a6c57923924cf96fa05c Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Tue, 3 Nov 2026 15:28:10 +0800 Subject: [PATCH 01/10] Fix table model UDF classloader lifecycle --- .../fragment/FragmentInstanceContext.java | 20 ++++++++++++ .../fragment/FragmentInstanceManager.java | 31 +++++++++++++++---- .../operator/process/TransformOperator.java | 4 +-- .../plan/execution/QueryExecution.java | 31 ++++++++++++++++++- .../plan/udf/UDFManagementService.java | 30 +++++++++++------- 5 files changed, 96 insertions(+), 20 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index a060bd5b91b51..84775f51132bf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -33,7 +33,9 @@ import org.apache.iotdb.commons.path.AlignedFullPath; import org.apache.iotdb.commons.path.IFullPath; import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; 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; @@ -116,6 +118,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 boolean tableModelUDFQueryInitialized; private final FragmentInstanceStateMachine stateMachine; @@ -1348,6 +1351,7 @@ private void releaseTVListOwnedByQuery() { * be decreased. */ public synchronized void releaseResource() { + finalizeTableModelUDFQuery(); try { getQueryChunkReaders().close(); } catch (IOException e) { @@ -1493,6 +1497,22 @@ private void releaseDataNodeQueryContext() { } } + public synchronized void initializeTableModelUDFQuery() { + if (!tableModelUDFQueryInitialized + && sessionInfo != null + && sessionInfo.getSqlDialect() == SqlDialect.TABLE) { + UDFClassLoaderManager.getInstance().initializeUDFQuery(id.getFullId()); + tableModelUDFQueryInitialized = true; + } + } + + public synchronized void finalizeTableModelUDFQuery() { + if (tableModelUDFQueryInitialized) { + UDFClassLoaderManager.getInstance().finalizeUDFQuery(id.getFullId()); + tableModelUDFQueryInitialized = false; + } + } + public void setMayHaveTmpFile(boolean mayHaveTmpFile) { this.mayHaveTmpFile = mayHaveTmpFile; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java index fa6a45e837496..b267bdba44d5e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -28,6 +28,8 @@ import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.exception.QueryTimeoutException; import org.apache.iotdb.commons.exception.SemanticException; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; +import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; @@ -179,12 +181,28 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( context.setHighestPriority(instance.isHighestPriority()); try { - List driverFactories = - planner.plan( - instance.getFragment().getPlanNodeTree(), - instance.getFragment().getTypeProvider(), - context, - dataNodeQueryContexts[0]); + List driverFactories; + if (instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) { + UDFManagementService.getInstance().acquireReadLock(); + try { + context.initializeTableModelUDFQuery(); + driverFactories = + planner.plan( + instance.getFragment().getPlanNodeTree(), + instance.getFragment().getTypeProvider(), + context, + dataNodeQueryContexts[0]); + } finally { + UDFManagementService.getInstance().releaseReadLock(); + } + } else { + driverFactories = + planner.plan( + instance.getFragment().getPlanNodeTree(), + instance.getFragment().getTypeProvider(), + context, + dataNodeQueryContexts[0]); + } List drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); @@ -235,6 +253,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.finalizeTableModelUDFQuery(); clearFIRelatedResources(instanceId); return null; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java index 9a4526c5bfe5f..4d5724c436390 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java @@ -136,7 +136,7 @@ protected void initTransformers( Map> inputLocations, Expression[] outputExpressions, Map, TSDataType> expressionTypes) { - UDFManagementService.getInstance().acquireLock(); + UDFManagementService.getInstance().acquireReadLock(); try { // This statement must be surrounded by the registration lock. UDFClassLoaderManager.getInstance().initializeUDFQuery(udtfQueryId); @@ -155,7 +155,7 @@ protected void initTransformers( .buildResultColumnPointReaders() .getOutputReaders(); } finally { - UDFManagementService.getInstance().releaseLock(); + UDFManagementService.getInstance().releaseReadLock(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 2324bea787789..528f7593f4bb0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -25,7 +25,9 @@ import org.apache.iotdb.commons.exception.IoTDBException; import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.queryengine.common.SqlDialect; +import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService; import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.query.KilledByOthersException; @@ -49,6 +51,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySpillManager; +import org.apache.iotdb.db.queryengine.plan.relational.planner.TableModelPlanner; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -109,6 +112,7 @@ public class QueryExecution implements IQueryExecution { private final AtomicBoolean resultHandleCleanUp; private final AtomicBoolean stopped; + private final Object tableUdfClassLoaderLock = new Object(); // cost time in ns of finished rpc private long totalExecutionTime = 0; @@ -164,6 +168,7 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService // preventing a premature close raced with late-scheduled FragmentInstances. context.releaseExternalTsFileQueryResources(); this.stop(cause); + finalizeTableUDFClassLoader(); } }); this.stopped = new AtomicBoolean(false); @@ -332,7 +337,24 @@ private void constructResultForMemorySource() { // Analyze the statement in QueryContext. Generate the analysis this query need private IAnalysis analyze(MPPQueryContext context) { - return planner.analyze(context); + if (!(planner instanceof TableModelPlanner)) { + return planner.analyze(context); + } + String queryId = context.getQueryId().getId(); + UDFManagementService.getInstance().acquireReadLock(); + try { + synchronized (tableUdfClassLoaderLock) { + UDFClassLoaderManager.getInstance().initializeUDFQuery(queryId); + try { + return planner.analyze(context); + } catch (RuntimeException | Error e) { + finalizeTableUDFClassLoader(); + throw e; + } + } + } finally { + UDFManagementService.getInstance().releaseReadLock(); + } } private void schedule() { @@ -436,6 +458,7 @@ public void stopAndCleanup(Throwable t) { /** Release the resources that current QueryExecution hold with a specified exception */ private void releaseResource(Throwable t) { + finalizeTableUDFClassLoader(); context.recordDeviceEntryDiskIOMetricsOnRelease(); // close ResultHandle to unblock client's getResult request // Actually, we should not close the ResultHandle when the QueryExecution is Finished. @@ -466,6 +489,12 @@ private void releaseResource(Throwable t) { } } + private void finalizeTableUDFClassLoader() { + synchronized (tableUdfClassLoaderLock) { + UDFClassLoaderManager.getInstance().finalizeUDFQuery(context.getQueryId().getId()); + } + } + /** * Clear up Coordinator.queryExecutionMap by calling Coordinator.cleanupQueryExecution if there is * no RPC in progress for this query (that is, the current RPC has finished and {@code diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java index fd5dc23d77dfa..bd931fb9324d7 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java @@ -42,17 +42,17 @@ import java.io.IOException; import java.lang.reflect.InvocationTargetException; import java.nio.ByteBuffer; -import java.util.concurrent.locks.ReentrantLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; public class UDFManagementService { private static final Logger LOGGER = LoggerFactory.getLogger(UDFManagementService.class); - private final ReentrantLock lock; + private final ReentrantReadWriteLock lock; private final UDFTable udfTable; private UDFManagementService() { - lock = new ReentrantLock(); + lock = new ReentrantReadWriteLock(); udfTable = new UDFTable(); // register tree model built-in functions for (BuiltinTimeSeriesGeneratingFunction builtinTimeSeriesGeneratingFunction : @@ -69,23 +69,31 @@ private UDFManagementService() { } } - public void acquireLock() { - lock.lock(); + public void acquireReadLock() { + lock.readLock().lock(); } - public void releaseLock() { - lock.unlock(); + public void releaseReadLock() { + lock.readLock().unlock(); + } + + public void acquireWriteLock() { + lock.writeLock().lock(); + } + + public void releaseWriteLock() { + lock.writeLock().unlock(); } public void register(UDFInformation udfInformation, ByteBuffer jarFile) throws Exception { Model model = checkAndGetModel(udfInformation); try { - acquireLock(); + acquireWriteLock(); checkIfRegistered(model, udfInformation); saveJarFile(udfInformation.getJarName(), jarFile); doRegister(model, udfInformation); } finally { - releaseLock(); + releaseWriteLock(); } } @@ -206,7 +214,7 @@ private void updateAllRegisteredClasses(Model model, UDFClassLoader activeClassL public void deregister(Model model, String functionName, boolean needToDeleteJar) throws Exception { try { - acquireLock(); + acquireWriteLock(); UDFInformation information = udfTable.getUDFInformation(model, functionName); if (information == null) { return; @@ -219,7 +227,7 @@ public void deregister(Model model, String functionName, boolean needToDeleteJar .removeFileUnderTemporaryRoot(functionName.toUpperCase() + ".txt"); } } finally { - releaseLock(); + releaseWriteLock(); } } From aa62c2fe2c93e846644cc76d5158d3e3a5f4965d Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Tue, 3 Nov 2026 16:10:11 +0800 Subject: [PATCH 02/10] Fix table UDF lifecycle cleanup guards --- .../execution/fragment/FragmentInstanceManager.java | 3 ++- .../iotdb/db/queryengine/plan/execution/QueryExecution.java | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java index b267bdba44d5e..10f5d5801a5b4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -182,7 +182,8 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( try { List driverFactories; - if (instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) { + if (instance.getSessionInfo() != null + && instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) { UDFManagementService.getInstance().acquireReadLock(); try { context.initializeTableModelUDFQuery(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 528f7593f4bb0..19924f49e889e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -490,6 +490,9 @@ private void releaseResource(Throwable t) { } private void finalizeTableUDFClassLoader() { + if (!(planner instanceof TableModelPlanner)) { + return; + } synchronized (tableUdfClassLoaderLock) { UDFClassLoaderManager.getInstance().finalizeUDFQuery(context.getQueryId().getId()); } From 104cde7975301c609e9d7a2ac96b5202871b378f Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 14:14:21 +0800 Subject: [PATCH 03/10] Fix UDF classloader lifecycle with query scopes --- .../statemachine/BaseStateMachine.java | 26 ++- .../dataregion/DataRegionStateMachine.java | 6 +- .../SchemaRegionStateMachine.java | 6 +- .../impl/DataNodeInternalRPCServiceImpl.java | 73 +++++-- .../fragment/FragmentInstanceContext.java | 32 ++-- .../fragment/FragmentInstanceManager.java | 74 +++----- .../operator/process/TransformOperator.java | 37 ++-- .../plan/execution/QueryExecution.java | 88 +++++---- .../plan/planner/LocalExecutionPlanner.java | 31 +++ .../plan/planner/plan/FragmentInstance.java | 28 ++- .../transformation/dag/udf/UDTFContext.java | 2 - .../planner/FragmentInstanceSerdeTest.java | 6 +- .../iotdb/commons/i18n/QueryMessages.java | 9 + .../iotdb/commons/i18n/QueryMessages.java | 9 + .../plan/udf/UDFManagementService.java | 87 ++++++--- .../apache/iotdb/commons/udf/UDFTable.java | 81 +++++--- .../commons/udf/service/UDFClassLoader.java | 18 +- .../udf/service/UDFClassLoaderManager.java | 179 ++++++++++++++++-- .../service/UDFClassLoaderManagerTest.java | 91 +++++++++ 19 files changed, 664 insertions(+), 219 deletions(-) create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java index bed0d16a049cb..569272c40d7a3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java @@ -22,14 +22,18 @@ 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); @@ -37,7 +41,27 @@ public abstract class BaseStateMachine implements IStateMachine, IStateMachine.E 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) { + if (classLoaderManager != null && queryScope != null) { + classLoaderManager.finalizeUDFQuery(queryScope); + } + throw e; + } } else if (request instanceof FragmentInstance) { instance = (FragmentInstance) request; } else { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java index 3d2fb3d2006b3..ce9663bc886c4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java @@ -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(); + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java index 9396d8cc35438..947c64aeb26b7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java @@ -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(); + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 4b1e05f4e857d..07303d69223e1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -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; @@ -528,9 +529,9 @@ 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); @@ -538,27 +539,57 @@ public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanc 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 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 84775f51132bf..7428876a943ef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -33,7 +33,6 @@ import org.apache.iotdb.commons.path.AlignedFullPath; import org.apache.iotdb.commons.path.IFullPath; import org.apache.iotdb.commons.queryengine.common.SessionInfo; -import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils; import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.commons.utils.TestOnly; @@ -118,7 +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 boolean tableModelUDFQueryInitialized; + private UDFClassLoaderManager.QueryScope udfQueryScope; private final FragmentInstanceStateMachine stateMachine; @@ -1351,7 +1350,7 @@ private void releaseTVListOwnedByQuery() { * be decreased. */ public synchronized void releaseResource() { - finalizeTableModelUDFQuery(); + finalizeUDFQuery(); try { getQueryChunkReaders().close(); } catch (IOException e) { @@ -1497,19 +1496,26 @@ private void releaseDataNodeQueryContext() { } } - public synchronized void initializeTableModelUDFQuery() { - if (!tableModelUDFQueryInitialized - && sessionInfo != null - && sessionInfo.getSqlDialect() == SqlDialect.TABLE) { - UDFClassLoaderManager.getInstance().initializeUDFQuery(id.getFullId()); - tableModelUDFQueryInitialized = true; + synchronized void adoptUDFQueryScope(UDFClassLoaderManager.QueryScope transferredScope) { + if (udfQueryScope == null) { + udfQueryScope = transferredScope; } } - public synchronized void finalizeTableModelUDFQuery() { - if (tableModelUDFQueryInitialized) { - UDFClassLoaderManager.getInstance().finalizeUDFQuery(id.getFullId()); - tableModelUDFQueryInitialized = false; + 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); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java index 10f5d5801a5b4..634a908bcd202 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -28,8 +28,6 @@ import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.exception.QueryTimeoutException; import org.apache.iotdb.commons.exception.SemanticException; -import org.apache.iotdb.commons.queryengine.common.SqlDialect; -import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; @@ -166,45 +164,30 @@ 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()); try { - List driverFactories; - if (instance.getSessionInfo() != null - && instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) { - UDFManagementService.getInstance().acquireReadLock(); - try { - context.initializeTableModelUDFQuery(); - driverFactories = - planner.plan( - instance.getFragment().getPlanNodeTree(), - instance.getFragment().getTypeProvider(), - context, - dataNodeQueryContexts[0]); - } finally { - UDFManagementService.getInstance().releaseReadLock(); - } - } else { - driverFactories = - planner.plan( - instance.getFragment().getPlanNodeTree(), - instance.getFragment().getTypeProvider(), - context, - dataNodeQueryContexts[0]); - } - + List driverFactories = + planner.plan( + instance.getFragment().getPlanNodeTree(), + instance.getFragment().getTypeProvider(), + context, + dataNodeQueryContexts[0]); List drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); // For ShowQueries related instances, isHighestPriority == true @@ -254,7 +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.finalizeTableModelUDFQuery(); + context.finalizeUDFQuery(); clearFIRelatedResources(instanceId); return null; } @@ -338,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()); @@ -355,7 +341,6 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( instance.getFragment().getTypeProvider(), context, schemaRegion); - List drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); context.initializeNumOfDrivers(drivers.size()); @@ -390,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; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java index 4d5724c436390..e8b7bd3a465e2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/TransformOperator.java @@ -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; @@ -136,27 +134,20 @@ protected void initTransformers( Map> inputLocations, Expression[] outputExpressions, Map, TSDataType> expressionTypes) { - UDFManagementService.getInstance().acquireReadLock(); - 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().releaseReadLock(); - } + // 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 { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 19924f49e889e..2ee3a6609647a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -25,7 +25,6 @@ import org.apache.iotdb.commons.exception.IoTDBException; import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.queryengine.common.SqlDialect; -import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService; import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics; import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.conf.IoTDBConfig; @@ -51,7 +50,6 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil; import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySpillManager; -import org.apache.iotdb.db.queryengine.plan.relational.planner.TableModelPlanner; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -112,7 +110,7 @@ public class QueryExecution implements IQueryExecution { private final AtomicBoolean resultHandleCleanUp; private final AtomicBoolean stopped; - private final Object tableUdfClassLoaderLock = new Object(); + private final UDFClassLoaderManager.QueryScope udfQueryScope; // cost time in ns of finished rpc private long totalExecutionTime = 0; @@ -133,9 +131,14 @@ public class QueryExecution implements IQueryExecution { public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService executor) { this.context = context; this.planner = planner; + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + this.udfQueryScope = + classLoaderManager == null + ? null + : classLoaderManager.openQueryScope(context.getQueryId().getId()); planner.beginAnalysisAttempt(); try { - this.analysis = analyze(context); + this.analysis = analyzeWithCleanupOnFailure(context); } catch (RuntimeException | Error e) { planner.rollbackAnalysisAttempt(); throw e; @@ -168,7 +171,9 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService // preventing a premature close raced with late-scheduled FragmentInstances. context.releaseExternalTsFileQueryResources(); this.stop(cause); - finalizeTableUDFClassLoader(); + if (context.isQuery()) { + finalizeUDFClassLoader(); + } } }); this.stopped = new AtomicBoolean(false); @@ -182,9 +187,21 @@ interface ISourceHandleSupplier { @Override public void start() { + boolean failed = false; try { - startInternal(); + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager == null || udfQueryScope == null) { + startInternalWithPlanCleanup(); + } else { + classLoaderManager.withQueryScope( + udfQueryScope, + () -> { + startInternalWithPlanCleanup(); + return null; + }); + } } catch (RuntimeException e) { + failed = true; finishAnalysisAttemptForCurrentState(); if (stateMachine.getState().isDone()) { // The query has already reached a terminal state (e.g. killed by others), and the status @@ -199,12 +216,17 @@ public void start() { } throw e; } catch (Error e) { + failed = true; finishAnalysisAttemptForCurrentState(); throw e; + } finally { + if (failed || !context.isQuery()) { + finalizeUDFClassLoader(); + } } } - private void startInternal() { + private void startInternalWithPlanCleanup() { final long startTime = System.nanoTime(); if (skipExecute()) { LOGGER.debug(DataNodeQueryMessages.SKIP_EXECUTE); @@ -273,14 +295,14 @@ private ExecutionResult retry() { LOGGER.warn(DataNodeQueryMessages.REACHMAXRETRYCOUNT); // Keep the same ownership order as an ordinary retry: plans and scheduler stop using the // analysis working state before the journal restores parser-owned objects. - this.stopAndCleanup(stateMachine.getFailureException()); + cleanupForRetry(stateMachine.getFailureException()); planner.rollbackAnalysisAttempt(); stateMachine.transitionToFailed(); return getStatus(); } LOGGER.warn(DataNodeQueryMessages.ERROR_WHEN_EXECUTING_QUERY, stateMachine.getFailureMessage()); // stop and clean up resources the QueryExecution used - this.stopAndCleanup(stateMachine.getFailureException()); + cleanupForRetry(stateMachine.getFailureException()); // The generated plans no longer read the statement after cleanup. Restore the parser-owned // state before the next analysis attempt so retry observes the same SQL as the first attempt. planner.rollbackAnalysisAttempt(); @@ -301,14 +323,20 @@ private ExecutionResult retry() { // re-stop this.stopped.compareAndSet(true, false); this.resultHandleCleanUp.compareAndSet(true, false); + if (stateMachine.getState().isDone()) { + return getStatus(); + } // re-analyze the query planner.beginAnalysisAttempt(); try { - this.analysis = analyze(context); + this.analysis = analyzeWithCleanupOnFailure(context); } catch (RuntimeException | Error e) { planner.rollbackAnalysisAttempt(); throw e; } + if (stateMachine.getState().isDone()) { + return getStatus(); + } // re-start the QueryExecution this.start(); return getStatus(); @@ -336,24 +364,16 @@ private void constructResultForMemorySource() { } // Analyze the statement in QueryContext. Generate the analysis this query need - private IAnalysis analyze(MPPQueryContext context) { - if (!(planner instanceof TableModelPlanner)) { - return planner.analyze(context); - } - String queryId = context.getQueryId().getId(); - UDFManagementService.getInstance().acquireReadLock(); + private IAnalysis analyzeWithCleanupOnFailure(MPPQueryContext context) { try { - synchronized (tableUdfClassLoaderLock) { - UDFClassLoaderManager.getInstance().initializeUDFQuery(queryId); - try { - return planner.analyze(context); - } catch (RuntimeException | Error e) { - finalizeTableUDFClassLoader(); - throw e; - } + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager == null || udfQueryScope == null) { + return planner.analyze(context); } - } finally { - UDFManagementService.getInstance().releaseReadLock(); + return classLoaderManager.withQueryScope(udfQueryScope, () -> planner.analyze(context)); + } catch (RuntimeException | Error e) { + finalizeUDFClassLoader(); + throw e; } } @@ -456,9 +476,13 @@ public void stopAndCleanup(Throwable t) { releaseResource(t); } + private void cleanupForRetry(Throwable t) { + stop(t); + releaseResource(t); + } + /** Release the resources that current QueryExecution hold with a specified exception */ private void releaseResource(Throwable t) { - finalizeTableUDFClassLoader(); context.recordDeviceEntryDiskIOMetricsOnRelease(); // close ResultHandle to unblock client's getResult request // Actually, we should not close the ResultHandle when the QueryExecution is Finished. @@ -489,12 +513,10 @@ private void releaseResource(Throwable t) { } } - private void finalizeTableUDFClassLoader() { - if (!(planner instanceof TableModelPlanner)) { - return; - } - synchronized (tableUdfClassLoaderLock) { - UDFClassLoaderManager.getInstance().finalizeUDFQuery(context.getQueryId().getId()); + private void finalizeUDFClassLoader() { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager != null && udfQueryScope != null) { + classLoaderManager.finalizeUDFQuery(udfQueryScope); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index c600b4eb9016d..92534af6b0e88 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -28,6 +28,7 @@ import org.apache.iotdb.commons.path.IFullPath; import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.conf.DataNodeMemoryConfig; import org.apache.iotdb.db.conf.IoTDBConfig; @@ -105,6 +106,21 @@ public List plan( FragmentInstanceContext instanceContext, DataNodeQueryContext dataNodeQueryContext) throws MemoryNotEnoughException { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + UDFClassLoaderManager.QueryScope queryScope = instanceContext.getOrCreateUDFQueryScope(); + if (classLoaderManager == null || queryScope == null) { + return planInternal(plan, types, instanceContext, dataNodeQueryContext); + } + return classLoaderManager.withQueryScope( + queryScope, () -> planInternal(plan, types, instanceContext, dataNodeQueryContext)); + } + + private List planInternal( + PlanNode plan, + TypeProvider types, + FragmentInstanceContext instanceContext, + DataNodeQueryContext dataNodeQueryContext) + throws MemoryNotEnoughException { if (Objects.isNull(plan)) { throw new IoTDBRuntimeException( DataNodeQueryMessages @@ -145,6 +161,21 @@ public List plan( FragmentInstanceContext instanceContext, ISchemaRegion schemaRegion) throws MemoryNotEnoughException { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + UDFClassLoaderManager.QueryScope queryScope = instanceContext.getOrCreateUDFQueryScope(); + if (classLoaderManager == null || queryScope == null) { + return planInternal(plan, types, instanceContext, schemaRegion); + } + return classLoaderManager.withQueryScope( + queryScope, () -> planInternal(plan, types, instanceContext, schemaRegion)); + } + + private List planInternal( + PlanNode plan, + TypeProvider types, + FragmentInstanceContext instanceContext, + ISchemaRegion schemaRegion) + throws MemoryNotEnoughException { if (Objects.isNull(plan)) { throw new IoTDBRuntimeException( DataNodeQueryMessages diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java index 463aeb5413196..10bba4daedc13 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.partition.ExecutorType; import org.apache.iotdb.commons.queryengine.common.SessionInfo; import org.apache.iotdb.commons.request.IConsensusRequest; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.commons.utils.ThriftCommonsSerDeUtils; import org.apache.iotdb.db.conf.IoTDBConfig; @@ -73,6 +74,8 @@ public class FragmentInstance implements IConsensusRequest { private boolean isHighestPriority; + private transient UDFClassLoaderManager.QueryScope udfQueryScope; + // indicate which index we are retrying private transient int nextRetryIndex = 0; @@ -242,7 +245,11 @@ public String toString() { } public static FragmentInstance deserializeFrom(ByteBuffer buffer) { - FragmentInstanceId id = FragmentInstanceId.deserialize(buffer); + return deserializeFrom(FragmentInstanceId.deserialize(buffer), buffer); + } + + /** Deserialize the remaining fragment payload after the caller has read its ID. */ + public static FragmentInstance deserializeFrom(FragmentInstanceId id, ByteBuffer buffer) { PlanFragment planFragment = PlanFragment.deserialize(buffer); long timeOut = ReadWriteIOUtils.readLong(buffer); boolean hasSessionInfo = ReadWriteIOUtils.readBool(buffer); @@ -342,6 +349,25 @@ public SessionInfo getSessionInfo() { return sessionInfo; } + public void attachUDFQueryScope(UDFClassLoaderManager.QueryScope queryScope) { + udfQueryScope = queryScope; + } + + /** Transfer query scope ownership out of this deserialization result. */ + public UDFClassLoaderManager.QueryScope claimUDFQueryScope() { + UDFClassLoaderManager.QueryScope queryScope = udfQueryScope; + udfQueryScope = null; + return queryScope; + } + + public void closeUDFQueryScopeIfPresent() { + UDFClassLoaderManager.QueryScope queryScope = claimUDFQueryScope(); + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (queryScope != null && classLoaderManager != null) { + classLoaderManager.finalizeUDFQuery(queryScope); + } + } + public boolean isExplainAnalyze() { return isExplainAnalyze; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/transformation/dag/udf/UDTFContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/transformation/dag/udf/UDTFContext.java index c9dd4bb0620a6..66876ec17d83b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/transformation/dag/udf/UDTFContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/transformation/dag/udf/UDTFContext.java @@ -21,7 +21,6 @@ import org.apache.iotdb.calc.service.AbstractTemporaryQueryDataFileService; import org.apache.iotdb.calc.transformation.dag.udf.UDTFExecutor; -import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.queryengine.plan.expression.Expression; import org.apache.iotdb.db.queryengine.plan.expression.multi.FunctionExpression; @@ -51,7 +50,6 @@ public void finalizeUDFExecutors(String queryId) { executor.beforeDestroy(); } } finally { - UDFClassLoaderManager.getInstance().finalizeUDFQuery(queryId); // close and delete UDF temp files AbstractTemporaryQueryDataFileService.getInstance().deregister(queryId); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java index b387baddff6b2..f54fbb111db79 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java @@ -33,6 +33,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.expression.ExpressionFactory; @@ -86,7 +87,10 @@ public void testSerializeAndDeserializeForTree1() throws IllegalPathException { ImmutableList.of(dataNodeLocation)); fragmentInstance.setExecutorAndHost(new StorageExecutor(regionReplicaSet)); ByteBuffer byteBuffer = fragmentInstance.serializeToByteBuffer(); - FragmentInstance deserializeFragmentInstance = FragmentInstance.deserializeFrom(byteBuffer); + FragmentInstanceId fragmentInstanceId = FragmentInstanceId.deserialize(byteBuffer); + FragmentInstance deserializeFragmentInstance = + FragmentInstance.deserializeFrom(fragmentInstanceId, byteBuffer); + assertNull(deserializeFragmentInstance.claimUDFQueryScope()); assertNull(deserializeFragmentInstance.getExecutorType()); // Because the ExecutorType won't be considered in serialization, we need to set it // from original object before comparison. diff --git a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/QueryMessages.java b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/QueryMessages.java index fc06f96c70271..ed20c1dfdfc8a 100644 --- a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/QueryMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/QueryMessages.java @@ -202,6 +202,15 @@ public final class QueryMessages { "Failed to reflect UDF instance, because UDF %s has not been registered."; public static final String UDF_REFLECT_INSTANCE_FAILED = "Failed to reflect UDF %s(%s) instance, because %s"; + public static final String EXCEPTION_UDF_CLASSLOADER_REQUIRES_QUERY_SCOPE_ARG_A39FB39E = + "Cannot pin UDF ClassLoader for class %s because no UDF query scope is active."; + public static final String + EXCEPTION_FAILED_TO_PIN_UDF_CLASSLOADER_QUERY_SCOPE_ARG_CLASS_ARG_275E1E20 = + "Failed to pin UDFClassLoader for query scope %s and class %s because the classloader is deprecated or closed."; + public static final String EXCEPTION_UDF_QUERY_SCOPE_IS_CLOSED_ARG_7381D949 = + "Cannot execute action in query scope %s because the scope has been closed."; + public static final String EXCEPTION_UDF_CLASSLOADER_MANAGER_NOT_INITIALIZED_ARG_DF66329C = + "Cannot pin UDF ClassLoader for class %s because UDFClassLoaderManager is not initialized."; public static final String AI_NODE_SERVICE_NOT_AVAILABLE = "Table function AINode service is not available in current node"; // ======================== UDTF Forecast ======================== diff --git a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/QueryMessages.java b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/QueryMessages.java index e08448271bbb2..6f429ac1102fc 100644 --- a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/QueryMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/QueryMessages.java @@ -202,6 +202,15 @@ public final class QueryMessages { "反射 UDF 实例失败，因为 UDF %s 尚未注册。"; public static final String UDF_REFLECT_INSTANCE_FAILED = "反射 UDF %s(%s) 实例失败，原因：%s"; + public static final String EXCEPTION_UDF_CLASSLOADER_REQUIRES_QUERY_SCOPE_ARG_A39FB39E = + "无法为类 %s pin UDF ClassLoader，因为当前没有活动的 UDF query scope。"; + public static final String + EXCEPTION_FAILED_TO_PIN_UDF_CLASSLOADER_QUERY_SCOPE_ARG_CLASS_ARG_275E1E20 = + "无法为 query scope %s 和类 %s pin UDFClassLoader，因为 classloader 已废弃或已关闭。"; + public static final String EXCEPTION_UDF_QUERY_SCOPE_IS_CLOSED_ARG_7381D949 = + "无法在 query scope %s 中执行操作，因为该 scope 已关闭。"; + public static final String EXCEPTION_UDF_CLASSLOADER_MANAGER_NOT_INITIALIZED_ARG_DF66329C = + "无法为类 %s pin UDF ClassLoader，因为 UDFClassLoaderManager 尚未初始化。"; public static final String AI_NODE_SERVICE_NOT_AVAILABLE = "当前节点不支持表函数 AINode 服务"; // ======================== UDTF Forecast ======================== diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java index bd931fb9324d7..d0e59ccb0c7f7 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java @@ -42,6 +42,11 @@ import java.io.IOException; import java.lang.reflect.InvocationTargetException; import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.concurrent.locks.ReentrantReadWriteLock; public class UDFManagementService { @@ -178,17 +183,36 @@ private void saveJarFile(String jarName, ByteBuffer byteBuffer) throws IOExcepti public void doRegister(Model model, UDFInformation udfInformation) throws UDFManagementException { String functionName = udfInformation.getFunctionName(); String className = udfInformation.getClassName(); + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + UDFClassLoader candidateClassLoader = null; + acquireWriteLock(); try { - UDFClassLoader currentActiveClassLoader = - UDFClassLoaderManager.getInstance().updateAndGetActiveClassLoader(); - updateAllRegisteredClasses(model, currentActiveClassLoader); - - Class functionClass = Class.forName(className, true, currentActiveClassLoader); + candidateClassLoader = classLoaderManager.createUDFClassLoader(); + List candidateInformation = new ArrayList<>(udfTable.getAllInformationList()); + candidateInformation.add(udfInformation); + Map>> candidateFunctionClasses = new EnumMap<>(Model.class); + candidateFunctionClasses.put(Model.TREE, new HashMap<>()); + candidateFunctionClasses.put(Model.TABLE, new HashMap<>()); + Class registeredFunctionClass = null; + for (UDFInformation information : candidateInformation) { + Class functionClass = + Class.forName(information.getClassName(), true, candidateClassLoader); + Model functionModel = information.getUdfType().isTreeModel() ? Model.TREE : Model.TABLE; + candidateFunctionClasses + .get(functionModel) + .put(information.getFunctionName().toUpperCase(), functionClass); + if (information == udfInformation) { + registeredFunctionClass = functionClass; + } + } - // ensure that it is a UDF class - getBaseClass(model).cast(functionClass.getDeclaredConstructor().newInstance()); - udfTable.addUDFInformation(functionName, udfInformation); - udfTable.addFunctionAndClass(model, functionName, functionClass); + // ensure that the new class is a valid UDF implementation + getBaseClass(model).cast(registeredFunctionClass.getDeclaredConstructor().newInstance()); + UDFClassLoader publishClassLoader = candidateClassLoader; + classLoaderManager.publishUDFClassLoader( + publishClassLoader, + () -> udfTable.publishUDF(functionName, udfInformation, candidateFunctionClasses)); + candidateClassLoader = null; } catch (IOException | InstantiationException | InvocationTargetException @@ -201,13 +225,11 @@ public void doRegister(Model model, UDFInformation udfInformation) throws UDFMan QueryMessages.UDF_REGISTER_INSTANCE_FAILED, functionName.toUpperCase(), className, e); LOGGER.warn(errorMessage, e); throw new UDFManagementException(errorMessage); - } - } - - private void updateAllRegisteredClasses(Model model, UDFClassLoader activeClassLoader) - throws ClassNotFoundException { - for (UDFInformation information : getUDFInformation(model)) { - udfTable.updateFunctionClass(information, activeClassLoader); + } finally { + if (candidateClassLoader != null) { + classLoaderManager.discardUDFClassLoader(candidateClassLoader); + } + releaseWriteLock(); } } @@ -240,17 +262,34 @@ public T reflect(String functionName, Class clazz) { } else { throw new UDFManagementException(QueryMessages.UDF_CLASS_TYPE_UNSUPPORTED); } - UDFInformation information = udfTable.getUDFInformation(model, functionName); - if (information == null) { - String errorMessage = - String.format(QueryMessages.UDF_REFLECT_NOT_REGISTERED, functionName.toUpperCase()); - LOGGER.warn(errorMessage); - throw new UDFException(errorMessage); + UDFInformation information; + Class functionClass; + try { + acquireReadLock(); + information = udfTable.getUDFInformation(model, functionName); + if (information == null) { + String errorMessage = + String.format(QueryMessages.UDF_REFLECT_NOT_REGISTERED, functionName.toUpperCase()); + LOGGER.warn(errorMessage); + throw new UDFException(errorMessage); + } + functionClass = udfTable.getFunctionClass(model, functionName); + if (functionClass.getClassLoader() instanceof UDFClassLoader) { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager == null) { + throw new IllegalStateException( + String.format( + QueryMessages.EXCEPTION_UDF_CLASSLOADER_MANAGER_NOT_INITIALIZED_ARG_DF66329C, + functionClass.getName())); + } + classLoaderManager.pinClassLoaderForCurrentQuery(functionClass); + } + } finally { + releaseReadLock(); } try { - return clazz.cast( - udfTable.getFunctionClass(model, functionName).getDeclaredConstructor().newInstance()); + return clazz.cast(functionClass.getDeclaredConstructor().newInstance()); } catch (InstantiationException | InvocationTargetException | NoSuchMethodException diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/UDFTable.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/UDFTable.java index 73f0705cc1f4d..c640a6edd6b7c 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/UDFTable.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/UDFTable.java @@ -20,7 +20,6 @@ package org.apache.iotdb.commons.udf; import org.apache.iotdb.common.rpc.thrift.Model; -import org.apache.iotdb.commons.udf.service.UDFClassLoader; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.tsfile.utils.ReadWriteIOUtils; @@ -29,6 +28,8 @@ import java.io.InputStream; import java.io.OutputStream; import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -44,18 +45,23 @@ public class UDFTable { private final Map> udfInformationMap; /** maintain a map for creating instance, model -> functionName -> class */ - private final Map>> functionToClassMap; + private volatile Map>> functionToClassMap; public UDFTable() { udfInformationMap = new ConcurrentHashMap<>(); - functionToClassMap = new ConcurrentHashMap<>(); + functionToClassMap = emptyFunctionClassMap(); udfInformationMap.put(Model.TREE, new ConcurrentHashMap<>()); udfInformationMap.put(Model.TABLE, new ConcurrentHashMap<>()); - functionToClassMap.put(Model.TREE, new ConcurrentHashMap<>()); - functionToClassMap.put(Model.TABLE, new ConcurrentHashMap<>()); } - public void addUDFInformation(String functionName, UDFInformation udfInformation) { + private static Map>> emptyFunctionClassMap() { + Map>> functionClassMap = new HashMap<>(); + functionClassMap.put(Model.TREE, Collections.emptyMap()); + functionClassMap.put(Model.TABLE, Collections.emptyMap()); + return Collections.unmodifiableMap(functionClassMap); + } + + public synchronized void addUDFInformation(String functionName, UDFInformation udfInformation) { if (udfInformation.getUdfType().isTreeModel()) { udfInformationMap.get(Model.TREE).put(functionName.toUpperCase(), udfInformation); } else { @@ -63,51 +69,65 @@ public void addUDFInformation(String functionName, UDFInformation udfInformation } } - public void removeUDFInformation(Model model, String functionName) { + public synchronized void removeUDFInformation(Model model, String functionName) { udfInformationMap.get(model).remove(functionName.toUpperCase()); } - public UDFInformation getUDFInformation(Model model, String functionName) { + public synchronized UDFInformation getUDFInformation(Model model, String functionName) { return udfInformationMap.get(model).get(functionName.toUpperCase()); } - public void addFunctionAndClass(Model model, String functionName, Class clazz) { - functionToClassMap.get(model).put(functionName.toUpperCase(), clazz); + public synchronized void addFunctionAndClass(Model model, String functionName, Class clazz) { + Map>> newFunctionClassMap = new HashMap<>(functionToClassMap); + Map> modelFunctionClassMap = new HashMap<>(functionToClassMap.get(model)); + modelFunctionClassMap.put(functionName.toUpperCase(), clazz); + newFunctionClassMap.put(model, Collections.unmodifiableMap(modelFunctionClassMap)); + functionToClassMap = Collections.unmodifiableMap(newFunctionClassMap); } - public Class getFunctionClass(Model model, String functionName) { + public synchronized Class getFunctionClass(Model model, String functionName) { return functionToClassMap.get(model).get(functionName.toUpperCase()); } - public void removeFunctionClass(Model model, String functionName) { - functionToClassMap.get(model).remove(functionName.toUpperCase()); + public synchronized void removeFunctionClass(Model model, String functionName) { + Map>> newFunctionClassMap = new HashMap<>(functionToClassMap); + Map> modelFunctionClassMap = new HashMap<>(functionToClassMap.get(model)); + modelFunctionClassMap.remove(functionName.toUpperCase()); + newFunctionClassMap.put(model, Collections.unmodifiableMap(modelFunctionClassMap)); + functionToClassMap = Collections.unmodifiableMap(newFunctionClassMap); } - public void updateFunctionClass(UDFInformation udfInformation, UDFClassLoader classLoader) - throws ClassNotFoundException { - Class functionClass = Class.forName(udfInformation.getClassName(), true, classLoader); - if (udfInformation.getUdfType().isTreeModel()) { - functionToClassMap - .get(Model.TREE) - .put(udfInformation.getFunctionName().toUpperCase(), functionClass); - } else { - functionToClassMap - .get(Model.TABLE) - .put(udfInformation.getFunctionName().toUpperCase(), functionClass); + private static Map>> createFunctionClassSnapshot( + Map>> newFunctionClassMap) { + Map>> snapshot = new HashMap<>(); + for (Model model : Model.values()) { + snapshot.put( + model, Collections.unmodifiableMap(new HashMap<>(newFunctionClassMap.get(model)))); } + return Collections.unmodifiableMap(snapshot); + } + + public synchronized void publishUDF( + String functionName, + UDFInformation udfInformation, + Map>> newFunctionClassMap) { + Map>> functionClassSnapshot = + createFunctionClassSnapshot(newFunctionClassMap); + addUDFInformation(functionName, udfInformation); + functionToClassMap = functionClassSnapshot; } - public List getUDFInformationList(Model model) { + public synchronized List getUDFInformationList(Model model) { return new ArrayList<>(udfInformationMap.get(model).values()); } - public List getAllInformationList() { + public synchronized List getAllInformationList() { return udfInformationMap.values().stream() .flatMap(map -> map.values().stream()) .collect(Collectors.toList()); } - public boolean containsUDF(Model model, String udfName) { + public synchronized boolean containsUDF(Model model, String udfName) { return udfInformationMap.get(model).containsKey(udfName.toUpperCase()); } @@ -124,7 +144,7 @@ public void serializeUDFTable(OutputStream outputStream) throws IOException { } } - public void deserializeUDFTable(InputStream inputStream) throws IOException { + public synchronized void deserializeUDFTable(InputStream inputStream) throws IOException { int size = ReadWriteIOUtils.readInt(inputStream); while (size > 0) { UDFInformation udfInformation = UDFInformation.deserialize(inputStream); @@ -137,10 +157,9 @@ public void deserializeUDFTable(InputStream inputStream) throws IOException { } } - public void clear() { + public synchronized void clear() { udfInformationMap.get(Model.TREE).clear(); udfInformationMap.get(Model.TABLE).clear(); - functionToClassMap.get(Model.TREE).clear(); - functionToClassMap.get(Model.TABLE).clear(); + functionToClassMap = emptyFunctionClassMap(); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoader.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoader.java index df5dd977b5dd5..72a8bd09ec67a 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoader.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoader.java @@ -46,11 +46,14 @@ public class UDFClassLoader extends URLClassLoader { */ private volatile boolean deprecated; + private boolean closed; + public UDFClassLoader(String libRoot) throws IOException { super(new URL[0]); this.libRoot = libRoot; activeQueriesCount = new AtomicLong(0); deprecated = false; + closed = false; addURLs(); } @@ -65,23 +68,28 @@ private void addURLs() throws IOException { } } - public void acquire() { + public synchronized boolean tryAcquire() { + if (deprecated || closed) { + return false; + } activeQueriesCount.incrementAndGet(); + return true; } - public void release() throws IOException { + public synchronized void release() throws IOException { activeQueriesCount.decrementAndGet(); closeIfPossible(); } - public void markAsDeprecated() throws IOException { + public synchronized void markAsDeprecated() throws IOException { deprecated = true; closeIfPossible(); } - public void closeIfPossible() throws IOException { - if (deprecated && activeQueriesCount.get() == 0) { + public synchronized void closeIfPossible() throws IOException { + if (!closed && deprecated && activeQueriesCount.get() == 0) { close(); + closed = true; } } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManager.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManager.java index 19e06c4f7342d..118d00f0339cb 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManager.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManager.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.exception.StartupException; import org.apache.iotdb.commons.file.SystemFileFactory; import org.apache.iotdb.commons.i18n.CommonMessages; +import org.apache.iotdb.commons.i18n.QueryMessages; import org.apache.iotdb.commons.service.IService; import org.apache.iotdb.commons.service.ServiceType; import org.apache.iotdb.commons.utils.TestOnly; @@ -30,9 +31,14 @@ import org.slf4j.LoggerFactory; import java.io.IOException; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; +import java.util.ArrayDeque; +import java.util.Collections; +import java.util.Deque; +import java.util.IdentityHashMap; +import java.util.Objects; +import java.util.Set; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; public class UDFClassLoaderManager implements IService { @@ -40,8 +46,7 @@ public class UDFClassLoaderManager implements IService { private final String libRoot; - /** The keys in the map are the query IDs of the UDF queries being executed. */ - private final Map queryIdToUDFClassLoaderMap; + private final ThreadLocal> queryScopeStack; /** * activeClassLoader is used to load all classes under libRoot. libRoot may be updated before the @@ -50,30 +55,135 @@ public class UDFClassLoaderManager implements IService { */ private final AtomicReference activeClassLoader = new AtomicReference<>(); + public static final class QueryScope { + private final String queryId; + private final Set pinnedClassLoaders = + Collections.newSetFromMap(new IdentityHashMap<>()); + private int activeActions; + private boolean finalizeRequested; + private boolean finalized; + + private QueryScope(String queryId) { + this.queryId = queryId; + } + } + private UDFClassLoaderManager(String libRoot) { this.libRoot = libRoot; LOGGER.info(CommonMessages.UDF_LIB_ROOT, libRoot); - queryIdToUDFClassLoaderMap = new ConcurrentHashMap<>(); + queryScopeStack = ThreadLocal.withInitial(ArrayDeque::new); } @TestOnly private UDFClassLoaderManager() throws IOException { this.libRoot = null; - queryIdToUDFClassLoaderMap = new ConcurrentHashMap<>(); + queryScopeStack = ThreadLocal.withInitial(ArrayDeque::new); activeClassLoader.set(new UDFClassLoader("")); } - public void initializeUDFQuery(String queryId) { - activeClassLoader.get().acquire(); - queryIdToUDFClassLoaderMap.put(queryId, activeClassLoader.get()); + public QueryScope openQueryScope(String queryId) { + return new QueryScope(Objects.requireNonNull(queryId)); } - public void finalizeUDFQuery(String queryId) { - UDFClassLoader classLoader = queryIdToUDFClassLoaderMap.remove(queryId); + /** + * Bind an existing logical query scope to the current thread. Pins acquired by reflect() remain + * owned by the scope after this method returns and must be released by finalizeUDFQuery(). + */ + public T withQueryScope(QueryScope queryScope, Supplier action) { + Objects.requireNonNull(queryScope); + Objects.requireNonNull(action); + synchronized (this) { + if (queryScope.finalizeRequested || queryScope.finalized) { + throw new IllegalStateException( + String.format( + QueryMessages.EXCEPTION_UDF_QUERY_SCOPE_IS_CLOSED_ARG_7381D949, + queryScope.queryId)); + } + queryScope.activeActions++; + } + + Deque scopeStack = queryScopeStack.get(); + scopeStack.push(queryScope); try { - if (classLoader != null) { - classLoader.release(); + return action.get(); + } finally { + scopeStack.pop(); + if (scopeStack.isEmpty()) { + queryScopeStack.remove(); } + synchronized (this) { + queryScope.activeActions--; + closeScopeIfPossible(queryScope); + } + } + } + + /** + * Pin the exact class loader that defined the reflected UDF class. Built-in functions loaded by + * the parent class loader do not create a pin. + */ + public void pinClassLoaderForCurrentQuery(Class udfClass) { + ClassLoader classLoader = udfClass.getClassLoader(); + if (!(classLoader instanceof UDFClassLoader)) { + return; + } + + Deque scopeStack = queryScopeStack.get(); + QueryScope queryScope = scopeStack.peek(); + if (queryScope == null) { + throw new IllegalStateException( + String.format( + QueryMessages.EXCEPTION_UDF_CLASSLOADER_REQUIRES_QUERY_SCOPE_ARG_A39FB39E, + udfClass.getName())); + } + + UDFClassLoader udfClassLoader = (UDFClassLoader) classLoader; + synchronized (this) { + if (queryScope.finalizeRequested || queryScope.finalized) { + throw new IllegalStateException( + String.format( + QueryMessages + .EXCEPTION_FAILED_TO_PIN_UDF_CLASSLOADER_QUERY_SCOPE_ARG_CLASS_ARG_275E1E20, + queryScope.queryId, + udfClass.getName())); + } + if (queryScope.pinnedClassLoaders.contains(udfClassLoader)) { + return; + } + if (!udfClassLoader.tryAcquire()) { + throw new IllegalStateException( + String.format( + QueryMessages + .EXCEPTION_FAILED_TO_PIN_UDF_CLASSLOADER_QUERY_SCOPE_ARG_CLASS_ARG_275E1E20, + queryScope.queryId, + udfClass.getName())); + } + queryScope.pinnedClassLoaders.add(udfClassLoader); + } + } + + public synchronized void finalizeUDFQuery(QueryScope queryScope) { + if (queryScope == null || queryScope.finalized) { + return; + } + queryScope.finalizeRequested = true; + closeScopeIfPossible(queryScope); + } + + private void closeScopeIfPossible(QueryScope queryScope) { + if (!queryScope.finalizeRequested || queryScope.activeActions > 0 || queryScope.finalized) { + return; + } + for (UDFClassLoader classLoader : queryScope.pinnedClassLoaders) { + releaseClassLoaderPin(queryScope.queryId, classLoader); + } + queryScope.pinnedClassLoaders.clear(); + queryScope.finalized = true; + } + + private void releaseClassLoaderPin(Object queryId, UDFClassLoader classLoader) { + try { + classLoader.release(); } catch (IOException e) { LOGGER.warn( CommonMessages.LOG_FAILED_CLOSE_UDFCLASSLOADER_QUERYID_ARG_BECAUSE_ARG_8B1C3739, @@ -82,13 +192,36 @@ public void finalizeUDFQuery(String queryId) { } } - public UDFClassLoader updateAndGetActiveClassLoader() throws IOException { + public UDFClassLoader createUDFClassLoader() throws IOException { + return new UDFClassLoader(libRoot); + } + + public void discardUDFClassLoader(UDFClassLoader classLoader) { + try { + classLoader.close(); + } catch (IOException e) { + LOGGER.warn( + CommonMessages.LOG_FAILED_CLOSE_UDFCLASSLOADER_QUERYID_ARG_BECAUSE_ARG_8B1C3739, + System.identityHashCode(classLoader), + e.toString()); + } + } + + public synchronized void publishUDFClassLoader( + UDFClassLoader newClassLoader, Runnable publishCatalog) { UDFClassLoader deprecatedClassLoader = activeClassLoader.get(); - activeClassLoader.set(new UDFClassLoader(libRoot)); + publishCatalog.run(); + activeClassLoader.set(newClassLoader); if (deprecatedClassLoader != null) { - deprecatedClassLoader.markAsDeprecated(); + try { + deprecatedClassLoader.markAsDeprecated(); + } catch (IOException e) { + LOGGER.warn( + CommonMessages.LOG_FAILED_CLOSE_UDFCLASSLOADER_QUERYID_ARG_BECAUSE_ARG_8B1C3739, + System.identityHashCode(deprecatedClassLoader), + e.toString()); + } } - return activeClassLoader.get(); } ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -99,7 +232,17 @@ public UDFClassLoader updateAndGetActiveClassLoader() throws IOException { public void start() throws StartupException { try { SystemFileFactory.INSTANCE.makeDirIfNecessary(libRoot); - activeClassLoader.set(new UDFClassLoader(libRoot)); + if (activeClassLoader.get() != null) { + return; + } + UDFClassLoader newClassLoader = createUDFClassLoader(); + synchronized (this) { + if (activeClassLoader.get() == null) { + activeClassLoader.set(newClassLoader); + } else { + newClassLoader.close(); + } + } } catch (IOException e) { throw new StartupException(this.getID().getName(), e.getMessage()); } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java new file mode 100644 index 0000000000000..d2bb5e9adbf2c --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java @@ -0,0 +1,91 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.udf.service; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.lang.reflect.Field; +import java.util.Set; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class UDFClassLoaderManagerTest { + + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void testFinalizeWaitsForAllNestedActions() throws Exception { + UDFClassLoaderManager manager = UDFClassLoaderManager.setupAndGetInstance(); + UDFClassLoaderManager.QueryScope queryScope = manager.openQueryScope("nested-actions"); + UDFClassLoader classLoader = createPinnedDeprecatedClassLoader(queryScope); + + manager.withQueryScope( + queryScope, + () -> { + manager.withQueryScope( + queryScope, + () -> { + manager.finalizeUDFQuery(queryScope); + return null; + }); + assertFalse(isClosed(classLoader)); + return null; + }); + + assertTrue(isClosed(classLoader)); + } + + @Test + public void testFinalizedScopeCannotBeBoundAgain() throws Exception { + UDFClassLoaderManager manager = UDFClassLoaderManager.setupAndGetInstance(); + UDFClassLoaderManager.QueryScope queryScope = manager.openQueryScope("closed-scope"); + + manager.finalizeUDFQuery(queryScope); + + assertThrows(IllegalStateException.class, () -> manager.withQueryScope(queryScope, () -> null)); + } + + private UDFClassLoader createPinnedDeprecatedClassLoader( + UDFClassLoaderManager.QueryScope queryScope) throws Exception { + UDFClassLoader classLoader = new UDFClassLoader(temporaryFolder.newFolder().getAbsolutePath()); + assertTrue(classLoader.tryAcquire()); + classLoader.markAsDeprecated(); + getPinnedClassLoaders(queryScope).add(classLoader); + return classLoader; + } + + @SuppressWarnings("unchecked") + private static Set getPinnedClassLoaders( + UDFClassLoaderManager.QueryScope queryScope) throws Exception { + Field field = UDFClassLoaderManager.QueryScope.class.getDeclaredField("pinnedClassLoaders"); + field.setAccessible(true); + return (Set) field.get(queryScope); + } + + private static boolean isClosed(UDFClassLoader classLoader) throws Exception { + Field field = UDFClassLoader.class.getDeclaredField("closed"); + field.setAccessible(true); + return field.getBoolean(classLoader); + } +} From 145c528d9144828a95ca613263052ed7e4193682 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 14:33:44 +0800 Subject: [PATCH 04/10] Add UDF query scope lifecycle regression coverage --- .../plan/execution/QueryExecution.java | 6 +- .../execution/QueryExecutionRetryTest.java | 46 ++++- .../planner/FragmentInstanceSerdeTest.java | 37 ++++ .../plan/udf/UDFManagementServiceTest.java | 180 ++++++++++++++++++ 4 files changed, 262 insertions(+), 7 deletions(-) create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 2ee3a6609647a..6090f0ad703ae 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -171,7 +171,7 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService // preventing a premature close raced with late-scheduled FragmentInstances. context.releaseExternalTsFileQueryResources(); this.stop(cause); - if (context.isQuery()) { + if (context.isQuery() || state != QueryState.FINISHED) { finalizeUDFClassLoader(); } } @@ -220,7 +220,9 @@ public void start() { finishAnalysisAttemptForCurrentState(); throw e; } finally { - if (failed || !context.isQuery()) { + if (failed + || (!context.isQuery() + && stateMachine.getState() != QueryState.PENDING_RETRY)) { finalizeUDFClassLoader(); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java index 8761346df5e6c..d893302ec29a4 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java @@ -21,6 +21,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; @@ -51,6 +52,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -60,7 +62,8 @@ public class QueryExecutionRetryTest { private ExecutorService stateMachineExecutor; @Before - public void setUp() { + public void setUp() throws Exception { + UDFClassLoaderManager.setupAndGetInstance(); stateMachineExecutor = Executors.newSingleThreadExecutor(); } @@ -70,7 +73,7 @@ public void tearDown() { } @Test - public void testAnalysisAttemptIsRolledBackOnlyWhenRetryStarts() { + public void testAnalysisAttemptIsRolledBackOnlyWhenRetryStarts() throws Exception { RecordingPlanner planner = new RecordingPlanner(false); QueryExecution execution = new QueryExecution(planner, createQueryContext("retry_succeeds"), stateMachineExecutor); @@ -91,6 +94,19 @@ public void testAnalysisAttemptIsRolledBackOnlyWhenRetryStarts() { assertEquals( Arrays.asList("begin", "analyze", "rollback", "begin", "analyze", "commit"), planner.getAnalysisEvents()); + assertUDFQueryScopeClosed(execution); + } + + @Test + public void testNonQueryScopeIsClosedAfterPlanningBeforeAsyncCompletion() throws Exception { + RecordingPlanner planner = new RecordingPlanner(false, false); + QueryExecution execution = + new QueryExecution(planner, createQueryContext("non_query_running"), stateMachineExecutor); + + execution.start(); + + assertEquals(1, planner.getScheduleAttempts()); + assertUDFQueryScopeClosed(execution); } @Test @@ -108,10 +124,11 @@ public void testAnalysisAttemptIsRolledBackWhenRetryLimitIsReached() throws Exce assertEquals(TSStatusCode.DISPATCH_ERROR.getStatusCode(), result.status.code); assertEquals(Arrays.asList("begin", "analyze", "rollback"), planner.getAnalysisEvents()); assertFalse(planner.getAnalysisEvents().contains("commit")); + assertUDFQueryScopeClosed(execution); } @Test - public void testPlanningFailureIsNotSwallowedWhileNotDone() { + public void testPlanningFailureIsNotSwallowedWhileNotDone() throws Exception { RecordingPlanner planner = new RecordingPlanner(false) { @Override @@ -124,10 +141,11 @@ public LogicalQueryPlan doLogicalPlan(IAnalysis analysis, MPPQueryContext contex IllegalStateException thrown = assertThrows(IllegalStateException.class, execution::start); assertEquals("planning failed", thrown.getMessage()); + assertUDFQueryScopeClosed(execution); } @Test - public void testCleanupFailureAfterQueryCanceledIsIgnoredDuringStart() { + public void testCleanupFailureAfterQueryCanceledIsIgnoredDuringStart() throws Exception { QueryExecution[] self = new QueryExecution[1]; RecordingPlanner planner = new RecordingPlanner(false) { @@ -152,6 +170,7 @@ public LogicalQueryPlan doLogicalPlan(IAnalysis analysis, MPPQueryContext contex assertEquals(Arrays.asList("begin", "analyze", "rollback"), planner.getAnalysisEvents()); ExecutionResult result = execution.getStatus(); assertEquals(TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), result.status.code); + assertUDFQueryScopeClosed(execution); } private static MPPQueryContext createQueryContext(String queryId) { @@ -174,6 +193,17 @@ private static void setRetryCount(QueryExecution execution, int retryCount) thro retryCountField.setInt(execution, retryCount); } + private static void assertUDFQueryScopeClosed(QueryExecution execution) throws Exception { + Field queryScopeField = QueryExecution.class.getDeclaredField("udfQueryScope"); + queryScopeField.setAccessible(true); + UDFClassLoaderManager.QueryScope queryScope = + (UDFClassLoaderManager.QueryScope) queryScopeField.get(execution); + assertNotNull(queryScope); + assertThrows( + IllegalStateException.class, + () -> UDFClassLoaderManager.getInstance().withQueryScope(queryScope, () -> null)); + } + private static class RecordingPlanner implements IPlanner { private final IAnalysis analysis = mock(IAnalysis.class); @@ -182,11 +212,17 @@ private static class RecordingPlanner implements IPlanner { private final DistributedQueryPlan distributedPlan = mock(DistributedQueryPlan.class); private final List analysisEvents = new ArrayList<>(); private final boolean alwaysFailDispatch; + private final boolean retryOnFirstDispatch; private int scheduleAttempts; private RecordingPlanner(boolean alwaysFailDispatch) { + this(alwaysFailDispatch, true); + } + + private RecordingPlanner(boolean alwaysFailDispatch, boolean retryOnFirstDispatch) { this.alwaysFailDispatch = alwaysFailDispatch; + this.retryOnFirstDispatch = retryOnFirstDispatch; when(distributedPlan.getInstances()).thenReturn(Collections.emptyList()); } @@ -230,7 +266,7 @@ public IScheduler doSchedule( QueryStateMachine stateMachine) { scheduleAttempts++; stateMachine.transitionToDispatching(); - if (alwaysFailDispatch || scheduleAttempts == 1) { + if (alwaysFailDispatch || (retryOnFirstDispatch && scheduleAttempts == 1)) { TSStatus dispatchFailure = RpcUtils.getStatus(TSStatusCode.DISPATCH_ERROR); stateMachine.transitionToPendingRetry(dispatchFailure); } else { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java index f54fbb111db79..45b63abae2c94 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java @@ -31,6 +31,7 @@ import org.apache.iotdb.commons.queryengine.common.SessionInfo; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; @@ -54,6 +55,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; public class FragmentInstanceSerdeTest { private static final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); @@ -166,4 +169,38 @@ private PlanNode constructPlanNodeTree() throws IllegalPathException { return offsetNode; } + + @Test + public void testUDFQueryScopeCanBeTransferredAndClaimed() throws Exception { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.setupAndGetInstance(); + PlanFragmentId planFragmentId = new PlanFragmentId("scope_transfer", 0); + FragmentInstance fragmentInstance = + new FragmentInstance( + new PlanFragment(planFragmentId, constructPlanNodeTree()), + planFragmentId.genFragmentInstanceId(), + null, + QueryType.READ, + config.getQueryTimeoutThreshold(), + sessionInfo, + false, + false); + ByteBuffer byteBuffer = fragmentInstance.serializeToByteBuffer(); + FragmentInstanceId fragmentInstanceId = FragmentInstanceId.deserialize(byteBuffer); + UDFClassLoaderManager.QueryScope queryScope = + classLoaderManager.openQueryScope(fragmentInstanceId.getFullId()); + + FragmentInstance deserializedFragmentInstance = + classLoaderManager.withQueryScope( + queryScope, + () -> FragmentInstance.deserializeFrom(fragmentInstanceId, byteBuffer)); + deserializedFragmentInstance.attachUDFQueryScope(queryScope); + + assertSame(queryScope, deserializedFragmentInstance.claimUDFQueryScope()); + assertNull(deserializedFragmentInstance.claimUDFQueryScope()); + classLoaderManager.finalizeUDFQuery(queryScope); + assertThrows( + IllegalStateException.class, + () -> classLoaderManager.withQueryScope(queryScope, () -> null)); + } + } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java new file mode 100644 index 0000000000000..08771e1f180e3 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.queryengine.plan.udf; + +import org.apache.iotdb.common.rpc.thrift.Model; +import org.apache.iotdb.commons.udf.UDFInformation; +import org.apache.iotdb.commons.udf.UDFTable; +import org.apache.iotdb.commons.udf.UDFType; +import org.apache.iotdb.commons.udf.service.UDFClassLoader; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; +import org.apache.iotdb.udf.api.relational.ScalarFunction; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import javax.tools.JavaCompiler; +import javax.tools.ToolProvider; + +import java.io.File; +import java.lang.reflect.Field; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.jar.JarEntry; +import java.util.jar.JarOutputStream; +import java.util.stream.Stream; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +public class UDFManagementServiceTest { + + private static final String PACKAGE_NAME = "org.apache.iotdb.test.generated.lazy"; + private static final String FUNCTION_CLASS = PACKAGE_NAME + ".LazyScalarFunction"; + + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void testReflectPinsDefiningClassLoaderForLazyDependency() throws Exception { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.setupAndGetInstance(); + Path jarPath = buildUDFJar(); + UDFClassLoader classLoader = new UDFClassLoader(jarPath.toString()); + Class functionClass = Class.forName(FUNCTION_CLASS, true, classLoader); + assertSame(classLoader, functionClass.getClassLoader()); + + UDFManagementService service = UDFManagementService.getInstance(); + UDFTable udfTable = getUDFTable(service); + String functionName = "lazy_scope_probe"; + UDFInformation information = + new UDFInformation(functionName, FUNCTION_CLASS, UDFType.TABLE_AVAILABLE_SCALAR); + udfTable.addUDFInformation(functionName, information); + udfTable.addFunctionAndClass(Model.TABLE, functionName, functionClass); + + UDFClassLoaderManager.QueryScope queryScope = + classLoaderManager.openQueryScope("lazy-scope-probe"); + try { + classLoaderManager.withQueryScope( + queryScope, + () -> { + ScalarFunction function = service.reflect(functionName, ScalarFunction.class); + classLoader.markAsDeprecated(); + classLoaderManager.finalizeUDFQuery(queryScope); + + assertFalse(isClosed(classLoader)); + assertNotNull(function.analyze(null)); + assertFalse(isClosed(classLoader)); + return null; + }); + assertTrue(isClosed(classLoader)); + } finally { + classLoaderManager.finalizeUDFQuery(queryScope); + udfTable.removeUDFInformation(Model.TABLE, functionName); + udfTable.removeFunctionClass(Model.TABLE, functionName); + } + } + + private Path buildUDFJar() throws Exception { + Path sourceRoot = temporaryFolder.newFolder("source").toPath(); + Path sourceDir = sourceRoot.resolve(PACKAGE_NAME.replace('.', File.separatorChar)); + Files.createDirectories(sourceDir); + + Path functionSource = sourceDir.resolve("LazyScalarFunction.java"); + Files.writeString( + functionSource, + "package " + + PACKAGE_NAME + + ";\n" + + "import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;\n" + + "import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;\n" + + "import org.apache.iotdb.udf.api.relational.ScalarFunction;\n" + + "import org.apache.iotdb.udf.api.relational.access.Record;\n" + + "import org.apache.iotdb.udf.api.type.Type;\n" + + "public class LazyScalarFunction implements ScalarFunction {\n" + + " public ScalarFunctionAnalysis analyze(FunctionArguments arguments) {\n" + + " Class.forName(\"" + + PACKAGE_NAME + + ".LazyHelper\", true, LazyScalarFunction.class.getClassLoader());\n" + + " return new ScalarFunctionAnalysis.Builder().outputDataType(Type.INT32).build();\n" + + " }\n" + + " public Object evaluate(Record input) { return null; }\n" + + "}\n", + StandardCharsets.UTF_8); + + Path helperSource = sourceDir.resolve("LazyHelper.java"); + Files.writeString( + helperSource, + "package " + PACKAGE_NAME + ";\npublic class LazyHelper {}\n", + StandardCharsets.UTF_8); + + Path classesDir = temporaryFolder.newFolder("classes").toPath(); + JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); + assertNotNull(compiler); + assertEquals( + 0, + compiler.run( + null, + null, + null, + "-classpath", + System.getProperty("java.class.path"), + "-d", + classesDir.toString(), + functionSource.toString(), + helperSource.toString())); + + Path jarPath = temporaryFolder.newFile("lazy-udf.jar").toPath(); + try (JarOutputStream outputStream = + new JarOutputStream(Files.newOutputStream(jarPath)); + Stream paths = Files.walk(classesDir)) { + paths + .filter(Files::isRegularFile) + .forEach( + path -> { + try { + String entryName = + classesDir.relativize(path).toString().replace(File.separatorChar, '/'); + outputStream.putNextEntry(new JarEntry(entryName)); + Files.copy(path, outputStream); + outputStream.closeEntry(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + return jarPath; + } + + private static UDFTable getUDFTable(UDFManagementService service) throws Exception { + Field field = UDFManagementService.class.getDeclaredField("udfTable"); + field.setAccessible(true); + return (UDFTable) field.get(service); + } + + private static boolean isClosed(UDFClassLoader classLoader) throws Exception { + Field field = UDFClassLoader.class.getDeclaredField("closed"); + field.setAccessible(true); + return field.getBoolean(classLoader); + } +} From 358408c8187b65caa12212435b562c6c2e11825c Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 15:05:19 +0800 Subject: [PATCH 05/10] spotless --- .../iotdb/db/queryengine/plan/execution/QueryExecution.java | 4 +--- .../queryengine/plan/planner/FragmentInstanceSerdeTest.java | 4 +--- .../queryengine/plan/udf/UDFManagementServiceTest.java | 3 +-- 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 6090f0ad703ae..838b4fa2dfe79 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -220,9 +220,7 @@ public void start() { finishAnalysisAttemptForCurrentState(); throw e; } finally { - if (failed - || (!context.isQuery() - && stateMachine.getState() != QueryState.PENDING_RETRY)) { + if (failed || (!context.isQuery() && stateMachine.getState() != QueryState.PENDING_RETRY)) { finalizeUDFClassLoader(); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java index 45b63abae2c94..bacb2f0c240bf 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/FragmentInstanceSerdeTest.java @@ -191,8 +191,7 @@ public void testUDFQueryScopeCanBeTransferredAndClaimed() throws Exception { FragmentInstance deserializedFragmentInstance = classLoaderManager.withQueryScope( - queryScope, - () -> FragmentInstance.deserializeFrom(fragmentInstanceId, byteBuffer)); + queryScope, () -> FragmentInstance.deserializeFrom(fragmentInstanceId, byteBuffer)); deserializedFragmentInstance.attachUDFQueryScope(queryScope); assertSame(queryScope, deserializedFragmentInstance.claimUDFQueryScope()); @@ -202,5 +201,4 @@ public void testUDFQueryScopeCanBeTransferredAndClaimed() throws Exception { IllegalStateException.class, () -> classLoaderManager.withQueryScope(queryScope, () -> null)); } - } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index 08771e1f180e3..90e02de31e154 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -145,8 +145,7 @@ private Path buildUDFJar() throws Exception { helperSource.toString())); Path jarPath = temporaryFolder.newFile("lazy-udf.jar").toPath(); - try (JarOutputStream outputStream = - new JarOutputStream(Files.newOutputStream(jarPath)); + try (JarOutputStream outputStream = new JarOutputStream(Files.newOutputStream(jarPath)); Stream paths = Files.walk(classesDir)) { paths .filter(Files::isRegularFile) From 437d6e9258c0cfa417befb785fbafc8298892498 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 15:45:42 +0800 Subject: [PATCH 06/10] Fix checked exceptions in UDF scope tests --- .../plan/udf/UDFManagementServiceTest.java | 23 +++++++++++++++---- .../service/UDFClassLoaderManagerTest.java | 12 ++++++---- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index 90e02de31e154..303e90349d8f5 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -35,6 +35,7 @@ import javax.tools.ToolProvider; import java.io.File; +import java.io.IOException; import java.lang.reflect.Field; import java.nio.charset.StandardCharsets; import java.nio.file.Files; @@ -79,7 +80,7 @@ public void testReflectPinsDefiningClassLoaderForLazyDependency() throws Excepti queryScope, () -> { ScalarFunction function = service.reflect(functionName, ScalarFunction.class); - classLoader.markAsDeprecated(); + markAsDeprecated(classLoader); classLoaderManager.finalizeUDFQuery(queryScope); assertFalse(isClosed(classLoader)); @@ -171,9 +172,21 @@ private static UDFTable getUDFTable(UDFManagementService service) throws Excepti return (UDFTable) field.get(service); } - private static boolean isClosed(UDFClassLoader classLoader) throws Exception { - Field field = UDFClassLoader.class.getDeclaredField("closed"); - field.setAccessible(true); - return field.getBoolean(classLoader); + private static void markAsDeprecated(UDFClassLoader classLoader) { + try { + classLoader.markAsDeprecated(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + private static boolean isClosed(UDFClassLoader classLoader) { + try { + Field field = UDFClassLoader.class.getDeclaredField("closed"); + field.setAccessible(true); + return field.getBoolean(classLoader); + } catch (ReflectiveOperationException e) { + throw new IllegalStateException(e); + } } } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java index d2bb5e9adbf2c..6262e12484980 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java @@ -83,9 +83,13 @@ private static Set getPinnedClassLoaders( return (Set) field.get(queryScope); } - private static boolean isClosed(UDFClassLoader classLoader) throws Exception { - Field field = UDFClassLoader.class.getDeclaredField("closed"); - field.setAccessible(true); - return field.getBoolean(classLoader); + private static boolean isClosed(UDFClassLoader classLoader) { + try { + Field field = UDFClassLoader.class.getDeclaredField("closed"); + field.setAccessible(true); + return field.getBoolean(classLoader); + } catch (ReflectiveOperationException e) { + throw new IllegalStateException(e); + } } } From a4b1455a8478bb5b04bc2b751776f1aa490b6074 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 16:32:22 +0800 Subject: [PATCH 07/10] Fix UDF metadata function type classification --- .../it/udf/IoTDBSQLFunctionManagementIT.java | 19 +++++++++++++++++++ .../config/metadata/ShowFunctionsTask.java | 16 +++++++--------- .../db/queryengine/plan/udf/TreeUDFUtils.java | 17 +++++------------ .../plan/udf/UDFManagementService.java | 12 ++++++++++++ .../plan/udf/UDFManagementServiceTest.java | 13 +++++++++++++ 5 files changed, 56 insertions(+), 21 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java index f56da20b05176..ff5cee4a3dd24 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java @@ -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; @@ -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_name = '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(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowFunctionsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowFunctionsTask.java index 9866b77cf13ca..3369b6233250b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowFunctionsTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowFunctionsTask.java @@ -19,12 +19,12 @@ package org.apache.iotdb.db.queryengine.plan.execution.config.metadata; +import org.apache.iotdb.common.rpc.thrift.FunctionType; import org.apache.iotdb.common.rpc.thrift.Model; import org.apache.iotdb.commons.queryengine.plan.relational.function.TableBuiltinTableFunction; import org.apache.iotdb.commons.queryengine.plan.udf.BuiltinAggregationFunction; import org.apache.iotdb.commons.queryengine.plan.udf.BuiltinScalarFunction; import org.apache.iotdb.commons.queryengine.plan.udf.BuiltinTimeSeriesGeneratingFunction; -import org.apache.iotdb.commons.queryengine.plan.udf.TableUDFUtils; import org.apache.iotdb.commons.schema.column.ColumnHeader; import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; import org.apache.iotdb.commons.udf.UDFInformation; @@ -272,14 +272,12 @@ public static Binary getFunctionType(UDFInformation udfInformation) { } else if (TreeUDFUtils.isUDAF(udfInformation.getFunctionName())) { return BINARY_MAP.get(FUNCTION_TYPE_EXTERNAL_UDAF); } - } else { - if (TableUDFUtils.isScalarFunction(udfInformation.getFunctionName())) { - return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC); - } else if (TableUDFUtils.isAggregateFunction(udfInformation.getFunctionName())) { - return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_AGG_FUNC); - } else if (TableUDFUtils.isTableFunction(udfInformation.getFunctionName())) { - return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_TABLE_FUNC); - } + } else if (FunctionType.SCALAR.equals(type.getType())) { + return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC); + } else if (FunctionType.AGGREGATE.equals(type.getType())) { + return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_AGG_FUNC); + } else if (FunctionType.TABLE.equals(type.getType())) { + return BINARY_MAP.get(FUNCTION_TYPE_USER_DEFINED_TABLE_FUNC); } } return BINARY_MAP.get(FUNCTION_TYPE_UNKNOWN); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/TreeUDFUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/TreeUDFUtils.java index fcfe39cd1b37b..95918ecbaff97 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/TreeUDFUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/TreeUDFUtils.java @@ -19,26 +19,19 @@ package org.apache.iotdb.db.queryengine.plan.udf; +import org.apache.iotdb.common.rpc.thrift.Model; import org.apache.iotdb.commons.queryengine.plan.udf.UDFManagementService; import org.apache.iotdb.udf.api.UDAF; import org.apache.iotdb.udf.api.UDTF; public class TreeUDFUtils { public static boolean isUDTF(String functionName) { - try { - UDFManagementService.getInstance().reflect(functionName, UDTF.class); - return true; - } catch (Throwable e) { - return false; - } + return UDFManagementService.getInstance() + .isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDTF.class); } public static boolean isUDAF(String functionName) { - try { - UDFManagementService.getInstance().reflect(functionName, UDAF.class); - return true; - } catch (Throwable e) { - return false; - } + return UDFManagementService.getInstance() + .isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDAF.class); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java index d0e59ccb0c7f7..f36adb004c6fb 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementService.java @@ -305,6 +305,18 @@ public T reflect(String functionName, Class clazz) { } } + /** Check a registered function class without instantiating it or acquiring a query scope. */ + public boolean isRegisteredFunctionAssignableTo( + Model model, String functionName, Class expectedType) { + try { + acquireReadLock(); + Class functionClass = udfTable.getFunctionClass(model, functionName); + return functionClass != null && expectedType.isAssignableFrom(functionClass); + } finally { + releaseReadLock(); + } + } + public UDFInformation[] getUDFInformation(Model model) { return udfTable.getUDFInformationList(model).toArray(new UDFInformation[0]); } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index 303e90349d8f5..d29d148288cdf 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -25,6 +25,8 @@ import org.apache.iotdb.commons.udf.UDFType; import org.apache.iotdb.commons.udf.service.UDFClassLoader; import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; +import org.apache.iotdb.udf.api.UDAF; +import org.apache.iotdb.udf.api.UDTF; import org.apache.iotdb.udf.api.relational.ScalarFunction; import org.junit.Rule; @@ -57,6 +59,17 @@ public class UDFManagementServiceTest { @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + @Test + public void testRegisteredFunctionTypeCheckDoesNotRequireQueryScope() { + UDFManagementService service = UDFManagementService.getInstance(); + String functionName = BuiltinTimeSeriesGeneratingFunction.values()[0].getFunctionName(); + + assertTrue( + service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDTF.class)); + assertFalse( + service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDAF.class)); + } + @Test public void testReflectPinsDefiningClassLoaderForLazyDependency() throws Exception { UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.setupAndGetInstance(); From b27cace414a09e7bc36d6c902bf93d7a7ba61152 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 09:41:20 +0800 Subject: [PATCH 08/10] spotless --- .../queryengine/plan/udf/UDFManagementServiceTest.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index d29d148288cdf..871dd914d7f2d 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -64,10 +64,8 @@ public void testRegisteredFunctionTypeCheckDoesNotRequireQueryScope() { UDFManagementService service = UDFManagementService.getInstance(); String functionName = BuiltinTimeSeriesGeneratingFunction.values()[0].getFunctionName(); - assertTrue( - service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDTF.class)); - assertFalse( - service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDAF.class)); + assertTrue(service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDTF.class)); + assertFalse(service.isRegisteredFunctionAssignableTo(Model.TREE, functionName, UDAF.class)); } @Test From 3ccfa984f429030390e7b3e080130cdcffc020bd Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 10:21:41 +0800 Subject: [PATCH 09/10] Fix UDF lifecycle test validation --- .../db/it/udf/IoTDBSQLFunctionManagementIT.java | 2 +- .../plan/udf/UDFManagementServiceTest.java | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java index ff5cee4a3dd24..333262b869809 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java @@ -144,7 +144,7 @@ public void testInformationSchemaFunctionsShowsExternalScalarType() { "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_name = 'UDSF'")) { + "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()); diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index 871dd914d7f2d..cbc82c4bb27ef 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -36,6 +36,7 @@ import javax.tools.JavaCompiler; import javax.tools.ToolProvider; +import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; import java.lang.reflect.Field; @@ -143,18 +144,21 @@ private Path buildUDFJar() throws Exception { Path classesDir = temporaryFolder.newFolder("classes").toPath(); JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); assertNotNull(compiler); - assertEquals( - 0, + String testClassPath = + System.getProperty("surefire.test.class.path", System.getProperty("java.class.path")); + ByteArrayOutputStream compilerOutput = new ByteArrayOutputStream(); + int compilerExitCode = compiler.run( null, - null, - null, + compilerOutput, + compilerOutput, "-classpath", - System.getProperty("java.class.path"), + testClassPath, "-d", classesDir.toString(), functionSource.toString(), - helperSource.toString())); + helperSource.toString()); + assertEquals(compilerOutput.toString(StandardCharsets.UTF_8), 0, compilerExitCode); Path jarPath = temporaryFolder.newFile("lazy-udf.jar").toPath(); try (JarOutputStream outputStream = new JarOutputStream(Files.newOutputStream(jarPath)); From 3f6dd878127b0a2ad5bf27e05c0dc0916f27175f Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 10:40:08 +0800 Subject: [PATCH 10/10] Fix generated UDF test compilation --- .../queryengine/plan/udf/UDFManagementServiceTest.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java index cbc82c4bb27ef..be0d98bfb4131 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -126,9 +126,13 @@ private Path buildUDFJar() throws Exception { + "import org.apache.iotdb.udf.api.type.Type;\n" + "public class LazyScalarFunction implements ScalarFunction {\n" + " public ScalarFunctionAnalysis analyze(FunctionArguments arguments) {\n" - + " Class.forName(\"" + + " try {\n" + + " Class.forName(\"" + PACKAGE_NAME + ".LazyHelper\", true, LazyScalarFunction.class.getClassLoader());\n" + + " } catch (ClassNotFoundException e) {\n" + + " throw new RuntimeException(e);\n" + + " }\n" + " return new ScalarFunctionAnalysis.Builder().outputDataType(Type.INT32).build();\n" + " }\n" + " public Object evaluate(Record input) { return null; }\n"