Pipe: fixed the bug that processor dies when encountered memory short…

…age for tablets in a period of time (#12042) Co-authored-by: Steve Yurong Su <[email protected]>
apache · Feb 19, 2024 · 1e7c9c0 · 1e7c9c0
1 parent 89b9804
commit 1e7c9c0
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 20 deletions.
diff --git a/...org/apache/iotdb/db/pipe/connector/protocol/thrift/sync/IoTDBThriftSyncClientManager.java b/...org/apache/iotdb/db/pipe/connector/protocol/thrift/sync/IoTDBThriftSyncClientManager.java
@@ -38,7 +38,6 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -74,7 +73,7 @@ public IoTDBThriftSyncClientManager(
     }
   }
 
-  public void checkClientStatusAndTryReconstructIfNecessary() throws IOException {
+  public void checkClientStatusAndTryReconstructIfNecessary() {
     // reconstruct all dead clients
     for (final Map.Entry<TEndPoint, Pair<IoTDBThriftSyncConnectorClient, Boolean>> entry :
         endPoint2ClientAndStatus.entrySet()) {

diff --git a/...c/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeRawTabletInsertionEvent.java b/...c/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeRawTabletInsertionEvent.java
@@ -25,7 +25,7 @@
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.pipe.resource.PipeResourceManager;
-import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock;
+import org.apache.iotdb.db.pipe.resource.memory.PipeTabletMemoryBlock;
 import org.apache.iotdb.pipe.api.access.Row;
 import org.apache.iotdb.pipe.api.collector.RowCollector;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
@@ -42,7 +42,7 @@ public class PipeRawTabletInsertionEvent extends EnrichedEvent implements Tablet
   private final EnrichedEvent sourceEvent;
   private boolean needToReport;
 
-  private PipeMemoryBlock allocatedMemoryBlock;
+  private PipeTabletMemoryBlock allocatedMemoryBlock;
 
   private TabletInsertionDataContainer dataContainer;
 

diff --git a/...re/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java b/...re/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java
@@ -61,7 +61,7 @@ public class PipeMemoryManager {
   // threshold, allocations of memory block for tablets will be rejected.
   private static final double TABLET_MEMORY_REJECT_THRESHOLD =
       PipeConfig.getInstance().getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold();
-  private long usedMemorySizeInBytesOfTablets;
+  private volatile long usedMemorySizeInBytesOfTablets;
 
   private final Set<PipeMemoryBlock> allocatedBlocks = new HashSet<>();
 
@@ -78,8 +78,13 @@ public synchronized PipeMemoryBlock forceAllocate(long sizeInBytes)
     return forceAllocate(sizeInBytes, false);
   }
 
-  public PipeMemoryBlock forceAllocateWithRetry(Tablet tablet)
+  public PipeTabletMemoryBlock forceAllocateWithRetry(Tablet tablet)
       throws PipeRuntimeOutOfMemoryCriticalException {
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      // No need to calculate the tablet size, skip it to save time
+      return new PipeTabletMemoryBlock(0);
+    }
+
     for (int i = 1; i <= MEMORY_ALLOCATE_MAX_RETRIES; i++) {
       if ((double) usedMemorySizeInBytesOfTablets / TOTAL_MEMORY_SIZE_IN_BYTES
           < TABLET_MEMORY_REJECT_THRESHOLD) {
@@ -104,7 +109,8 @@ public PipeMemoryBlock forceAllocateWithRetry(Tablet tablet)
     }
 
     synchronized (this) {
-      final PipeMemoryBlock block = forceAllocate(calculateTabletSizeInBytes(tablet), true);
+      final PipeTabletMemoryBlock block =
+          (PipeTabletMemoryBlock) forceAllocate(calculateTabletSizeInBytes(tablet), true);
       usedMemorySizeInBytesOfTablets += block.getMemoryUsageInBytes();
       return block;
     }
@@ -113,8 +119,9 @@ public PipeMemoryBlock forceAllocateWithRetry(Tablet tablet)
   private PipeMemoryBlock forceAllocate(long sizeInBytes, boolean isForTablet)
       throws PipeRuntimeOutOfMemoryCriticalException {
     if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
-      // No need to consider isForTablet, for memory control is disabled
-      return new PipeMemoryBlock(sizeInBytes);
+      return isForTablet
+          ? new PipeTabletMemoryBlock(sizeInBytes)
+          : new PipeMemoryBlock(sizeInBytes);
     }
 
     for (int i = 1; i <= MEMORY_ALLOCATE_MAX_RETRIES; i++) {
@@ -155,6 +162,11 @@ public synchronized PipeMemoryBlock forceAllocateIfSufficient(
     if (usedThreshold < 0.0f || usedThreshold > 1.0f) {
       return null;
     }
+
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      return new PipeMemoryBlock(sizeInBytes);
+    }
+
     if (TOTAL_MEMORY_SIZE_IN_BYTES - usedMemorySizeInBytes >= sizeInBytes
         && (float) usedMemorySizeInBytes / TOTAL_MEMORY_SIZE_IN_BYTES < usedThreshold) {
       return forceAllocate(sizeInBytes);

diff --git a/...e/datanode/src/main/java/org/apache/iotdb/db/pipe/task/connection/PipeEventCollector.java b/...e/datanode/src/main/java/org/apache/iotdb/db/pipe/task/connection/PipeEventCollector.java
@@ -70,6 +70,8 @@ public synchronized void collect(Event event) {
       } else {
         collectEvent(event);
       }
+    } catch (PipeException e) {
+      throw e;
     } catch (Exception e) {
       throw new PipeException("Error occurred when collecting events from processor.", e);
     }

diff --git a/...ore/datanode/src/main/java/org/apache/iotdb/db/pipe/task/subtask/PipeDataNodeSubtask.java b/...ore/datanode/src/main/java/org/apache/iotdb/db/pipe/task/subtask/PipeDataNodeSubtask.java
@@ -70,8 +70,8 @@ private void onEnrichedEventFailure(@NotNull Throwable throwable) {
           throwable);
     }
 
-    if (retryCount.get() < MAX_RETRY_TIMES) {
-      retryCount.incrementAndGet();
+    retryCount.incrementAndGet();
+    if (retryCount.get() <= MAX_RETRY_TIMES) {
       LOGGER.warn(
           "Retry executing subtask {} (creation time: {}, simple class: {}), retry count [{}/{}]",
           taskID,

diff --git a/...e/src/main/java/org/apache/iotdb/db/pipe/task/subtask/processor/PipeProcessorSubtask.java b/...e/src/main/java/org/apache/iotdb/db/pipe/task/subtask/processor/PipeProcessorSubtask.java
@@ -19,8 +19,10 @@
 
 package org.apache.iotdb.db.pipe.task.subtask.processor;
 
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
 import org.apache.iotdb.commons.pipe.execution.scheduler.PipeSubtaskScheduler;
 import org.apache.iotdb.commons.pipe.task.EventSupplier;
+import org.apache.iotdb.db.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.pipe.event.UserDefinedEnrichedEvent;
 import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
 import org.apache.iotdb.db.pipe.metric.PipeProcessorMetrics;
@@ -139,6 +141,11 @@ protected boolean executeOnce() throws Exception {
         }
       }
       releaseLastEvent(!isClosed.get() && outputEventCollector.hasNoCollectInvocationAfterReset());
+    } catch (PipeRuntimeOutOfMemoryCriticalException e) {
+      LOGGER.info(
+          "Temporarily out of memory in pipe event processing, will wait for the memory to release.",
+          e);
+      return false;
     } catch (Exception e) {
       if (!isClosed.get()) {
         throw new PipeException(
@@ -162,6 +169,10 @@ public void submitSelf() {
     // and the worker will be submitted to the executor
   }
 
+  public boolean isStoppedByException() {
+    return lastEvent instanceof EnrichedEvent && retryCount.get() > MAX_RETRY_TIMES;
+  }
+
   @Override
   public void close() {
     PipeProcessorMetrics.getInstance().deregister(taskID);

diff --git a/...main/java/org/apache/iotdb/db/pipe/task/subtask/processor/PipeProcessorSubtaskWorker.java b/...main/java/org/apache/iotdb/db/pipe/task/subtask/processor/PipeProcessorSubtaskWorker.java
@@ -28,8 +28,6 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.stream.Collectors;
 
-import static org.apache.iotdb.db.pipe.task.subtask.PipeDataNodeSubtask.MAX_RETRY_TIMES;
-
 public class PipeProcessorSubtaskWorker extends WrappedRunnable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(PipeProcessorSubtaskWorker.class);
@@ -71,9 +69,7 @@ private boolean runSubtasks() {
     boolean canSleepBeforeNextRound = true;
 
     for (final PipeProcessorSubtask subtask : subtasks.keySet()) {
-      if (subtask.isClosed()
-          || !subtask.isSubmittingSelf()
-          || MAX_RETRY_TIMES <= subtask.getRetryCount()) {
+      if (subtask.isClosed() || !subtask.isSubmittingSelf() || subtask.isStoppedByException()) {
         continue;
       }
 

diff --git a/...re/node-commons/src/main/java/org/apache/iotdb/commons/pipe/task/subtask/PipeSubtask.java b/...re/node-commons/src/main/java/org/apache/iotdb/commons/pipe/task/subtask/PipeSubtask.java
@@ -161,8 +161,4 @@ public String getTaskID() {
   public long getCreationTime() {
     return creationTime;
   }
-
-  public int getRetryCount() {
-    return retryCount.get();
-  }
 }