Subscription: fix consumer infinite pulling event & fully managed tsfile parsing process & increase the reference count for subscribed parsed raw tablet event & disrupt parsing requests through the introduction of randomness & disable prefetch by default (#14856) (#14886)

VGalaxies · web-flow · commit f0fa1fd801a4 · 2025-02-19T10:03:20.000+08:00
diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/config/TopicConfig.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/config/TopicConfig.java
@@ -53,6 +53,8 @@ public TopicConfig(final Map<String, String> attributes) {
 
   private static final Map<String, String> SINK_TABLET_FORMAT_CONFIG =
       Collections.singletonMap("format", "tablet");
+  private static final Map<String, String> SINK_TS_FILE_FORMAT_CONFIG =
+      Collections.singletonMap("format", "tsfile");
 
   private static final Map<String, String> SNAPSHOT_MODE_CONFIG =
       Collections.singletonMap("mode", MODE_SNAPSHOT_VALUE);
@@ -149,7 +151,12 @@ public Map<String, String> getAttributesWithProcessorPrefix() {
   }
 
   public Map<String, String> getAttributesWithSinkFormat() {
-    return Collections.emptyMap(); // default to hybrid
+    // refer to
+    // org.apache.iotdb.db.pipe.agent.task.connection.PipeEventCollector.parseAndCollectEvent(org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent)
+    return TopicConstant.FORMAT_TS_FILE_HANDLER_VALUE.equalsIgnoreCase(
+            attributes.getOrDefault(TopicConstant.FORMAT_KEY, TopicConstant.FORMAT_DEFAULT_VALUE))
+        ? SINK_TS_FILE_FORMAT_CONFIG
+        : SINK_TABLET_FORMAT_CONFIG;
   }
 
   public Map<String, String> getAttributesWithSinkPrefix() {
diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/payload/poll/TabletsPayload.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/subscription/payload/poll/TabletsPayload.java
@@ -40,8 +40,8 @@ public class TabletsPayload implements SubscriptionPollPayload {
    * <ul>
    *   <li>If nextOffset is 1, it indicates that the current payload is the first payload (its
    *       tablets are empty) and the fetching should continue.
-   *   <li>If nextOffset is negative, it indicates all tablets have been fetched, and -nextOffset
-   *       represents the total number of tablets.
+   *   <li>If nextOffset is negative (or zero), it indicates all tablets have been fetched, and
+   *       -nextOffset represents the total number of tablets.
    * </ul>
    */
   private transient int nextOffset;
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionConsumer.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/subscription/consumer/SubscriptionConsumer.java
@@ -926,7 +926,7 @@ private Optional<SubscriptionMessage> pollTabletsInternal(
 
     int nextOffset = ((TabletsPayload) initialResponse.getPayload()).getNextOffset();
     while (true) {
-      if (nextOffset < 0) {
+      if (nextOffset <= 0) {
         if (!Objects.equals(tablets.size(), -nextOffset)) {
           final String errorMessage =
               String.format(
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java
@@ -49,6 +49,7 @@
 import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_FORMAT_HYBRID_VALUE;
 import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_FORMAT_KEY;
 import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_FORMAT_TABLET_VALUE;
+import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.CONNECTOR_FORMAT_TS_FILE_VALUE;
 import static org.apache.iotdb.commons.pipe.config.constant.PipeConnectorConstant.SINK_FORMAT_KEY;
 
 public class PipeDataNodeTaskBuilder {
@@ -139,7 +140,16 @@ public PipeDataNodeTask build() {
                 .getStringOrDefault(
                     Arrays.asList(CONNECTOR_FORMAT_KEY, SINK_FORMAT_KEY),
                     CONNECTOR_FORMAT_HYBRID_VALUE)
-                .equals(CONNECTOR_FORMAT_TABLET_VALUE));
+                .equals(CONNECTOR_FORMAT_TABLET_VALUE),
+            PipeType.SUBSCRIPTION.equals(pipeType)
+                &&
+                // should not skip parsing when the format is tsfile
+                !pipeStaticMeta
+                    .getConnectorParameters()
+                    .getStringOrDefault(
+                        Arrays.asList(CONNECTOR_FORMAT_KEY, SINK_FORMAT_KEY),
+                        CONNECTOR_FORMAT_HYBRID_VALUE)
+                    .equals(CONNECTOR_FORMAT_TS_FILE_VALUE));
 
     return new PipeDataNodeTask(
         pipeStaticMeta.getPipeName(), regionId, extractorStage, processorStage, connectorStage);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/connection/PipeEventCollector.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/connection/PipeEventCollector.java
@@ -54,6 +54,8 @@ public class PipeEventCollector implements EventCollector {
 
   private final boolean forceTabletFormat;
 
+  private final boolean skipParseTsFile;
+
   private final AtomicInteger collectInvocationCount = new AtomicInteger(0);
   private boolean hasNoGeneratedEvent = true;
   private boolean isFailedToIncreaseReferenceCount = false;
@@ -62,11 +64,13 @@ public PipeEventCollector(
       final UnboundedBlockingPendingQueue<Event> pendingQueue,
       final long creationTime,
       final int regionId,
-      final boolean forceTabletFormat) {
+      final boolean forceTabletFormat,
+      final boolean skipParseTsFile) {
     this.pendingQueue = pendingQueue;
     this.creationTime = creationTime;
     this.regionId = regionId;
     this.forceTabletFormat = forceTabletFormat;
+    this.skipParseTsFile = skipParseTsFile;
   }
 
   @Override
@@ -120,6 +124,11 @@ private void parseAndCollectEvent(final PipeTsFileInsertionEvent sourceEvent) th
       return;
     }
 
+    if (skipParseTsFile) {
+      collectEvent(sourceEvent);
+      return;
+    }
+
     if (!forceTabletFormat && !sourceEvent.shouldParseTimeOrPattern()) {
       collectEvent(sourceEvent);
       return;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/stage/PipeTaskProcessorStage.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/stage/PipeTaskProcessorStage.java
@@ -68,7 +68,8 @@ public PipeTaskProcessorStage(
       final UnboundedBlockingPendingQueue<Event> pipeConnectorOutputPendingQueue,
       final PipeProcessorSubtaskExecutor executor,
       final PipeTaskMeta pipeTaskMeta,
-      final boolean forceTabletFormat) {
+      final boolean forceTabletFormat,
+      final boolean skipParseTsFile) {
     final PipeProcessorRuntimeConfiguration runtimeConfiguration =
         new PipeTaskRuntimeConfiguration(
             new PipeTaskProcessorRuntimeEnvironment(
@@ -100,7 +101,11 @@ public PipeTaskProcessorStage(
     final String taskId = pipeName + "_" + regionId + "_" + creationTime;
     final PipeEventCollector pipeConnectorOutputEventCollector =
         new PipeEventCollector(
-            pipeConnectorOutputPendingQueue, creationTime, regionId, forceTabletFormat);
+            pipeConnectorOutputPendingQueue,
+            creationTime,
+            regionId,
+            forceTabletFormat,
+            skipParseTsFile);
     this.pipeProcessorSubtask =
         new PipeProcessorSubtask(
             taskId,
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java
@@ -197,7 +197,8 @@ public SubscriptionEvent pollInternal(final String consumerId) {
               return null;
             },
             SubscriptionAgent.receiver().remainingMs());
-      } catch (final Exception ignored) {
+      } catch (final Exception e) {
+        LOGGER.warn("Exception {} occurred when {} execute receiver subtask", this, e, e);
       }
     }
 
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingTabletQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingTabletQueue.java
@@ -159,7 +159,7 @@ public SubscriptionEvent pollTablets(
                 String.format(
                     "exception occurred when fetching next response: %s, consumer id: %s, commit context: %s, offset: %s, prefetching queue: %s",
                     e, consumerId, commitContext, offset, this);
-            LOGGER.warn(errorMessage);
+            LOGGER.warn(errorMessage, e);
             eventRef.set(generateSubscriptionPollErrorResponse(errorMessage));
           }
 
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingTsFileQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingTsFileQueue.java
@@ -211,7 +211,7 @@ public SubscriptionEvent pollTsFile(
                 String.format(
                     "exception occurred when fetching next response: %s, consumer id: %s, commit context: %s, writing offset: %s, prefetching queue: %s",
                     e, consumerId, commitContext, writingOffset, this);
-            LOGGER.warn(errorMessage);
+            LOGGER.warn(errorMessage, e);
             eventRef.set(generateSubscriptionPollErrorResponse(errorMessage));
           }
 
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java
@@ -37,7 +37,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -61,7 +60,7 @@ public class SubscriptionPipeTabletEventBatch extends SubscriptionPipeEventBatch
   private final Meter insertNodeTabletInsertionEventSizeEstimator;
   private final Meter rawTabletInsertionEventSizeEstimator;
 
-  private volatile List<EnrichedEvent> iteratedEnrichedEvents;
+  private volatile SubscriptionPipeTabletIterationSnapshot iterationSnapshot;
   private final AtomicInteger referenceCount = new AtomicInteger();
 
   private static final long ITERATED_COUNT_REPORT_FREQ =
@@ -88,7 +87,17 @@ public SubscriptionPipeTabletEventBatch(
   @Override
   public synchronized void ack() {
     referenceCount.decrementAndGet();
-    // do nothing for iterated enriched events, see SubscriptionPipeTabletBatchEvents
+
+    // we decrease the reference count of events if and only if when the whole batch is consumed
+    if (!hasNext() && referenceCount.get() == 0) {
+      for (final EnrichedEvent enrichedEvent : enrichedEvents) {
+        if (enrichedEvent instanceof PipeTsFileInsertionEvent) {
+          // close data container in tsfile event
+          ((PipeTsFileInsertionEvent) enrichedEvent).close();
+        }
+        enrichedEvent.decreaseReferenceCount(this.getClass().getName(), true);
+      }
+    }
   }
 
   @Override
@@ -200,26 +209,29 @@ private void updateEstimatedRawTabletInsertionEventSize(final long size) {
 
   /////////////////////////////// iterator ///////////////////////////////
 
-  public List<EnrichedEvent> sendIterationSnapshot() {
-    final List<EnrichedEvent> result = Collections.unmodifiableList(iteratedEnrichedEvents);
-    iteratedEnrichedEvents = new ArrayList<>();
+  public synchronized SubscriptionPipeTabletIterationSnapshot sendIterationSnapshot() {
+    final SubscriptionPipeTabletIterationSnapshot result = iterationSnapshot;
+    iterationSnapshot = new SubscriptionPipeTabletIterationSnapshot();
     referenceCount.incrementAndGet();
     return result;
   }
 
-  public void resetForIteration() {
+  public synchronized void resetForIteration() {
     currentEnrichedEventsIterator = enrichedEvents.iterator();
     currentTabletInsertionEventsIterator = null;
     currentTsFileInsertionEvent = null;
 
-    iteratedEnrichedEvents = new ArrayList<>();
+    if (Objects.nonNull(iterationSnapshot)) {
+      iterationSnapshot.clear();
+    }
+    iterationSnapshot = new SubscriptionPipeTabletIterationSnapshot();
     referenceCount.set(0);
 
     iteratedCount.set(0);
   }
 
   @Override
-  public boolean hasNext() {
+  public synchronized boolean hasNext() {
     if (Objects.nonNull(currentTabletInsertionEventsIterator)) {
       if (currentTabletInsertionEventsIterator.hasNext()) {
         return true;
@@ -245,7 +257,7 @@ public boolean hasNext() {
   }
 
   @Override
-  public List<Tablet> next() {
+  public synchronized List<Tablet> next() {
     final List<Tablet> tablets = nextInternal();
     if (Objects.isNull(tablets)) {
       return null;
@@ -267,8 +279,16 @@ private List<Tablet> nextInternal() {
       if (currentTabletInsertionEventsIterator.hasNext()) {
         final TabletInsertionEvent tabletInsertionEvent =
             currentTabletInsertionEventsIterator.next();
+        if (!((EnrichedEvent) tabletInsertionEvent)
+            .increaseReferenceCount(this.getClass().getName())) {
+          LOGGER.warn(
+              "SubscriptionPipeTabletEventBatch: Failed to increase the reference count of event {}, skipping it.",
+              ((EnrichedEvent) tabletInsertionEvent).coreReportMessage());
+        } else {
+          iterationSnapshot.addParsedEnrichedEvent((EnrichedEvent) tabletInsertionEvent);
+        }
         if (!currentTabletInsertionEventsIterator.hasNext()) {
-          iteratedEnrichedEvents.add((EnrichedEvent) currentTsFileInsertionEvent);
+          iterationSnapshot.addIteratedEnrichedEvent((EnrichedEvent) currentTsFileInsertionEvent);
         }
         return convertToTablets(tabletInsertionEvent);
       } else {
@@ -297,11 +317,13 @@ private List<Tablet> nextInternal() {
       currentTsFileInsertionEvent = tsFileInsertionEvent;
       currentTabletInsertionEventsIterator =
           tsFileInsertionEvent
-              .toTabletInsertionEvents(SubscriptionAgent.receiver().remainingMs())
+              .toTabletInsertionEvents(
+                  // disrupt parsing requests through the introduction of randomness
+                  (long) ((1 + Math.random()) * SubscriptionAgent.receiver().remainingMs()))
               .iterator();
       return next();
     } else if (enrichedEvent instanceof TabletInsertionEvent) {
-      iteratedEnrichedEvents.add(enrichedEvent);
+      iterationSnapshot.addIteratedEnrichedEvent(enrichedEvent);
       return convertToTablets((TabletInsertionEvent) enrichedEvent);
     } else {
       LOGGER.warn(
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletIterationSnapshot.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletIterationSnapshot.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.subscription.event.batch;
+
+import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
+import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
+import org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+public class SubscriptionPipeTabletIterationSnapshot {
+
+  private final List<EnrichedEvent> iteratedEnrichedEvents = new ArrayList<>();
+  private final List<EnrichedEvent> parsedEnrichedEvents = new ArrayList<>();
+
+  public List<EnrichedEvent> getIteratedEnrichedEvents() {
+    return Collections.unmodifiableList(iteratedEnrichedEvents);
+  }
+
+  public void addIteratedEnrichedEvent(final EnrichedEvent enrichedEvent) {
+    iteratedEnrichedEvents.add(enrichedEvent);
+  }
+
+  public void addParsedEnrichedEvent(final EnrichedEvent enrichedEvent) {
+    parsedEnrichedEvents.add(enrichedEvent);
+  }
+
+  public void clear() {
+    for (final EnrichedEvent enrichedEvent : iteratedEnrichedEvents) {
+      if (enrichedEvent instanceof PipeTsFileInsertionEvent) {
+        // close data container in tsfile event
+        ((PipeTsFileInsertionEvent) enrichedEvent).close();
+      }
+    }
+
+    for (final EnrichedEvent enrichedEvent : parsedEnrichedEvents) {
+      if (enrichedEvent instanceof PipeRawTabletInsertionEvent) {
+        // decrease reference count in raw tablet event
+        enrichedEvent.decreaseReferenceCount(this.getClass().getName(), true);
+      }
+    }
+  }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/pipe/SubscriptionPipeTabletBatchEvents.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/pipe/SubscriptionPipeTabletBatchEvents.java
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java

Original file line number	Diff line number	Diff line change
`@@ -197,7 +197,8 @@ public SubscriptionEvent pollInternal(final String consumerId) {`
`197`	`197`	`return null;`
`198`	`198`	`},`
`199`	`199`	`SubscriptionAgent.receiver().remainingMs());`
`200`		`- } catch (final Exception ignored) {`
	`200`	`+ } catch (final Exception e) {`
	`201`	`+ LOGGER.warn("Exception {} occurred when {} execute receiver subtask", this, e, e);`
`201`	`202`	`}`
`202`	`203`	`}`
`203`	`204`
Original file line number	Diff line number	Diff line change
`@@ -159,7 +159,7 @@ public SubscriptionEvent pollTablets(`
`159`	`159`	`String.format(`
`160`	`160`	`"exception occurred when fetching next response: %s, consumer id: %s, commit context: %s, offset: %s, prefetching queue: %s",`
`161`	`161`	`e, consumerId, commitContext, offset, this);`
`162`		`- LOGGER.warn(errorMessage);`
	`162`	`+ LOGGER.warn(errorMessage, e);`
`163`	`163`	`eventRef.set(generateSubscriptionPollErrorResponse(errorMessage));`
`164`	`164`	`}`
`165`	`165`
Original file line number	Diff line number	Diff line change
`@@ -211,7 +211,7 @@ public SubscriptionEvent pollTsFile(`
`211`	`211`	`String.format(`
`212`	`212`	`"exception occurred when fetching next response: %s, consumer id: %s, commit context: %s, writing offset: %s, prefetching queue: %s",`
`213`	`213`	`e, consumerId, commitContext, writingOffset, this);`
`214`		`- LOGGER.warn(errorMessage);`
	`214`	`+ LOGGER.warn(errorMessage, e);`
`215`	`215`	`eventRef.set(generateSubscriptionPollErrorResponse(errorMessage));`
`216`	`216`	`}`
`217`	`217`