[To dev/1.3] Fix same-region flush stalls during write admission

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

@@ -1201,6 +1201,10 @@
         if (!SystemInfo.getInstance().reportStorageGroupStatus(dataRegionInfo, this)) {
           long startTime = System.currentTimeMillis();
           while (SystemInfo.getInstance().isRejected()) {
+            // The asynchronous submitter needs the region write lock held by this insertion.
+            // Submit other selected MemTables before checking whether this processor was selected
+            // too, so they can release memory even when this insertion can now proceed.
+            submitOtherSelectedMemTablesToFlush();
             if (workMemTable.shouldFlush()) {
               break;
             }
@@ -1227,6 +1231,19 @@
     workMemTable.addTextDataSize(textDataIncrement);
   }
 
+  private void submitOtherSelectedMemTablesToFlush() {
+    for (TsFileProcessor processor : dataRegionInfo.getAllReportedTsp()) {
+      // This processor still owns the pending write and its memory reservations.
+      if (processor == this) {
+        continue;
+      }
+      IMemTable memTable = processor.getWorkMemTable();
+      if (memTable != null && memTable.shouldFlush()) {
+        processor.submitAFlushTask();
+      }
+    }
+  }
+
   private void rollbackMemoryInfo(long[] memIncrements) {
     long memTableIncrement = memIncrements[0];
     long textDataIncrement = memIncrements[1];
@@ -2304,7 +2321,7 @@
 
   private long getQueryTimeLowerBound(String[] device) {
     long deviceTTL = DataNodeTTLCache.getInstance().getTTL(device);
     return deviceTTL != Long.MAX_VALUE
         ? CommonDateTimeUtils.currentTime() - deviceTTL
         : Long.MIN_VALUE;
   }

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

@@ -58,6 +58,7 @@
 import org.apache.iotdb.db.storageengine.dataregion.compaction.utils.CompactionConfigRestorer;
 import org.apache.iotdb.db.storageengine.dataregion.flush.FlushManager;
 import org.apache.iotdb.db.storageengine.dataregion.flush.TsFileFlushPolicy;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.ReadOnlyMemChunk;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor;
 import org.apache.iotdb.db.storageengine.dataregion.read.QueryDataSource;
@@ -76,6 +77,12 @@
 import org.apache.tsfile.file.metadata.enums.CompressionType;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.read.TimeValuePair;
+import org.apache.tsfile.read.TsFileReader;
+import org.apache.tsfile.read.TsFileSequenceReader;
+import org.apache.tsfile.read.common.Path;
+import org.apache.tsfile.read.common.RowRecord;
+import org.apache.tsfile.read.expression.QueryExpression;
+import org.apache.tsfile.read.query.dataset.QueryDataSet;
 import org.apache.tsfile.read.reader.IPointReader;
 import org.apache.tsfile.utils.BitMap;
 import org.apache.tsfile.write.record.TSRecord;
@@ -93,6 +100,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -181,6 +189,141 @@
     return insertRowNode;
   }
 
+  @Test
+  public void testRejectedRowFlushesOnlySelectedOtherMemtable() throws Exception {
+    assertRejectedInsertionFlushesOnlySelectedOtherMemtable(false);
+  }
+
+  @Test
+  public void testRejectedTabletFlushesOnlySelectedOtherMemtable() throws Exception {
+    assertRejectedInsertionFlushesOnlySelectedOtherMemtable(true);
+  }
+
+  private void assertRejectedInsertionFlushesOnlySelectedOtherMemtable(boolean tablet)
+      throws Exception {
+    final SystemInfo systemInfo = SystemInfo.getInstance();
+    final int originalWait = config.getMaxWaitingTimeWhenInsertBlocked();
+    final int originalPeriod = config.getCheckPeriodWhenInsertBlocked();
+    final boolean originalSeparateData = config.isEnableSeparateData();
+    final long partitionInterval = TimePartitionUtils.getTimePartitionInterval();
+    try {
+      config.setEnableSeparateData(true);
+      for (int i = 1; i <= 384; i++) {
+        insertMemoryAdmissionTestRow(i, i);
+      }
+      final TsFileProcessor selected =
+          dataRegion.getWorkSequenceTsFileProcessors().stream()
+              .max(
+                  (left, right) ->
+                      Long.compare(left.getWorkMemTableRamCost(), right.getWorkMemTableRamCost()))
+              .get();
+      // Fill the first primitive arrays so that the rejected write owns a new reservation.
+      for (int i = 1; i <= 128; i++) {
+        insertMemoryAdmissionTestRow(partitionInterval + i, i);
+      }
+      final TsFileProcessor current =
+          dataRegion.getWorkSequenceTsFileProcessors().stream()
+              .filter(processor -> processor != selected)
+              .findFirst()
+              .get();
+      insertMemoryAdmissionTestRow(2 * partitionInterval + 1, 1);
+      final TsFileProcessor unselected =
+          dataRegion.getWorkSequenceTsFileProcessors().stream()
+              .filter(processor -> processor != selected && processor != current)
+              .findFirst()
+              .get();
+      final IMemTable currentMemTable = current.getWorkMemTable();
+      final IMemTable unselectedMemTable = unselected.getWorkMemTable();
+      final IMemTable selectedMemTable = selected.getWorkMemTable();
+
+      final Field regionInfoField = DataRegion.class.getDeclaredField("dataRegionInfo");
+      regionInfoField.setAccessible(true);
+      final DataRegionInfo regionInfo = (DataRegionInfo) regionInfoField.get(dataRegion);
+      final long memoryBefore = regionInfo.getMemCost();
+      setMemoryAdmissionTestThreshold("memorySizeForMemtable", 2 * memoryBefore);
+      // One new row pushes memory above this threshold. Flushing the largest (selected) memtable
+      // brings it back below the threshold, so the current processor remains unselected.
+      setMemoryAdmissionTestThreshold("FLUSH_THRESHOLD", (double) memoryBefore);
+      setMemoryAdmissionTestThreshold("REJECT_THRESHOLD", (double) memoryBefore - 1);
+      config.setMaxWaitingTimeWhenInsertBlocked(5_000);
+      config.setCheckPeriodWhenInsertBlocked(10);
+      regionInfo.setNeedToReportToSystem(true);
+
+      // Keep the region lock held until the write succeeds. The asynchronous flush submitter
+      // cannot acquire it, so the insertion must submit the selected processor itself.
+      dataRegion.writeLock("testMemoryAdmission");
+      try {
+        if (tablet) {
+          dataRegion.insertTablet(
+              new InsertTabletNode(
+                  new PlanNodeId("memoryAdmission"),
+                  new PartialPath(deviceId),
+                  false,
+                  new String[] {measurementId},
+                  new TSDataType[] {TSDataType.INT32},
+                  new MeasurementSchema[] {new MeasurementSchema(measurementId, TSDataType.INT32)},
+                  new long[] {partitionInterval + 129},
+                  null,
+                  new Object[] {new int[] {129}},
+                  1));
+        } else {
+          insertMemoryAdmissionTestRow(partitionInterval + 129, 129);
+        }
+        Assert.assertTrue(selectedMemTable.shouldFlush());
+        Assert.assertNull(selected.getWorkMemTable());
+        Assert.assertSame(currentMemTable, current.getWorkMemTable());
+        Assert.assertSame(unselectedMemTable, unselected.getWorkMemTable());
+        Assert.assertFalse(currentMemTable.shouldFlush());
+        Assert.assertFalse(unselectedMemTable.shouldFlush());
+        Assert.assertEquals(129, currentMemTable.getTotalPointsNum());
+        Assert.assertFalse(systemInfo.isRejected());
+      } finally {
+        dataRegion.writeUnlock();
+      }
+
+      // A later write still reuses the small current MemTable instead of creating another file.
+      insertMemoryAdmissionTestRow(partitionInterval + 130, 130);
+      Assert.assertSame(currentMemTable, current.getWorkMemTable());
+      Assert.assertEquals(3, dataRegion.getTsFileManager().getTsFileList(true).size());
+      dataRegion.syncCloseAllWorkingTsFileProcessors();
+      int points = 0;
+      for (TsFileResource resource : dataRegion.getTsFileManager().getTsFileList(true)) {
+        try (TsFileSequenceReader reader = new TsFileSequenceReader(resource.getTsFilePath());
+            TsFileReader tsFileReader = new TsFileReader(reader)) {
+          final QueryDataSet dataSet =
+              tsFileReader.query(
+                  QueryExpression.create(
+                      Collections.singletonList(new Path(deviceId, measurementId, false)), null));
+          while (dataSet.hasNext()) {
+            final RowRecord row = dataSet.next();
+            Assert.assertEquals(
+                row.getTimestamp() % partitionInterval, row.getFields().get(0).getIntV());
+            points++;
+          }
+        }
+      }
+      Assert.assertEquals(515, points);
+    } finally {
+      systemInfo.allocateWriteMemory();
+      config.setMaxWaitingTimeWhenInsertBlocked(originalWait);
+      config.setCheckPeriodWhenInsertBlocked(originalPeriod);
+      config.setEnableSeparateData(originalSeparateData);
+    }
+  }
+
+  private void insertMemoryAdmissionTestRow(long time, int value) throws Exception {
+    final TSRecord record = new TSRecord(time, deviceId);
+    record.addTuple(
+        DataPoint.getDataPoint(TSDataType.INT32, measurementId, Integer.toString(value)));
+    dataRegion.insert(buildInsertRowNodeByTSRecord(record));
+  }
+
+  private void setMemoryAdmissionTestThreshold(String name, Object value) throws Exception {
+    final Field field = SystemInfo.class.getDeclaredField(name);
+    field.setAccessible(true);
+    field.set(SystemInfo.getInstance(), value);
+  }
+
   @Test
   public void testUnseqUnsealedDelete()
       throws WriteProcessException, IOException, MetadataException {
@@ -969,7 +1112,7 @@
           DataRegionException,
           TsFileProcessorException {
     DataRegion dataRegion1 = new DummyDataRegion(systemDir, "root.Rows");
     long[] time = new long[] {3, 4, 1, 2};
     List<Integer> indexList = new ArrayList<>();
     List<InsertRowNode> nodes = new ArrayList<>();
     for (int i = 0; i < 4; i++) {
@@ -1130,8 +1273,8 @@
 
     final List<Integer> indexList = Arrays.asList(0, 1);
     final List<InsertRowNode> nodes = new ArrayList<>();
     final long[] times = new long[] {1, failingTime};
     final int[] values = new int[] {10, 20};
     for (int i = 0; i < times.length; i++) {
       final long time = times[i];
       final TSRecord record = new TSRecord(time, "root.cache_rows");
@@ -1186,14 +1329,14 @@
     DataNodeSchemaCache.getInstance()
         .declareLastCache(dataRegion1.getDatabaseName(), lastCachePath);
 
     final String[] measurements = new String[] {measurementId};
     final TSDataType[] dataTypes = new TSDataType[] {TSDataType.INT32};
     final MeasurementSchema[] measurementSchemas =
         new MeasurementSchema[] {
           new MeasurementSchema(measurementId, TSDataType.INT32, TSEncoding.PLAIN)
         };
     final long[] times = new long[] {1, 2};
     final Object[] columns = new Object[] {new int[] {10, 20}};
     final InsertTabletNode insertTabletNode =
         new InsertTabletNode(
             new QueryId("test_write").genPlanNodeId(),