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]

Reply via email to