Skip to content

Protect UDF classloader lifetime with query scopes - #19107

Open
shuwenwei wants to merge 10 commits into
masterfrom
fix/table-model-udf-classloader-lifecycle
Open

shuwenwei wants to merge 10 commits into
masterfrom
fix/table-model-udf-classloader-lifecycle

Conversation

@shuwenwei

@shuwenwei shuwenwei commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

Summary

Protect custom UDF classloaders by pinning the exact loader that defined the reflected UDF class, for the lifetime of an explicit query or fragment scope.

The previous implementation pinned the active loader at scope entry and broadcast every loader generation to all active queries. That could pin unrelated loaders and still miss classes restored on older generations. This change moves pinning to UDFManagementService.reflect() and keeps the loader pinned only until the owning query or fragment scope closes.

Implementation

  • Add explicit UDFClassLoaderManager.QueryScope handles with active-action tracking and deferred close.
  • Pin functionClass.getClassLoader() inside the UDF registry read lock during reflect().
  • Keep a scope usable across retries, but reject any attempt to bind a finalized scope.
  • Publish a complete tree/table catalog snapshot before switching the active loader, then deprecate the old loader.
  • Bind query scopes across analysis, logical/distributed planning, and scheduling; non-query scopes close after planning.
  • Bind fragment scopes around LocalExecutionPlanner.plan() and transfer deserialization ownership into FragmentInstanceContext.
  • Apply scope handling to RPC and consensus fragment deserialization, including cleanup when fragment construction does not claim ownership.
  • Remove TransformOperator-local classloader pin management.
  • Classify tree UDF metadata through the registered Class hierarchy without instantiating the UDF or requiring a QueryScope.
  • Classify table UDF metadata directly from UDFInformation function type.
  • Add localized errors for missing query scopes and closed scopes.

Tests

Added regression coverage for:

  • finalize requests waiting for all nested actions;
  • finalized query scopes rejecting a second bind;
  • reflect() pinning the defining UDF classloader while a helper class is loaded lazily;
  • non-query, error, cancel, and retry query scope cleanup;
  • fragment scope handle transfer and no implicit scope from plain deserialization;
  • UDF type classification without a QueryScope;
  • information_schema.functions reporting an external table scalar UDF correctly.

Local validation is limited to static inspection and git diff --check; Maven tests were not run locally.

@codecov

codecov Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 52.33918% with 163 lines in your changes missing coverage. Please review.
✅ Project coverage is 48.53%. Comparing base (e6b44b5) to head (3f6dd87).
⚠️ Report is 35 commits behind head on master.

Files with missing lines Patch % Lines
...ons/queryengine/plan/udf/UDFManagementService.java 38.33% 37 Missing ⚠️
...tdb/commons/udf/service/UDFClassLoaderManager.java 59.09% 36 Missing ⚠️
...ol/thrift/impl/DataNodeInternalRPCServiceImpl.java 0.00% 33 Missing ⚠️
...db/db/consensus/statemachine/BaseStateMachine.java 0.00% 16 Missing ⚠️
...in/java/org/apache/iotdb/commons/udf/UDFTable.java 64.00% 9 Missing ⚠️
.../db/queryengine/plan/execution/QueryExecution.java 82.50% 7 Missing ⚠️
...n/execution/config/metadata/ShowFunctionsTask.java 0.00% 6 Missing ⚠️
...ueryengine/plan/planner/LocalExecutionPlanner.java 50.00% 6 Missing ⚠️
...ueryengine/plan/planner/plan/FragmentInstance.java 54.54% 5 Missing ⚠️
...tatemachine/dataregion/DataRegionStateMachine.java 0.00% 2 Missing ⚠️
... and 4 more
Additional details and impacted files
@@             Coverage Diff              @@
##             master   #19107      +/-   ##
============================================
+ Coverage     48.13%   48.53%   +0.40%     
- Complexity      625      635      +10     
============================================
  Files          5572     5573       +1     
  Lines        411555   412445     +890     
  Branches      54193    54373     +180     
============================================
+ Hits         198084   200177    +2093     
+ Misses       213471   212268    -1203     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@JackieTien97 JackieTien97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for tackling this. The UDF classloader does need to stay alive for the whole table-model query lifecycle, but I have a few concerns with the current approach (roughly in order of severity; details inline):

  1. Possible deadlock between the new UDF read lock and instanceExecution.computeIfAbsent(...) while a CREATE/DROP FUNCTION is waiting for the write lock (FragmentInstanceManager L187).
  2. The pin targets activeClassLoader, not the loader that defined the table UDF class. doRegister only reloads classes of the model being registered. With both tree and table UDFs (including right after a DataNode restart), table UDF classes can sit on an already-closed loader, and this PR does not protect them (QueryExecution L347).
  3. Pin leak for non-query table statements that finish during analysis and never reach a terminal state, e.g. LOAD ... WITH ('async'='true') (QueryExecution L171).
  4. The read lock now spans the entire table-model analysis, including every table insert/LOAD/DELETE/UPDATE. One CREATE FUNCTION can stall all table-model reads and writes on a DataNode behind the slowest in-flight analysis (QueryExecution L344).

A few other paths still instantiate table UDFs without a pin (schema FIs, FI deserialization), and there are some minor points inline.

On the overall approach: pinning and locking at each entry point is easy to miss for the next caller. Schema FIs, plan deserialization, schema-region UPDATE, and tree-model analysis are still uncovered, and the global read/write lock is what introduces the stall and the deadlock above. It may be simpler and more robust to fix this inside the loader management itself, for example:

  • when doRegister rotates the active loader, reload the registered classes of all models (and have UDFClassLoaderManager.start() deprecate the loader it replaces), so UDFTable never references a deprecated loader;
  • make UDFClassLoaderManager#initializeUDFQuery atomic on its own: read activeClassLoader once and tryAcquire() it with a CAS that fails once the loader is closed, retrying on failure. Query paths then don't need the global lock at all;
  • or treat the classes still referenced by UDFTable as pins, so a loader is never closed while it still defines registered functions.

Tests: there is no UT/IT for the new lifecycle. Two tests would catch issues 2 and 3:

  • a UT asserting that queryIdToUDFClassLoaderMap is empty after an async LOAD or a no-match UPDATE;
  • an IT that runs CREATE FUNCTION for a tree UDF and then a table UDF that lazily loads a helper class.

The failing CI jobs (IoTDBMQTTServiceIT hanging on Windows, IoTDBSessionSimpleIT.testWriteRestartAndDeleteDB, and the PageCacheDeletionBufferTest timeout) look unrelated at first glance; a re-run would help confirm.

List<PipelineDriverFactory> driverFactories;
if (instance.getSessionInfo() != null
&& instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) {
UDFManagementService.getInstance().acquireReadLock();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This blocks on the UDF read lock inside the instanceExecution.computeIfAbsent(...) mapping function, i.e. while holding that CHM bin. Together with the read lock that QueryExecution#analyze now holds for the entire table analysis, I think this can deadlock:

  1. Thread A analyzes a table-model insert with new devices. It holds the read lock (QueryExecution L344). TableDeviceSchemaValidator → TableDeviceSchemaFetcher#fetchMissingDeviceSchemaForDataInsertion runs a nested FetchDevice query with timeout Long.MAX_VALUE, whose local schema FI goes through execSchemaQueryFragmentInstance → instanceExecution.computeIfAbsent(...).
  2. A CREATE FUNCTION arrives, and register() → acquireWriteLock() waits for A.
  3. Thread T dispatches a table data FI, enters this computeIfAbsent lambda (holding bin B), and blocks here. The non-fair ReentrantReadWriteLock blocks new readers once a writer is queued.
  4. If A's schema FI lands in bin B, A blocks on that bin.

That makes the cycle A → T → writer → A. None of lock(), synchronized, or the nested fetch time out, so it never resolves. Every later table analysis or FI planning on that DataNode (and tree UDF queries via TransformOperator) then queues behind the writer, and the ConfigNode keeps holding its udfInfo lock until its DataNode requests time out. The more table FIs arrive while the writer is queued, the more bins are held, so a collision gets more likely under load.

Could we take the read lock before entering computeIfAbsent instead of blocking inside it? It would also help not to hold the read lock across the whole analysis (see the comment on QueryExecution L344).

UDFManagementService.getInstance().acquireReadLock();
try {
synchronized (tableUdfClassLoaderLock) {
UDFClassLoaderManager.getInstance().initializeUDFQuery(queryId);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This pins activeClassLoader, which is not necessarily the loader that defined the table UDF classes stored in UDFTable. UDFManagementService#doRegister rotates the global active loader on every registration, but updateAllRegisteredClasses(model, ...) only reloads classes of the model being registered. The previous loader is deprecated and, if nothing pins it, closed right away in markAsDeprecated().

Example:

  1. CREATE FUNCTION my_scalar (table): the active loader becomes L1, and table classes are loaded by L1.
  2. CREATE FUNCTION my_udtf (tree): the active loader becomes L2, and only tree classes are reloaded. L1 is deprecated and closed immediately.
  3. SELECT my_scalar(s1) FROM t pins L2, here and in FragmentInstanceContext#initializeTableModelUDFQuery, but my_scalar's class still comes from the closed L1. The first lazy load of a helper class, resource, or ServiceLoader entry fails with NoClassDefFoundError/ClassNotFoundException.

The same happens after a DataNode restart when both models have UDFs. prepareUDFResources() calls doRegister once per UDF, rotating each time, and UDFClassLoaderManager.start() then replaces the active loader again. So the failure this PR targets can still happen in mixed tree/table deployments. Reloading all models' classes on rotation, or pinning clazz.getClassLoader() instead of the active loader, would close this gap.

// preventing a premature close raced with late-scheduled FragmentInstances.
context.releaseExternalTsFileQueryResources();
this.stop(cause);
finalizeTableUDFClassLoader();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For non-query statements this listener is the only place that releases the analysis pin. releaseResource() is reached through Coordinator#cleanupQueryExecution, but Coordinator#execution only puts isQuery() executions into queryExecutionMap. So a non-query table statement that never reaches a terminal state leaks its pin:

  • LOAD '/x.tsfile' WITH ('async'='true') in the table model: visitLoadTsFile sets QueryType.OTHER, and a successful doAsyncLoad sets finishQueryAfterAnalyze without a fail status. startInternal() then calls constructResultForMemorySource() and transitionToRunning() and returns, so the state stays RUNNING forever.
  • UPDATE t SET a = 'v' WHERE <predicate that matches no device>: QueryType.WRITE, and parseWhere returns false so the analysis is marked finished.
  • Any non-query statement whose start() throws from doLogicalPlan/doDistributedPlan/schedule().

Each of these leaves an entry in queryIdToUDFClassLoaderMap and a permanent activeQueriesCount on the loader. After the next CREATE FUNCTION, the deprecated loader is never closed: jar handles and metaspace leak, and on Windows the jar can't be deleted. Could the pin be released at the end of start() for non-query statements, or could only QueryType.READ statements be pinned?

return planner.analyze(context);
}
String queryId = context.getQueryId().getId();
UDFManagementService.getInstance().acquireReadLock();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This holds the UDF read lock for the whole planner.analyze(context), and TableModelPlanner analysis covers every table-model statement, including inserts, LOAD, DELETE, and UPDATE. That analysis can block for a long time:

  • the insert-time FetchDevice uses Long.MAX_VALUE as its timeout;
  • schema auto-creation goes through the ConfigNode;
  • LOAD analyzes whole TsFiles.

Once register()/deregister() is queued for the write lock, the non-fair ReentrantReadWriteLock blocks every new reader.

So one CREATE FUNCTION issued while a slow LOAD or insert is in analysis blocks all new table-model inserts and queries on that DataNode until that analysis finishes. That covers this line and FI planning in FragmentInstanceManager. Meanwhile UDFManager#createFunction on the ConfigNode keeps holding the udfInfo lock while it waits, which blocks other CREATE/DROP FUNCTION cluster-wide. If the ConfigNode-side request times out, the DataNode may still finish the registration afterwards. Before this PR, table writes were not affected by UDF DDL at all.

Could the lock be narrowed to resolving UDF classes and taking the pin? Or could it be removed from the query path with an atomic initializeUDFQuery (see the review summary)?

dataNodeQueryContexts[0]);
List<PipelineDriverFactory> driverFactories;
if (instance.getSessionInfo() != null
&& instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only the data-query path is covered. Table-model schema FIs also instantiate scalar UDFs, but they go through execSchemaQueryFragmentInstance with neither the lock nor a pin:

  • DataNodeTableOperatorGenerator#visitTableDeviceQueryCount builds a DevicePredicateFilter with ColumnTransformerBuilder for tagFuzzyPredicate, e.g. COUNT DEVICES FROM t WHERE my_udf(tag1) = 'x'.
  • The schema-region attribute update does the same: SchemaRegionMemoryImpl#updateTableDeviceAttribute → constructDevicePredicateUpdater, e.g. UPDATE t SET attr = my_udf(attr).

On the DataNode that hosts the schema region, nothing pins the defining loader there, and reflect() can still race with deregister().

&& instance.getSessionInfo().getSqlDialect() == SqlDialect.TABLE) {
UDFManagementService.getInstance().acquireReadLock();
try {
context.initializeTableModelUDFQuery();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For remote FIs, user table-function code has already run before this pin is taken. DataNodeInternalRPCServiceImpl#sendFragmentInstance → FragmentInstance.deserializeFrom → TableFunctionProcessorNode.deserialize calls TableFunctionFactory.getTableFunction(name).createTableFunctionHandle() and handle.deserialize(bytes). If a CREATE FUNCTION rotates the loader between deserialization and this line, the handle class belongs to the old, possibly already closed loader, while this FI pins the new active one. The window is narrow, but it is the same kind of race this PR is trying to close.

if (!(planner instanceof TableModelPlanner)) {
return;
}
synchronized (tableUdfClassLoaderLock) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

analyze() holds tableUdfClassLoaderLock for the entire planner.analyze(context), and this method needs the same monitor. Suppose the query is killed or fails while retry() is re-analyzing. The state-change listener runs on the shared fixed-size MPP-Coordinator-Executor pool (coordinator_read_executor_size: 20 by default, 2 in the Edge config). It reaches releaseResource() → here and blocks for the rest of that analysis, which can be unbounded for an insert waiting on FetchDevice. That delays state notifications for other queries.

Also, as soon as the listener gets the monitor it drops the pin, while start() keeps planning; CTE materialization and scalar subqueries can still instantiate UDFs at that point. Could the monitor guard only the initialize/finalize bookkeeping (e.g. a pinned flag) instead of the whole analysis?

// Analyze the statement in QueryContext. Generate the analysis this query need
private IAnalysis analyze(MPPQueryContext context) {
return planner.analyze(context);
if (!(planner instanceof TableModelPlanner)) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor/perf: this condition applies to every table-model statement. High-rate insertTablet/insertRecords, LOAD, and DELETE, which never resolve UDFs, now pay for:

  • the global read lock and a synchronized block;
  • a put/remove on queryIdToUDFClassLoaderMap;
  • an increment/decrement on the single AtomicLong of the active loader.

All writer threads contend on those cache lines. Could pinning be limited to statements that can reference functions (e.g. QueryType.READ and UPDATE), or done lazily on the first UDF resolution?

UDFManagementService.getInstance().releaseReadLock();
}
} else {
driverFactories =

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the two branches duplicate the planner.plan(...) call, and the dialect check repeats the one inside FragmentInstanceContext#initializeTableModelUDFQuery(). Two options would keep the paths from drifting apart:

  • a boolean flag around a single planner.plan call;
  • a small helper that does "read lock → pin → run", shared with QueryExecution#analyze.

@shuwenwei shuwenwei changed the title Fix table model UDF classloader lifecycle Protect UDF classloader lifetime with query scopes Nov 5, 2026
@sonarqubecloud

sonarqubecloud Bot commented Nov 6, 2026

Copy link
Copy Markdown

Quality Gate Failed Quality Gate failed

Failed conditions
C Reliability Rating on New Code (required ≥ A)

See analysis details on SonarQube Cloud

Catch issues before they fail your Quality Gate with our IDE extension SonarQube for IDE

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants