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..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 @@ -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_table = 'UDSF'")) { + assertTrue(resultSet.next()); + assertEquals(FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC, resultSet.getString(1)); + assertFalse(resultSet.next()); + } + } catch (SQLException throwable) { + fail(throwable.getMessage()); + } + } + @Test public void testCreateShowDropAggregateFunction() { try (Connection connection = EnvFactory.getEnv().getTableConnection(); 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 a060bd5b91b51..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 @@ -34,6 +34,7 @@ import org.apache.iotdb.commons.path.IFullPath; import org.apache.iotdb.commons.queryengine.common.SessionInfo; import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; @@ -116,6 +117,7 @@ public class FragmentInstanceContext extends QueryContext { private static final long LONG_WAIT_DURATION = 5_000_000_000L; private static final int MODS_MEMORY_ESTIMATE_READ_INTERVAL = 10_000; private final FragmentInstanceId id; + private UDFClassLoaderManager.QueryScope udfQueryScope; private final FragmentInstanceStateMachine stateMachine; @@ -1348,6 +1350,7 @@ private void releaseTVListOwnedByQuery() { * be decreased. */ public synchronized void releaseResource() { + finalizeUDFQuery(); try { getQueryChunkReaders().close(); } catch (IOException e) { @@ -1493,6 +1496,29 @@ private void releaseDataNodeQueryContext() { } } + synchronized void adoptUDFQueryScope(UDFClassLoaderManager.QueryScope transferredScope) { + if (udfQueryScope == null) { + udfQueryScope = transferredScope; + } + } + + public synchronized UDFClassLoaderManager.QueryScope getOrCreateUDFQueryScope() { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (udfQueryScope == null && classLoaderManager != null) { + udfQueryScope = classLoaderManager.openQueryScope(id.getFullId()); + } + return udfQueryScope; + } + + synchronized void finalizeUDFQuery() { + UDFClassLoaderManager.QueryScope queryScope = udfQueryScope; + udfQueryScope = null; + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (queryScope != null && classLoaderManager != null) { + classLoaderManager.finalizeUDFQuery(queryScope); + } + } + public void setMayHaveTmpFile(boolean mayHaveTmpFile) { this.mayHaveTmpFile = mayHaveTmpFile; } 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..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 @@ -164,16 +164,19 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( dataNodeQueryContexts[0] = getOrCreateDataNodeQueryContext( instanceId.getQueryId(), instance.getDataNodeFINum()); - return createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - dataRegion, - instance.getGlobalTimePredicate(), - dataNodeQueryContextMap, - DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), - instance.isDebug(), - instance.isVerbose()); + FragmentInstanceContext createdContext = + createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + dataRegion, + instance.getGlobalTimePredicate(), + dataNodeQueryContextMap, + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), + instance.isDebug(), + instance.isVerbose()); + createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); + return createdContext; }); rejectIfRepeatedDispatch(contextCreated[0], instanceId); context.setHighestPriority(instance.isHighestPriority()); @@ -185,7 +188,6 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( instance.getFragment().getTypeProvider(), context, dataNodeQueryContexts[0]); - List drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); // For ShowQueries related instances, isHighestPriority == true @@ -235,6 +237,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( // still being created. Capture the failure result from the local context before // returning from computeIfAbsent instead of looking it up from the map later. failedInstanceInfo.set(context.getInstanceInfo()); + context.finalizeUDFQuery(); clearFIRelatedResources(instanceId); return null; } @@ -318,12 +321,15 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( instanceId, fragmentInstanceId -> { contextCreated[0] = true; - return createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - instance.isDebug(), - instance.isVerbose()); + FragmentInstanceContext createdContext = + createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + instance.isDebug(), + instance.isVerbose()); + createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); + return createdContext; }); rejectIfRepeatedDispatch(contextCreated[0], instanceId); context.setHighestPriority(instance.isHighestPriority()); @@ -335,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()); @@ -370,6 +375,7 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( // See execDataQueryFragmentInstance for why this result must not be fetched from // instanceContext after computeIfAbsent returns. failedInstanceInfo.set(context.getInstanceInfo()); + context.finalizeUDFQuery(); clearFIRelatedResources(instanceId); return null; } 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..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().acquireLock(); - try { - // This statement must be surrounded by the registration lock. - UDFClassLoaderManager.getInstance().initializeUDFQuery(udtfQueryId); - // UDF executors will be initialized at the same time - transformers = - new EvaluationDAGBuilder( - udtfQueryId, - inputLayer, - inputLocations, - outputExpressions, - expressionTypes, - udtfContext, - udfTransformerMemoryBudgetInMB + udfCollectorMemoryBudgetInMB) - .buildLayerMemoryAssigner() - .bindInputLayerColumnIndexWithExpression() - .buildResultColumnPointReaders() - .getOutputReaders(); - } finally { - UDFManagementService.getInstance().releaseLock(); - } + // UDF executors will be initialized at the same time + transformers = + new EvaluationDAGBuilder( + udtfQueryId, + inputLayer, + inputLocations, + outputExpressions, + expressionTypes, + udtfContext, + udfTransformerMemoryBudgetInMB + udfCollectorMemoryBudgetInMB) + .buildLayerMemoryAssigner() + .bindInputLayerColumnIndexWithExpression() + .buildResultColumnPointReaders() + .getOutputReaders(); } protected YieldableState iterateAllColumnsToNextValid() throws Exception { 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..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 @@ -26,6 +26,7 @@ import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.queryengine.common.SqlDialect; 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; @@ -109,6 +110,7 @@ public class QueryExecution implements IQueryExecution { private final AtomicBoolean resultHandleCleanUp; private final AtomicBoolean stopped; + private final UDFClassLoaderManager.QueryScope udfQueryScope; // cost time in ns of finished rpc private long totalExecutionTime = 0; @@ -129,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; @@ -164,6 +171,9 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService // preventing a premature close raced with late-scheduled FragmentInstances. context.releaseExternalTsFileQueryResources(); this.stop(cause); + if (context.isQuery() || state != QueryState.FINISHED) { + finalizeUDFClassLoader(); + } } }); this.stopped = new AtomicBoolean(false); @@ -177,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 @@ -194,12 +216,17 @@ public void start() { } throw e; } catch (Error e) { + failed = true; finishAnalysisAttemptForCurrentState(); throw e; + } finally { + if (failed || (!context.isQuery() && stateMachine.getState() != QueryState.PENDING_RETRY)) { + finalizeUDFClassLoader(); + } } } - private void startInternal() { + private void startInternalWithPlanCleanup() { final long startTime = System.nanoTime(); if (skipExecute()) { LOGGER.debug(DataNodeQueryMessages.SKIP_EXECUTE); @@ -268,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(); @@ -296,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(); @@ -331,8 +364,17 @@ private void constructResultForMemorySource() { } // Analyze the statement in QueryContext. Generate the analysis this query need - private IAnalysis analyze(MPPQueryContext context) { - return planner.analyze(context); + private IAnalysis analyzeWithCleanupOnFailure(MPPQueryContext context) { + try { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager == null || udfQueryScope == null) { + return planner.analyze(context); + } + return classLoaderManager.withQueryScope(udfQueryScope, () -> planner.analyze(context)); + } catch (RuntimeException | Error e) { + finalizeUDFClassLoader(); + throw e; + } } private void schedule() { @@ -434,6 +476,11 @@ 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) { context.recordDeviceEntryDiskIOMetricsOnRelease(); @@ -466,6 +513,13 @@ private void releaseResource(Throwable t) { } } + private void finalizeUDFClassLoader() { + UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance(); + if (classLoaderManager != null && udfQueryScope != null) { + classLoaderManager.finalizeUDFQuery(udfQueryScope); + } + } + /** * 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/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/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/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/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/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 b387baddff6b2..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 @@ -31,8 +31,10 @@ 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; 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; @@ -53,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(); @@ -86,7 +90,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. @@ -162,4 +169,36 @@ 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/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 fd5dc23d77dfa..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 @@ -42,17 +42,22 @@ import java.io.IOException; import java.lang.reflect.InvocationTargetException; import java.nio.ByteBuffer; -import java.util.concurrent.locks.ReentrantLock; +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 { 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 +74,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(); } } @@ -170,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 @@ -193,20 +225,18 @@ 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(); } } 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 +249,7 @@ public void deregister(Model model, String functionName, boolean needToDeleteJar .removeFileUnderTemporaryRoot(functionName.toUpperCase() + ".txt"); } } finally { - releaseLock(); + releaseWriteLock(); } } @@ -232,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 @@ -258,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/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/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..be0d98bfb4131 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/queryengine/plan/udf/UDFManagementServiceTest.java @@ -0,0 +1,211 @@ +/* + * 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.UDAF; +import org.apache.iotdb.udf.api.UDTF; +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.ByteArrayOutputStream; +import java.io.File; +import java.io.IOException; +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 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(); + 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); + markAsDeprecated(classLoader); + 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" + + " 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" + + "}\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); + String testClassPath = + System.getProperty("surefire.test.class.path", System.getProperty("java.class.path")); + ByteArrayOutputStream compilerOutput = new ByteArrayOutputStream(); + int compilerExitCode = + compiler.run( + null, + compilerOutput, + compilerOutput, + "-classpath", + testClassPath, + "-d", + classesDir.toString(), + functionSource.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)); + 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 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 new file mode 100644 index 0000000000000..6262e12484980 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/udf/service/UDFClassLoaderManagerTest.java @@ -0,0 +1,95 @@ +/* + * 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) { + try { + Field field = UDFClassLoader.class.getDeclaredField("closed"); + field.setAccessible(true); + return field.getBoolean(classLoader); + } catch (ReflectiveOperationException e) { + throw new IllegalStateException(e); + } + } +}