[Performance] Reduce temporary allocations on insert path (#18240)

* Reduce temporary allocations on insert path

* Add opt-in write path performance benchmarks

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/InsertNodeMemoryEstimator.java

@@ -28,6 +28,7 @@
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
+import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertMultiTabletsNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
@@ -155,6 +156,12 @@ public class InsertNodeMemoryEstimator {
   // from the actual result because the properties of the parent class are not added.
   private static final double INSERT_ROW_NODE_EXPANSION_FACTOR = 1.3;
 
+  // Composite insert nodes are estimated on write threads. Reuse the identity set between events,
+  // but discard unusually large sets to avoid retaining a large table on every write thread.
+  private static final int MAX_RETAINED_DEDUPLICATED_OBJECTS = 1024;
+  private static final ThreadLocal<Set<Object>> REUSABLE_DEDUPLICATED_OBJECTS =
+      ThreadLocal.withInitial(InsertNodeMemoryEstimator::newDeduplicatedObjectSet);
+
   public static long sizeOf(final InsertNode insertNode) {
     try {
       final String className = insertNode.getClass().getSimpleName();
@@ -248,46 +255,62 @@ private static long sizeOfInsertRowNode(
 
   private static long sizeOfInsertRowsNode(final InsertRowsNode node) {
     final Set<Object> deduplicatedObjects =
-        newDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
-    long size = INSERT_ROWS_NODE_SIZE;
-    size += calculateFullInsertNodeSize(node, deduplicatedObjects);
-    size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
-    size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
-    size += sizeOfResults(node.getResults());
-    return size;
+        acquireDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
+    try {
+      long size = INSERT_ROWS_NODE_SIZE;
+      size += calculateFullInsertNodeSize(node, deduplicatedObjects);
+      size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
+      size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
+      size += sizeOfResults(node.getResults());
+      return size;
+    } finally {
+      releaseDeduplicatedObjectSet(deduplicatedObjects);
+    }
   }
 
   private static long sizeOfInsertRowsOfOneDeviceNode(final InsertRowsOfOneDeviceNode node) {
     final Set<Object> deduplicatedObjects =
-        newDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
-    long size = INSERT_ROWS_OF_ONE_DEVICE_NODE_SIZE;
-    size += calculateFullInsertNodeSize(node, deduplicatedObjects);
-    size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
-    size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
-    size += sizeOfResults(node.getResults());
-    return size;
+        acquireDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
+    try {
+      long size = INSERT_ROWS_OF_ONE_DEVICE_NODE_SIZE;
+      size += calculateFullInsertNodeSize(node, deduplicatedObjects);
+      size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
+      size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
+      size += sizeOfResults(node.getResults());
+      return size;
+    } finally {
+      releaseDeduplicatedObjectSet(deduplicatedObjects);
+    }
   }
 
   private static long sizeOfInsertMultiTabletsNode(final InsertMultiTabletsNode node) {
     final Set<Object> deduplicatedObjects =
-        newDeduplicatedObjectSetIfNeeded(node.getInsertTabletNodeList());
-    long size = INSERT_MULTI_TABLETS_NODE_SIZE;
-    size += calculateFullInsertNodeSize(node, deduplicatedObjects);
-    size += sizeOfInsertTabletNodeList(node.getInsertTabletNodeList(), deduplicatedObjects);
-    size += sizeOfIntegerList(node.getParentInsertTabletNodeIndexList());
-    size += sizeOfResults(node.getResults());
-    return size;
+        acquireDeduplicatedObjectSetIfNeeded(node.getInsertTabletNodeList());
+    try {
+      long size = INSERT_MULTI_TABLETS_NODE_SIZE;
+      size += calculateFullInsertNodeSize(node, deduplicatedObjects);
+      size += sizeOfInsertTabletNodeList(node.getInsertTabletNodeList(), deduplicatedObjects);
+      size += sizeOfIntegerList(node.getParentInsertTabletNodeIndexList());
+      size += sizeOfResults(node.getResults());
+      return size;
+    } finally {
+      releaseDeduplicatedObjectSet(deduplicatedObjects);
+    }
   }
 
   private static long sizeOfRelationalInsertRowsNode(final RelationalInsertRowsNode node) {
     final Set<Object> deduplicatedObjects =
-        newDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
-    long size = RELATIONAL_INSERT_ROWS_NODE_SIZE;
-    size += calculateFullInsertNodeSize(node, deduplicatedObjects);
-    size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
-    size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
-    // ignore deviceIDs
-    return size;
+        acquireDeduplicatedObjectSetIfNeeded(node.getInsertRowNodeList());
+    try {
+      long size = RELATIONAL_INSERT_ROWS_NODE_SIZE;
+      size += calculateFullInsertNodeSize(node, deduplicatedObjects);
+      size += sizeOfInsertRowNodeList(node.getInsertRowNodeList(), deduplicatedObjects);
+      size += sizeOfIntegerList(node.getInsertRowNodeIndexList());
+      // ignore deviceIDs
+      return size;
+    } finally {
+      releaseDeduplicatedObjectSet(deduplicatedObjects);
+    }
   }
 
   private static long sizeOfRelationalInsertRowNode(final RelationalInsertRowNode node) {
@@ -775,8 +798,24 @@ private static Set<Object> newDeduplicatedObjectSet() {
     return Collections.newSetFromMap(new IdentityHashMap<>());
   }
 
-  private static Set<Object> newDeduplicatedObjectSetIfNeeded(final List<?> children) {
-    return children != null && children.size() > 1 ? newDeduplicatedObjectSet() : null;
+  private static Set<Object> acquireDeduplicatedObjectSetIfNeeded(final List<?> children) {
+    return children != null && children.size() > 1 ? REUSABLE_DEDUPLICATED_OBJECTS.get() : null;
+  }
+
+  private static void releaseDeduplicatedObjectSet(final Set<Object> deduplicatedObjects) {
+    if (deduplicatedObjects == null) {
+      return;
+    }
+    final boolean oversized = deduplicatedObjects.size() > MAX_RETAINED_DEDUPLICATED_OBJECTS;
+    deduplicatedObjects.clear();
+    if (oversized) {
+      REUSABLE_DEDUPLICATED_OBJECTS.remove();
+    }
+  }
+
+  @TestOnly
+  static void clearReusableDeduplicatedObjectsForTest() {
+    REUSABLE_DEDUPLICATED_OBJECTS.remove();
   }
 
   private static boolean shouldCountObject(

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java

@@ -33,6 +33,7 @@
 import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.WritePlanNode;
+import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.LastCacheUpdateSource;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.fetcher.cache.TreeDeviceSchemaCacheManager;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.AbstractMemTable;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunkGroup;
@@ -60,7 +61,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class InsertRowNode extends InsertNode implements WALEntryValue {
+public class InsertRowNode extends InsertNode implements WALEntryValue, LastCacheUpdateSource {
 
   private static final byte TYPE_RAW_STRING = -1;
 
@@ -1010,6 +1011,21 @@ public TimeValuePair composeTimeValuePair(int columnIndex) {
     return new TimeValuePair(time, TsPrimitiveType.getByType(dataTypes[columnIndex], value));
   }
 
+  @Override
+  public long getLastCacheTimestamp() {
+    return time;
+  }
+
+  @Override
+  public boolean hasLastCacheValue(final int index) {
+    return canComposeTimeValuePair(index);
+  }
+
+  @Override
+  public TimeValuePair getLastCacheValue(final int index) {
+    return composeTimeValuePair(index);
+  }
+
   private boolean canComposeTimeValuePair(final int columnIndex) {
     return measurements != null
         && columnIndex >= 0
@@ -1026,18 +1042,9 @@ private boolean canComposeTimeValuePair(final int columnIndex) {
   }
 
   public void updateLastCache(String databaseName) {
-    TimeValuePair[] timeValuePairs = new TimeValuePair[measurements.length];
-    for (int i = 0; i < measurements.length; i++) {
-      timeValuePairs[i] = composeTimeValuePair(i);
-    }
     TreeDeviceSchemaCacheManager.getInstance()
         .updateLastCacheIfExists(
-            databaseName,
-            getDeviceID(),
-            measurements,
-            timeValuePairs,
-            isAligned,
-            measurementSchemas);
+            databaseName, getDeviceID(), measurements, this, isAligned, measurementSchemas);
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowNode.java

@@ -35,7 +35,6 @@
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.IDeviceID.Factory;
-import org.apache.tsfile.read.TimeValuePair;
 import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.apache.tsfile.write.schema.MeasurementSchema;
 
@@ -281,13 +280,9 @@ protected PartialPath readTargetPath(DataInputStream stream)
 
   @Override
   public void updateLastCache(String databaseName) {
-    TimeValuePair[] timeValuePairs = new TimeValuePair[measurements.length];
-    for (int i = 0; i < measurements.length; i++) {
-      timeValuePairs[i] = composeTimeValuePair(i);
-    }
     TableDeviceSchemaCache.getInstance()
         .updateLastCacheIfExists(
-            databaseName, getDeviceID(), measurements, measurementSchemas, timeValuePairs);
+            databaseName, getDeviceID(), measurements, measurementSchemas, this);
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/cache/LastCacheUpdateSource.java

@@ -0,0 +1,37 @@
+/*
+ * 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.queryengine.plan.relational.metadata.fetcher.cache;
+
+import org.apache.tsfile.read.TimeValuePair;
+
+/**
+ * Provides row values lazily when updating last cache on the write path.
+ *
+ * <p>{@link #getLastCacheValue(int)} is called only when the corresponding cache entry exists and
+ * its timestamp is eligible for update.
+ */
+public interface LastCacheUpdateSource {
+
+  long getLastCacheTimestamp();
+
+  boolean hasLastCacheValue(int index);
+
+  TimeValuePair getLastCacheValue(int index);
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/metadata/fetcher/cache/TableDeviceCacheEntry.java

@@ -140,13 +140,22 @@ int setMeasurementSchema(
       return 0;
     }
     // Safe here because tree schema is invalidated by the whole entry
-    final int result =
-        (deviceSchema.compareAndSet(null, new TreeDeviceNormalSchema(database, isAligned))
-            ? TreeDeviceNormalSchema.INSTANCE_SIZE
-            : 0);
-    return deviceSchema.get() instanceof TreeDeviceNormalSchema
-        ? result + ((TreeDeviceNormalSchema) deviceSchema.get()).update(measurements, schemas)
-        : 0;
+    IDeviceSchema schema = deviceSchema.get();
+    int result = 0;
+    if (schema == null) {
+      final TreeDeviceNormalSchema newSchema = new TreeDeviceNormalSchema(database, isAligned);
+      if (deviceSchema.compareAndSet(null, newSchema)) {
+        schema = newSchema;
+        result = TreeDeviceNormalSchema.INSTANCE_SIZE;
+      } else {
+        schema = deviceSchema.get();
+      }
+    }
+    if (!(schema instanceof TreeDeviceNormalSchema)) {
+      return 0;
+    }
+    result += ((TreeDeviceNormalSchema) schema).update(measurements, schemas);
+    return deviceSchema.get() == schema ? result : 0;
   }
 
   IDeviceSchema getDeviceSchema() {
@@ -203,6 +212,18 @@ int tryUpdateLastCache(
     return Objects.nonNull(lastCache.get()) ? result : 0;
   }
 
+  int tryUpdateLastCache(
+      final String[] measurements,
+      final @Nullable IMeasurementSchema[] measurementSchemas,
+      final LastCacheUpdateSource updateSource) {
+    final TableDeviceLastCache cache = lastCache.get();
+    final int result =
+        Objects.nonNull(cache)
+            ? cache.tryUpdate(measurements, measurementSchemas, updateSource)
+            : 0;
+    return Objects.nonNull(lastCache.get()) ? result : 0;
+  }
+
   int tryUpdateLastCache(final String[] measurements, final TimeValuePair[] timeValuePairs) {
     return tryUpdateLastCache(measurements, timeValuePairs, false);
   }