Fix pipe event accounting and aligned tablet failure reporting

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

@@ -448,7 +448,10 @@ protected void extractProgressReportEvent(final PipeRealtimeEvent event) {
       if (lastEvent == null || !(lastEvent.getEvent() instanceof PipeHeartbeatEvent)) {
         break;
       }
-      pendingQueue.pollLast();
+      final PipeRealtimeEvent droppedEvent = (PipeRealtimeEvent) pendingQueue.pollLast();
+      if (droppedEvent != null) {
+        droppedEvent.decreaseReferenceCount(PipeRealtimeDataRegionSource.class.getName(), false);
+      }
     }
     final Event last = pendingQueue.peekLast();
     if (last instanceof PipeRealtimeEvent
@@ -459,6 +462,7 @@ protected void extractProgressReportEvent(final PipeRealtimeEvent event) {
           oldEvent
               .getProgressIndex()
               .updateToMinimumEqualOrIsAfterProgressIndex(event.getProgressIndex()));
+      event.decreaseReferenceCount(PipeRealtimeDataRegionSource.class.getName(), false);
       return;
     }
     pendingQueue.offer(event);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java

@@ -491,7 +491,6 @@ private long[] scheduleMemoryBlock(
       InsertTabletNode insertTabletNode,
       List<int[]> rangeList,
       TSStatus[] results,
-      boolean noFailure,
       long[] infoForMetrics)
       throws WriteProcessException {
     long memControlStartTime = System.nanoTime();
@@ -500,7 +499,7 @@ private long[] scheduleMemoryBlock(
       int start = range[0];
       int end = range[1];
       try {
-        long[] memIncrements = checkMemCost(insertTabletNode, start, end, noFailure, results);
+        long[] memIncrements = checkMemCost(insertTabletNode, start, end, results);
         for (int i = 0; i < memIncrements.length; i++) {
           totalMemIncrements[i] += memIncrements[i];
         }
@@ -518,11 +517,11 @@ private long[] scheduleMemoryBlock(
   }
 
   private long[] checkMemCost(
-      InsertTabletNode insertTabletNode, int start, int end, boolean noFailure, TSStatus[] results)
+      InsertTabletNode insertTabletNode, int start, int end, TSStatus[] results)
       throws WriteProcessException {
     long[] memIncrements;
     if (insertTabletNode.isAligned()) {
-      memIncrements = checkAlignedMemCost(insertTabletNode, start, end, noFailure, results);
+      memIncrements = checkAlignedMemCost(insertTabletNode, start, end, results);
     } else {
       memIncrements =
           checkMemCostAndAddToTspInfoForTablet(
@@ -538,7 +537,7 @@ private long[] checkMemCost(
   }
 
   private long[] checkAlignedMemCost(
-      InsertTabletNode insertTabletNode, int start, int end, boolean noFailure, TSStatus[] results)
+      InsertTabletNode insertTabletNode, int start, int end, TSStatus[] results)
       throws WriteProcessException {
     List<Pair<IDeviceID, Integer>> deviceEndPosList = insertTabletNode.splitByDevice(start, end);
     long[] memIncrements = new long[NUM_MEM_TO_ESTIMATE];
@@ -555,7 +554,6 @@ private long[] checkAlignedMemCost(
               insertTabletNode.getColumnCategories(),
               splitStart,
               splitEnd,
-              noFailure,
               results);
       for (int i = 0; i < NUM_MEM_TO_ESTIMATE; i++) {
         memIncrements[i] += splitMemIncrements[i];
@@ -586,7 +584,7 @@ public void insertTablet(
     workMemTable.checkDataType(insertTabletNode);
 
     long[] memIncrements =
-        scheduleMemoryBlock(insertTabletNode, rangeList, results, noFailure, infoForMetrics);
+        scheduleMemoryBlock(insertTabletNode, rangeList, results, infoForMetrics);
 
     long startTime = System.nanoTime();
     WALFlushListener walFlushListener;
@@ -648,7 +646,10 @@ public void insertTablet(
         throw new WriteProcessException(e);
       }
       for (int i = start; i < end; i++) {
-        results[i] = RpcUtils.SUCCESS_STATUS;
+        if (results[i] == null
+            || results[i].getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+          results[i] = RpcUtils.SUCCESS_STATUS;
+        }
       }
 
       final List<Pair<IDeviceID, Integer>> deviceEndOffsetPairs =
@@ -977,7 +978,6 @@ private long[] checkAlignedMemCostAndAddToTspForTablet(
       TsTableColumnCategory[] columnCategories,
       int start,
       int end,
-      boolean noFailure,
       TSStatus[] results)
       throws WriteProcessException {
     if (start >= end) {
@@ -994,7 +994,6 @@ private long[] checkAlignedMemCostAndAddToTspForTablet(
         memIncrements,
         columns,
         columnCategories,
-        noFailure,
         results);
     long memTableIncrement = memIncrements[0];
     long textDataIncrement = memIncrements[1];
@@ -1051,19 +1050,8 @@ private void updateAlignedMemCost(
       long[] memIncrements,
       Object[] columns,
       TsTableColumnCategory[] columnCategories,
-      boolean noFailure,
       TSStatus[] results) {
-    int incomingPointNum;
-    if (noFailure) {
-      incomingPointNum = end - start;
-    } else {
-      incomingPointNum = end - start;
-      for (TSStatus result : results) {
-        if (result != null && result.code != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-          incomingPointNum--;
-        }
-      }
-    }
+    int incomingPointNum = end - start;
 
     TSDataType[] writableFieldDataTypes =
         getWritableFieldDataTypes(measurementIds, dataTypes, columns, columnCategories);

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionSourceTest.java

@@ -0,0 +1,113 @@
+/*
+ * 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.pipe.source.dataregion.realtime;
+
+import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
+import org.apache.iotdb.commons.pipe.event.ProgressReportEvent;
+import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
+import org.apache.iotdb.db.pipe.event.realtime.PipeRealtimeEvent;
+import org.apache.iotdb.db.pipe.event.realtime.PipeRealtimeEventFactory;
+import org.apache.iotdb.pipe.api.event.Event;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class PipeRealtimeDataRegionSourceTest {
+
+  private static final String TEST_REFERENCE_HOLDER =
+      PipeRealtimeDataRegionSourceTest.class.getName();
+
+  @Test
+  public void progressReportEventReleasesDroppedHeartbeatEvent() throws Exception {
+    try (final ProgressReportTestSource source = new ProgressReportTestSource()) {
+      final PipeRealtimeEvent heartbeatEvent = createHeartbeatEvent();
+      source.extract(heartbeatEvent);
+      Assert.assertEquals(1, source.getEventCount());
+      Assert.assertEquals(1, source.getPipeHeartbeatEventCount());
+
+      final PipeRealtimeEvent progressReportEvent = createProgressReportEvent();
+      source.extract(progressReportEvent);
+
+      Assert.assertTrue(heartbeatEvent.getEvent().isReleased());
+      Assert.assertEquals(0, heartbeatEvent.getEvent().getReferenceCount());
+      Assert.assertEquals(1, source.getEventCount());
+      Assert.assertEquals(0, source.getPipeHeartbeatEventCount());
+      Assert.assertFalse(progressReportEvent.getEvent().isReleased());
+    }
+  }
+
+  @Test
+  public void mergedProgressReportEventReleasesNewEvent() throws Exception {
+    try (final ProgressReportTestSource source = new ProgressReportTestSource()) {
+      final PipeRealtimeEvent firstProgressReportEvent = createProgressReportEvent();
+      final PipeRealtimeEvent secondProgressReportEvent = createProgressReportEvent();
+
+      source.extract(firstProgressReportEvent);
+      source.extract(secondProgressReportEvent);
+
+      Assert.assertFalse(firstProgressReportEvent.getEvent().isReleased());
+      Assert.assertTrue(secondProgressReportEvent.getEvent().isReleased());
+      Assert.assertEquals(1, source.getEventCount());
+    }
+  }
+
+  private static PipeRealtimeEvent createHeartbeatEvent() {
+    final PipeRealtimeEvent event = PipeRealtimeEventFactory.createRealtimeEvent(1, false);
+    Assert.assertTrue(event.increaseReferenceCount(TEST_REFERENCE_HOLDER));
+    return event;
+  }
+
+  private static PipeRealtimeEvent createProgressReportEvent() {
+    final ProgressReportEvent progressReportEvent = new ProgressReportEvent("pipe", 1L, null);
+    progressReportEvent.bindProgressIndex(MinimumProgressIndex.INSTANCE);
+
+    final PipeRealtimeEvent event =
+        PipeRealtimeEventFactory.createRealtimeEvent(progressReportEvent);
+    Assert.assertTrue(event.increaseReferenceCount(TEST_REFERENCE_HOLDER));
+    return event;
+  }
+
+  private static class ProgressReportTestSource extends PipeRealtimeDataRegionSource {
+
+    @Override
+    protected void doExtract(final PipeRealtimeEvent event) {
+      if (event.getEvent() instanceof PipeHeartbeatEvent) {
+        extractHeartbeat(event);
+        return;
+      }
+      pendingQueue.offer(event);
+    }
+
+    @Override
+    public Event supply() {
+      return pendingQueue.directPoll();
+    }
+
+    @Override
+    public boolean isNeedListenToTsFile() {
+      return false;
+    }
+
+    @Override
+    public boolean isNeedListenToInsertNode() {
+      return false;
+    }
+  }
+}

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessorTest.java

@@ -41,9 +41,12 @@
 import org.apache.iotdb.db.storageengine.dataregion.DataRegionInfo;
 import org.apache.iotdb.db.storageengine.dataregion.DataRegionTest;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.rescon.memory.PrimitiveArrayManager;
 import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo;
 import org.apache.iotdb.db.utils.EnvironmentUtils;
 import org.apache.iotdb.db.utils.constant.TestConstant;
+import org.apache.iotdb.rpc.RpcUtils;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.external.commons.io.FileUtils;
@@ -73,6 +76,7 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -560,6 +564,33 @@ public void alignedTvListRamCostTest()
     Assert.assertEquals(721560, memTable.memSize());
   }
 
+  @Test
+  public void alignedTabletKeepsFailedStatusesAndCountsWrittenRows()
+      throws MetadataException, WriteProcessException, IOException, IllegalPathException {
+    final int rowCount = PrimitiveArrayManager.ARRAY_SIZE + 2;
+    final List<int[]> rangeList = Collections.singletonList(new int[] {0, rowCount - 1});
+
+    final TsFileProcessor expectedProcessor = newTestProcessor(filePath + ".expected");
+    final TSStatus[] expectedResults = new TSStatus[rowCount];
+    Arrays.fill(expectedResults, RpcUtils.SUCCESS_STATUS);
+    expectedProcessor.insertTablet(
+        genSingleMeasurementTablet(rowCount, true), rangeList, expectedResults, false, new long[5]);
+
+    final TsFileProcessor actualProcessor = newTestProcessor(filePath + ".actual");
+    final TSStatus[] actualResults = new TSStatus[rowCount];
+    Arrays.fill(actualResults, RpcUtils.SUCCESS_STATUS);
+    final int failedIndex = rowCount - 2;
+    actualResults[failedIndex] = RpcUtils.getStatus(TSStatusCode.OUT_OF_TTL, "failed row");
+    actualProcessor.insertTablet(
+        genSingleMeasurementTablet(rowCount, true), rangeList, actualResults, false, new long[5]);
+
+    Assert.assertEquals(
+        expectedProcessor.getWorkMemTable().getTVListsRamCost(),
+        actualProcessor.getWorkMemTable().getTVListsRamCost());
+    Assert.assertEquals(
+        TSStatusCode.OUT_OF_TTL.getStatusCode(), actualResults[failedIndex].getCode());
+  }
+
   @Test
   public void alignedTvListRamCostTest2()
       throws MetadataException, WriteProcessException, IOException {
@@ -1261,6 +1292,49 @@ private void closeTsFileProcessor(TsFileProcessor unsealedTsFileProcessor)
     }
   }
 
+  private TsFileProcessor newTestProcessor(String path) throws IOException, WriteProcessException {
+    TsFileProcessor newProcessor =
+        new TsFileProcessor(
+            storageGroup,
+            SystemFileFactory.INSTANCE.getFile(path),
+            sgInfo,
+            this::closeTsFileProcessor,
+            (tsFileProcessor, updateMap, systemFlushTime) -> {},
+            true);
+    TsFileProcessorInfo tsFileProcessorInfo = new TsFileProcessorInfo(sgInfo);
+    newProcessor.setTsFileProcessorInfo(tsFileProcessorInfo);
+    this.sgInfo.initTsFileProcessorInfo(newProcessor);
+    SystemInfo.getInstance().reportStorageGroupStatus(sgInfo, newProcessor);
+    return newProcessor;
+  }
+
+  private InsertTabletNode genSingleMeasurementTablet(int rowCount, boolean isAligned)
+      throws IllegalPathException {
+    String[] measurements = new String[] {measurementId};
+    TSDataType[] dataTypes = new TSDataType[] {dataType};
+    MeasurementSchema[] schemas =
+        new MeasurementSchema[] {new MeasurementSchema(measurementId, dataType, encoding)};
+    long[] times = new long[rowCount];
+    Object[] columns = new Object[] {new int[rowCount]};
+
+    for (int i = 0; i < rowCount; i++) {
+      times[i] = i;
+      ((int[]) columns[0])[i] = i;
+    }
+
+    return new InsertTabletNode(
+        new QueryId("test_write").genPlanNodeId(),
+        new PartialPath(deviceId),
+        isAligned,
+        measurements,
+        dataTypes,
+        schemas,
+        times,
+        null,
+        columns,
+        rowCount);
+  }
+
   private InsertTabletNode genInsertTableNode(long startTime, boolean isAligned)
       throws IllegalPathException {
     String deviceId = "root.sg.device5";

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

@@ -39,6 +39,8 @@ public E peekLast() {
   }
 
   public E pollLast() {
-    return pendingDeque.pollLast();
+    final E event = pendingDeque.pollLast();
+    eventCounter.decreaseEventCount(event);
+    return event;
   }
 }