[Pipe] Fix batch and async resource lifecycle leaks

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java

@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.db.pipe.event.common.tsfile.parser;
 
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
 import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
 import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock;
 
@@ -64,7 +65,17 @@ public long getMemoryUsageInBytes() {
 
     @Override
     public void forceResize(final long newSizeInBytes) {
-      PipeDataNodeResourceManager.memory().forceResize(delegate, newSizeInBytes);
+      final long oldSize = delegate.getMemoryUsageInBytes();
+      if (!PipeDataNodeResourceManager.memory().tryResize(delegate, newSizeInBytes)) {
+        throw new PipeRuntimeOutOfMemoryCriticalException(
+            String.format(
+                "forceResize: failed to allocate memory after %d retries, total memory size %d "
+                    + "bytes, used memory size %d bytes, requested memory size %d bytes",
+                0,
+                PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes(),
+                PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(),
+                newSizeInBytes - oldSize));
+      }
     }
 
     @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java

@@ -690,66 +690,43 @@ public void forceResize(final PipeMemoryBlock block, final long targetSize) {
     resize(block, targetSize, true);
   }
 
+  /**
+   * Attempts a single resize without waiting for other pipe tasks to release memory.
+   *
+   * <p>This is intended for callers that hold payload/batch locks and can actively release memory
+   * after a failed attempt. Waiting in that situation can prevent the caller itself from making
+   * forward progress.
+   */
+  public synchronized boolean tryResize(final PipeMemoryBlock block, final long targetSize) {
+    if (targetSize < 0) {
+      return false;
+    }
+    if (block == null || block.isReleased()) {
+      LOGGER.warn("forceResize: cannot resize a null or released memory block");
+      return false;
+    }
+    return tryResizeInternal(block, targetSize);
+  }
+
   public synchronized void resize(
       final PipeMemoryBlock block, final long targetSize, final boolean force) {
     if (block == null || block.isReleased()) {
       LOGGER.warn("forceResize: cannot resize a null or released memory block");
       return;
     }
 
-    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
-      block.setMemoryUsageInBytes(targetSize);
+    if (tryResizeInternal(block, targetSize)) {
       return;
     }
 
-    final long oldSize = block.getMemoryUsageInBytes();
-    if (oldSize >= targetSize) {
-      usedMemorySizeInBytes -= oldSize - targetSize;
-      if (block instanceof PipeTabletMemoryBlock) {
-        usedMemorySizeInBytesOfTablets -= oldSize - targetSize;
-      }
-      if (block instanceof PipeTsFileMemoryBlock) {
-        usedMemorySizeInBytesOfTsFiles -= oldSize - targetSize;
-      }
-      block.setMemoryUsageInBytes(targetSize);
-
-      // If no memory is used in the block, we can remove it from the allocated blocks.
-      if (targetSize == 0) {
-        allocatedBlocks.remove(block);
-      }
-
-      notifyNextTsFileParserMemoryReservationInternal();
-      this.notifyAll();
-      return;
-    }
-
-    long sizeInBytes = targetSize - oldSize;
+    final long sizeInBytes = targetSize - block.getMemoryUsageInBytes();
     final int memoryAllocateMaxRetries = PipeConfig.getInstance().getPipeMemoryAllocateMaxRetries();
     for (int i = 1; i <= memoryAllocateMaxRetries; i++) {
-      // Dynamically resized data-structure blocks must obey the same admission thresholds as
-      // blocks allocated with a non-zero initial size. Otherwise they can exhaust the pool and
-      // prevent downstream consumers from allocating the memory needed to release them.
-      if (isHardEnoughForResizing(block, sizeInBytes)
-          && getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes) {
-        usedMemorySizeInBytes += sizeInBytes;
-        if (oldSize == 0) {
-          // If the memory block is not registered, we need to register it first.
-          // Otherwise, the memory usage will be inconsistent.
-          // See registerMemoryBlock for more details.
-          allocatedBlocks.add(block);
-        }
-        if (block instanceof PipeTabletMemoryBlock) {
-          usedMemorySizeInBytesOfTablets += sizeInBytes;
-        }
-        if (block instanceof PipeTsFileMemoryBlock) {
-          usedMemorySizeInBytesOfTsFiles += sizeInBytes;
-        }
-        block.setMemoryUsageInBytes(targetSize);
-        return;
-      }
-
       try {
         tryShrinkUntilFreeMemorySatisfy(sizeInBytes);
+        if (tryResizeInternal(block, targetSize)) {
+          return;
+        }
         this.wait(PipeConfig.getInstance().getPipeMemoryAllocateRetryIntervalInMs());
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
@@ -770,6 +747,57 @@ && getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes
     }
   }
 
+  private boolean tryResizeInternal(final PipeMemoryBlock block, final long targetSize) {
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      block.setMemoryUsageInBytes(targetSize);
+      return true;
+    }
+
+    final long oldSize = block.getMemoryUsageInBytes();
+    if (oldSize >= targetSize) {
+      final long releasedSize = oldSize - targetSize;
+      usedMemorySizeInBytes -= releasedSize;
+      if (block instanceof PipeTabletMemoryBlock) {
+        usedMemorySizeInBytesOfTablets -= releasedSize;
+      }
+      if (block instanceof PipeTsFileMemoryBlock) {
+        usedMemorySizeInBytesOfTsFiles -= releasedSize;
+      }
+      block.setMemoryUsageInBytes(targetSize);
+
+      if (targetSize == 0) {
+        allocatedBlocks.remove(block);
+      }
+
+      notifyNextTsFileParserMemoryReservationInternal();
+      this.notifyAll();
+      return true;
+    }
+
+    final long sizeInBytes = targetSize - oldSize;
+    // Dynamically resized data-structure blocks must obey the same admission thresholds as blocks
+    // allocated with a non-zero initial size. Otherwise they can exhaust the pool and prevent
+    // downstream consumers from allocating the memory needed to release them.
+    if (!isHardEnoughForResizing(block, sizeInBytes)
+        || getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes < sizeInBytes) {
+      return false;
+    }
+
+    usedMemorySizeInBytes += sizeInBytes;
+    if (oldSize == 0) {
+      // Zero-sized blocks are registered lazily on their first successful expansion.
+      allocatedBlocks.add(block);
+    }
+    if (block instanceof PipeTabletMemoryBlock) {
+      usedMemorySizeInBytesOfTablets += sizeInBytes;
+    }
+    if (block instanceof PipeTsFileMemoryBlock) {
+      usedMemorySizeInBytesOfTsFiles += sizeInBytes;
+    }
+    block.setMemoryUsageInBytes(targetSize);
+    return true;
+  }
+
   /**
    * Allocate a {@link PipeMemoryBlock} for pipe only if memory used after allocation is less than
    * the specified threshold.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFilePublicResource.java

@@ -37,7 +37,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 
 public class PipeTsFilePublicResource extends PipeTsFileResource {
   private static final Logger LOGGER = LoggerFactory.getLogger(PipeTsFilePublicResource.class);
@@ -52,24 +51,15 @@ public PipeTsFilePublicResource(File hardlinkOrCopiedFile) {
   }
 
   @Override
-  public void close() {
+  public synchronized void close() {
     super.close();
-
-    if (deviceMeasurementsMap != null) {
-      deviceMeasurementsMap = null;
-    }
-
-    if (deviceIsAlignedMap != null) {
-      deviceIsAlignedMap = null;
-    }
-
-    if (measurementDataTypeMap != null) {
-      measurementDataTypeMap = null;
-    }
-
-    if (allocatedMemoryBlock != null) {
-      allocatedMemoryBlock.close();
-      allocatedMemoryBlock = null;
+    deviceMeasurementsMap = null;
+    deviceIsAlignedMap = null;
+    measurementDataTypeMap = null;
+    final PipeMemoryBlock block = allocatedMemoryBlock;
+    allocatedMemoryBlock = null;
+    if (block != null) {
+      block.close();
     }
   }
 
@@ -113,46 +103,53 @@ synchronized boolean cacheDeviceIsAlignedMapIfAbsent(final File tsFile) throws I
     // See if pipe memory is sufficient to be allocated for TsFileSequenceReader.
     // Only allocate when pipe memory used is less than 50%, because memory here
     // is hard to shrink and may consume too much memory.
-    allocatedMemoryBlock =
+    final PipeMemoryBlock readerMemoryBlock =
         PipeDataNodeResourceManager.memory()
             .forceAllocateIfSufficient(
                 PipeConfig.getInstance().getPipeMemoryAllocateForTsFileSequenceReaderInBytes(),
                 MEMORY_SUFFICIENT_THRESHOLD);
-    if (allocatedMemoryBlock == null) {
+    if (readerMemoryBlock == null) {
       LOGGER.info(
           "Failed to cacheDeviceIsAlignedMapIfAbsent for tsfile {}, because memory usage is high",
           tsFile.getPath());
       return false;
     }
 
+    final Map<IDeviceID, Boolean> cachedDeviceIsAlignedMap = new HashMap<>();
     long memoryRequiredInBytes = 0L;
-    try (TsFileSequenceReader sequenceReader =
-        new TsFileSequenceReader(tsFile.getPath(), true, false)) {
-      deviceIsAlignedMap = new HashMap<>();
-      final TsFileDeviceIterator deviceIsAlignedIterator =
-          sequenceReader.getAllDevicesIteratorWithIsAligned();
-      while (deviceIsAlignedIterator.hasNext()) {
-        final Pair<IDeviceID, Boolean> deviceIsAlignedPair = deviceIsAlignedIterator.next();
-        deviceIsAlignedMap.put(deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight());
+    try {
+      try (TsFileSequenceReader sequenceReader =
+          new TsFileSequenceReader(tsFile.getPath(), true, false)) {
+        final TsFileDeviceIterator deviceIsAlignedIterator =
+            sequenceReader.getAllDevicesIteratorWithIsAligned();
+        while (deviceIsAlignedIterator.hasNext()) {
+          final Pair<IDeviceID, Boolean> deviceIsAlignedPair = deviceIsAlignedIterator.next();
+          cachedDeviceIsAlignedMap.put(
+              deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight());
+        }
       }
-      memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(deviceIsAlignedMap);
+      memoryRequiredInBytes +=
+          PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(cachedDeviceIsAlignedMap);
+    } finally {
+      // The reader block is temporary and must never become the persistent metadata block.
+      readerMemoryBlock.close();
     }
-    // Release memory of TsFileSequenceReader.
-    allocatedMemoryBlock.close();
-    allocatedMemoryBlock = null;
 
     // Allocate again for the cached objects.
-    allocatedMemoryBlock =
+    final PipeMemoryBlock cachedMemoryBlock =
         PipeDataNodeResourceManager.memory()
             .forceAllocateIfSufficient(memoryRequiredInBytes, MEMORY_SUFFICIENT_THRESHOLD);
-    if (allocatedMemoryBlock == null) {
+    if (cachedMemoryBlock == null) {
       LOGGER.info(
           "PipeTsFileResource: Failed to cache objects for tsfile {} in cache, because memory usage is high",
           tsFile.getPath());
-      deviceIsAlignedMap = null;
       return false;
     }
 
+    // Publish the map only after its accounting block has been acquired.  Readers never observe
+    // a partially built map or a map without a corresponding memory reservation.
+    deviceIsAlignedMap = cachedDeviceIsAlignedMap;
+    allocatedMemoryBlock = cachedMemoryBlock;
     LOGGER.info("PipeTsFileResource: Cached deviceIsAlignedMap for tsfile {}.", tsFile.getPath());
     return true;
   }
@@ -163,65 +160,75 @@ synchronized boolean cacheObjectsIfAbsent(final File tsFile) throws IOException
         return true;
       } else {
         // Recalculate it again because only deviceIsAligned map is cached
-        allocatedMemoryBlock.close();
+        final PipeMemoryBlock oldMemoryBlock = allocatedMemoryBlock;
         allocatedMemoryBlock = null;
+        deviceIsAlignedMap = null;
+        oldMemoryBlock.close();
       }
     }
 
     // See if pipe memory is sufficient to be allocated for TsFileSequenceReader.
     // Only allocate when pipe memory used is less than 50%, because memory here
     // is hard to shrink and may consume too much memory.
-    allocatedMemoryBlock =
+    final PipeMemoryBlock readerMemoryBlock =
         PipeDataNodeResourceManager.memory()
             .forceAllocateIfSufficient(
                 PipeConfig.getInstance().getPipeMemoryAllocateForTsFileSequenceReaderInBytes(),
                 MEMORY_SUFFICIENT_THRESHOLD);
-    if (allocatedMemoryBlock == null) {
+    if (readerMemoryBlock == null) {
       LOGGER.info(
           "Failed to cacheObjectsIfAbsent for tsfile {}, because memory usage is high",
           tsFile.getPath());
       return false;
     }
 
+    Map<IDeviceID, List<String>> cachedDeviceMeasurementsMap = null;
+    Map<IDeviceID, Boolean> cachedDeviceIsAlignedMap = null;
+    Map<String, TSDataType> cachedMeasurementDataTypeMap = null;
     long memoryRequiredInBytes = 0L;
-    try (TsFileSequenceReader sequenceReader =
-        new TsFileSequenceReader(tsFile.getPath(), true, true)) {
-      deviceMeasurementsMap = sequenceReader.getDeviceMeasurementsMap();
-      memoryRequiredInBytes +=
-          PipeMemoryWeightUtil.memoryOfIDeviceID2StrList(deviceMeasurementsMap);
-
-      if (Objects.isNull(deviceIsAlignedMap)) {
-        deviceIsAlignedMap = new HashMap<>();
+    try {
+      try (TsFileSequenceReader sequenceReader =
+          new TsFileSequenceReader(tsFile.getPath(), true, true)) {
+        cachedDeviceMeasurementsMap = sequenceReader.getDeviceMeasurementsMap();
+        memoryRequiredInBytes +=
+            PipeMemoryWeightUtil.memoryOfIDeviceID2StrList(cachedDeviceMeasurementsMap);
+
+        cachedDeviceIsAlignedMap = new HashMap<>();
         final TsFileDeviceIterator deviceIsAlignedIterator =
             sequenceReader.getAllDevicesIteratorWithIsAligned();
         while (deviceIsAlignedIterator.hasNext()) {
           final Pair<IDeviceID, Boolean> deviceIsAlignedPair = deviceIsAlignedIterator.next();
-          deviceIsAlignedMap.put(deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight());
+          cachedDeviceIsAlignedMap.put(
+              deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight());
         }
-      }
-      memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(deviceIsAlignedMap);
+        memoryRequiredInBytes +=
+            PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(cachedDeviceIsAlignedMap);
 
-      measurementDataTypeMap = sequenceReader.getFullPathDataTypeMap();
-      memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfStr2TSDataType(measurementDataTypeMap);
+        cachedMeasurementDataTypeMap = sequenceReader.getFullPathDataTypeMap();
+        memoryRequiredInBytes +=
+            PipeMemoryWeightUtil.memoryOfStr2TSDataType(cachedMeasurementDataTypeMap);
+      }
+    } finally {
+      // The reader block is temporary and must be released even when metadata traversal fails.
+      readerMemoryBlock.close();
     }
-    // Release memory of TsFileSequenceReader.
-    allocatedMemoryBlock.close();
-    allocatedMemoryBlock = null;
 
     // Allocate again for the cached objects.
-    allocatedMemoryBlock =
+    final PipeMemoryBlock cachedMemoryBlock =
         PipeDataNodeResourceManager.memory()
             .forceAllocateIfSufficient(memoryRequiredInBytes, MEMORY_SUFFICIENT_THRESHOLD);
-    if (allocatedMemoryBlock == null) {
+    if (cachedMemoryBlock == null) {
       LOGGER.info(
           "PipeTsFileResource: Failed to cache objects for tsfile {} in cache, because memory usage is high",
           tsFile.getPath());
-      deviceIsAlignedMap = null;
-      deviceMeasurementsMap = null;
-      measurementDataTypeMap = null;
       return false;
     }
 
+    // Publish all metadata only after the persistent accounting block is ready.
+    deviceMeasurementsMap = cachedDeviceMeasurementsMap;
+    deviceIsAlignedMap = cachedDeviceIsAlignedMap;
+    measurementDataTypeMap = cachedMeasurementDataTypeMap;
+    allocatedMemoryBlock = cachedMemoryBlock;
     LOGGER.info("PipeTsFileResource: Cached objects for tsfile {}.", tsFile.getPath());
     return true;
   }