Allow heartbeat retries while a realtime TsFile is in flight

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionTsFileSource.java

@@ -42,6 +42,7 @@
 import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.Deque;
+import java.util.Iterator;
 import java.util.PriorityQueue;
 import java.util.Queue;
 
@@ -177,7 +178,7 @@ public void close() throws Exception {
   public Event supply() {
     if (isRegionLevelDowngradingEnabled) {
       synchronized (tsFileOrderingLock) {
-        if (isClosed.get() || inFlightTsFileEpoch != null) {
+        if (isClosed.get()) {
           return null;
         }
 
@@ -191,21 +192,37 @@ public Event supply() {
             bufferedEvents.offerLast(pendingEvent);
           }
         }
-        return supplyInternal();
+        // A heartbeat may trigger retry of the in-flight TsFile at the sink. Keep files and other
+        // events ordered behind it, but do not block those retries while waiting for its commit.
+        return supplyInternal(inFlightTsFileEpoch != null);
       }
     }
-    return supplyInternal();
+    return supplyInternal(false);
   }
 
-  private PipeRealtimeEvent pollNextEvent() {
+  private PipeRealtimeEvent pollNextEvent(final boolean heartbeatOnly) {
     if (!isRegionLevelDowngradingEnabled) {
       return (PipeRealtimeEvent) pendingQueue.directPoll();
     }
+    if (heartbeatOnly) {
+      final Iterator<PipeRealtimeEvent> iterator = bufferedEvents.iterator();
+      while (iterator.hasNext()) {
+        final PipeRealtimeEvent event = iterator.next();
+        if (event.getEvent() instanceof PipeHeartbeatEvent) {
+          if (((PipeHeartbeatEvent) event.getEvent()).isCompletionBarrier()) {
+            return null;
+          }
+          iterator.remove();
+          return event;
+        }
+      }
+      return null;
+    }
     return orderedTsFileEvents.isEmpty() ? bufferedEvents.pollFirst() : orderedTsFileEvents.poll();
   }
 
-  private Event supplyInternal() {
-    PipeRealtimeEvent realtimeEvent = pollNextEvent();
+  private Event supplyInternal(final boolean heartbeatOnly) {
+    PipeRealtimeEvent realtimeEvent = pollNextEvent(heartbeatOnly);
 
     while (realtimeEvent != null) {
       Event suppliedEvent = null;
@@ -255,7 +272,7 @@ private Event supplyInternal() {
         return suppliedEvent;
       }
 
-      realtimeEvent = pollNextEvent();
+      realtimeEvent = pollNextEvent(heartbeatOnly);
     }
 
     // means the pending queue is empty.

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/source/PipeRealtimeExtractTest.java

@@ -1333,6 +1333,60 @@ public void testTsFileSourceRegionLevelDowngradingUsesPipeParameters() throws Ex
     }
   }
 
+  @Test
+  public void testTsFileSourceRegionLevelDowngradingSuppliesHeartbeatWhileTsFileIsInFlight()
+      throws Exception {
+    final PipeTaskMeta pipeTaskMeta = new PipeTaskMeta(MinimumProgressIndex.INSTANCE, 1);
+    try (final PipeRealtimeDataRegionTsFileSource extractor =
+        createTsFileSourceForTest(
+            PipeSourceConstant.SOURCE_REALTIME_REGION_LEVEL_DOWNGRADING_KEY, pipeTaskMeta)) {
+      final PipeRealtimeEvent firstEvent =
+          extractTsFileForTest(
+              extractor,
+              pipeTaskMeta,
+              createTsFileResource(dataRegion1, "123-1-0-0.tsfile"),
+              false);
+      final PipeRealtimeEvent secondEvent =
+          extractTsFileForTest(
+              extractor,
+              pipeTaskMeta,
+              createTsFileResource(dataRegion1, "123-2-0-0.tsfile"),
+              false);
+      final PipeRealtimeEvent heartbeat =
+          PipeRealtimeEventFactory.createRealtimeEvent(dataRegion1, false);
+      Assert.assertTrue(heartbeat.increaseReferenceCount(TEST_REFERENCE_HOLDER));
+      extractor.extract(heartbeat);
+      final PipeRealtimeEvent completionBarrier =
+          PipeRealtimeEventFactory.createCompletionBarrierEvent(dataRegion1);
+      Assert.assertTrue(completionBarrier.increaseReferenceCount(TEST_REFERENCE_HOLDER));
+      extractor.extract(completionBarrier);
+
+      // Drain the pending queue before adding another heartbeat so heartbeat coalescing does not
+      // discard it behind the completion barrier.
+      Assert.assertSame(firstEvent.getEvent(), extractor.supply());
+      final PipeRealtimeEvent heartbeatAfterCompletionBarrier =
+          PipeRealtimeEventFactory.createRealtimeEvent(dataRegion1, false);
+      Assert.assertTrue(
+          heartbeatAfterCompletionBarrier.increaseReferenceCount(TEST_REFERENCE_HOLDER));
+      extractor.extract(heartbeatAfterCompletionBarrier);
+
+      Assert.assertSame(heartbeat.getEvent(), extractor.supply());
+      Assert.assertNull(extractor.supply());
+      Assert.assertFalse(secondEvent.getEvent().isReleased());
+      Assert.assertFalse(completionBarrier.getEvent().isReleased());
+      Assert.assertFalse(heartbeatAfterCompletionBarrier.getEvent().isReleased());
+
+      releaseSuppliedEvent(heartbeat.getEvent());
+      releaseSuppliedEvent(firstEvent.getEvent());
+      Assert.assertSame(secondEvent.getEvent(), extractor.supply());
+      releaseSuppliedEvent(secondEvent.getEvent());
+      Assert.assertSame(completionBarrier.getEvent(), extractor.supply());
+      releaseSuppliedEvent(completionBarrier.getEvent());
+      Assert.assertSame(heartbeatAfterCompletionBarrier.getEvent(), extractor.supply());
+      releaseSuppliedEvent(heartbeatAfterCompletionBarrier.getEvent());
+    }
+  }
+
   @Test
   public void testTsFileSourceRegionLevelDowngradingFollowsQueryPriority() throws Exception {
     final PipeEventCommitManager commitManager = PipeEventCommitManager.getInstance();