This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new ab881319f22 Fix fragment failure race with concurrent cancellation
(#18439)
ab881319f22 is described below
commit ab881319f22df935abb6a4185037c41c8b3e3575
Author: Jackie Tien <[email protected]>
AuthorDate: Tue Aug 11 14:09:18 2026 +0800
Fix fragment failure race with concurrent cancellation (#18439)
---
.../fragment/FragmentInstanceManager.java | 19 +--
.../fragment/FragmentInstanceManagerTest.java | 131 +++++++++++++++++++++
2 files changed, 143 insertions(+), 7 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
index 5ef70a99354..342dec0a7ba 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
@@ -64,6 +64,7 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
import static java.util.Objects.requireNonNull;
import static
org.apache.iotdb.calc.execution.schedule.queue.IndexedBlockingQueue.TOO_MANY_CONCURRENT_QUERIES_ERROR_MSG;
@@ -141,6 +142,7 @@ public class FragmentInstanceManager {
FragmentInstanceId instanceId = instance.getId();
AtomicLong driversCount = new AtomicLong();
try (SetThreadName fragmentInstanceName = new
SetThreadName(instanceId.getFullId())) {
+ AtomicReference<FragmentInstanceInfo> failedInstanceInfo = new
AtomicReference<>();
FragmentInstanceExecution execution =
instanceExecution.computeIfAbsent(
instanceId,
@@ -228,6 +230,10 @@ public class FragmentInstanceManager {
DataNodeQueryMessages.ERROR_WHEN_CREATE_FRAGMENTINSTANCEEXECUTION, t);
stateMachine.failed(t);
}
+ // cancelTask may remove the context from instanceContext
while this execution is
+ // still being created. Capture the failure result from the
local context before
+ // returning from computeIfAbsent instead of looking it up
from the map later.
+ failedInstanceInfo.set(context.getInstanceInfo());
clearFIRelatedResources(instanceId);
return null;
}
@@ -244,7 +250,7 @@ public class FragmentInstanceManager {
});
return execution.getInstanceInfo();
} else {
- return createFailedInstanceInfo(instanceId);
+ return failedInstanceInfo.get();
}
} finally {
QueryRelatedResourceMetricSet.getInstance()
@@ -297,6 +303,7 @@ public class FragmentInstanceManager {
public FragmentInstanceInfo execSchemaQueryFragmentInstance(
FragmentInstance instance, ISchemaRegion schemaRegion) {
FragmentInstanceId instanceId = instance.getId();
+ AtomicReference<FragmentInstanceInfo> failedInstanceInfo = new
AtomicReference<>();
FragmentInstanceExecution execution =
instanceExecution.computeIfAbsent(
instanceId,
@@ -359,6 +366,9 @@ public class FragmentInstanceManager {
logger.warn(DataNodeQueryMessages.EXECUTE_ERROR_CAUSED_BY,
t);
stateMachine.failed(t);
}
+ // See execDataQueryFragmentInstance for why this result must
not be fetched from
+ // instanceContext after computeIfAbsent returns.
+ failedInstanceInfo.set(context.getInstanceInfo());
clearFIRelatedResources(instanceId);
return null;
}
@@ -374,7 +384,7 @@ public class FragmentInstanceManager {
});
return execution.getInstanceInfo();
} else {
- return createFailedInstanceInfo(instanceId);
+ return failedInstanceInfo.get();
}
}
@@ -447,11 +457,6 @@ public class FragmentInstanceManager {
return statisticsResp == null ? new TFetchFragmentInstanceStatisticsResp()
: statisticsResp;
}
- private FragmentInstanceInfo createFailedInstanceInfo(FragmentInstanceId
instanceId) {
- FragmentInstanceContext context = instanceContext.get(instanceId);
- return context.getInstanceInfo();
- }
-
private void removeOldInstances() {
long oldestAllowedInstance = System.currentTimeMillis() -
infoCacheTime.toMillis();
instanceContext
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java
new file mode 100644
index 00000000000..d04fb263a31
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java
@@ -0,0 +1,131 @@
+/*
+ * 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.execution.fragment;
+
+import org.apache.iotdb.commons.exception.MetadataLeaseFencedException;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment;
+import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion;
+import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import static
org.apache.iotdb.commons.exception.MetadataLeaseFencedException.LeaseFencedRetryPolicy.RETRY_UNTIL_SUCCESS;
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class FragmentInstanceManagerTest {
+
+ private static final String FAILURE_MESSAGE = "metadata lease fenced";
+
+ @BeforeClass
+ public static void setUpClass() {
+ IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1);
+ }
+
+ @Test
+ public void testDataQueryFailureInfoAfterConcurrentContextRemoval() throws
Exception {
+ assertFailureInfoAfterConcurrentContextRemoval(true);
+ }
+
+ @Test
+ public void testSchemaQueryFailureInfoAfterConcurrentContextRemoval() throws
Exception {
+ assertFailureInfoAfterConcurrentContextRemoval(false);
+ }
+
+ private void assertFailureInfoAfterConcurrentContextRemoval(boolean
dataQuery) throws Exception {
+ String queryType = dataQuery ? "data" : "schema";
+ FragmentInstanceId instanceId =
+ new FragmentInstanceId(
+ new PlanFragmentId(new QueryId(queryType + "_context_removal"),
0), "0");
+ FragmentInstance instance = mock(FragmentInstance.class);
+ PlanFragment fragment = mock(PlanFragment.class);
+ CountDownLatch planningStarted = new CountDownLatch(1);
+ CountDownLatch failPlanning = new CountDownLatch(1);
+
+ when(instance.getId()).thenReturn(instanceId);
+ when(instance.getFragment()).thenReturn(fragment);
+ when(instance.getDataNodeFINum()).thenReturn(1);
+ when(instance.getTimeOut()).thenReturn(TimeUnit.SECONDS.toMillis(30));
+ when(fragment.getPlanNodeTree())
+ .thenAnswer(
+ ignored -> {
+ planningStarted.countDown();
+ if (!failPlanning.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError("Timed out waiting to fail planning");
+ }
+ throw new MetadataLeaseFencedException(FAILURE_MESSAGE,
RETRY_UNTIL_SUCCESS);
+ });
+
+ FragmentInstanceManager manager = FragmentInstanceManager.getInstance();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ try {
+ Future<FragmentInstanceInfo> executionFuture =
+ executor.submit(
+ () ->
+ dataQuery
+ ? manager.execDataQueryFragmentInstance(
+ instance, mock(IDataRegionForQuery.class))
+ : manager.execSchemaQueryFragmentInstance(
+ instance, mock(ISchemaRegion.class)));
+
+ assertTrue(planningStarted.await(10, TimeUnit.SECONDS));
+ Future<FragmentInstanceInfo> cancellationFuture =
+ executor.submit(() -> manager.cancelTask(instanceId, true));
+
+ // cancelTask removes the context before it waits for the in-progress
computeIfAbsent on the
+ // execution map. This recreates the race that previously made the
failure path look up a
+ // null context.
+ await().atMost(10, TimeUnit.SECONDS).until(() ->
manager.getInstanceInfo(instanceId) == null);
+ failPlanning.countDown();
+
+ FragmentInstanceInfo failureInfo = executionFuture.get(10,
TimeUnit.SECONDS);
+ assertNotNull(failureInfo);
+ assertTrue(failureInfo.getState().isFailed());
+ assertEquals(FAILURE_MESSAGE, failureInfo.getMessage());
+ assertTrue(failureInfo.getErrorCode().isPresent());
+ assertEquals(
+ TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode(),
+ failureInfo.getErrorCode().get().getCode());
+ assertNotNull(cancellationFuture.get(10, TimeUnit.SECONDS));
+ } finally {
+ failPlanning.countDown();
+ executor.shutdownNow();
+ assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+ manager.cancelTask(instanceId, true);
+ }
+ }
+}
