[To dev/1.3] Load: Fixed multiple bugs (#17565) (#17623)

* Load: Fixed multiple bugs (#17565)

* fix

* test-del

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java

@@ -128,6 +128,34 @@ public LoadTsFileManager() {
     activeLoadAgent.start();
   }
 
+  private long getCleanupTaskDelayInMs() {
+    return CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000L;
+  }
+
+  private void createCleanupTaskIfAbsent(final String uuid) {
+    synchronized (uuid2CleanupTask) {
+      if (uuid2CleanupTask.containsKey(uuid)) {
+        return;
+      }
+
+      final CleanupTask cleanupTask = new CleanupTask(uuid, getCleanupTaskDelayInMs());
+      uuid2CleanupTask.put(uuid, cleanupTask);
+      cleanupTaskQueue.add(cleanupTask);
+    }
+  }
+
+  private void rescheduleCleanupTask(final CleanupTask cleanupTask) {
+    synchronized (uuid2CleanupTask) {
+      if (uuid2CleanupTask.get(cleanupTask.uuid) != cleanupTask) {
+        return;
+      }
+
+      cleanupTaskQueue.remove(cleanupTask);
+      cleanupTask.resetScheduledTime();
+      cleanupTaskQueue.add(cleanupTask);
+    }
+  }
+
   private void registerCleanupTaskExecutor() {
     PipeDataNodeAgent.runtime()
         .registerPeriodicalJob(
@@ -199,26 +227,13 @@ private void recover() {
               uuid2WriterManager.put(uuid, writerManager);
               writerManager.close();
 
-              synchronized (uuid2CleanupTask) {
-                final CleanupTask cleanupTask =
-                    new CleanupTask(
-                        uuid, CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000);
-                uuid2CleanupTask.put(uuid, cleanupTask);
-                cleanupTaskQueue.add(cleanupTask);
-              }
+              createCleanupTaskIfAbsent(uuid);
             });
   }
 
   public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNode, String uuid)
       throws IOException, PageException {
-    if (!uuid2WriterManager.containsKey(uuid)) {
-      synchronized (uuid2CleanupTask) {
-        final CleanupTask cleanupTask =
-            new CleanupTask(uuid, CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000);
-        uuid2CleanupTask.put(uuid, cleanupTask);
-        cleanupTaskQueue.add(cleanupTask);
-      }
-    }
+    createCleanupTaskIfAbsent(uuid);
 
     final Optional<CleanupTask> cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid));
     cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning);
@@ -293,6 +308,8 @@ public boolean loadAll(
       return false;
     }
 
+    createCleanupTaskIfAbsent(uuid);
+
     final Optional<CleanupTask> cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid));
     cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning);
     try {
@@ -315,9 +332,10 @@ public boolean deleteAll(String uuid) {
 
   private void clean(String uuid) {
     synchronized (uuid2CleanupTask) {
-      final CleanupTask cleanupTask = uuid2CleanupTask.get(uuid);
+      final CleanupTask cleanupTask = uuid2CleanupTask.remove(uuid);
       if (cleanupTask != null) {
         cleanupTask.cancel();
+        cleanupTaskQueue.remove(cleanupTask);
       }
     }
 
@@ -744,12 +762,12 @@ private CleanupTask(String uuid, long delayInMs) {
 
     public void markLoadTaskRunning() {
       isLoadTaskRunning = true;
-      resetScheduledTime();
+      rescheduleCleanupTask(this);
     }
 
     public void markLoadTaskNotRunning() {
       isLoadTaskRunning = false;
-      resetScheduledTime();
+      rescheduleCleanupTask(this);
     }
 
     public void resetScheduledTime() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManager.java

@@ -126,7 +126,18 @@ public synchronized LoadTsFileDataCacheMemoryBlock allocateDataCacheMemoryBlock(
     if (dataCacheMemoryBlock == null) {
       final long actuallyAllocateMemoryInBytes =
           tryAllocateFromQuery(MEMORY_TOTAL_SIZE_FROM_QUERY_IN_BYTES >> 2);
-      dataCacheMemoryBlock = new LoadTsFileDataCacheMemoryBlock(actuallyAllocateMemoryInBytes);
+      try {
+        dataCacheMemoryBlock = new LoadTsFileDataCacheMemoryBlock(actuallyAllocateMemoryInBytes);
+      } catch (RuntimeException e) {
+        if (actuallyAllocateMemoryInBytes > 0) {
+          try {
+            releaseToQuery(actuallyAllocateMemoryInBytes);
+          } catch (RuntimeException releaseException) {
+            e.addSuppressed(releaseException);
+          }
+        }
+        throw e;
+      }
       LOGGER.info(
           "Create Data Cache Memory Block {}, allocate memory {}",
           dataCacheMemoryBlock,

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java

@@ -0,0 +1,106 @@
+/*
+ * 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.storageengine.load.memory;
+
+import org.apache.iotdb.db.exception.load.LoadRuntimeOutOfMemoryException;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
+
+public class LoadTsFileMemoryManagerTest {
+
+  @Test
+  public void testAllocateDataCacheMemoryBlockDoesNotDoubleCountMemory() throws Exception {
+    final long allocatedMemoryInBytes = 2L * 1024 * 1024;
+    final LoadTsFileMemoryManager manager = spy(newMemoryManager());
+
+    doAnswer(
+            invocation -> {
+              setUsedMemorySize(manager, allocatedMemoryInBytes);
+              return allocatedMemoryInBytes;
+            })
+        .when(manager)
+        .tryAllocateFromQuery(anyLong());
+
+    manager.allocateDataCacheMemoryBlock();
+
+    Assert.assertEquals(allocatedMemoryInBytes, manager.getUsedMemorySizeInBytes());
+    Assert.assertNotNull(getDataCacheMemoryBlock(manager));
+  }
+
+  @Test
+  public void testAllocateDataCacheMemoryBlockRollsBackPartialAllocationOnFailure()
+      throws Exception {
+    final long allocatedMemoryInBytes = 512L;
+    final LoadTsFileMemoryManager manager = spy(newMemoryManager());
+
+    doAnswer(
+            invocation -> {
+              setUsedMemorySize(manager, allocatedMemoryInBytes);
+              return allocatedMemoryInBytes;
+            })
+        .when(manager)
+        .tryAllocateFromQuery(anyLong());
+    doAnswer(
+            invocation -> {
+              setUsedMemorySize(manager, 0L);
+              return null;
+            })
+        .when(manager)
+        .releaseToQuery(anyLong());
+
+    try {
+      manager.allocateDataCacheMemoryBlock();
+      Assert.fail("Expected LoadRuntimeOutOfMemoryException");
+    } catch (LoadRuntimeOutOfMemoryException e) {
+      Assert.assertEquals(0L, manager.getUsedMemorySizeInBytes());
+      Assert.assertNull(getDataCacheMemoryBlock(manager));
+    }
+  }
+
+  private static LoadTsFileMemoryManager newMemoryManager() throws Exception {
+    final Constructor<LoadTsFileMemoryManager> constructor =
+        LoadTsFileMemoryManager.class.getDeclaredConstructor();
+    constructor.setAccessible(true);
+    return constructor.newInstance();
+  }
+
+  private static void setUsedMemorySize(
+      final LoadTsFileMemoryManager manager, final long usedMemorySizeInBytes) throws Exception {
+    final Field field = LoadTsFileMemoryManager.class.getDeclaredField("usedMemorySizeInBytes");
+    field.setAccessible(true);
+    ((AtomicLong) field.get(manager)).set(usedMemorySizeInBytes);
+  }
+
+  private static LoadTsFileDataCacheMemoryBlock getDataCacheMemoryBlock(
+      final LoadTsFileMemoryManager manager) throws Exception {
+    final Field field = LoadTsFileMemoryManager.class.getDeclaredField("dataCacheMemoryBlock");
+    field.setAccessible(true);
+    return (LoadTsFileDataCacheMemoryBlock) field.get(manager);
+  }
+}