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..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 @@ -72,16 +72,20 @@ 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( "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) { @@ -127,17 +131,27 @@ 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); - // 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..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 @@ -46,6 +46,31 @@ 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() { + statementContext = statementContext.createNextExecuteContext(); + return statementContext; + } + public void setStartTime() { startTime = System.currentTimeMillis(); } 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 ca0d8306f535c9..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 @@ -17,8 +17,11 @@ 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.TableIf; import org.apache.doris.datasource.CatalogIf; import org.apache.doris.datasource.ExternalScanTaskCacheKey; @@ -30,11 +33,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; @@ -168,10 +175,21 @@ 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"); + // 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, - 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()); } @@ -199,16 +217,25 @@ public void testExternalScanTasksUseANewGenerationForEveryExecute() throws Excep () -> 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 = - statementContext.getExternalScanTaskCache(); + 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 = - statementContext.getExternalScanTaskCache(); + secondExecuteContext.getExternalScanTaskCache(); Assertions.assertNotSame(firstExecuteGeneration, secondExecuteGeneration); Assertions.assertEquals(Collections.singletonList("execute-3"), secondExecuteGeneration.getOrLoad(key, @@ -255,7 +282,71 @@ public void testIcebergWriteSchemaContextIsResetForEveryExecute() throws Excepti org.apache.doris.datasource.iceberg.IcebergWriteSchemaContext.class))); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - Assertions.assertFalse(statementContext.getIcebergWriteSchemaContext().isPresent()); + // 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 + // 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"); + + // 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.getOlapTable()).thenReturn(table); + 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); + + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + 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); + + 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) { 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()