From 11ec38ec7e7d9ac454cf82adf5cef5a065931d93 Mon Sep 17 00:00:00 2001 From: lancelly <1435078631@qq.com> Date: Wed, 19 Jun 2024 16:53:26 +0800 Subject: [PATCH 1/4] fix --- .../exchange/MPPDataExchangeManager.java | 21 +++++++++++++++---- .../fragment/FragmentInstanceManager.java | 21 +++++++++---------- 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index a38ef923bbb38..49c17b54d7482 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -585,11 +585,13 @@ public MPPDataExchangeManager( ExecutorService executorService, IClientManager mppDataExchangeServiceClientManager) { - this.localMemoryManager = Validate.notNull(localMemoryManager); - this.tsBlockSerdeFactory = Validate.notNull(tsBlockSerdeFactory); - this.executorService = Validate.notNull(executorService); + this.localMemoryManager = Validate.notNull(localMemoryManager, "localMemoryManager is null."); + this.tsBlockSerdeFactory = + Validate.notNull(tsBlockSerdeFactory, "tsBlockSerdeFactory is null."); + this.executorService = Validate.notNull(executorService, "executorService is null."); this.mppDataExchangeServiceClientManager = - Validate.notNull(mppDataExchangeServiceClientManager); + Validate.notNull( + mppDataExchangeServiceClientManager, "mppDataExchangeServiceClientManager is null."); sourceHandles = new ConcurrentHashMap<>(); shuffleSinkHandles = new ConcurrentHashMap<>(); } @@ -607,6 +609,17 @@ public void deRegisterFragmentInstanceFromMemoryPool(String queryId, String frag .deRegisterFragmentInstanceFromQueryMemoryMap(queryId, fragmentInstanceId); } + public void removeSourceHandlesAndShuffleSinkHandlesOfFI(TFragmentInstanceId fragmentInstanceId) { + Map removedSourceHandles = sourceHandles.remove(fragmentInstanceId); + if (removedSourceHandles != null) { + removedSourceHandles.values().forEach(ISourceHandle::abort); + } + ISinkHandle shuffleSinkHandle = shuffleSinkHandles.remove(fragmentInstanceId); + if (shuffleSinkHandle != null) { + shuffleSinkHandle.abort(); + } + } + public LocalMemoryManager getLocalMemoryManager() { return localMemoryManager; } 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 bea961c7525a9..71cb0ab0f98fd 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 @@ -177,6 +177,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( instance.isExplainAnalyze(), exchangeManager); } catch (Throwable t) { + clearFIRelatedResources(instanceId); logger.warn("error when create FragmentInstanceExecution.", t); stateMachine.failed(t); return null; @@ -205,6 +206,14 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( } } + private void clearFIRelatedResources(FragmentInstanceId instanceId) { + // close and remove all the handles of the fragment instance + exchangeManager.removeSourceHandlesAndShuffleSinkHandlesOfFI(instanceId.toThrift()); + // clear MemoryPool + exchangeManager.deRegisterFragmentInstanceFromMemoryPool( + instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId()); + } + private DataNodeQueryContext getOrCreateDataNodeQueryContext(QueryId queryId, int dataNodeFINum) { return dataNodeQueryContextMap.computeIfAbsent( queryId, queryId1 -> new DataNodeQueryContext(dataNodeFINum)); @@ -249,6 +258,7 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( false, exchangeManager); } catch (Throwable t) { + clearFIRelatedResources(instanceId); logger.warn("Execute error caused by ", t); stateMachine.failed(t); return null; @@ -269,17 +279,6 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( } } - /** Aborts a FragmentInstance. keep FragmentInstanceContext for later state tracking */ - public FragmentInstanceInfo abortFragmentInstance(FragmentInstanceId fragmentInstanceId) { - instanceExecution.remove(fragmentInstanceId); - FragmentInstanceContext context = instanceContext.get(fragmentInstanceId); - if (context != null) { - context.abort(); - return context.getInstanceInfo(); - } - return null; - } - /** Cancels a FragmentInstance. */ public FragmentInstanceInfo cancelTask(FragmentInstanceId instanceId, boolean hasThrowable) { logger.debug("[CancelFI]"); From 77ec32f397e0e04aefff1c5d834cc2571ac7a68c Mon Sep 17 00:00:00 2001 From: lancelly <1435078631@qq.com> Date: Wed, 19 Jun 2024 17:39:17 +0800 Subject: [PATCH 2/4] fix --- .../execution/exchange/MPPDataExchangeManager.java | 11 ----------- .../execution/fragment/FragmentInstanceManager.java | 2 +- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 49c17b54d7482..8f46e1945bca2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -609,17 +609,6 @@ public void deRegisterFragmentInstanceFromMemoryPool(String queryId, String frag .deRegisterFragmentInstanceFromQueryMemoryMap(queryId, fragmentInstanceId); } - public void removeSourceHandlesAndShuffleSinkHandlesOfFI(TFragmentInstanceId fragmentInstanceId) { - Map removedSourceHandles = sourceHandles.remove(fragmentInstanceId); - if (removedSourceHandles != null) { - removedSourceHandles.values().forEach(ISourceHandle::abort); - } - ISinkHandle shuffleSinkHandle = shuffleSinkHandles.remove(fragmentInstanceId); - if (shuffleSinkHandle != null) { - shuffleSinkHandle.abort(); - } - } - public LocalMemoryManager getLocalMemoryManager() { return localMemoryManager; } 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 71cb0ab0f98fd..62e8e939f5822 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 @@ -208,7 +208,7 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( private void clearFIRelatedResources(FragmentInstanceId instanceId) { // close and remove all the handles of the fragment instance - exchangeManager.removeSourceHandlesAndShuffleSinkHandlesOfFI(instanceId.toThrift()); + exchangeManager.forceDeregisterFragmentInstance(instanceId.toThrift()); // clear MemoryPool exchangeManager.deRegisterFragmentInstanceFromMemoryPool( instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId()); From 3a2ff6d920effb71add7af06b0182146bac6b481 Mon Sep 17 00:00:00 2001 From: lancelly <1435078631@qq.com> Date: Wed, 19 Jun 2024 18:20:23 +0800 Subject: [PATCH 3/4] add forceDeregister --- .../execution/exchange/MPPDataExchangeManager.java | 5 +++-- .../execution/fragment/FragmentInstanceExecution.java | 2 +- .../execution/fragment/FragmentInstanceManager.java | 2 +- .../db/queryengine/execution/memory/MemoryPool.java | 9 ++++++++- .../db/queryengine/plan/execution/QueryExecution.java | 3 ++- 5 files changed, 15 insertions(+), 6 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 8f46e1945bca2..ac0b411109089 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -603,10 +603,11 @@ public MPPDataExchangeServiceImpl getOrCreateMPPDataExchangeServiceImpl() { return mppDataExchangeService; } - public void deRegisterFragmentInstanceFromMemoryPool(String queryId, String fragmentInstanceId) { + public void deRegisterFragmentInstanceFromMemoryPool( + String queryId, String fragmentInstanceId, boolean forceDeregister) { localMemoryManager .getQueryPool() - .deRegisterFragmentInstanceFromQueryMemoryMap(queryId, fragmentInstanceId); + .deRegisterFragmentInstanceFromQueryMemoryMap(queryId, fragmentInstanceId, forceDeregister); } public LocalMemoryManager getLocalMemoryManager() { 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 d4ba3abd9697c..b31d4e13eeb94 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 @@ -309,7 +309,7 @@ private void initialize(IDriverScheduler scheduler, boolean isExplainAnalyze) { // release memory exchangeManager.deRegisterFragmentInstanceFromMemoryPool( - instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId()); + instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId(), true); if (newState.isFailed()) { scheduler.abortFragmentInstance(instanceId); 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 62e8e939f5822..5280dcc404e37 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 @@ -211,7 +211,7 @@ private void clearFIRelatedResources(FragmentInstanceId instanceId) { exchangeManager.forceDeregisterFragmentInstance(instanceId.toThrift()); // clear MemoryPool exchangeManager.deRegisterFragmentInstanceFromMemoryPool( - instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId()); + instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId(), false); } private DataNodeQueryContext getOrCreateDataNodeQueryContext(QueryId queryId, int dataNodeFINum) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 188833e329ebc..a9ab6ed5d8120 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -181,7 +181,7 @@ public void registerPlanNodeIdToQueryMemoryMap( * @throws MemoryLeakException throw {@link MemoryLeakException} */ public void deRegisterFragmentInstanceFromQueryMemoryMap( - String queryId, String fragmentInstanceId) { + String queryId, String fragmentInstanceId, boolean forceDeregister) { Map> queryRelatedMemory = queryMemoryReservations.get(queryId); if (queryRelatedMemory != null) { Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); @@ -192,6 +192,13 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( hasPotentialMemoryLeak = fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L); } + if (!forceDeregister && hasPotentialMemoryLeak) { + // If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles + // have not been closed. + // We should wait for them to be closed. Make sure this method is called again with + // forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed. + return; + } synchronized (queryMemoryReservations) { queryRelatedMemory.remove(fragmentInstanceId); if (queryRelatedMemory.isEmpty()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 40a188e3d6a7b..8cb8ea48ea827 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -365,7 +365,8 @@ private void cleanUpResultHandle() { .deRegisterFragmentInstanceFromMemoryPool( fragmentInstanceId.queryId, FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId( - fragmentInstanceId)); + fragmentInstanceId), + true); } } From 9391d781fe5bfa5dec768b21dbf1b94886613320 Mon Sep 17 00:00:00 2001 From: lancelly <1435078631@qq.com> Date: Wed, 19 Jun 2024 20:17:36 +0800 Subject: [PATCH 4/4] add back deleted method --- .../execution/fragment/FragmentInstanceManager.java | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 5280dcc404e37..900865dc88701 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 @@ -279,6 +279,17 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( } } + /** Aborts a FragmentInstance. keep FragmentInstanceContext for later state tracking */ + public FragmentInstanceInfo abortFragmentInstance(FragmentInstanceId fragmentInstanceId) { + instanceExecution.remove(fragmentInstanceId); + FragmentInstanceContext context = instanceContext.get(fragmentInstanceId); + if (context != null) { + context.abort(); + return context.getInstanceInfo(); + } + return null; + } + /** Cancels a FragmentInstance. */ public FragmentInstanceInfo cancelTask(FragmentInstanceId instanceId, boolean hasThrowable) { logger.debug("[CancelFI]");