From ba964d8dcd386650eddaffd14d3c33f142ecdb16 Mon Sep 17 00:00:00 2001 From: Christofer Dutz Date: Mon, 2 Nov 2026 11:24:16 +0100 Subject: [PATCH] [Query] Cancel dispatched fragment instances right away when the dispatch fails When the dispatch of a query fails part way, for example because one fragment instance can't get enough operator memory, ClusterScheduler fails the query and SimpleQueryTerminator cancels the fragment instances that were already dispatched. It waits for a grace period of one second before it does so. Those fragment instances can't finish on their own: their consumer is gone, so they keep the operator memory they reserved in LocalExecutionPlanner.checkMemory() until they are cancelled. If more queries arrive within that second, they find the memory taken, fail their dispatch too and hold more memory for another second. With a small memory pool this cascades. Remember in ClusterScheduler that the dispatch failed, and terminate such a query without the grace period. Queries that ran normally keep the grace period. --- .../plan/scheduler/ClusterScheduler.java | 12 +++- .../plan/scheduler/IQueryTerminator.java | 8 +++ .../plan/scheduler/SimpleQueryTerminator.java | 12 +++- .../scheduler/SimpleQueryTerminatorTest.java | 72 +++++++++++++++++++ 4 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java index 4c828c3808cbf..03fa0b666a071 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java @@ -65,6 +65,8 @@ public class ClusterScheduler implements IScheduler { private final IFragInstanceDispatcher dispatcher; private IFragInstanceStateTracker stateTracker; private IQueryTerminator queryTerminator; + // Set when the dispatch of the fragment instances failed (see start()). + private volatile boolean dispatchFailed = false; private static final QueryExecutionMetricSet QUERY_EXECUTION_METRICS = QueryExecutionMetricSet.getInstance(); @@ -122,6 +124,7 @@ public void start() { try { FragInstanceDispatchResult result = dispatchResultFuture.get(); if (!result.isSuccessful()) { + dispatchFailed = true; if (needRetry(result.getFailureStatus())) { stateMachine.transitionToPendingRetry(result.getFailureStatus()); } else { @@ -132,6 +135,7 @@ public void start() { } } catch (InterruptedException | ExecutionException e) { // If the dispatch request cannot be sent or TException is caught, we will retry this query. + dispatchFailed = true; if (e instanceof InterruptedException) { Thread.currentThread().interrupt(); } @@ -162,7 +166,13 @@ public void stop(Throwable t) { stateTracker.abort(); } if (queryTerminator != null) { - queryTerminator.terminate(t); + if (dispatchFailed) { + // The FragmentInstances dispatched before the failure can't finish on their own and hold + // operator memory until they are cancelled, so don't wait for the grace period. + queryTerminator.terminateImmediately(t); + } else { + queryTerminator.terminate(t); + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java index b7be0efdc247c..3be8958684732 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java @@ -22,5 +22,13 @@ import java.util.concurrent.Future; public interface IQueryTerminator { + /** Cancels the unfinished FragmentInstances of the query after a grace period. */ Future terminate(Throwable t); + + /** + * Cancels the unfinished FragmentInstances of the query without a grace period. Used when the + * dispatch failed: the dispatched FragmentInstances can't finish on their own and hold their + * resources until they are cancelled. + */ + Future terminateImmediately(Throwable t); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java index f0e91ae8075f0..8af9117fe0546 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java @@ -86,6 +86,15 @@ private void calculateParameter(final List instances) { @Override public Future terminate(final Throwable t) { + return terminate(t, TERMINATION_GRACE_PERIOD_IN_MS); + } + + @Override + public Future terminateImmediately(final Throwable t) { + return terminate(t, 0L); + } + + private Future terminate(final Throwable t, final long delayInMs) { // For the failure dispatch, the termination should not be triggered because of connection issue this.relatedHost = this.relatedHost.stream() @@ -93,8 +102,7 @@ public Future terminate(final Throwable t) { .collect(Collectors.toList()); this.hasThrowable = Objects.nonNull(t); - return scheduledExecutor.schedule( - this::syncTerminate, TERMINATION_GRACE_PERIOD_IN_MS, TimeUnit.MILLISECONDS); + return scheduledExecutor.schedule(this::syncTerminate, delayInMs, TimeUnit.MILLISECONDS); } public Boolean syncTerminate() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java new file mode 100644 index 0000000000000..911074f2eca35 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.scheduler; + +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.QueryId; + +import org.junit.Before; +import org.junit.Test; + +import java.util.Collections; +import java.util.concurrent.Callable; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class SimpleQueryTerminatorTest { + + private ScheduledExecutorService scheduledExecutor; + private SimpleQueryTerminator terminator; + + @Before + public void setUp() { + scheduledExecutor = mock(ScheduledExecutorService.class); + final MPPQueryContext queryContext = mock(MPPQueryContext.class); + when(queryContext.getQueryId()).thenReturn(new QueryId("test_query")); + when(queryContext.getEndPointBlackList()).thenReturn(Collections.emptySet()); + terminator = + new SimpleQueryTerminator( + scheduledExecutor, + queryContext, + Collections.emptyList(), + null, + mock(IFragInstanceStateTracker.class)); + } + + @Test + public void terminateWaitsForTheGracePeriod() { + terminator.terminate(new RuntimeException("query failed")); + + verify(scheduledExecutor).schedule(any(Callable.class), eq(1000L), eq(TimeUnit.MILLISECONDS)); + } + + @Test + public void terminateImmediatelySkipsTheGracePeriod() { + terminator.terminateImmediately(new RuntimeException("dispatch failed")); + + verify(scheduledExecutor).schedule(any(Callable.class), eq(0L), eq(TimeUnit.MILLISECONDS)); + } +}