[Storage] Cap the temporary flush memory estimate at the memtable size

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java

@@ -141,15 +141,44 @@ public MemTableFlushTask(
         writer.getFile().getName());
   }
 
+  /**
+   * Headroom of the queued chunk buffers over the raw data counted by {@code memSize}, as a
+   * fraction (5/2 = 2.5): the encoded data can be up to 1.25 times the raw values (variable-length
+   * integers take 5 bytes for an {@code int} and 10 for a {@code long}; the other encodings stay
+   * within about 6%), and the buffers grow by doubling, so their capacity can reach twice the
+   * encoded data.
+   */
+  private static final long CHUNK_BUFFER_HEADROOM_NUMERATOR = 5;
+
+  private static final long CHUNK_BUFFER_HEADROOM_DENOMINATOR = 2;
+
+  /**
+   * Temporary memory of a flush: the bounded IO queue holds up to {@code ioTaskQueueSize} encoded
+   * chunks of average size. With no more series than queue slots the queue can't hold more than the
+   * memtable, so the raw share is capped at {@code memSize}; without the cap a memtable of 3 series
+   * reserved 3.3 times its size, which can exceed the whole MemTable budget and make the reject
+   * threshold negative. The share is multiplied by 2.5 for the encoding expansion and the growth
+   * slack of the chunk buffers.
+   */
+  static long estimateTemporaryMemory(
+      long memSize, int seriesNumber, int ioTaskQueueSize, boolean encodingFasterThanIo) {
+    if (!encodingFasterThanIo || seriesNumber == 0) {
+      return 0;
+    }
+    long queuedShare =
+        seriesNumber <= ioTaskQueueSize ? memSize : memSize / seriesNumber * ioTaskQueueSize;
+    return queuedShare * CHUNK_BUFFER_HEADROOM_NUMERATOR / CHUNK_BUFFER_HEADROOM_DENOMINATOR;
+  }
+
   /** the function for flushing memtable. */
   @SuppressWarnings("squid:S3776")
   public void syncFlushMemTable() throws ExecutionException, InterruptedException {
     long estimatedBytes =
-        encodingFasterThanIo && memTable.getSeriesNumber() != 0
-            ? memTable.memSize()
-                / memTable.getSeriesNumber()
-                * config.getIoTaskQueueSizeForFlushing()
-            : 0;
+        estimateTemporaryMemory(
+            memTable.memSize(),
+            memTable.getSeriesNumber(),
+            config.getIoTaskQueueSizeForFlushing(),
+            encodingFasterThanIo);
     try (SystemInfo.FlushMemoryReservation ignored =
         SystemInfo.getInstance()
             .reserveTemporaryMemoryForFlushing(

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

@@ -0,0 +1,53 @@
+/*
+ * 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.dataregion.flush;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/** The temporary memory reserved for a flush (reduces the MemTable budget while it runs). */
+public class FlushMemoryEstimateTest {
+
+  private static final long MB = 1024 * 1024;
+
+  @Test
+  public void manySeriesReserveTheQueuedShareWithHeadroom() {
+    // 300 series, a queue of 10 encoded chunks: 10 series' share of the memtable, times 2.5 for the
+    // encoding expansion and the growth slack of the chunk buffers
+    assertEquals(
+        30 * MB / 300 * 10 * 5 / 2,
+        MemTableFlushTask.estimateTemporaryMemory(30 * MB, 300, 10, true));
+  }
+
+  @Test
+  public void fewSeriesReserveAtMostTwoAndAHalfTimesTheMemTable() {
+    // 3 series cannot fill 10 queue slots: at most the memtable itself is queued
+    assertEquals(130 * MB, MemTableFlushTask.estimateTemporaryMemory(52 * MB, 3, 10, true));
+    assertEquals(130 * MB, MemTableFlushTask.estimateTemporaryMemory(52 * MB, 1, 10, true));
+    assertEquals(130 * MB, MemTableFlushTask.estimateTemporaryMemory(52 * MB, 10, 10, true));
+  }
+
+  @Test
+  public void noReservationWithoutABoundedQueueOrSeries() {
+    assertEquals(0, MemTableFlushTask.estimateTemporaryMemory(52 * MB, 3, 10, false));
+    assertEquals(0, MemTableFlushTask.estimateTemporaryMemory(52 * MB, 0, 10, true));
+  }
+}

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

@@ -20,15 +20,20 @@
 package org.apache.iotdb.db.storageengine.dataregion.memtable;
 
 import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.db.exception.WriteProcessException;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode;
 import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask;
 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.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.file.metadata.ChunkMetadata;
 import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.fileSystem.FSFactoryProducer;
 import org.apache.tsfile.read.TsFileReader;
@@ -38,6 +43,7 @@
 import org.apache.tsfile.read.expression.QueryExpression;
 import org.apache.tsfile.read.query.dataset.QueryDataSet;
 import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.PublicBAOS;
 import org.apache.tsfile.write.chunk.IChunkWriter;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.apache.tsfile.write.schema.MeasurementSchema;
@@ -49,9 +55,14 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Array;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CountDownLatch;
@@ -672,4 +683,201 @@ private boolean isColumnMappingBuilt() {
       return columnMappingBuilt;
     }
   }
+
+  /**
+   * With the IO consumer stalled the encoder fills the IO queue; the temporary reservation must
+   * cover the capacity of the queued chunk buffers, which grow by doubling and can exceed the raw
+   * data counted by memSize (up to 2x), and whose encoded data can itself exceed the raw values
+   * (variable-length integers: up to 1.25x).
+   */
+  @Test(timeout = 60000)
+  public void testReservationCoversQueuedChunkBuffersWithStalledIo() throws Exception {
+    String oldTimeEncoder = TSFileDescriptor.getInstance().getConfig().getTimeEncoder();
+    TSFileDescriptor.getInstance().getConfig().setTimeEncoder("PLAIN");
+    try {
+      // fewer series than queue slots
+      assertReservationCoversQueuedBuffers(9, 60_000, TSDataType.DOUBLE, TSEncoding.PLAIN);
+      assertReservationCoversQueuedBuffers(1, 60_000, TSDataType.DOUBLE, TSEncoding.PLAIN);
+      // more series than queue slots
+      assertReservationCoversQueuedBuffers(300, 6_000, TSDataType.DOUBLE, TSEncoding.PLAIN);
+      // the encoded values are larger than the raw ones (5 bytes per PLAIN INT32), and the buffer
+      // has just grown past 512 KiB to 1 MiB: 1,048,912 B for a memSize of 492,000 B
+      assertReservationCoversQueuedBuffers(1, 41_000, TSDataType.INT32, TSEncoding.PLAIN);
+    } finally {
+      TSFileDescriptor.getInstance().getConfig().setTimeEncoder(oldTimeEncoder);
+    }
+  }
+
+  private void assertReservationCoversQueuedBuffers(
+      int seriesNumber, int rowCount, TSDataType dataType, TSEncoding encoding) throws Exception {
+    IMemTable table = new PrimitiveMemTable(storageGroup, dataRegionId);
+    String[] measurements = new String[seriesNumber];
+    TSDataType[] dataTypes = new TSDataType[seriesNumber];
+    MeasurementSchema[] schemas = new MeasurementSchema[seriesNumber];
+    Object[] columns = new Object[seriesNumber];
+    long[] times = new long[rowCount];
+    for (int row = 0; row < rowCount; row++) {
+      times[row] = row;
+    }
+    for (int i = 0; i < seriesNumber; i++) {
+      measurements[i] = String.format("s%03d", i);
+      dataTypes[i] = dataType;
+      schemas[i] =
+          new MeasurementSchema(measurements[i], dataType, encoding, CompressionType.UNCOMPRESSED);
+      columns[i] = largeValues(dataType, rowCount, i);
+    }
+    table.insertTablet(
+        new InsertTabletNode(
+            new PlanNodeId("stalled-io"),
+            new PartialPath("root.sg.d1"),
+            false,
+            measurements,
+            dataTypes,
+            schemas,
+            times,
+            null,
+            columns,
+            rowCount),
+        0,
+        rowCount);
+
+    SystemInfo system = SystemInfo.getInstance();
+    boolean old = system.isEncodingFasterThanIo();
+    long before = system.getFlushMemoryManager().getMemoryUsageInBytes();
+    CountDownLatch stalled = new CountDownLatch(1);
+    CountDownLatch release = new CountDownLatch(1);
+    RestorableTsFileIOWriter stalledWriter = mock(RestorableTsFileIOWriter.class);
+    when(stalledWriter.getFile()).thenReturn(new File(filePath));
+    doAnswer(
+            invocation -> {
+              stalled.countDown();
+              release.await();
+              throw new IOException("stalled IO released");
+            })
+        .when(stalledWriter)
+        .startChunkGroup(any(IDeviceID.class));
+    system.setEncodingFasterThanIo(true);
+    MemTableFlushTask task =
+        new MemTableFlushTask(table, stalledWriter, storageGroup, dataRegionId, false);
+    Field queueField = MemTableFlushTask.class.getDeclaredField("ioTaskQueue");
+    queueField.setAccessible(true);
+    BlockingQueue<?> ioQueue = (BlockingQueue<?>) queueField.get(task);
+    Thread flusher =
+        new Thread(
+            () -> {
+              try {
+                task.syncFlushMemTable();
+              } catch (ExecutionException expected) {
+                // the stalled writer fails once released
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              }
+            });
+    try {
+      flusher.start();
+      assertTrue(stalled.await(10, TimeUnit.SECONDS));
+      // the encoder either blocks on the full queue or has queued every series, the end of the
+      // chunk group and the end of the task
+      int expectedQueued = Math.min(config().getIoTaskQueueSizeForFlushing(), seriesNumber + 2);
+      long deadline = System.currentTimeMillis() + 20_000;
+      while (ioQueue.size() < expectedQueued && System.currentTimeMillis() < deadline) {
+        Thread.sleep(10);
+      }
+      assertEquals(expectedQueued, ioQueue.size());
+      long reserved = system.getFlushMemoryManager().getMemoryUsageInBytes() - before;
+      long buffers = 0;
+      for (Object queued : new ArrayList<>(ioQueue)) {
+        if (queued instanceof IChunkWriter) {
+          buffers += retainedArrayBytes(queued);
+        }
+      }
+      assertTrue(buffers > table.memSize() / seriesNumber); // the probe found the buffers
+      assertTrue(
+          String.format(
+              "%d series: queued chunk buffers %d B exceed the reservation %d B (memSize %d B)",
+              seriesNumber, buffers, reserved, table.memSize()),
+          buffers <= reserved);
+    } finally {
+      release.countDown();
+      flusher.join(10_000);
+      system.setEncodingFasterThanIo(old);
+    }
+    assertEquals(before, system.getFlushMemoryManager().getMemoryUsageInBytes());
+  }
+
+  /** Values of large magnitude, which take the most bytes in variable-length encodings. */
+  private static Object largeValues(TSDataType dataType, int rowCount, int series) {
+    switch (dataType) {
+      case INT32:
+        int[] ints = new int[rowCount];
+        for (int row = 0; row < rowCount; row++) {
+          ints[row] = Integer.MAX_VALUE - row - series;
+        }
+        return ints;
+      case DOUBLE:
+        double[] doubles = new double[rowCount];
+        for (int row = 0; row < rowCount; row++) {
+          doubles[row] = row * 1.5 + series;
+        }
+        return doubles;
+      default:
+        throw new IllegalArgumentException(dataType.toString());
+    }
+  }
+
+  private static org.apache.iotdb.db.conf.IoTDBConfig config() {
+    return org.apache.iotdb.db.conf.IoTDBDescriptor.getInstance().getConfig();
+  }
+
+  /** Capacity of the byte buffers and primitive arrays a queued chunk writer keeps alive. */
+  private static long retainedArrayBytes(Object root) throws IllegalAccessException {
+    IdentityHashMap<Object, Boolean> seen = new IdentityHashMap<>();
+    ArrayDeque<Object> todo = new ArrayDeque<>();
+    todo.add(root);
+    long bytes = 0;
+    while (!todo.isEmpty()) {
+      Object o = todo.poll();
+      if (seen.put(o, Boolean.TRUE) != null) {
+        continue;
+      }
+      Class<?> c = o.getClass();
+      if (o instanceof PublicBAOS) {
+        bytes += ((PublicBAOS) o).getBuf().length;
+      } else if (c.isArray()) {
+        if (c.getComponentType().isPrimitive()) {
+          bytes += (long) Array.getLength(o) * primitiveBytes(c.getComponentType());
+        } else {
+          for (int i = 0; i < Array.getLength(o); i++) {
+            Object element = Array.get(o, i);
+            if (element != null) {
+              todo.add(element);
+            }
+          }
+        }
+      } else if (!c.getName().startsWith("java.")) {
+        for (Class<?> k = c; k != null && !k.getName().startsWith("java."); k = k.getSuperclass()) {
+          for (Field f : k.getDeclaredFields()) {
+            if (!Modifier.isStatic(f.getModifiers()) && !f.getType().isPrimitive()) {
+              f.setAccessible(true);
+              Object value = f.get(o);
+              if (value != null) {
+                todo.add(value);
+              }
+            }
+          }
+        }
+      }
+    }
+    return bytes;
+  }
+
+  private static int primitiveBytes(Class<?> type) {
+    if (type == long.class || type == double.class) {
+      return 8;
+    }
+    if (type == int.class || type == float.class) {
+      return 4;
+    }
+    return type == short.class || type == char.class ? 2 : 1;
+  }
 }