[Pipe] Restore sink progress under memory pressure

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java

@@ -24,8 +24,8 @@
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
+import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock;
 import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager;
-import org.apache.iotdb.db.pipe.resource.memory.PipeTabletMemoryBlock;
 import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink;
 import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException;
 import org.apache.iotdb.pipe.api.event.Event;
@@ -51,7 +51,7 @@ public abstract class PipeTabletEventBatch implements AutoCloseable {
   private long firstEventProcessingTime = Long.MIN_VALUE;
 
   protected long totalBufferSize = 0;
-  private final PipeTabletMemoryBlock allocatedMemoryBlock;
+  private final PipeMemoryBlock allocatedMemoryBlock;
   private boolean shouldEmitOnMemoryPressure = false;
 
   protected volatile boolean isClosed = false;
@@ -64,8 +64,7 @@ protected PipeTabletEventBatch(
 
     // limit in buffer size
     this.maxBatchSizeInBytes = requestMaxBatchSizeInBytes;
-    this.allocatedMemoryBlock =
-        PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0);
+    this.allocatedMemoryBlock = PipeDataNodeResourceManager.memory().forceAllocate(0);
     if (recordMetric != null) {
       this.recordMetric = recordMetric;
     } else {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java

@@ -22,6 +22,8 @@
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
+import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
+import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager;
 import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
 
@@ -94,6 +96,29 @@ protected boolean constructBatch(final TabletInsertionEvent event) {
     }
   }
 
+  @Test
+  public void testBatchMemoryIsNotCountedAsTabletMemory() throws Exception {
+    final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory();
+    final PipeRawTabletInsertionEvent event = createEvent(1);
+
+    try {
+      Assert.assertTrue(event.increaseReferenceCount(getClass().getName()));
+      final long tabletMemoryBeforeBatch = memoryManager.getUsedMemorySizeInBytesOfTablets();
+      final long totalMemoryBeforeBatch = memoryManager.getUsedMemorySizeInBytes();
+
+      try (final PipeTabletEventBatch batch =
+          new PipeTabletEventPlainBatch(Integer.MAX_VALUE, Long.MAX_VALUE, null)) {
+        batch.onEvent(event);
+
+        Assert.assertEquals(
+            tabletMemoryBeforeBatch, memoryManager.getUsedMemorySizeInBytesOfTablets());
+        Assert.assertTrue(memoryManager.getUsedMemorySizeInBytes() > totalMemoryBeforeBatch);
+      }
+    } finally {
+      event.clearReferenceCount(getClass().getName());
+    }
+  }
+
   private static PipeRawTabletInsertionEvent createEvent(final int value) {
     final Tablet tablet =
         new Tablet(

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/subtask/PipeAbstractSinkSubtask.java

@@ -361,9 +361,8 @@ protected void handleException(final Event event, final Exception e) {
 
       if (failureType == PipeResourceFailureType.MEMORY_TIMEOUT) {
         PipeLogger.log(LOGGER::info, e, PipeMessages.TEMPORARILY_OUT_OF_MEMORY);
-      } else {
-        sleep4NonReportException();
       }
+      sleep4NonReportException();
     } else if (e instanceof PipeRuntimeSinkNonReportTimeConfigurableException) {
       if (lastExceptionTime == Long.MAX_VALUE) {
         lastExceptionTime = System.currentTimeMillis();

iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/task/PipeSleepIntervalTest.java

@@ -22,9 +22,11 @@
 import org.apache.iotdb.commons.conf.CommonConfig;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeException;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
 import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask;
 import org.apache.iotdb.commons.pipe.config.PipeConfig;
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
+import org.apache.iotdb.pipe.api.event.Event;
 import org.apache.iotdb.pipe.api.exception.PipeConnectionException;
 
 import org.junit.After;
@@ -65,6 +67,10 @@ boolean isAuthenticationFailureException(final Throwable throwable) {
     void sleepWithoutHighPriorityTask(final long sleepMillis) throws InterruptedException {
       sleepIfNoHighPriorityTask(sleepMillis);
     }
+
+    void handle(final Event event, final Exception exception) {
+      handleException(event, exception);
+    }
   }
 
   private long oldPipeSinkSubtaskSleepIntervalInitMs;
@@ -136,4 +142,15 @@ public void testSleepIfNoHighPriorityTaskWaits() throws Exception {
       Assert.assertTrue(System.currentTimeMillis() - startTime >= 15L);
     }
   }
+
+  @Test
+  public void testMemoryTimeoutRetryWaits() {
+    try (final TestSinkSubtask subtask = new TestSinkSubtask()) {
+      final long startTime = System.currentTimeMillis();
+      subtask.handle(null, new PipeRuntimeOutOfMemoryCriticalException("memory unavailable"));
+      Assert.assertTrue(
+          System.currentTimeMillis() - startTime
+              >= PipeConfig.getInstance().getPipeSinkSubtaskSleepIntervalInitMs());
+    }
+  }
 }