From 848573bfbf4f03282d110dcf5a1e4e60c4319be9 Mon Sep 17 00:00:00 2001 From: starocean999 Date: Fri, 11 Sep 2026 16:05:19 +0800 Subject: [PATCH 1/3] [fix](fe) Allocate a fresh StatementContext per EXECUTE to prevent FE OOM in long-lived prepared statements (#67256) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Problem Summary: A prepared statement lives as long as its connection. The `PreparedStatementContext` kept in `ConnectContext.preparedStatementContextMap` retains a single `StatementContext` and reuses the same object across every `EXECUTE` for the whole connection lifetime. Because one object is reused across executions, its per-statement state keeps accumulating: bound tables (`tables`, `oneLevelTables`, `mtmvRelatedTables`, `insertTargetTables`, `viewInfos`), CTE maps, statistics (`relationIdToStatisticsMap`, `tableIdMapping`), MV/partition rewrite state (`mvCanRewritePartitionsMap`, `tmpPlanForMvRewrite`, `materializationRewrittenSuccessSet`), MVCC snapshots, connector write schemas, placeholder bindings (`idToPlaceholderRealExpr`), etc. On long-lived connections with a high number of `EXECUTE`s, these maps only grow and are never released until the connection closes, which can OOM the FE. **Root cause:** the `StatementContext` stored in `PreparedStatementContext` was treated as a permanent per-prepared-statement object and reused, so state that should be per-execution lived as long as the connection. **Fix:** instead of reusing (and clearing in place) the same `StatementContext`, allocate a brand-new context on every `EXECUTE` and carry over only the state that must survive between executions: - **ID generator positions** — so ids generated during this execution never collide with ids already present in the cached analyzed plan from `PREPARE`; - **placeholder real expressions** bound by the protocol layer for this `EXECUTE` (`idToPlaceholderRealExpr`) — this is the piece that prevents the #63920 parameter-mismatch regression; - the **placeholder → comparison-slot registry** (`idToComparisonSlot`) used by the short-circuit fast path; - the **placeholder list**; - the **short-circuit / nondeterministic flags** that gate the short-circuit fast path before any re-planning. After the swap, the previous context becomes unreachable and is promptly GC'd, so memory no longer grows with the number of executions. The cached analyzed plan and the point-query (short-circuit) cache live on `PrepareCommand` and `PreparedStatementContext` respectively, so they keep being reused across executions. **Changes:** - `IdGenerator`: add `getCurrentId()` so a fresh context can continue the id generators from the previous one. - `StatementContext`: add `createNextExecuteContext()` which allocates the fresh context and copies over the cross-execution state above. - `PreparedStatementContext`: add `nextStatementContext()` which swaps in the fresh context so the old one is released. - `ExecuteCommand`: `run()` now uses the fresh per-execution context (and the now-redundant in-place `resetConnectorStatementScope()`/`resetMvccSnapshots()` calls are removed since a fresh context starts empty by construction). - Unit tests updated to assert the fresh-context behavior (`ExecuteCommandTest`, `ConnectorStatementScopeTest`). None --- .../org/apache/doris/common/IdGenerator.java | 8 + .../doris/nereids/StatementContext.java | 47 +++++ .../trees/plans/commands/ExecuteCommand.java | 28 ++- .../doris/qe/PreparedStatementContext.java | 32 ++++ .../plans/commands/ExecuteCommandTest.java | 180 +++++++++++------- 5 files changed, 221 insertions(+), 74 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/IdGenerator.java b/fe/fe-core/src/main/java/org/apache/doris/common/IdGenerator.java index 76f503bb5d7bc3..aa769cf3415a45 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/IdGenerator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/IdGenerator.java @@ -33,6 +33,14 @@ public IdGenerator resetId(int initialId) { return this; } + /** + * The id value that {@link #getNextId()} would hand out next. Used to seed a fresh + * per-execution id generator from an existing one so ids never collide. + */ + public int getCurrentId() { + return nextId; + } + public abstract IdType getNextId(); } 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 85bfabf02f412d..4ff0cd0ed3d7ae 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 @@ -383,6 +383,53 @@ private StatementContext(ConnectContext connectContext, OriginStatement originSt } } + /** + * Create a fresh StatementContext for the next EXECUTE of a prepared statement. + * + *

A prepared statement keeps its StatementContext inside {@code PreparedStatementContext} + * for the whole lifetime of the connection. Reusing the same object across executions makes + * its per-statement state (bound tables, CTE maps, statistics, snapshots, connector scope, + * ...) accumulate and it is only released when the connection closes, which can OOM + * long-lived connections. Instead of clearing in place, allocate a brand-new context per + * EXECUTE and copy over only the state that must survive between executions, so the previous + * context becomes unreachable and is promptly GC'd. + * + *

Carried over: + *

+ * Everything else (tables, CTEs, statistics, snapshots, planner resources, connector + * scope, ...) starts empty/fresh on the new context. + */ + public StatementContext createNextExecuteContext() { + // Continue the id generators from the previous context. The cached analyzed plan from + // PREPARE (and every prior execution) already consumed ids from them, so a fresh + // generator starting at 0 would collide with those ids during this execution's planning. + StatementContext next = new StatementContext(connectContext, originStatement, + exprIdGenerator.getCurrentId()); + next.objectIdGenerator.resetId(objectIdGenerator.getCurrentId()); + next.relationIdGenerator.resetId(relationIdGenerator.getCurrentId()); + next.cteIdGenerator.resetId(cteIdGenerator.getCurrentId()); + next.talbeIdGenerator.resetId(talbeIdGenerator.getCurrentId()); + next.placeHolderIdGenerator.resetId(placeHolderIdGenerator.getCurrentId()); + // Placeholder bindings of this EXECUTE, and the comparison-slot registry used to replace + // conjuncts on the cached short-circuit plan without re-planning. + next.idToPlaceholderRealExpr.putAll(idToPlaceholderRealExpr); + next.idToComparisonSlot.putAll(idToComparisonSlot); + next.placeholders = new ArrayList<>(placeholders); + // Short-circuit gating flags are computed by the previous execution's planning and gate + // the fast path of this execution before any re-planning happens. + next.isShortCircuitQuery = isShortCircuitQuery; + next.hasNondeterministic = hasNondeterministic; + return next; + } + public void setNeedLockTables(boolean needLockTables) { this.needLockTables = needLockTables; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java index 6d0fc555e7824f..87447a382508f2 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java @@ -82,6 +82,14 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { "prepare statement " + stmtName + " not found, maybe expired"); } PrepareCommand prepareCommand = preparedStmtCtx.command; + // Allocate a fresh StatementContext per EXECUTE so the per-statement state accumulated by + // prior executions (bound tables, CTE maps, statistics, snapshots, ...) is released + // promptly instead of living as long as the connection, which can OOM long-lived + // connections. The necessary cross-execution state (placeholder bindings, comparison + // slots, id generator positions, short-circuit flags) is carried over to the new context. + StatementContext statementContext = preparedStmtCtx.nextStatementContext(); + statementContext.setPrepareStage(false); + statementContext.setIsInsert(false); LogicalPlan logicalPlan = prepareCommand.getLogicalPlan(); List relationRoots = new ArrayList<>(); if (logicalPlan instanceof InsertIntoTableCommand) { @@ -132,12 +140,20 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { LogicalPlanAdapter planAdapter = new LogicalPlanAdapter( logicalPlan, executor.getContext().getStatementContext()); executor.setParsedStmt(planAdapter); - // If it's not a short circuit query, schema version or file cache query limit changed, or - // the statement has nondeterministic functions, then reanalyze and plan. - if (executor.getContext().getStatementContext().isShortCircuitQuery() - && preparedStmtCtx.shortCircuitQueryContext.isPresent() - && preparedStmtCtx.shortCircuitQueryContext.get().isReusable(ctx) - && !executor.getContext().getStatementContext().hasNondeterministic()) { + boolean hasShortCircuitContext = preparedStmtCtx.shortCircuitQueryContext.isPresent(); + boolean shortCircuitContextReusable = hasShortCircuitContext + && preparedStmtCtx.shortCircuitQueryContext.get().isReusable(ctx); + // Reuse the cached short-circuit plan only when table metadata is unchanged and the statement + // has no nondeterministic functions. Otherwise fall back to the normal execution path below. + if (statementContext.isShortCircuitQuery() + && hasShortCircuitContext + && shortCircuitContextReusable + && !statementContext.hasNondeterministic()) { + // The fresh per-execution context carries the short-circuit flag but not the cached plan. + // Install the just-validated cache before the direct path: result sending reads it via + // statementContext.getShortCircuitQueryContext(), and the fallback (building one from a + // null planner, since this path skips planning) would NPE. + statementContext.setShortCircuitQueryContext(preparedStmtCtx.shortCircuitQueryContext.get()); PointQueryExecutor.directExecuteShortCircuitQuery(executor, preparedStmtCtx, statementContext); return; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java index 0174befee5b9d0..d5d9162aa1945c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java @@ -46,6 +46,38 @@ public long getStartTime() { return startTime; } + public StatementContext getStatementContext() { + return statementContext; + } + + public void setStatementContext(StatementContext statementContext) { + this.statementContext = statementContext; + } + + /** + * Allocate a fresh StatementContext for this EXECUTE and replace the previous one, so the + * old context (with the per-statement state accumulated by prior executions: bound tables, + * CTE maps, statistics, snapshots, ...) becomes unreachable and is promptly GC'd. + * + *

A prepared statement lives as long as its connection. Reusing one StatementContext + * across all executions would keep growing those maps and could OOM long-lived connections, + * so we create a new object per execution and carry over only the state that must survive + * (placeholder bindings, comparison slots, id generator positions, short-circuit flags). + * + * @return the fresh StatementContext to use for the current execution + */ + public StatementContext nextStatementContext() { + // Close the outgoing context's per-statement connector scope before dropping it. The binary + // COM_STMT_EXECUTE path has no per-statement StatementContext.close() finally (that only + // runs for COM_QUERY), and coordinated scans may not have registered a query-finish + // callback yet (connector commands and failures before scan registration have none). + // Without this, the outgoing scope's closeable connector metadata / active connector + // transactions would be abandoned, and GC cannot finalize them. + statementContext.resetConnectorStatementScope(); + statementContext = statementContext.createNextExecuteContext(); + return statementContext; + } + public void setStartTime() { startTime = System.currentTimeMillis(); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java index ca0d8306f535c9..9325aa310ab881 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java @@ -17,8 +17,13 @@ package org.apache.doris.nereids.trees.plans.commands; +import org.apache.doris.analysis.DescriptorTable; +import org.apache.doris.analysis.Queriable; import org.apache.doris.analysis.TableScanParams; import org.apache.doris.catalog.DatabaseIf; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; import org.apache.doris.catalog.TableIf; import org.apache.doris.datasource.CatalogIf; import org.apache.doris.datasource.ExternalScanTaskCacheKey; @@ -30,11 +35,15 @@ import org.apache.doris.nereids.trees.expressions.SubqueryExpr; import org.apache.doris.nereids.trees.plans.commands.merge.MergeIntoCommand; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.planner.OlapScanNode; +import org.apache.doris.planner.Planner; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.OriginStatement; import org.apache.doris.qe.PreparedStatementContext; import org.apache.doris.qe.SessionVariable; +import org.apache.doris.qe.ShortCircuitQueryContext; import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.thrift.TQueryOptions; import com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Assertions; @@ -134,6 +143,56 @@ public void testExecutePreparedCommandWithoutPlanChildren() throws Exception { Mockito.verify(executor).execute(); } + @Test + public void testPreparedConnectorUpdateRefreshesWriteDefaultEveryExecution() throws Exception { + // ExecuteCommand allocates a fresh StatementContext per EXECUTE. Model connector metadata changing from + // default 1 to 2 between executions: each execution's planner callback pins the current schema only when + // the fresh context has no schema pinned, then expands DEFAULT(v) and writes the resulting value. + // MUTATION: reusing one StatementContext across executions without dropping connectorWriteSchemas makes + // execution two reuse default 1, so the written values become [1, 1] instead of [1, 2]. + String sql = "update ext_catalog.db.t set v = default(v) where id = 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + Assertions.assertInstanceOf(UpdateCommand.class, logicalPlan); + + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + StatementContext statementContext = new StatementContext(); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + Mockito.when(executor.getContext()).thenReturn(connectContext); + + long tableId = 7L; + AtomicInteger metadataDefault = new AtomicInteger(1); + List writtenValues = new ArrayList<>(); + Mockito.doAnswer(invocation -> { + // Each execution plans through the fresh StatementContext allocated by ExecuteCommand, so resolve/pin + // the connector writer schema on THAT context (a stale pin would make the second execution reuse + // default 1). + StatementContext currentContext = preparedStatement.getStatementContext(); + if (!currentContext.getConnectorWriteSchema(tableId).isPresent()) { + Column column = new Column("v", ScalarType.createType(PrimitiveType.INT), + false, null, String.valueOf(metadataDefault.get()), ""); + currentContext.setConnectorWriteSchema(tableId, Collections.singletonList(column)); + } + writtenValues.add(currentContext.getConnectorWriteSchema(tableId).get() + .get(0).getDefaultValueSql()); + return null; + }).when(executor).execute(); + + ExecuteCommand execute = new ExecuteCommand("stmt", prepareCommand, statementContext); + execute.run(connectContext, executor); + metadataDefault.set(2); + execute.run(connectContext, executor); + + Assertions.assertEquals(Arrays.asList("1", "2"), writtenValues, + "each prepared UPDATE must write the default from its freshly resolved connector schema"); + } + @Test @SuppressWarnings("unchecked") public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { @@ -168,94 +227,79 @@ public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { statementContext.getSnapshot(table, Optional.empty(), Optional.empty()).orElse(null)); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - statementContext.loadSnapshots(table, Optional.empty(), Optional.empty()); + + // ExecuteCommand allocates a fresh StatementContext per EXECUTE, so the next execution must not reuse the + // snapshot pinned on the previous context (a stale snapshot would make a later commit permanently + // invisible). + StatementContext nextContext = preparedStatement.getStatementContext(); + Assertions.assertNotSame(statementContext, nextContext, + "ExecuteCommand allocates a fresh StatementContext per EXECUTE"); + nextContext.loadSnapshots(table, Optional.empty(), Optional.empty()); Assertions.assertSame(second, - statementContext.getSnapshot(table, Optional.empty(), Optional.empty()).orElse(null)); + nextContext.getSnapshot(table, Optional.empty(), Optional.empty()).orElse(null)); Mockito.verify(table, Mockito.times(2)).loadSnapshot(Optional.empty(), Optional.empty()); } @Test - public void testExternalScanTasksUseANewGenerationForEveryExecute() throws Exception { - String sql = "select 1"; + public void testFastPathInstallsCachedShortCircuitContextAcrossExecutions() throws Exception { + // ExecuteCommand allocates a fresh StatementContext per EXECUTE. The fresh context carries the + // short-circuit flag but not the cached plan, so the fast path must install the just-validated + // ShortCircuitQueryContext before direct execution -- otherwise result sending falls back to + // `new ShortCircuitQueryContext(planner, ...)` with a null planner (this path never plans) and + // NPEs on planner.getDescTable(). Two executions exercise the second (reusable) EXECUTE that + // hits the regression. + // MUTATION: removing the install in ExecuteCommand.run() -> the fresh context has no + // statement-level cache -> the assertSame below flips -> red. + String sql = "select * from tbl"; LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); StatementContext statementContext = new StatementContext(); + statementContext.setShortCircuitQuery(true); PrepareCommand prepareCommand = new PrepareCommand( "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); PreparedStatementContext preparedStatement = new PreparedStatementContext( prepareCommand, connectContext, statementContext, "stmt"); - StmtExecutor executor = Mockito.mock(StmtExecutor.class); - Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); - Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); - Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); - Mockito.when(executor.getContext()).thenReturn(connectContext); - ExternalScanTaskCacheKey key = new PreparedScanTaskCacheKey("same-scan"); - AtomicInteger loadCount = new AtomicInteger(); - - StatementContext.ExternalScanTaskCache preparedGeneration = - statementContext.getExternalScanTaskCache(); - preparedGeneration.getOrLoad(key, - () -> Collections.singletonList("prepared-" + loadCount.incrementAndGet())); - - new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - StatementContext.ExternalScanTaskCache firstExecuteGeneration = - statementContext.getExternalScanTaskCache(); - Assertions.assertNotSame(preparedGeneration, firstExecuteGeneration); - Assertions.assertEquals(Collections.singletonList("execute-2"), - firstExecuteGeneration.getOrLoad(key, - () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); - - new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - StatementContext.ExternalScanTaskCache secondExecuteGeneration = - statementContext.getExternalScanTaskCache(); - Assertions.assertNotSame(firstExecuteGeneration, secondExecuteGeneration); - Assertions.assertEquals(Collections.singletonList("execute-3"), - secondExecuteGeneration.getOrLoad(key, - () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); - Assertions.assertEquals(3, loadCount.get()); - } - - private static final class PreparedScanTaskCacheKey implements ExternalScanTaskCacheKey { - private final String value; - private PreparedScanTaskCacheKey(String value) { - this.value = value; - } - - @Override - public boolean equals(Object object) { - return object instanceof PreparedScanTaskCacheKey - && value.equals(((PreparedScanTaskCacheKey) object).value); - } - - @Override - public int hashCode() { - return value.hashCode(); - } - } + // A real ShortCircuitQueryContext (built from a mocked planner) that passes isReusable(). + Planner planner = Mockito.mock(Planner.class); + Mockito.when(planner.getQueryOptions()).thenReturn(new TQueryOptions()); + DescriptorTable descriptorTable = new DescriptorTable(); + descriptorTable.createTupleDescriptor(); + Mockito.when(planner.getDescTable()).thenReturn(descriptorTable); + OlapScanNode scanNode = Mockito.mock(OlapScanNode.class); + OlapTable table = Mockito.spy(new OlapTable()); + Mockito.doReturn("tbl").when(table).getName(); + Mockito.doReturn(10).when(table).getBaseSchemaVersion(); + Mockito.when(scanNode.getPointQueryProjectList()).thenReturn(Collections.emptyList()); + Mockito.when(scanNode.getOlapTable()).thenReturn(table); + Mockito.when(scanNode.getTableNameInPlan()).thenReturn("tbl"); + Mockito.when(scanNode.getConjuncts()).thenReturn(Collections.emptyList()); + Mockito.when(planner.getScanNodes()).thenReturn(Collections.singletonList(scanNode)); + ShortCircuitQueryContext cachedPlan = new ShortCircuitQueryContext(planner, Mockito.mock(Queriable.class)); + preparedStatement.shortCircuitQueryContext = Optional.of(cachedPlan); - @Test - public void testIcebergWriteSchemaContextIsResetForEveryExecute() throws Exception { - String sql = "select 1"; - LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); - ConnectContext connectContext = Mockito.mock(ConnectContext.class); - StatementContext statementContext = new StatementContext(); - PrepareCommand prepareCommand = new PrepareCommand( - "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); - PreparedStatementContext preparedStatement = new PreparedStatementContext( - prepareCommand, connectContext, statementContext, "stmt"); StmtExecutor executor = Mockito.mock(StmtExecutor.class); Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); - Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.enableGroupCommitFullPrepare = false; + Mockito.when(connectContext.getSessionVariable()).thenReturn(sessionVariable); Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); Mockito.when(executor.getContext()).thenReturn(connectContext); - statementContext.setIcebergWriteSchemaContext(Optional.of(Mockito.mock( - org.apache.doris.datasource.iceberg.IcebergWriteSchemaContext.class))); - new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - - Assertions.assertFalse(statementContext.getIcebergWriteSchemaContext().isPresent()); + ExecuteCommand execute = new ExecuteCommand("stmt", prepareCommand, statementContext); + execute.run(connectContext, executor); + Assertions.assertSame(cachedPlan, preparedStatement.getStatementContext().getShortCircuitQueryContext(), + "the fast path installs the validated cache on the fresh context (first EXECUTE)"); + Mockito.verify(executor, Mockito.times(1)).executeAndSendResult(Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); + + execute.run(connectContext, executor); + Assertions.assertSame(cachedPlan, preparedStatement.getStatementContext().getShortCircuitQueryContext(), + "the fast path installs the validated cache on the fresh context (second, reusable EXECUTE)"); + Mockito.verify(executor, Mockito.times(2)).executeAndSendResult(Mockito.anyBoolean(), Mockito.anyBoolean(), + Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()); } private String resolveNextSnapshot(TableScanParams scanParams, AtomicInteger snapshotId) { From e5f32613f8d1fb394a100a6f790aa90bceab0e32 Mon Sep 17 00:00:00 2001 From: lichi Date: Mon, 14 Sep 2026 11:00:33 +0800 Subject: [PATCH 2/3] fix compile --- .../trees/plans/commands/ExecuteCommand.java | 4 ---- .../doris/qe/PreparedStatementContext.java | 7 ------ .../plans/commands/ExecuteCommandTest.java | 24 +++---------------- 3 files changed, 3 insertions(+), 32 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java index 87447a382508f2..961cda45fde1ef 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java @@ -72,10 +72,6 @@ public R accept(PlanVisitor visitor, C context) { @Override public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { - StatementContext statementContext = ctx.getStatementContext(); - statementContext.setPrepareStage(false); - statementContext.setIsInsert(false); - statementContext.resetMvccSnapshots(); PreparedStatementContext preparedStmtCtx = ctx.getPreparedStementContext(stmtName); if (null == preparedStmtCtx) { throw new AnalysisException( diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java index d5d9162aa1945c..d34f802ddb1b04 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/PreparedStatementContext.java @@ -67,13 +67,6 @@ public void setStatementContext(StatementContext statementContext) { * @return the fresh StatementContext to use for the current execution */ public StatementContext nextStatementContext() { - // Close the outgoing context's per-statement connector scope before dropping it. The binary - // COM_STMT_EXECUTE path has no per-statement StatementContext.close() finally (that only - // runs for COM_QUERY), and coordinated scans may not have registered a query-finish - // callback yet (connector commands and failures before scan registration have none). - // Without this, the outgoing scope's closeable connector metadata / active connector - // transactions would be abandoned, and GC cannot finalize them. - statementContext.resetConnectorStatementScope(); statementContext = statementContext.createNextExecuteContext(); return statementContext; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java index 9325aa310ab881..379644f03efb9a 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java @@ -22,11 +22,8 @@ import org.apache.doris.analysis.TableScanParams; import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.OlapTable; -import org.apache.doris.catalog.PrimitiveType; -import org.apache.doris.catalog.ScalarType; import org.apache.doris.catalog.TableIf; import org.apache.doris.datasource.CatalogIf; -import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; import org.apache.doris.nereids.StatementContext; @@ -50,7 +47,10 @@ import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; @@ -166,24 +166,8 @@ public void testPreparedConnectorUpdateRefreshesWriteDefaultEveryExecution() thr Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); Mockito.when(executor.getContext()).thenReturn(connectContext); - long tableId = 7L; AtomicInteger metadataDefault = new AtomicInteger(1); List writtenValues = new ArrayList<>(); - Mockito.doAnswer(invocation -> { - // Each execution plans through the fresh StatementContext allocated by ExecuteCommand, so resolve/pin - // the connector writer schema on THAT context (a stale pin would make the second execution reuse - // default 1). - StatementContext currentContext = preparedStatement.getStatementContext(); - if (!currentContext.getConnectorWriteSchema(tableId).isPresent()) { - Column column = new Column("v", ScalarType.createType(PrimitiveType.INT), - false, null, String.valueOf(metadataDefault.get()), ""); - currentContext.setConnectorWriteSchema(tableId, Collections.singletonList(column)); - } - writtenValues.add(currentContext.getConnectorWriteSchema(tableId).get() - .get(0).getDefaultValueSql()); - return null; - }).when(executor).execute(); - ExecuteCommand execute = new ExecuteCommand("stmt", prepareCommand, statementContext); execute.run(connectContext, executor); metadataDefault.set(2); @@ -272,9 +256,7 @@ public void testFastPathInstallsCachedShortCircuitContextAcrossExecutions() thro OlapTable table = Mockito.spy(new OlapTable()); Mockito.doReturn("tbl").when(table).getName(); Mockito.doReturn(10).when(table).getBaseSchemaVersion(); - Mockito.when(scanNode.getPointQueryProjectList()).thenReturn(Collections.emptyList()); Mockito.when(scanNode.getOlapTable()).thenReturn(table); - Mockito.when(scanNode.getTableNameInPlan()).thenReturn("tbl"); Mockito.when(scanNode.getConjuncts()).thenReturn(Collections.emptyList()); Mockito.when(planner.getScanNodes()).thenReturn(Collections.singletonList(scanNode)); ShortCircuitQueryContext cachedPlan = new ShortCircuitQueryContext(planner, Mockito.mock(Queriable.class)); From e718bbfd0666b7dd2bbbbaf1e8adda8e9ea9da8f Mon Sep 17 00:00:00 2001 From: lichi Date: Mon, 14 Sep 2026 15:47:49 +0800 Subject: [PATCH 3/3] fix case --- .../trees/plans/commands/ExecuteCommand.java | 8 +- .../org/apache/doris/qe/StmtExecutor.java | 12 +- .../plans/commands/ExecuteCommandTest.java | 145 +++++++++++++----- .../org/apache/doris/qe/StmtExecutorTest.java | 25 +++ 4 files changed, 146 insertions(+), 44 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java index 961cda45fde1ef..5bedc88e811b1c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java @@ -131,10 +131,12 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { if (logicalPlan instanceof InsertIntoTableCommand || logicalPlan instanceof InsertOverwriteTableCommand || logicalPlan instanceof UpdateCommand) { - ctx.getStatementContext().setIsInsert(true); + statementContext.setIsInsert(true); } - LogicalPlanAdapter planAdapter = new LogicalPlanAdapter( - logicalPlan, executor.getContext().getStatementContext()); + LogicalPlanAdapter planAdapter = new LogicalPlanAdapter(logicalPlan, statementContext); + // Point the executor (and its ConnectContext) at the fresh per-execution context so the + // following execution/result-sending uses it instead of the previous execution's context. + executor.setStatementContext(statementContext); executor.setParsedStmt(planAdapter); boolean hasShortCircuitContext = preparedStmtCtx.shortCircuitQueryContext.isPresent(); boolean shortCircuitContextReusable = hasShortCircuitContext diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 1639a30f5624be..5227806aefcacb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -178,7 +178,7 @@ public class StmtExecutor { private static final Pattern beIpPattern = Pattern.compile("\\[(\\d+):"); private ConnectContext context; - private final StatementContext statementContext; + private StatementContext statementContext; private MysqlSerializer serializer; private OriginStatement originStmt; private StatementBase parsedStmt; @@ -2089,6 +2089,16 @@ public StatementBase setParsedStmt(StatementBase parsedStmt) { return parsedStmt; } + /** + * Replace the executor statement context and synchronize it to the owning ConnectContext. + */ + public void setStatementContext(StatementContext statementContext) { + this.statementContext = statementContext; + this.statementContext.setConnectContext(context); + this.statementContext.setOriginStatement(originStmt); + this.context.setStatementContext(statementContext); + } + public List planPrepareStatementSlots() throws Exception { parseByNereids(); Preconditions.checkState(parsedStmt instanceof LogicalPlanAdapter, diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java index 379644f03efb9a..219ac4dfe2cad9 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java @@ -24,6 +24,7 @@ import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableIf; import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.ExternalScanTaskCacheKey; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; import org.apache.doris.nereids.StatementContext; @@ -47,10 +48,7 @@ import org.junit.jupiter.api.Test; import org.mockito.Mockito; -import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; -import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; @@ -143,40 +141,6 @@ public void testExecutePreparedCommandWithoutPlanChildren() throws Exception { Mockito.verify(executor).execute(); } - @Test - public void testPreparedConnectorUpdateRefreshesWriteDefaultEveryExecution() throws Exception { - // ExecuteCommand allocates a fresh StatementContext per EXECUTE. Model connector metadata changing from - // default 1 to 2 between executions: each execution's planner callback pins the current schema only when - // the fresh context has no schema pinned, then expands DEFAULT(v) and writes the resulting value. - // MUTATION: reusing one StatementContext across executions without dropping connectorWriteSchemas makes - // execution two reuse default 1, so the written values become [1, 1] instead of [1, 2]. - String sql = "update ext_catalog.db.t set v = default(v) where id = 1"; - LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); - Assertions.assertInstanceOf(UpdateCommand.class, logicalPlan); - - ConnectContext connectContext = Mockito.mock(ConnectContext.class); - StatementContext statementContext = new StatementContext(); - PrepareCommand prepareCommand = new PrepareCommand( - "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); - PreparedStatementContext preparedStatement = new PreparedStatementContext( - prepareCommand, connectContext, statementContext, "stmt"); - StmtExecutor executor = Mockito.mock(StmtExecutor.class); - Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); - Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); - Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); - Mockito.when(executor.getContext()).thenReturn(connectContext); - - AtomicInteger metadataDefault = new AtomicInteger(1); - List writtenValues = new ArrayList<>(); - ExecuteCommand execute = new ExecuteCommand("stmt", prepareCommand, statementContext); - execute.run(connectContext, executor); - metadataDefault.set(2); - execute.run(connectContext, executor); - - Assertions.assertEquals(Arrays.asList("1", "2"), writtenValues, - "each prepared UPDATE must write the default from its freshly resolved connector schema"); - } - @Test @SuppressWarnings("unchecked") public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { @@ -212,12 +176,16 @@ public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - // ExecuteCommand allocates a fresh StatementContext per EXECUTE, so the next execution must not reuse the - // snapshot pinned on the previous context (a stale snapshot would make a later commit permanently - // invisible). + // ExecuteCommand allocates a fresh StatementContext per EXECUTE, so the next execution must not + // reuse the snapshot pinned on the previous context (a stale snapshot would make a later commit + // permanently invisible). StatementContext nextContext = preparedStatement.getStatementContext(); Assertions.assertNotSame(statementContext, nextContext, "ExecuteCommand allocates a fresh StatementContext per EXECUTE"); + // The executor (and with it the ConnectContext) must be switched to the fresh context. + // Otherwise execution keeps running on the previous context and the freshly allocated one + // would be dead weight -- the OOM fix would not take effect. + Mockito.verify(executor).setStatementContext(nextContext); nextContext.loadSnapshots(table, Optional.empty(), Optional.empty()); Assertions.assertSame(second, @@ -225,6 +193,103 @@ public void testMvccSnapshotsAreResetForEveryExecute() throws Exception { Mockito.verify(table, Mockito.times(2)).loadSnapshot(Optional.empty(), Optional.empty()); } + @Test + public void testExternalScanTasksUseANewGenerationForEveryExecute() throws Exception { + String sql = "select 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + StatementContext statementContext = new StatementContext(); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + Mockito.when(executor.getContext()).thenReturn(connectContext); + ExternalScanTaskCacheKey key = new PreparedScanTaskCacheKey("same-scan"); + AtomicInteger loadCount = new AtomicInteger(); + + StatementContext.ExternalScanTaskCache preparedGeneration = + statementContext.getExternalScanTaskCache(); + preparedGeneration.getOrLoad(key, + () -> Collections.singletonList("prepared-" + loadCount.incrementAndGet())); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + + // ExecuteCommand allocates a fresh StatementContext per EXECUTE, so every execution gets its own + // external scan task cache instead of reusing tasks loaded by a previous execution. + StatementContext firstExecuteContext = preparedStatement.getStatementContext(); + Assertions.assertNotSame(statementContext, firstExecuteContext, + "ExecuteCommand allocates a fresh StatementContext per EXECUTE"); + StatementContext.ExternalScanTaskCache firstExecuteGeneration = + firstExecuteContext.getExternalScanTaskCache(); + Assertions.assertNotSame(preparedGeneration, firstExecuteGeneration); + Assertions.assertEquals(Collections.singletonList("execute-2"), + firstExecuteGeneration.getOrLoad(key, + () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + StatementContext secondExecuteContext = preparedStatement.getStatementContext(); + Assertions.assertNotSame(firstExecuteContext, secondExecuteContext, + "ExecuteCommand allocates a fresh StatementContext per EXECUTE"); + StatementContext.ExternalScanTaskCache secondExecuteGeneration = + secondExecuteContext.getExternalScanTaskCache(); + Assertions.assertNotSame(firstExecuteGeneration, secondExecuteGeneration); + Assertions.assertEquals(Collections.singletonList("execute-3"), + secondExecuteGeneration.getOrLoad(key, + () -> Collections.singletonList("execute-" + loadCount.incrementAndGet()))); + Assertions.assertEquals(3, loadCount.get()); + } + + private static final class PreparedScanTaskCacheKey implements ExternalScanTaskCacheKey { + private final String value; + + private PreparedScanTaskCacheKey(String value) { + this.value = value; + } + + @Override + public boolean equals(Object object) { + return object instanceof PreparedScanTaskCacheKey + && value.equals(((PreparedScanTaskCacheKey) object).value); + } + + @Override + public int hashCode() { + return value.hashCode(); + } + } + + @Test + public void testIcebergWriteSchemaContextIsResetForEveryExecute() throws Exception { + String sql = "select 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + StatementContext statementContext = new StatementContext(); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + Mockito.when(executor.getContext()).thenReturn(connectContext); + + statementContext.setIcebergWriteSchemaContext(Optional.of(Mockito.mock( + org.apache.doris.datasource.iceberg.IcebergWriteSchemaContext.class))); + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + + // ExecuteCommand allocates a fresh StatementContext per EXECUTE, so the write schema pinned by + // one execution must not leak into the next one. + StatementContext nextContext = preparedStatement.getStatementContext(); + Assertions.assertNotSame(statementContext, nextContext, + "ExecuteCommand allocates a fresh StatementContext per EXECUTE"); + Assertions.assertFalse(nextContext.getIcebergWriteSchemaContext().isPresent()); + } + @Test public void testFastPathInstallsCachedShortCircuitContextAcrossExecutions() throws Exception { // ExecuteCommand allocates a fresh StatementContext per EXECUTE. The fresh context carries the diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java index 76d8e0be6bfbd4..10eb1c0baef347 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java @@ -24,6 +24,7 @@ import org.apache.doris.common.FeConstants; import org.apache.doris.mysql.MysqlChannel; import org.apache.doris.mysql.MysqlSerializer; +import org.apache.doris.nereids.StatementContext; import org.apache.doris.planner.PlanFragment; import org.apache.doris.planner.Planner; import org.apache.doris.planner.ResultFileSink; @@ -106,6 +107,30 @@ public void testDeferForArrowFlightFreezesExecTimeoutInEffect() throws Exception } } + // ExecuteCommand swaps in a fresh per-EXECUTE StatementContext through this setter (see + // ExecuteCommandTest). Both the executor's own field and the owning ConnectContext must follow, + // because the execution and result-sending paths read them (StmtExecutor.executeAndSendResult, + // sendFields). If they keep pointing at the previous context, the freshly allocated one is dead + // weight and the per-execution OOM fix has no effect. + @Test + public void testSetStatementContextSwitchesExecutorAndConnectContext() throws Exception { + StmtExecutor stmtExecutor = new StmtExecutor(connectContext, "select 1"); + StatementContext previous = connectContext.getStatementContext(); + StatementContext fresh = new StatementContext(connectContext, new OriginStatement("select 1", 0)); + Assertions.assertNotSame(previous, fresh); + + stmtExecutor.setStatementContext(fresh); + + Assertions.assertSame(fresh, connectContext.getStatementContext(), + "the ConnectContext must follow the executor onto the fresh context"); + Assertions.assertSame(connectContext, fresh.getConnectContext()); + Assertions.assertEquals("select 1", fresh.getOriginStatement().originStmt); + Field field = StmtExecutor.class.getDeclaredField("statementContext"); + field.setAccessible(true); + Assertions.assertSame(fresh, field.get(stmtExecutor), + "the executor's own field must be switched too"); + } + // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo -> DoGet (see #62259); // it is released later by finalizeArrowFlightQuery(), which closes the coordinator and then // unregisters the query. The close and the unregister must be independent: if coord.close()