Fix NPE when fragment context is removed concurrently

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<FragmentInstanceInfo> 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<FragmentInstanceInfo> 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

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);
+    }
+  }
+}