This is an automated email from the ASF dual-hosted git repository.
morningman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 5b3d2244210 [fix](fe) Stop leaking connector scope state after a
statement ends (#69718)
5b3d2244210 is described below
commit 5b3d2244210d439d32aaefc95a06590b08c51fed
Author: Mingyu Chen (Rayner) <[email protected]>
AuthorDate: Thu Oct 29 04:58:19 2026 +0000
[fix](fe) Stop leaking connector scope state after a statement ends (#69718)
### What problem does this PR solve?
Issue Number: None
Related PR: #69298, #69566, #66530
Problem Summary:
#### In short
A statement's connector state could outlive the statement in two ways on
master. A split producer still inside a read when the statement ends
went on to load values into the already-closed connector statement
scope, such as an Iceberg table borrowed from the table cache, and
nothing ever closed them. `EXPLAIN REFRESH MATERIALIZED VIEW` plans on a
context of its own that has no query id, so its scans registered their
query-finish cleanup under `""` and its connector scope was never
closed: each such EXPLAIN over an external base table left one registry
entry and one open scope behind, without bound. This PR makes a closed
scope refuse to load, and finishes the explain context's own statement
when the EXPLAIN ends.
This was found while checking master against branch-4.x's #66473 (scan
reuse, ported as #66530). branch-4.x has neither the connector statement
scope nor the query-finish registry, and its statement resources already
refuse registration once closed.
#### Background
- Every statement has a connector statement scope
(`ConnectorStatementScopeImpl`, hung on the `StatementContext`).
Connectors keep what the statement shares in it: the `ConnectorMetadata`
of each catalog, Iceberg table leases borrowed from the table cache
(returning one is what lets the cache close the table's FileIO when it
evicts it), Iceberg tracked tables, write transactions, and the query's
task groups since #69298. `closeAll()` releases them when the statement
ends.
- A scan registers two query-finish callbacks under the query id when it
starts planning splits: one stops its split producers, the other
(`QueryFinishCleanup`) releases a read transaction and closes the scope.
#69298 runs them at the end of every attempt, also for EXPLAIN and for
plans that fail; #69566 does the same for prepared EXECUTE.
- Batch and streaming scans plan on split producers. A producer calls
into the connector off the statement thread, for example Iceberg's
`streamSplits`, which resolves its table through the scope.
#### The problem, and what it cost
1. **Late producers load into the closed scope.** Stopping a scan only
flags its producers (`SplitProducer.close()`), and closing the query's
task groups cancels queued work but deliberately does not interrupt
metadata IO. A producer already inside `getNextBatch` keeps going after
`closeAll()`. `closeAll()` had also cleared the scope, so the producer
found nothing and loaded again: Iceberg borrowed a new table lease, Hive
partition batches built a new `ConnectorMetadata`. Both stayed in the
closed scope forever, so the lease was never returned. A value whose
load was in flight while `closeAll()` ran was missed the same way. This
happens whenever a statement ends while a producer is still reading: an
EXPLAIN of a large table (EXPLAIN also starts batch scans), a LIMIT that
finishes early, a KILL, a timeout. Nothing fails and no result is wrong;
the FE just keeps the leases, metadata and FileIO alive.
2. **`EXPLAIN REFRESH MATERIALIZED VIEW` leaks its own statement.** It
plans on an MTMV context of its own
(`RefreshMTMVCommand.getExplainConnectContext`), and that context had no
query id, so the connector sessions of its scans had `""` as their query
id and registered their cleanup under `""`. The end of the outer
statement finishes only the outer query id; nothing finishes `""`, and
it could not be finished at the end of one statement anyway, since every
context without a query id shares it. The explain context's
`StatementContext` was never closed either.
| Before this PR | Effect |
|---|---|
| A producer reading when the statement ends | Borrows a new table lease
(or builds a new metadata) into the closed scope; never released |
| `EXPLAIN REFRESH MATERIALIZED VIEW` over an external base table | One
permanent registry entry under `""` and one open connector scope per
execution |
#### How this PR fixes it
- `ConnectorStatementScopeImpl`: after `closeAll()` the statement is
over, so `computeIfAbsent` (and with it `computeIfAbsentForPlanning` and
`getOrCreateMetadata`) does not run the loader and throws
`CancellationException`, as the query's task groups already do once
closed; Iceberg's planning paths treat it as "the query stopped" and do
not fall back to the SDK. A load that raced the close is released by its
caller, because `closeAll()` could not see it. `closeAll()` now removes
each entry before releasing its value, still finalizing transactions
first, so every value is released exactly once. `getIfPresent` returns
nothing once closed.
- `StatementContext.getConnectorStatementScopeForNewSession()`: a
session built on a statement whose scope is already closed (the lineage
of an INSERT, after the insert finished its query, say) captures
`ConnectorStatementScope.NONE` and the direct task scheduler, as offline
work does, instead of failing in or loading into the closed scope.
Sessions captured before the close keep the closed scope, so a producer
holding one stops at its next load.
- `ExplainCommand`: an explain context of its own gets a query id, and
when the EXPLAIN ends, whether planning succeeded or not, it runs that
query id's query-finish callbacks and closes the context's
`StatementContext`, as `MTMVTask` does for a refresh. The thread-local
context is still restored if that cleanup throws.
- Comments that said EXPLAIN never registers a query-finish cleanup and
that `closeAll()` runs after the producers have quiesced are corrected;
the SPI javadoc of `computeIfAbsent` states the closed-scope behavior.
#### Results
| | Before | After |
|---|---|---|
| Producer resolving its table after the statement closed the scope |
New lease borrowed, never returned | `CancellationException`, nothing
borrowed |
| Value whose load raced `closeAll()` | Kept in the closed scope, never
closed | Closed once by its caller |
| Session built after the scope closed | Loads task groups into the
closed scope | Offline session (no scope, direct scheduler) |
| `EXPLAIN REFRESH MATERIALIZED VIEW` | Entry under `""` and open scope
left behind | Its callbacks run once, its scope closed |
New tests:
- `ConnectorStatementScopeTest`:
`closedScopeRefusesToLoadAndHandsOutNothing`,
`valueLoadedWhileTheScopeClosesIsReleasedOnceByItsCaller` (the loader
blocks on another thread while `closeAll()` runs),
`sessionBuiltAfterTheScopeClosesRunsAsOfflineWork`.
-
`SplitProducerTest.readThatOutlivesItsStatementLoadsNothingIntoTheClosedScope`:
a producer is stopped inside `getNextBatch`, the scope closes, then the
read resolves its table.
-
`ExplainRefreshMTMVCommandTest.testExplainCommandFinishesTheStatementOfItsOwnContext`:
planning registers a callback under the explain context's query id and
keeps a closeable in its scope; after the EXPLAIN, both succeeding and
failing, the callback ran once and the value was closed.
Each new test was checked against a mutation of the code it guards;
every mutation turns its test red:
| Mutation | Red test |
|---|---|
| Drop the closed check before loading |
`closedScopeRefusesToLoadAndHandsOutNothing` (the loader ran twice),
`readThatOutlivesItsStatementLoadsNothingIntoTheClosedScope` (the late
read loaded once) |
| Drop the closed check after loading |
`valueLoadedWhileTheScopeClosesIsReleasedOnceByItsCaller` (the caller
got the value) |
| Let a new session capture the closed scope |
`sessionBuiltAfterTheScopeClosesRunsAsOfflineWork` (building the session
throws: the scope cannot load `engine:query-task-groups`) |
| Do not finish the explain context |
`testExplainCommandFinishesTheStatementOfItsOwnContext` (callback never
ran) |
| Give the explain context no query id |
`testExplainCommandFinishesTheStatementOfItsOwnContext` (planned without
a query id) |
---
.../connector/spi/ConnectorStatementScope.java | 5 +-
.../doris/connector/ConnectorSessionBuilder.java | 9 +-
.../connector/ConnectorStatementScopeImpl.java | 95 ++++++++++++++-----
.../org/apache/doris/nereids/StatementContext.java | 34 +++++--
.../trees/plans/commands/ExplainCommand.java | 40 ++++++--
.../connector/ConnectorStatementScopeTest.java | 102 +++++++++++++++++++++
.../doris/datasource/split/SplitProducerTest.java | 48 ++++++++++
.../commands/ExplainRefreshMTMVCommandTest.java | 62 +++++++++++++
8 files changed, 348 insertions(+), 47 deletions(-)
diff --git
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorStatementScope.java
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorStatementScope.java
index 2809da1cb1c..c95c3eb4142 100644
---
a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorStatementScope.java
+++
b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/ConnectorStatementScope.java
@@ -40,7 +40,10 @@ public interface ConnectorStatementScope {
/**
* Returns the value cached under {@code key}, computing it with {@code
loader} on first access and
* caching it for the rest of the statement. Within one statement the same
key returns the same
- * instance to every caller; under {@link #NONE} the loader runs on every
call.
+ * instance to every caller; under {@link #NONE} the loader runs on every
call. Once the engine has
+ * closed the scope ({@link #closeAll}) the statement is over: the
engine's scope throws
+ * {@link java.util.concurrent.CancellationException} instead of loading,
so work that outlives its
+ * statement stops rather than holding a value nobody closes.
*/
<T> T computeIfAbsent(String key, Supplier<T> loader);
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java
b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java
index aa866eba680..d2fd3e68fdc 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorSessionBuilder.java
@@ -214,9 +214,10 @@ public final class ConnectorSessionBuilder {
* otherwise read it off the originating {@link ConnectContext}'s
statement context (preferring the
* retained context from {@link #from} over the thread-local, so a session
built off the request thread
* still captures correctly). Two-level null (no context / no statement
context) yields
- * {@link ConnectorStatementScope#NONE} -- mirrors {@code
MvccUtil.getSnapshotFromContext}. Capturing the
- * reference here (not reading it live) lets off-thread scan pumps that
reuse this one session reach the
- * same scope even though they have no ConnectContext thread-local.
+ * {@link ConnectorStatementScope#NONE} -- mirrors {@code
MvccUtil.getSnapshotFromContext}, and so does a
+ * statement whose scope is already closed (see {@link
StatementContext#getConnectorStatementScopeForNewSession}).
+ * Capturing the reference here (not reading it live) lets off-thread scan
pumps that reuse this one session
+ * reach the same scope even though they have no ConnectContext
thread-local.
*/
private ConnectorStatementScope captureStatementScope() {
if (statementScope != null) {
@@ -230,7 +231,7 @@ public final class ConnectorSessionBuilder {
if (sc == null) {
return ConnectorStatementScope.NONE;
}
- return sc.getOrCreateConnectorStatementScope();
+ return sc.getConnectorStatementScopeForNewSession();
}
/**
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorStatementScopeImpl.java
b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorStatementScopeImpl.java
index e097e08d55c..40d72238246 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorStatementScopeImpl.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/connector/ConnectorStatementScopeImpl.java
@@ -25,7 +25,9 @@ import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Iterator;
+import java.util.Map;
import java.util.Set;
+import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;
@@ -54,18 +56,45 @@ public class ConnectorStatementScopeImpl implements
ConnectorStatementScope {
private final ConcurrentHashMap<String, Object> cache = new
ConcurrentHashMap<>();
// Keys of the values only planning reads; endPlanning drops them.
private final Set<String> planningKeys = ConcurrentHashMap.newKeySet();
- private boolean closed;
+ // Read without the lock by the off-thread pumps, which can still be
loading when the statement ends.
+ private volatile boolean closed;
+ /**
+ * {@inheritDoc}
+ *
+ * <p>After {@link #closeAll} the statement is over, so the loader does
not run and a
+ * {@link CancellationException} is thrown: a pump that outlives its
statement (stopping a split producer only
+ * flags it, and closing the query's task groups does not interrupt
metadata IO) would otherwise load a value,
+ * such as a borrowed table, that nothing closes any more. A value whose
load raced the close is released here,
+ * because closeAll may not have seen it.</p>
+ */
@Override
@SuppressWarnings("unchecked")
public <T> T computeIfAbsent(String key, Supplier<T> loader) {
- return (T) cache.computeIfAbsent(key, k -> loader.get());
+ checkOpen(key);
+ Object value = cache.computeIfAbsent(key, k -> {
+ checkOpen(k);
+ return loader.get();
+ });
+ if (closed) {
+ // Whoever removes the entry, closeAll or this caller, releases
the value, so it is released once.
+ if (cache.remove(key, value)) {
+ release(value);
+ }
+ throw closedScope(key);
+ }
+ return (T) value;
}
@Override
@SuppressWarnings("unchecked")
public <T> T getIfPresent(String key) {
- return (T) cache.get(key);
+ return closed ? null : (T) cache.get(key);
+ }
+
+ /** Whether {@link #closeAll} has run: the statement has ended and the
scope loads nothing any more. */
+ public boolean isClosed() {
+ return closed;
}
@Override
@@ -94,8 +123,9 @@ public class ConnectorStatementScopeImpl implements
ConnectorStatementScope {
* Idempotent: guarded by {@code closed} so a second trigger (the
query-finish callback vs. a reused prepared
* statement's per-execution reset) is a harmless no-op and no value is
double-closed. Best-effort per value —
* a failure on one does not abort the rest — mirroring the isolation of
the engine's query-finish callback
- * registry. Runs after the scan off-thread pumps have quiesced, so it
does not race a concurrent
- * {@code computeIfAbsent}.
+ * registry. Off-thread pumps may still be running: each entry is removed
before its value is released, and a
+ * pump that loads concurrently releases what it loaded itself (see {@link
#computeIfAbsent}), so every value
+ * is released exactly once and none is left behind.
*/
@Override
public synchronized void closeAll() {
@@ -106,30 +136,45 @@ public class ConnectorStatementScopeImpl implements
ConnectorStatementScope {
// Pass 1: finalize the statement's write transaction(s) first, so a
transaction aborted mid-flight is
// rolled back / released before pass 2 closes the shared metadata
instance it was minted from. On every
// normal path the executor already finished the transaction, so this
is a no-op.
- for (Object value : cache.values()) {
- if (value instanceof CatalogStatementTransaction) {
- try {
- ((CatalogStatementTransaction)
value).finalizeAtStatementEnd();
- } catch (Exception e) {
- LOG.warn("failed to finalize per-statement transaction;
continuing", e);
- }
+ for (Map.Entry<String, Object> entry : cache.entrySet()) {
+ if (entry.getValue() instanceof CatalogStatementTransaction
+ && cache.remove(entry.getKey(), entry.getValue())) {
+ release(entry.getValue());
+ }
+ }
+ // Pass 2: close every remaining AutoCloseable value once.
+ for (Map.Entry<String, Object> entry : cache.entrySet()) {
+ if (cache.remove(entry.getKey(), entry.getValue())) {
+ release(entry.getValue());
}
}
- // Pass 2: close every remaining AutoCloseable value once (the
transaction holders handled above are
- // skipped here).
- for (Object value : cache.values()) {
- if (value instanceof CatalogStatementTransaction) {
- continue;
+ }
+
+ private void checkOpen(String key) {
+ if (closed) {
+ throw closedScope(key);
+ }
+ }
+
+ private static CancellationException closedScope(String key) {
+ return new CancellationException("the statement has ended, its
connector scope cannot load " + key);
+ }
+
+ // Finalizes a transaction and closes any other AutoCloseable value; a
plain value holds nothing to release.
+ private static void release(Object value) {
+ if (value instanceof CatalogStatementTransaction) {
+ try {
+ ((CatalogStatementTransaction) value).finalizeAtStatementEnd();
+ } catch (Exception e) {
+ LOG.warn("failed to finalize per-statement transaction;
continuing", e);
}
- if (value instanceof AutoCloseable) {
- try {
- ((AutoCloseable) value).close();
- } catch (Exception e) {
- LOG.warn("failed to close per-statement scope value of
type {}; continuing",
- value.getClass().getName(), e);
- }
+ } else if (value instanceof AutoCloseable) {
+ try {
+ ((AutoCloseable) value).close();
+ } catch (Exception e) {
+ LOG.warn("failed to close per-statement scope value of type
{}; continuing",
+ value.getClass().getName(), e);
}
}
- cache.clear();
}
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
index 898f1bc12fa..c854a4c8f92 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java
@@ -233,7 +233,7 @@ public class StatementContext implements Closeable {
// NOT cleared on close()/releasePlannerResources -- it survives to
statement GC. A reused
// prepared-statement context drops it per execution via
resetConnectorStatementScope() (see
// ExecuteCommand) so one execution's cached tables never leak into the
next.
- private ConnectorStatementScope connectorStatementScope;
+ private ConnectorStatementScopeImpl connectorStatementScope;
// table locks
private final Stack<CloseableResource> plannerResources = new Stack<>();
@@ -871,6 +871,18 @@ public class StatementContext implements Closeable {
return this.connectorStatementScope;
}
+ /**
+ * The scope a connector session built now captures: {@link
#getOrCreateConnectorStatementScope}, or
+ * {@link ConnectorStatementScope#NONE} once that scope is closed. The
statement has ended by then, and what
+ * still builds a session on its context (the lineage of an INSERT, say)
runs as offline work does instead of
+ * failing in the closed scope. A session captured before the close keeps
the closed scope, which refuses to
+ * load, so a split producer that outlives the statement stops.
+ */
+ public synchronized ConnectorStatementScope
getConnectorStatementScopeForNewSession() {
+ getOrCreateConnectorStatementScope();
+ return connectorStatementScope.isClosed() ?
ConnectorStatementScope.NONE : connectorStatementScope;
+ }
+
/**
* Drops the values of this statement's connector scope that only planning
reads (see
* {@link ConnectorStatementScope#endPlanning}), when a connector has
opened the scope; unlike
@@ -896,7 +908,8 @@ public class StatementContext implements Closeable {
* this is called at the start of each, so a prior execution's / attempt's
cached tables, pinned connector
* writer schema, and closeable state are released (closeAll is idempotent
— a no-op if the query-finish
* callback already closed it) and never leak into the next. Callers
invoke this only after the prior
- * execution/attempt has finished (no off-thread pump still running), so
close-before-drop is safe.
+ * execution/attempt has finished; a split producer of it still running is
refused by the closed scope, so
+ * close-before-drop is safe.
*/
public synchronized void resetConnectorStatementScope() {
closeConnectorStatementScope();
@@ -1293,14 +1306,15 @@ public class StatementContext implements Closeable {
// After the table locks: stopping a remote Doris scan's node sends a
CloseSession to the other
// frontend, which must not be waited for under a lock.
stopScanNodesLeftBehind();
- // Fallback deterministic close of the per-statement connector scope,
for statements that never reach the
- // query-finish callback: external DDL / SHOW / DESCRIBE / EXPLAIN /
foreground ANALYZE run via Command.run
- // with no coordinator, so PluginDrivenScanNode.getSplits never
registers a primary close for them. close()
- // runs in ConnectProcessor's per-statement finally (direct
connection) and the forwarded-request finally
- // (master side), so it covers them. Guarded by
isReturnResultFromLocal so an arrow-flight statement (which
- // returns results asynchronously and defers cleanup to its own
query-finish close) is not closed early
- // here. Idempotent (closeAll is close-once), so a coordinated
statement whose scope the primary already
- // closed is a no-op. Command statements have no off-thread scan pump,
so close-after-use holds.
+ // Fallback deterministic close of the per-statement connector scope,
for statements that open it without
+ // planning a scan, so no query-finish callback closes it: external
DDL / SHOW / DESCRIBE run via
+ // Command.run and only read metadata. (Every scan registers that
callback when it starts planning splits,
+ // also under EXPLAIN or a plan that fails later; StmtExecutor runs it
when the attempt ends.) close() runs
+ // in ConnectProcessor's per-statement finally (direct connection) and
the forwarded-request finally (master
+ // side), so it covers them. Guarded by isReturnResultFromLocal so an
arrow-flight statement (which returns
+ // results asynchronously and defers cleanup to its own query-finish
close) is not closed early here.
+ // Idempotent (closeAll is close-once), so a statement whose scope the
callback already closed is a no-op;
+ // a split producer still running after the close is refused by the
closed scope.
if (connectorStatementScope != null && connectContext != null &&
connectContext.isReturnResultFromLocal()) {
connectorStatementScope.closeAll();
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExplainCommand.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExplainCommand.java
index 6c144b5879c..0d74eaec498 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExplainCommand.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExplainCommand.java
@@ -20,6 +20,7 @@ package org.apache.doris.nereids.trees.plans.commands;
import org.apache.doris.analysis.ExplainOptions;
import org.apache.doris.analysis.StmtType;
import org.apache.doris.common.AnalysisException;
+import org.apache.doris.common.util.UniqueIdUtils;
import org.apache.doris.nereids.NereidsPlanner;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import
org.apache.doris.nereids.rules.exploration.mv.InitMaterializationContextHook;
@@ -33,6 +34,7 @@ import
org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
import org.apache.doris.planner.ScanNode;
import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.QeProcessorImpl;
import org.apache.doris.qe.StmtExecutor;
import java.util.Optional;
@@ -91,6 +93,13 @@ public class ExplainCommand extends Command implements
NoForward {
boolean resetTargetTableId = false;
try {
explainCtx = explainable.getExplainConnectContext(ctx);
+ if (explainCtx != ctx) {
+ // A context of its own (EXPLAIN REFRESH MATERIALIZED VIEW
plans on an MTMV context) is a statement
+ // of its own, which the cleanup at the end of this one does
not reach. Give it a query id for its
+ // scans to register their query-finish cleanup under, which
the finally below runs: without one
+ // they register under "" (ConnectorSessionBuilder), a key
every such context shares.
+ explainCtx.setQueryId(UniqueIdUtils.fastUniqueId());
+ }
if (explainable instanceof InsertIntoTableCommand
|| explainable instanceof InsertOverwriteTableCommand
|| explainable instanceof UpdateCommand) {
@@ -141,18 +150,35 @@ public class ExplainCommand extends Command implements
NoForward {
scanNode.stop();
}
} finally {
- if (resetTargetTableId) {
-
explainCtx.setSyntheticWriteColTargetTableId(previousTargetTableId);
- }
- if (ConnectContext.get() != previousCtx) {
- ConnectContext.remove();
- if (previousCtx != null) {
- previousCtx.setThreadLocalInfo();
+ try {
+ if (resetTargetTableId) {
+
explainCtx.setSyntheticWriteColTargetTableId(previousTargetTableId);
+ }
+ if (explainCtx != null && explainCtx != ctx) {
+ closeExplainConnectContext(explainCtx);
+ }
+ } finally {
+ if (ConnectContext.get() != previousCtx) {
+ ConnectContext.remove();
+ if (previousCtx != null) {
+ previousCtx.setThreadLocalInfo();
+ }
}
}
}
}
+ // Ends the statement of an explain context of its own as MTMVTask ends a
refresh's: runs the query-finish
+ // cleanup its scans registered (which stops their split producers and
closes their connector scope), then
+ // closes its statement context. It has no coordinator, so only the
callbacks are left to finish.
+ private static void closeExplainConnectContext(ConnectContext explainCtx) {
+ try {
+
QeProcessorImpl.INSTANCE.finishQueryCallbacks(explainCtx.queryId());
+ } finally {
+ explainCtx.getStatementContext().close();
+ }
+ }
+
@Override
public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
return visitor.visitExplainCommand(this, context);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/connector/ConnectorStatementScopeTest.java
b/fe/fe-core/src/test/java/org/apache/doris/connector/ConnectorStatementScopeTest.java
index ada9a96ef20..1ad040c7c0c 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/connector/ConnectorStatementScopeTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/connector/ConnectorStatementScopeTest.java
@@ -20,6 +20,8 @@ package org.apache.doris.connector;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.connector.spi.ConnectorSession;
import org.apache.doris.connector.spi.ConnectorStatementScope;
+import org.apache.doris.connector.spi.ConnectorStatementScopes;
+import org.apache.doris.connector.spi.ConnectorTaskScheduler;
import org.apache.doris.connector.spi.ConnectorWriteOps;
import org.apache.doris.connector.spi.handle.ConnectorTableHandle;
import org.apache.doris.connector.spi.handle.ConnectorTransaction;
@@ -41,6 +43,12 @@ import org.mockito.Mockito;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
@@ -284,6 +292,100 @@ public class ConnectorStatementScopeTest {
"the transaction is finalized before the shared metadata is
closed");
}
+ @Test
+ public void closedScopeRefusesToLoadAndHandsOutNothing() {
+ // The statement is over once closeAll() ran. A pump that outlives it
(a split producer still listing
+ // files: stopping it only flags it) must not load a value, such as a
borrowed table, that nothing would
+ // close any more. MUTATION: dropping the closed check in
computeIfAbsent -> the loader runs -> red.
+ ConnectorStatementScopeImpl scope = new ConnectorStatementScopeImpl();
+ scope.computeIfAbsent("plain", Object::new);
+ scope.closeAll();
+
+ AtomicInteger loads = new AtomicInteger();
+ Supplier<AutoCloseable> loader = () -> {
+ loads.incrementAndGet();
+ return () -> { };
+ };
+ Assertions.assertTrue(scope.isClosed());
+ Assertions.assertThrows(CancellationException.class, () ->
scope.computeIfAbsent("lease", loader));
+ Assertions.assertThrows(CancellationException.class,
+ () -> scope.computeIfAbsentForPlanning("memo", () -> {
+ loads.incrementAndGet();
+ return new Object();
+ }));
+ Assertions.assertEquals(0, loads.get(), "a closed scope never runs a
loader");
+ Assertions.assertNull(scope.getIfPresent("plain"), "a closed scope
hands out nothing it held");
+ }
+
+ @Test
+ public void valueLoadedWhileTheScopeClosesIsReleasedOnceByItsCaller()
throws Exception {
+ // closeAll() can run while a pump is inside a loader. The value is
not in the scope yet, so closeAll cannot
+ // release it; it used to land in the closed scope afterwards and
never be closed. The caller that loaded it
+ // releases it now and gets CancellationException instead of the value.
+ // MUTATION: dropping the closed check after the load -> closes==0 and
the caller gets the value -> red.
+ ConnectorStatementScopeImpl scope = new ConnectorStatementScopeImpl();
+ CountDownLatch loading = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicInteger closes = new AtomicInteger();
+ ExecutorService pump = Executors.newSingleThreadExecutor();
+ try {
+ Future<AutoCloseable> load = pump.submit(() ->
scope.computeIfAbsent("lease", () -> {
+ loading.countDown();
+ try {
+ Assertions.assertTrue(release.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ return (AutoCloseable) closes::incrementAndGet;
+ }));
+ Assertions.assertTrue(loading.await(5, TimeUnit.SECONDS));
+ scope.closeAll();
+ release.countDown();
+
+ ExecutionException failure =
Assertions.assertThrows(ExecutionException.class,
+ () -> load.get(5, TimeUnit.SECONDS));
+ Assertions.assertInstanceOf(CancellationException.class,
failure.getCause());
+ Assertions.assertEquals(1, closes.get(), "the caller releases the
value closeAll could not see");
+ scope.closeAll();
+ Assertions.assertEquals(1, closes.get(), "and nothing releases it
a second time");
+ } finally {
+ release.countDown();
+ pump.shutdownNow();
+ }
+ }
+
+ @Test
+ public void sessionBuiltAfterTheScopeClosesRunsAsOfflineWork() {
+ // What still builds a session on a statement that has ended (the
lineage of an INSERT, say) must neither
+ // fail in nor memoize into the closed scope: it captures NONE and the
direct scheduler, as offline work
+ // does. A session captured before the close keeps the closed scope
and is refused, so a split producer
+ // holding it stops instead of borrowing a table nothing returns.
+ // MUTATION: capturing the closed scope for a new session -> building
it throws -> red.
+ ConnectContext ctx = new ConnectContext();
+ ctx.setThreadLocalInfo();
+ try {
+ StatementContext stmtCtx = new StatementContext();
+ ctx.setStatementContext(stmtCtx);
+ ConnectorSession before =
ConnectorSessionBuilder.from(ctx).withCatalogId(1L).build();
+ ConnectorStatementScope live =
stmtCtx.getOrCreateConnectorStatementScope();
+ Assertions.assertSame(live, before.getStatementScope());
+
+ live.closeAll();
+
+ ConnectorSession after =
ConnectorSessionBuilder.from(ctx).withCatalogId(1L).build();
+ Assertions.assertSame(ConnectorStatementScope.NONE,
after.getStatementScope());
+ Assertions.assertSame(ConnectorTaskScheduler.DIRECT,
+
after.getTaskScheduler(ConnectorTaskScheduler.Stage.SPLIT));
+ AtomicInteger loads = new AtomicInteger();
+ Assertions.assertThrows(CancellationException.class, () ->
ConnectorStatementScopes.resolveInStatement(
+ before, "test_table", "db", "tbl",
loads::incrementAndGet));
+ Assertions.assertEquals(0, loads.get(), "a session of the ended
statement loads nothing");
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+
@Test
public void resetClosesTheDroppedScopeBeforeStartingFresh() {
// A prepared EXECUTE / retry reuses one StatementContext and calls
resetConnectorStatementScope() at the
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
index fde370bc637..ce05d95cb95 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/datasource/split/SplitProducerTest.java
@@ -20,6 +20,7 @@ package org.apache.doris.datasource.split;
import org.apache.doris.common.UserException;
import org.apache.doris.common.util.FairTaskExecutor;
import org.apache.doris.common.util.LocationPath;
+import org.apache.doris.connector.ConnectorStatementScopeImpl;
import org.apache.doris.connector.spi.ConnectorTaskScheduler;
import org.apache.doris.datasource.scan.FederationBackendPolicy;
import org.apache.doris.spi.Split;
@@ -382,6 +383,53 @@ class SplitProducerTest {
}
}
+ @Test
+ void readThatOutlivesItsStatementLoadsNothingIntoTheClosedScope() throws
Exception {
+ // Stopping the scan only flags a producer that is already inside
getNextBatch, and the statement then
+ // closes its connector scope. When that read goes on to resolve its
table (Iceberg's streamSplits borrows
+ // one through the scope), the closed scope refuses it, so it borrows
nothing that would never be returned.
+ // MUTATION: a scope that still loads after closeAll -> loads==1 ->
red.
+ Fixture fixture = new Fixture(1, 1);
+ ConnectorStatementScopeImpl scope = new ConnectorStatementScopeImpl();
+ AtomicInteger loads = new AtomicInteger();
+ CountDownLatch reading = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ RecordingSource source = new RecordingSource(1, 0) {
+ @Override
+ public List<Split> getNextBatch() throws IOException {
+ reading.countDown();
+ try {
+ Assertions.assertTrue(release.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ scope.computeIfAbsent("iceberg:table:1:db:tbl", () -> {
+ loads.incrementAndGet();
+ return (AutoCloseable) () -> { };
+ });
+ return super.getNextBatch();
+ }
+ };
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ try {
+ fixture.start(executor, source);
+ Assertions.assertTrue(reading.await(5, TimeUnit.SECONDS));
+ fixture.assignment.stop();
+ scope.closeAll();
+ release.countDown();
+ Assertions.assertTrue(source.closed.await(5, TimeUnit.SECONDS));
+ Assertions.assertEquals(0, loads.get(), "the closed scope refuses
the late read's load");
+ Assertions.assertEquals(1, source.closeCalls.get());
+ Assertions.assertEquals(1, fixture.completions.get());
+ Mockito.verify(fixture.policy,
Mockito.never()).computeScanRangeAssignment(Mockito.any());
+ } finally {
+ release.countDown();
+ executor.shutdownNow();
+ Assertions.assertTrue(executor.awaitTermination(5,
TimeUnit.SECONDS));
+ }
+ }
+
@Test
void rejectedCancellationCleanupStillClosesWithoutAddingAnError() throws
Exception {
Fixture fixture = new Fixture(1, 1);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExplainRefreshMTMVCommandTest.java
b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExplainRefreshMTMVCommandTest.java
index 88268fa0f8e..3916c83f832 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExplainRefreshMTMVCommandTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExplainRefreshMTMVCommandTest.java
@@ -24,6 +24,7 @@ import org.apache.doris.catalog.MTMV;
import org.apache.doris.catalog.TableIf.TableType;
import org.apache.doris.catalog.info.TableNameInfo;
import org.apache.doris.common.MetaNotFoundException;
+import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.datasource.InternalCatalog;
import org.apache.doris.mtmv.MTMVRefreshEnum.RefreshMethod;
import org.apache.doris.mtmv.MTMVRefreshInfo;
@@ -42,7 +43,9 @@ import
org.apache.doris.nereids.trees.plans.commands.insert.InsertIntoTableComma
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.planner.ScanNode;
import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.QeProcessorImpl;
import org.apache.doris.qe.StmtExecutor;
+import org.apache.doris.thrift.TUniqueId;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
@@ -52,6 +55,8 @@ import org.mockito.Mockito;
import java.util.Collections;
import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
public class ExplainRefreshMTMVCommandTest {
private final NereidsParser parser = new NereidsParser();
@@ -177,6 +182,63 @@ public class ExplainRefreshMTMVCommandTest {
}
}
+ @Test
+ public void testExplainCommandFinishesTheStatementOfItsOwnContext() throws
Exception {
+ // EXPLAIN REFRESH MATERIALIZED VIEW plans on an MTMV context of its
own, whose cleanup the end of the outer
+ // statement (it runs only the outer query id) does not reach. What
its scans registered while planning,
+ // and its connector scope, must still be finished when the EXPLAIN
ends, whether planning succeeds or not.
+ // MUTATION: dropping closeExplainConnectContext -> cleanups==0 and
closes==0 -> red.
+ assertExplainFinishesItsOwnContext(false);
+ assertExplainFinishesItsOwnContext(true);
+ }
+
+ private void assertExplainFinishesItsOwnContext(boolean planningFails)
throws Exception {
+ NereidsPlanner planner = Mockito.mock(NereidsPlanner.class);
+
Mockito.when(planner.getExplainString(Mockito.any())).thenReturn("plan");
+
Mockito.when(planner.getScanNodes()).thenReturn(Collections.<ScanNode>emptyList());
+
Mockito.when(planner.getCascadesContext()).thenReturn(Mockito.mock(org.apache.doris.nereids.CascadesContext.class));
+ AtomicInteger cleanups = new AtomicInteger();
+ AtomicInteger closes = new AtomicInteger();
+ AtomicReference<TUniqueId> planningQueryId = new AtomicReference<>();
+ Mockito.doAnswer(invocation -> {
+ // What a scan does while it plans splits: register its cleanup
under the query id of the context it
+ // plans on, and keep a closeable value (a borrowed table) in that
statement's connector scope.
+ ConnectContext planningCtx = ConnectContext.get();
+ planningQueryId.set(planningCtx.queryId());
+ QeProcessorImpl.INSTANCE.registerQueryFinishCallback(
+ DebugUtil.printId(planningCtx.queryId()),
cleanups::incrementAndGet);
+
planningCtx.getStatementContext().getOrCreateConnectorStatementScope()
+ .computeIfAbsent("lease", () -> (AutoCloseable)
closes::incrementAndGet);
+ if (planningFails) {
+ throw new IllegalStateException("planning failed");
+ }
+ return null;
+ }).when(planner).plan(Mockito.any(LogicalPlanAdapter.class),
Mockito.any());
+ RecordingRefreshMTMVCommand command = new
RecordingRefreshMTMVCommand(planner);
+ ConnectContext originCtx = new ConnectContext();
+ originCtx.setQueryId(new TUniqueId(0x13cL, planningFails ? 1L : 2L));
+ originCtx.setThreadLocalInfo();
+ try {
+ ExplainCommand explain = new ExplainCommand(ExplainLevel.NORMAL,
command, false);
+ if (planningFails) {
+ Assertions.assertThrows(IllegalStateException.class,
+ () -> explain.run(originCtx,
Mockito.mock(StmtExecutor.class)));
+ } else {
+ explain.run(originCtx, Mockito.mock(StmtExecutor.class));
+ }
+ Assertions.assertNotNull(planningQueryId.get(), "the explain
context plans under a query id");
+ Assertions.assertSame(command.explainCtx.queryId(),
planningQueryId.get());
+ Assertions.assertNotEquals(originCtx.queryId(),
planningQueryId.get(), "a query id of its own");
+ Assertions.assertEquals(1, cleanups.get(), "the EXPLAIN runs what
its own context's scans registered");
+ Assertions.assertEquals(1, closes.get(), "and closes that
context's connector scope");
+
QeProcessorImpl.INSTANCE.finishQueryCallbacks(planningQueryId.get());
+ Assertions.assertEquals(1, cleanups.get(), "nothing is left
registered under it");
+ Assertions.assertSame(originCtx, ConnectContext.get());
+ } finally {
+ ConnectContext.remove();
+ }
+ }
+
@Test
public void testExplainCommandRestoresPreviousThreadLocalContext() throws
Exception {
StmtExecutor executor = Mockito.mock(StmtExecutor.class);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]