diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java index 0b6e376d3b8fe..b2fdeaac4f1f3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.execution.fragment; import org.apache.iotdb.calc.exception.MemoryNotEnoughException; +import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.utils.FileUtils; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; @@ -332,7 +333,9 @@ private void initialize(IDriverScheduler scheduler, boolean isExplainAnalyze) { try { clearShuffleSinkHandle(newState); } catch (Throwable t) { - LOGGER.error(DataNodeQueryMessages.ERRORS_RELEASING_SINK, t); + if (!CommonDescriptor.getInstance().getConfig().isStopping()) { + LOGGER.error(DataNodeQueryMessages.ERRORS_RELEASING_SINK, t); + } } // close the driver after sink is aborted or closed because in driver.close() it @@ -358,8 +361,10 @@ private void initialize(IDriverScheduler scheduler, boolean isExplainAnalyze) { exchangeManager.deRegisterFragmentInstanceFromMemoryPool( instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId(), true); } catch (Throwable t) { - LOGGER.error( - DataNodeQueryMessages.ERRORS_DEREGISTER_FI_FROM_MEMORY_POOL, newState, t); + if (!CommonDescriptor.getInstance().getConfig().isStopping()) { + LOGGER.error( + DataNodeQueryMessages.ERRORS_DEREGISTER_FI_FROM_MEMORY_POOL, newState, t); + } } try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java index 0485de570fe5b..7f84e3ee761a0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java @@ -23,6 +23,7 @@ import org.apache.iotdb.calc.execution.schedule.queue.IndexedBlockingQueue; import org.apache.iotdb.calc.execution.schedule.queue.IndexedBlockingReserveQueue; import org.apache.iotdb.commons.concurrent.ThreadName; +import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.commons.exception.StartupException; import org.apache.iotdb.commons.queryengine.common.SessionInfo; @@ -394,7 +395,9 @@ private void clearDriverTask(DriverTask task) { try { task.getDriver().failed(task.getAbortCause().get()); } catch (Exception e) { - logger.error(DataNodeQueryMessages.CLEAR_DRIVERTASK_FAILED, e); + if (!CommonDescriptor.getInstance().getConfig().isStopping()) { + logger.error(DataNodeQueryMessages.CLEAR_DRIVERTASK_FAILED, e); + } } } if (task.getStatus() == DriverTaskStatus.ABORTED) { @@ -405,7 +408,9 @@ private void clearDriverTask(DriverTask task) { task.getDriverTaskId().getFragmentId().getId(), task.getDriverTaskId().getFragmentInstanceId().getInstanceId())); } catch (Exception e) { - logger.error(DataNodeQueryMessages.CLEAR_DRIVERTASK_FAILED, e); + if (!CommonDescriptor.getInstance().getConfig().isStopping()) { + logger.error(DataNodeQueryMessages.CLEAR_DRIVERTASK_FAILED, e); + } } } } finally { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java index 16f17916eb0c7..828c4b811ac49 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java @@ -112,11 +112,11 @@ public void run() { DNAuditLogger.getInstance().log(fields, () -> logMessage); startWatcher(); + CommonDescriptor.getInstance().getConfig().setStopping(true); // Stop external rpc service firstly. ExternalRPCService.getInstance().stop(); // Reject write operations to make sure all tsfiles will be sealed - CommonDescriptor.getInstance().getConfig().setStopping(true); CommonDescriptor.getInstance() .getConfig() .setNodeStatusWithReason(NodeStatus.ReadOnly, NodeStatus.STOPPING); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index af96360b12e7d..2d9e5683cfdfb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient; import org.apache.iotdb.commons.memory.MemoryManager; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.exception.runtime.MemoryLeakException; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener; import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; @@ -47,10 +48,12 @@ import java.util.ArrayList; import java.util.List; 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 java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -83,6 +86,83 @@ public void testAbortStopsFragmentFetchAfterInFlightRpc() throws Exception { testCancellationStopsFragmentFetchAfterInFlightRpc(true); } + @Test + public void reproduceAbortAfterFragmentInstanceMemoryIsForceDeregistered() throws Exception { + final String queryId = "q0"; + final String fragmentInstanceId = "0.0"; + final String localPlanNodeId = "exchange_0"; + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + final MemoryPool memoryPool = + new MemoryPool("test", new MemoryManager(2 * MOCK_TSBLOCK_SIZE), 2 * MOCK_TSBLOCK_SIZE); + final LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); + Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); + final ExecutorService executorService = Mockito.mock(ExecutorService.class); + final SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + localPlanNodeId, + 0, + localMemoryManager, + executorService, + Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE), + Mockito.mock(SourceHandleListener.class), + Mockito.mock(IClientManager.class)); + + sourceHandle.isBlocked(); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + Assert.assertEquals(MOCK_TSBLOCK_SIZE, sourceHandle.getBufferRetainedSizeInBytes()); + Assert.assertEquals(MOCK_TSBLOCK_SIZE, memoryPool.getReservedBytes()); + + final CountDownLatch memoryDeregistered = new CountDownLatch(1); + final AtomicReference deregisterFailure = new AtomicReference<>(); + final AtomicReference abortFailure = new AtomicReference<>(); + final Thread fragmentInstanceCleanupThread = + new Thread( + () -> { + try { + memoryPool.deRegisterFragmentInstanceFromQueryMemoryMap( + queryId, fragmentInstanceId, true); + } catch (Throwable t) { + deregisterFailure.set(t); + } finally { + memoryDeregistered.countDown(); + } + }); + final Thread driverSchedulerCleanupThread = + new Thread( + () -> { + try { + if (!memoryDeregistered.await(10, TimeUnit.SECONDS)) { + abortFailure.set( + new AssertionError("Timed out waiting for FI memory deregistration")); + return; + } + sourceHandle.abort(); + } catch (Throwable t) { + abortFailure.set(t); + } + }); + + fragmentInstanceCleanupThread.start(); + driverSchedulerCleanupThread.start(); + fragmentInstanceCleanupThread.join(TimeUnit.SECONDS.toMillis(10)); + driverSchedulerCleanupThread.join(TimeUnit.SECONDS.toMillis(10)); + + Assert.assertFalse(fragmentInstanceCleanupThread.isAlive()); + Assert.assertFalse(driverSchedulerCleanupThread.isAlive()); + Assert.assertTrue(deregisterFailure.get() instanceof MemoryLeakException); + Assert.assertTrue(abortFailure.get() instanceof IllegalArgumentException); + Assert.assertEquals( + "RelatedMemoryReserved can't be null when freeing memory", abortFailure.get().getMessage()); + Assert.assertEquals(MOCK_TSBLOCK_SIZE, sourceHandle.getBufferRetainedSizeInBytes()); + Assert.assertEquals(MOCK_TSBLOCK_SIZE, memoryPool.getReservedBytes()); + } + private void testCancellationStopsFragmentFetchAfterInFlightRpc(boolean abort) throws Exception { final String queryId = "q0"; final TEndPoint remoteEndpoint =