Subscription: poll heartbeat event when there are no consumers (#15307) (#15325)

Author: VGalaxies

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/connector/PipeRealtimePriorityBlockingQueue.java

@@ -161,6 +161,15 @@ public Event waitedPoll() {
     return event;
   }
 
+  @Override
+  public Event peek() {
+    final Event event = pendingQueue.peek();
+    if (Objects.nonNull(event)) {
+      return event;
+    }
+    return tsfileInsertEventDeque.peek();
+  }
+
   @Override
   public void clear() {
     super.clear();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBlockingPendingQueue.java

@@ -33,6 +33,10 @@ public SubscriptionBlockingPendingQueue(
 
   public abstract Event waitedPoll();
 
+  public abstract Event peek();
+
+  public abstract void directOffer(final Event event);
+
   public int size() {
     return inputPendingQueue.size();
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
 import org.apache.iotdb.db.pipe.agent.task.execution.PipeSubtaskExecutorManager;
 import org.apache.iotdb.db.pipe.event.UserDefinedEnrichedEvent;
+import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
 import org.apache.iotdb.db.pipe.event.common.terminate.PipeTerminateEvent;
 import org.apache.iotdb.db.subscription.agent.SubscriptionAgent;
 import org.apache.iotdb.db.subscription.event.SubscriptionEvent;
@@ -269,6 +270,7 @@ public boolean executePrefetch() {
             committedCleaner, pollableNacker, responsePrefetcher, responseSerializer);
         return true;
       } else {
+        peekOnce();
         remapInFlightEventsSnapshot(committedCleaner, pollableNacker);
         return false;
       }
@@ -311,6 +313,30 @@ public void prefetchEvent(@NonNull final SubscriptionEvent thisEvent) {
     prefetchingQueue.add(thisEvent);
   }
 
+  private synchronized void peekOnce() {
+    final Event peekedEvent = UserDefinedEnrichedEvent.maybeOf(inputPendingQueue.peek());
+    if (Objects.isNull(peekedEvent)) {
+      return;
+    }
+
+    if (!(peekedEvent instanceof PipeHeartbeatEvent)) {
+      return;
+    }
+
+    final Event polledEvent = inputPendingQueue.waitedPoll();
+    if (!Objects.equals(peekedEvent, polledEvent)) {
+      LOGGER.warn(
+          "Subscription: inconsistent heartbeat event when {} peeking (broken invariant), expected {}, actual {}, offer back",
+          this,
+          peekedEvent,
+          polledEvent);
+      inputPendingQueue.directOffer(polledEvent);
+    } else {
+      ((PipeHeartbeatEvent) peekedEvent)
+          .decreaseReferenceCount(SubscriptionPrefetchingQueue.class.getName(), false);
+    }
+  }
+
   /**
    * Prefetch at most one {@link SubscriptionEvent} from {@link
    * SubscriptionPrefetchingQueue#inputPendingQueue} to {@link
@@ -319,7 +345,7 @@ public void prefetchEvent(@NonNull final SubscriptionEvent thisEvent) {
    * <p>It will continuously attempt to prefetch and generate a {@link SubscriptionEvent} until
    * {@link SubscriptionPrefetchingQueue#inputPendingQueue} is empty.
    */
-  private void tryPrefetch() {
+  private synchronized void tryPrefetch() {
     while (!inputPendingQueue.isEmpty()) {
       final Event event = UserDefinedEnrichedEvent.maybeOf(inputPendingQueue.waitedPoll());
       if (Objects.isNull(event)) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/TsFileDeduplicationBlockingPendingQueue.java

@@ -58,6 +58,16 @@ public Event waitedPoll() {
     return filter(inputPendingQueue.waitedPoll());
   }
 
+  @Override
+  public Event peek() {
+    return inputPendingQueue.peek();
+  }
+
+  @Override
+  public void directOffer(final Event event) {
+    inputPendingQueue.directOffer(event);
+  }
+
   private synchronized Event filter(final Event event) { // make it synchronized
     if (Objects.isNull(event)) {
       return null;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/connection/BlockingPendingQueue.java

@@ -107,6 +107,10 @@ public E waitedPoll() {
     return event;
   }
 
+  public E peek() {
+    return pendingQueue.peek();
+  }
+
   public void clear() {
     isClosed.set(true);
     pendingQueue.clear();