From c422ae85357c6c0cc1c765481e3b011804dd41bf Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Tue, 11 Aug 2026 11:24:56 +0800 Subject: [PATCH] Fix fragment failure race with concurrent cancellation --- .../fragment/FragmentInstanceManager.java | 19 ++- .../fragment/FragmentInstanceManagerTest.java | 131 ++++++++++++++++++ 2 files changed, 143 insertions(+), 7 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java 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 5ef70a99354f2..342dec0a7bab7 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.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 FragmentInstanceInfo execDataQueryFragmentInstance( FragmentInstanceId instanceId = instance.getId(); AtomicLong driversCount = new AtomicLong(); try (SetThreadName fragmentInstanceName = new SetThreadName(instanceId.getFullId())) { + AtomicReference failedInstanceInfo = new AtomicReference<>(); FragmentInstanceExecution execution = instanceExecution.computeIfAbsent( instanceId, @@ -228,6 +230,10 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( 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 FragmentInstanceInfo execDataQueryFragmentInstance( }); return execution.getInstanceInfo(); } else { - return createFailedInstanceInfo(instanceId); + return failedInstanceInfo.get(); } } finally { QueryRelatedResourceMetricSet.getInstance() @@ -297,6 +303,7 @@ private DataNodeQueryContext getOrCreateDataNodeQueryContext(QueryId queryId, in public FragmentInstanceInfo execSchemaQueryFragmentInstance( FragmentInstance instance, ISchemaRegion schemaRegion) { FragmentInstanceId instanceId = instance.getId(); + AtomicReference failedInstanceInfo = new AtomicReference<>(); FragmentInstanceExecution execution = instanceExecution.computeIfAbsent( instanceId, @@ -359,6 +366,9 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( 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 FragmentInstanceInfo execSchemaQueryFragmentInstance( }); return execution.getInstanceInfo(); } else { - return createFailedInstanceInfo(instanceId); + return failedInstanceInfo.get(); } } @@ -447,11 +457,6 @@ public TFetchFragmentInstanceStatisticsResp getFragmentInstanceStatistics( 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 0000000000000..d04fb263a31d6 --- /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 executionFuture = + executor.submit( + () -> + dataQuery + ? manager.execDataQueryFragmentInstance( + instance, mock(IDataRegionForQuery.class)) + : manager.execSchemaQueryFragmentInstance( + instance, mock(ISchemaRegion.class))); + + assertTrue(planningStarted.await(10, TimeUnit.SECONDS)); + Future 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); + } + } +}