Fix IoTConsensusV2 progress for split LOAD TsFiles (#18927)

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java

@@ -31,6 +31,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.exception.load.LoadFileException;
 import org.apache.iotdb.db.i18n.StorageEngineMessages;
+import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusOp;
 import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache;
@@ -532,7 +533,7 @@ void prepare(
 
   /** Formally imports sealed staged TsFiles into the target storage engine data regions. */
   void loadAll(final boolean isGeneratedByPipe, final boolean deleteStagedSource)
-      throws LoadFileException {
+      throws IOException, LoadFileException {
     for (final Map.Entry<TTimePartitionSlot, TsFileResource> entry :
         timePartition2Resource.entrySet()) {
       final DataRegion targetDataRegion = dataRegion;
@@ -547,6 +548,14 @@ void loadAll(final boolean isGeneratedByPipe, final boolean deleteStagedSource)
                 uuidDirName(tsFileResource.getTsFile())));
       }
 
+      if (LoadConsensusCompat.usesIoTConsensusV2()) {
+        // V2 only replicates files carrying this node's progress. The coordinator's index may
+        // belong to another node and precede writes between PREPARE and COMMIT. Assign progress
+        // at import, and persist it so recovery keeps the file eligible for replication.
+        PipeDataNodeAgent.runtime().assignProgressIndexForTsFileLoad(tsFileResource);
+        tsFileResource.serialize();
+      }
+
       targetDataRegion.loadNewTsFile(
           tsFileResource, deleteStagedSource, isGeneratedByPipe, false, Optional.empty());
 

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java

@@ -20,16 +20,22 @@
 package org.apache.iotdb.db.storageengine.load;
 
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.consensus.index.ProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.RecoverProgressIndex;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.consensus.iot.log.ConsensusReqReader;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.exception.load.LoadFileException;
+import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
+import org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent;
+import org.apache.iotdb.db.pipe.processor.iotconsensusv2.IoTConsensusV2Processor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusOp;
 import org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator;
 import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus;
 import org.apache.iotdb.db.storageengine.dataregion.wal.node.IWALNode;
 import org.apache.iotdb.db.storageengine.dataregion.wal.node.WALNode;
 import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALMode;
@@ -798,6 +804,70 @@ public void testStagedDirectoryWithoutTerminalRecordIsKeptOnRecovery() throws Ex
     assertFalse(manager.prepare(prepareNode("in-progress-load"), Collections.emptyMap()));
   }
 
+  @Test
+  public void testIoTConsensusV2CommitUsesReceivingNodeProgress() throws Exception {
+    assertCommittedProgress(true, config.getDataNodeId() + 1);
+  }
+
+  @Test
+  public void testIoTConsensusV2CommitAdvancesPreparedProgress() throws Exception {
+    assertCommittedProgress(true, config.getDataNodeId());
+  }
+
+  @Test
+  public void testOtherConsensusCommitPreservesCoordinatorProgress() throws Exception {
+    assertCommittedProgress(false, config.getDataNodeId() + 1);
+  }
+
+  private void assertCommittedProgress(final boolean useIoTConsensusV2, final int coordinatorNodeId)
+      throws Exception {
+    LoadConsensusCompat.setIoTConsensusV2ForTest(useIoTConsensusV2);
+    final String loadId = "load-progress";
+    final LoadTsFileManager manager = new LoadTsFileManager(dataRegion);
+    try {
+      emulateStagedFileImport();
+      final LoadTsFileConsensusNode piece = stagedPiece(loadId, laidOutChunk());
+      manager.writePiece(piece);
+      final File staged = new File(tempDir, piece.getPieceRefs().get(0).getRelativePath());
+      final RecoverProgressIndex preparedLocalProgress =
+          PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad();
+      final ProgressIndex coordinatorProgress =
+          new RecoverProgressIndex(
+              coordinatorNodeId,
+              preparedLocalProgress.getDataNodeId2LocalIndex().get(config.getDataNodeId()));
+      assertTrue(
+          manager.prepare(
+              prepareNode(loadId),
+              Collections.singletonMap(new TTimePartitionSlot(0L), coordinatorProgress)));
+
+      // A local write after PREPARE must precede the progress assigned when the file is imported.
+      final ProgressIndex interveningWrite =
+          PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad();
+      final LoadTsFileConsensusNode commit = commitNode(loadId);
+      assertTrue(manager.loadAll(commit, Collections.emptyMap()));
+
+      final TsFileResource imported =
+          new TsFileResource(new File(new File(tempDir, "data"), staged.getName()));
+      imported.deserialize();
+      imported.setStatus(TsFileResourceStatus.NORMAL);
+      final ProgressIndex persistedProgress = imported.getMaxProgressIndexAfterClose();
+      final PipeTsFileInsertionEvent event = Mockito.mock(PipeTsFileInsertionEvent.class);
+      Mockito.when(event.forceGetProgressIndex()).thenReturn(persistedProgress);
+      assertEquals(
+          "A V2 LOAD must remain eligible for consensus replication after resource recovery",
+          useIoTConsensusV2,
+          IoTConsensusV2Processor.isShouldReplicate(event));
+      if (useIoTConsensusV2) {
+        assertTrue(persistedProgress.isAfter(interveningWrite));
+      } else {
+        assertEquals(coordinatorProgress, persistedProgress);
+      }
+    } finally {
+      manager.stop();
+      LoadConsensusCompat.setIoTConsensusV2ForTest(null);
+    }
+  }
+
   @Test
   public void testCommitKeepsTheStagedBytesWhileAReplicaMayStillNeedThem() throws Exception {
     // The watermark lags behind the COMMIT command, so a replica that catches up from this node's