Pipe: Avoid invalid leader cache hints for table-model batches

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

@@ -1279,18 +1279,27 @@ private static long parsePipeCreationTime(final String pipeCreationTime) {
   }
 
   /**
-   * For {@link InsertRowsStatement} and {@link InsertMultiTabletsStatement}, the returned {@link
-   * TSStatus} will use sub-status to record the endpoint for redirection. Each sub-status records
-   * the redirection endpoint for one device path, and the order is the same as the order of the
-   * device paths in the statement. However, this order is not guaranteed to be the same as in the
-   * request. So for each sub-status which needs to redirect, we record the device path using the
-   * message field.
+   * For tree-model {@link InsertRowsStatement} and {@link InsertMultiTabletsStatement}, the
+   * returned {@link TSStatus} uses sub-statuses to record redirection endpoints. Their order is the
+   * same as the device paths in the statement, but not necessarily the request, so attach the
+   * device path to each redirected sub-status.
    */
   private TSStatus executeBatchStatementAndAddRedirectInfo(final InsertBaseStatement statement) {
     final TSStatus result = executeStatementAndClassifyExceptions(statement, 5);
+    return addRedirectInfoForBatch(statement, result, receiverId.get());
+  }
 
+  static TSStatus addRedirectInfoForBatch(
+      final InsertBaseStatement statement, final TSStatus result, final long receiverId) {
     if (result.getCode() == TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()
         && result.getSubStatusSize() > 0) {
+      // A table-model batch may contain rows for multiple devices. The pipe sink currently routes
+      // the entire event by its device ID, so caching a row's endpoint for the whole tablet/table
+      // could misroute later writes. Keep the successful write status without cache hints.
+      if (statement.isWriteToTable()) {
+        return result;
+      }
+
       final List<PartialPath> devicePaths;
       if (statement instanceof InsertRowsStatement) {
         devicePaths = ((InsertRowsStatement) statement).getDevicePaths();
@@ -1299,7 +1308,7 @@ private TSStatus executeBatchStatementAndAddRedirectInfo(final InsertBaseStateme
       } else {
         LOGGER.warn(
             DataNodePipeMessages.RECEIVER_ID_UNSUPPORTED_STATEMENT_TYPE_FOR_REDIRECTION,
-            receiverId.get(),
+            receiverId,
             statement);
         return result;
       }
@@ -1313,7 +1322,7 @@ private TSStatus executeBatchStatementAndAddRedirectInfo(final InsertBaseStateme
       } else {
         LOGGER.warn(
             DataNodePipeMessages.RECEIVER_ID_THE_NUMBER_OF_DEVICE_PATHS,
-            receiverId.get(),
+            receiverId,
             statement,
             result);
       }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/cacher/LeaderCacheUtils.java

@@ -45,19 +45,24 @@ public static List<Pair<String, TEndPoint>> parseRecommendedRedirections(TSStatu
     // requests may contain any number of statements because rows are grouped by database and table.
     final List<Pair<String, TEndPoint>> redirectList = new ArrayList<>();
 
-    if (!status.isSetSubStatus()) {
+    if (status == null || !status.isSetSubStatus()) {
       return redirectList;
     }
 
     for (final TSStatus subStatus : status.getSubStatus()) {
-      if (subStatus.getCode() != TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()) {
+      if (subStatus == null
+          || subStatus.getCode() != TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()
+          || !subStatus.isSetSubStatus()) {
         continue;
       }
 
       for (final TSStatus innerSubStatus : subStatus.getSubStatus()) {
-        if (innerSubStatus.isSetRedirectNode()) {
-          // We assume that innerSubStatus.getMessage() is a device path.
-          // The message field should be a device path.
+        if (innerSubStatus != null
+            && innerSubStatus.isSetRedirectNode()
+            && innerSubStatus.isSetMessage()
+            && !innerSubStatus.getMessage().isEmpty()) {
+          // The receiver sets the message to a device path only when it can safely associate the
+          // redirection with a single tree-model device.
           redirectList.add(
               new Pair<>(innerSubStatus.getMessage(), innerSubStatus.getRedirectNode()));
         }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java

@@ -19,7 +19,9 @@
 
 package org.apache.iotdb.db.pipe.receiver.protocol.thrift;
 
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry;
 import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
@@ -62,6 +64,74 @@ public void tearDown() {
     registry.clear();
   }
 
+  @Test
+  public void testTableModelTabletRedirectDoesNotCachePerRowLeader() {
+    final InsertTabletStatement statement = new InsertTabletStatement();
+    statement.setWriteToTable(true);
+    statement.setRowCount(2);
+
+    final TSStatus firstRow =
+        new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())
+            .setRedirectNode(new TEndPoint("127.0.0.2", 6667));
+    final TSStatus secondRow =
+        new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())
+            .setRedirectNode(new TEndPoint("127.0.0.3", 6667));
+    final TSStatus result =
+        new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode())
+            .setSubStatus(Arrays.asList(firstRow, secondRow));
+
+    Assert.assertSame(result, IoTDBDataNodeReceiver.addRedirectInfoForBatch(statement, result, 1));
+    Assert.assertFalse(firstRow.isSetMessage());
+    Assert.assertFalse(secondRow.isSetMessage());
+  }
+
+  @Test
+  public void testTableModelRowsRedirectDoesNotCachePerRowLeader() throws Exception {
+    final InsertRowStatement firstRow = new InsertRowStatement();
+    firstRow.setDevicePath(new PartialPath("table1"));
+    final InsertRowStatement secondRow = new InsertRowStatement();
+    secondRow.setDevicePath(new PartialPath("table1"));
+    final InsertRowsStatement statement = new InsertRowsStatement();
+    statement.setWriteToTable(true);
+    statement.setInsertRowStatementList(Arrays.asList(firstRow, secondRow));
+
+    final TSStatus firstRowStatus =
+        new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())
+            .setRedirectNode(new TEndPoint("127.0.0.2", 6667));
+    final TSStatus secondRowStatus =
+        new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())
+            .setRedirectNode(new TEndPoint("127.0.0.3", 6667));
+    final TSStatus result =
+        new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode())
+            .setSubStatus(Arrays.asList(firstRowStatus, secondRowStatus));
+
+    Assert.assertSame(result, IoTDBDataNodeReceiver.addRedirectInfoForBatch(statement, result, 1));
+    Assert.assertFalse(firstRowStatus.isSetMessage());
+    Assert.assertFalse(secondRowStatus.isSetMessage());
+  }
+
+  @Test
+  public void testTreeModelBatchRedirectHasDevicePath() throws Exception {
+    final InsertRowStatement firstRow = new InsertRowStatement();
+    firstRow.setDevicePath(new PartialPath("root.sg.d1"));
+    final InsertRowStatement secondRow = new InsertRowStatement();
+    secondRow.setDevicePath(new PartialPath("root.sg.d2"));
+    final InsertRowsStatement statement = new InsertRowsStatement();
+    statement.setInsertRowStatementList(Arrays.asList(firstRow, secondRow));
+
+    final TSStatus local = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+    final TSStatus redirected =
+        new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())
+            .setRedirectNode(new TEndPoint("127.0.0.2", 6667));
+    final TSStatus result =
+        new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode())
+            .setSubStatus(Arrays.asList(local, redirected));
+
+    Assert.assertSame(result, IoTDBDataNodeReceiver.addRedirectInfoForBatch(statement, result, 1));
+    Assert.assertFalse(local.isSetMessage());
+    Assert.assertEquals("root.sg.d2", redirected.getMessage());
+  }
+
   @Test
   public void testLoadTsFileSyncStatementUsesTreeDatabaseLevelFromDatabaseName() throws Exception {
     final Path tsFile = Files.createTempFile("pipe-load-tree-database-level", ".tsfile");

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/util/cacher/LeaderCacheUtilsTest.java

@@ -59,4 +59,46 @@ public void testParseRecommendedRedirectionsFromVariableStatementCount() {
     Assert.assertEquals("table1.device1", redirects.get(0).getLeft());
     Assert.assertEquals(redirectEndPoint, redirects.get(0).getRight());
   }
+
+  @Test
+  public void testIgnoreRedirectsWithoutDevicePath() {
+    final TEndPoint redirectEndPoint = new TEndPoint("127.0.0.2", 6667);
+    final TSStatus tableRowWithoutPath =
+        RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS).setRedirectNode(redirectEndPoint);
+    final TSStatus rowWithEmptyPath =
+        RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)
+            .setMessage("")
+            .setRedirectNode(redirectEndPoint);
+    final TSStatus treeRowWithPath =
+        RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)
+            .setMessage("root.sg.d1")
+            .setRedirectNode(redirectEndPoint);
+    final TSStatus batchStatus =
+        RpcUtils.getStatus(TSStatusCode.REDIRECTION_RECOMMEND)
+            .setSubStatus(
+                Arrays.asList(
+                    RpcUtils.getStatus(TSStatusCode.REDIRECTION_RECOMMEND)
+                        .setSubStatus(Arrays.asList(tableRowWithoutPath, rowWithEmptyPath)),
+                    RpcUtils.getStatus(TSStatusCode.REDIRECTION_RECOMMEND)
+                        .setSubStatus(Collections.singletonList(treeRowWithPath))));
+
+    final List<Pair<String, TEndPoint>> redirects =
+        LeaderCacheUtils.parseRecommendedRedirections(batchStatus);
+
+    Assert.assertEquals(1, redirects.size());
+    Assert.assertEquals("root.sg.d1", redirects.get(0).getLeft());
+    Assert.assertEquals(redirectEndPoint, redirects.get(0).getRight());
+  }
+
+  @Test
+  public void testIgnoreMalformedRedirectStatus() {
+    final TSStatus redirectWithoutSubStatus =
+        RpcUtils.getStatus(TSStatusCode.REDIRECTION_RECOMMEND);
+    final TSStatus batchStatus =
+        RpcUtils.getStatus(TSStatusCode.REDIRECTION_RECOMMEND)
+            .setSubStatus(Arrays.asList(null, redirectWithoutSubStatus));
+
+    Assert.assertTrue(LeaderCacheUtils.parseRecommendedRedirections(batchStatus).isEmpty());
+    Assert.assertTrue(LeaderCacheUtils.parseRecommendedRedirections(null).isEmpty());
+  }
 }