[Test] Use consistent time units for asynchronous assertion flushes

integration-test/src/test/java/org/apache/iotdb/db/it/utils/IoTDBTestUtilsIT.java

@@ -37,6 +37,7 @@
 import java.util.function.Consumer;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -66,6 +67,12 @@ public void testAssertAlwaysFail() {
         () -> TestUtils.assertAlwaysFail(EnvFactory.getEnv(), "show databases", 1));
   }
 
+  @Test
+  public void testHalfTimeoutElapsedUsesSeconds() {
+    assertFalse(TestUtils.isHalfTimeoutElapsed(299_999, 600));
+    assertTrue(TestUtils.isHalfTimeoutElapsed(300_001, 600));
+  }
+
   @Test
   public void testAssertDataAlwaysFailsOnMismatch() {
     final BaseEnv env = EnvFactory.getEnv();

integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java

@@ -85,6 +85,10 @@ public class TestUtils {
 
   public static final String TIME_PRECISION_IN_NS = "ns";
 
+  static boolean isHalfTimeoutElapsed(final long elapsedTimeInMillis, final long timeoutSeconds) {
+    return elapsedTimeInMillis > TimeUnit.SECONDS.toMillis(timeoutSeconds) / 2;
+  }
+
   /**
    * Creates a statement on the write node without comparing results across DataNodes.
    *
@@ -1702,7 +1706,9 @@ public static void assertDataEventuallyOnEnv(
                   // For IoTV2 batch mode, the pipe receiver may need to flush because the replica
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
-                  if (!flushed[0] && System.currentTimeMillis() - startTime > timeoutSeconds >> 1) {
+                  if (!flushed[0]
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, timeoutSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -1783,7 +1789,9 @@ public static void assertDataSizeEventuallyOnEnv(
                   // For IoTV2 batch mode, the pipe receiver may need to flush because the replica
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
-                  if (!flushed[0] && System.currentTimeMillis() - startTime > timeoutSeconds >> 1) {
+                  if (!flushed[0]
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, timeoutSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -1822,7 +1830,9 @@ public static void assertDataEventuallyOnEnv(
                   // For IoTV2 batch mode, the pipe receiver may need to flush because the replica
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
-                  if (!flushed[0] && System.currentTimeMillis() - startTime > timeoutSeconds >> 1) {
+                  if (!flushed[0]
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, timeoutSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -1904,7 +1914,9 @@ public static void assertDataEventuallyOnEnv(
                   // For IoTV2 batch mode, the pipe receiver may need to flush because the replica
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
-                  if (!flushed[0] && System.currentTimeMillis() - startTime > timeoutSeconds >> 1) {
+                  if (!flushed[0]
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, timeoutSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -2024,7 +2036,9 @@ public static void assertDataEventuallyOnEnv(
                   // For IoTV2 batch mode, the pipe receiver may need to flush because the replica
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
-                  if (!flushed[0] && System.currentTimeMillis() - startTime > timeoutSeconds >> 1) {
+                  if (!flushed[0]
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, timeoutSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -2095,8 +2109,8 @@ public static void assertDataAlwaysOnEnv(
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
                   if (!flushed[0]
-                      && System.currentTimeMillis() - startTime
-                          > TimeUnit.SECONDS.toMillis(consistentSeconds) / 2) {
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, consistentSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -2137,8 +2151,8 @@ public static void assertDataAlwaysOnEnv(
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
                   if (!flushed[0]
-                      && System.currentTimeMillis() - startTime
-                          > TimeUnit.SECONDS.toMillis(consistentSeconds) / 2) {
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, consistentSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }
@@ -2192,8 +2206,8 @@ public static void assertDataAlwaysOnEnv(
                   // sync requires tsFile to process. We flush in the middle of assertion because we
                   // don't know when the data reaches the receiver in general cases
                   if (!flushed[0]
-                      && System.currentTimeMillis() - startTime
-                          > TimeUnit.SECONDS.toMillis(consistentSeconds) / 2) {
+                      && isHalfTimeoutElapsed(
+                          System.currentTimeMillis() - startTime, consistentSeconds)) {
                     flushed[0] = true;
                     statement.execute("flush");
                   }

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

@@ -24,12 +24,16 @@
 import org.apache.iotdb.commons.consensus.index.ProgressIndex;
 import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
 import org.apache.iotdb.commons.file.SystemFileFactory;
+import org.apache.iotdb.commons.schema.table.TsFileTableSchemaUtil;
+import org.apache.iotdb.commons.schema.table.TsTable;
+import org.apache.iotdb.commons.utils.PathUtils;
 import org.apache.iotdb.commons.utils.RetryUtils;
 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.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;
 import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
 import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
@@ -502,6 +506,18 @@ void prepare(
                   fileLength));
         }
       }
+      // Rebuild from the staged devices so recovered writers also retain the table schemas that
+      // Pipe needs to enumerate table data in the sealed file.
+      if (PathUtils.isTableModelDatabase(dataRegion.getDatabaseName())) {
+        for (final String tableName : writer.getTableNames()) {
+          final TsTable table =
+              DataNodeTableCache.getInstance()
+                  .getTable(dataRegion.getDatabaseName(), tableName, false);
+          if (table != null) {
+            writer.registerTableSchema(TsFileTableSchemaUtil.toTsFileTableSchemaNoAttribute(table));
+          }
+        }
+      }
       writer.close();
 
       final TsFileResource tsFileResource = timePartition2Resource.get(entry.getKey());

iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java

@@ -31,6 +31,7 @@
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.MeasurementMetadataIndexEntry;
 import org.apache.tsfile.file.metadata.MetadataIndexNode;
+import org.apache.tsfile.file.metadata.TableSchema;
 import org.apache.tsfile.file.metadata.TimeseriesMetadata;
 import org.apache.tsfile.file.metadata.TsFileMetadata;
 import org.apache.tsfile.file.metadata.enums.MetadataIndexNodeType;
@@ -53,6 +54,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 
@@ -80,6 +82,7 @@ public class TsFilePrecalculatedChunkWriter implements AutoCloseable {
   private final File file;
   private final Map<IDeviceID, Map<String, List<IChunkMetadata>>> device2MetadataMap =
       new TreeMap<>();
+  private final Map<String, TableSchema> tableSchemaMap = new TreeMap<>();
 
   private long maxPhysicalFileSize;
   private boolean sealed;
@@ -385,6 +388,7 @@ public void close() throws IOException {
       }
 
       tsFileMetadata.setTableMetadataIndexNodeMap(tableIndexNodeMap);
+      tsFileMetadata.setTableSchemaMap(tableSchemaMap);
       tsFileMetadata.setMetaOffset(metaOffset);
       tsFileMetadata.setBloomFilter(bloomFilter);
 
@@ -446,6 +450,16 @@ public Map<IDeviceID, List<IChunkMetadata>> getChunkMetadataListMap() {
     return result;
   }
 
+  public Set<String> getTableNames() {
+    return device2MetadataMap.keySet().stream()
+        .map(IDeviceID::getTableName)
+        .collect(Collectors.toSet());
+  }
+
+  public void registerTableSchema(final TableSchema tableSchema) {
+    tableSchemaMap.put(tableSchema.getTableName(), tableSchema);
+  }
+
   private String getFileForLog() {
     return file == null ? "<memory-output>" : file.getAbsolutePath();
   }

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

@@ -0,0 +1,181 @@
+/*
+ * 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.storageengine.load;
+
+import org.apache.iotdb.commons.pipe.datastructure.pattern.TablePattern;
+import org.apache.iotdb.commons.schema.table.TsFileTableSchemaUtil;
+import org.apache.iotdb.commons.schema.table.TsTable;
+import org.apache.iotdb.commons.schema.table.column.FieldColumnSchema;
+import org.apache.iotdb.commons.schema.table.column.TagColumnSchema;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
+import org.apache.iotdb.db.pipe.event.common.tsfile.parser.table.TsFileInsertionEventTableParser;
+import org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator;
+import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
+import org.apache.iotdb.db.storageengine.load.splitter.ChunkData;
+import org.apache.iotdb.db.storageengine.load.splitter.TsFileData;
+import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter;
+import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
+
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.TableSchema;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.read.TsFileSequenceReader;
+import org.apache.tsfile.write.TsFileWriter;
+import org.apache.tsfile.write.record.Tablet;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.Mockito;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+
+public class LoadTsFileTableSchemaTest {
+
+  private static final String DATABASE = "load_table_schema_test";
+  private static final String TABLE = "measurements";
+  private static final int ROW_COUNT = 20;
+
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+  private String[] originalLoadDirs;
+  private DataRegion dataRegion;
+  private TsTable table;
+
+  @Before
+  public void setUp() throws Exception {
+    originalLoadDirs = IoTDBDescriptor.getInstance().getConfig().getLoadTsFileDirs();
+    IoTDBDescriptor.getInstance()
+        .getConfig()
+        .setLoadTsFileDirs(new String[] {temporaryFolder.getRoot().getAbsolutePath()});
+    dataRegion = Mockito.mock(DataRegion.class);
+    Mockito.when(dataRegion.getDatabaseName()).thenReturn(DATABASE);
+    Mockito.when(dataRegion.getDataRegionIdString()).thenReturn("0");
+    Mockito.when(dataRegion.getNonSystemDatabaseName()).thenReturn(Optional.of(DATABASE));
+
+    table = new TsTable(TABLE);
+    table.addColumnSchema(new TagColumnSchema("device", TSDataType.STRING));
+    table.addColumnSchema(
+        new FieldColumnSchema("value", TSDataType.INT32, TSEncoding.PLAIN, CompressionType.LZ4));
+    DataNodeTableCache.getInstance().preUpdateTable(DATABASE, table, null);
+    DataNodeTableCache.getInstance().commitUpdateTable(DATABASE, TABLE, null);
+  }
+
+  @After
+  public void tearDown() {
+    DataNodeTableCache.getInstance().invalid(DATABASE);
+    IoTDBDescriptor.getInstance().getConfig().setLoadTsFileDirs(originalLoadDirs);
+  }
+
+  @Test
+  public void testLoadedTableRemainsReadableByPipe() throws Exception {
+    assertLoadedTableReadable(false, false);
+  }
+
+  @Test
+  public void testRecoveredTableRemainsReadableByPipe() throws Exception {
+    assertLoadedTableReadable(true, false);
+  }
+
+  @Test
+  public void testLoadedTableWithNullTagsRemainsReadableByPipe() throws Exception {
+    assertLoadedTableReadable(false, true);
+  }
+
+  private void assertLoadedTableReadable(final boolean recover, final boolean nullTags)
+      throws Exception {
+    final File source = new File(temporaryFolder.getRoot(), "source.tsfile");
+    final TableSchema schema = TsFileTableSchemaUtil.toTsFileTableSchemaNoAttribute(table);
+    try (final TsFileWriter writer = new TsFileWriter(source)) {
+      writer.registerTableSchema(schema);
+      final Tablet tablet =
+          new Tablet(
+              TABLE,
+              Arrays.asList("device", "value"),
+              Arrays.asList(TSDataType.STRING, TSDataType.INT32),
+              Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD),
+              ROW_COUNT);
+      for (int i = 0; i < ROW_COUNT; i++) {
+        tablet.addTimestamp(i, i);
+        tablet.addValue("device", i, nullTags ? null : "d" + i % 2);
+        tablet.addValue("value", i, i);
+      }
+      writer.writeTable(tablet);
+    }
+
+    final File taskDir = temporaryFolder.newFolder("task");
+    TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir);
+    try {
+      final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator();
+      final List<TsFileData> pieces = new ArrayList<>();
+      new TsFileSplitter(
+              source,
+              data -> {
+                calculator.assign((ChunkData) data);
+                pieces.add(data);
+                return true;
+              })
+          .splitTsFileByDataPartition();
+      manager.writePiece(pieces);
+      if (recover) {
+        manager.close(true);
+        manager = new TsFileWriterManager(dataRegion, taskDir).recoverFromDisk();
+      }
+      manager.prepare(false, Collections.emptyMap(), Collections.emptyMap());
+
+      final File staged = new File(taskDir, DATABASE + "-0-0.tsfile");
+      int rowCount = 0;
+      try (final TsFileInsertionEventTableParser parser =
+          new TsFileInsertionEventTableParser(
+              staged,
+              new TablePattern(true, null, null),
+              Long.MIN_VALUE,
+              Long.MAX_VALUE,
+              null,
+              null,
+              null,
+              false)) {
+        for (final TabletInsertionEvent event : parser.toTabletInsertionEvents()) {
+          final PipeRawTabletInsertionEvent tabletEvent = (PipeRawTabletInsertionEvent) event;
+          rowCount += tabletEvent.convertToTablet().getRowSize();
+          tabletEvent.clearReferenceCount(getClass().getName());
+        }
+      }
+      assertEquals(ROW_COUNT, rowCount);
+      try (final TsFileSequenceReader reader = new TsFileSequenceReader(staged.getPath())) {
+        assertEquals(schema, reader.getTableSchemaMap().get(TABLE));
+      }
+    } finally {
+      manager.close();
+    }
+  }
+}