Fix RPC buffer memory control activation

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java

@@ -40,11 +40,18 @@ class AutoResizingBuffer implements AutoCloseable {
   private long lastShrinkTime;
   private int accountedCapacity;
   private boolean closed;
+  private AutoResizingBufferMemoryManager.MemoryAccount memoryAccount;
 
   public AutoResizingBuffer(int initialCapacity) throws IOException {
-    AutoResizingBufferMemoryManager.allocate(initialCapacity);
-    this.array = new byte[initialCapacity];
     this.initialCapacity = initialCapacity;
+    this.memoryAccount = AutoResizingBufferMemoryManager.createMemoryAccount();
+    memoryAccount.allocate(initialCapacity);
+    try {
+      this.array = new byte[initialCapacity];
+    } catch (RuntimeException | Error e) {
+      memoryAccount.release(initialCapacity);
+      throw e;
+    }
     this.accountedCapacity = initialCapacity;
   }
 
@@ -76,7 +83,7 @@ public byte[] array() {
   @Override
   public void close() {
     if (!closed) {
-      AutoResizingBufferMemoryManager.release(accountedCapacity);
+      memoryAccount.release(accountedCapacity);
       accountedCapacity = 0;
       closed = true;
     }
@@ -86,25 +93,28 @@ private void resize(int newCapacity) throws IOException {
     final int currentCapacity = array.length;
     if (newCapacity > currentCapacity) {
       final int delta = newCapacity - currentCapacity;
-      AutoResizingBufferMemoryManager.allocate(delta);
+      memoryAccount.allocate(delta);
       try {
         array = Arrays.copyOf(array, newCapacity);
         accountedCapacity += delta;
       } catch (RuntimeException | Error e) {
-        AutoResizingBufferMemoryManager.release(delta);
+        memoryAccount.release(delta);
         throw e;
       }
     } else if (newCapacity < currentCapacity) {
       array = Arrays.copyOf(array, newCapacity);
       final int delta = currentCapacity - newCapacity;
-      AutoResizingBufferMemoryManager.release(delta);
+      memoryAccount.release(delta);
       accountedCapacity -= delta;
     }
   }
 
   private void reserveCurrentCapacityIfReleased() throws IOException {
     if (closed) {
-      AutoResizingBufferMemoryManager.allocate(array.length);
+      AutoResizingBufferMemoryManager.MemoryAccount newMemoryAccount =
+          AutoResizingBufferMemoryManager.createMemoryAccount();
+      newMemoryAccount.allocate(array.length);
+      memoryAccount = newMemoryAccount;
       accountedCapacity = array.length;
       closed = false;
     }

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryManager.java

@@ -68,7 +68,12 @@ static void setMemoryAllocateRetryIntervalInMs(long memoryAllocateRetryIntervalI
         memoryAllocateRetryIntervalInMs;
   }
 
-  static void allocate(long sizeInBytes) throws IOException {
+  static MemoryAccount createMemoryAccount() {
+    return new MemoryAccount(memoryControl);
+  }
+
+  private static void allocate(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes)
+      throws IOException {
     if (sizeInBytes <= 0) {
       return;
     }
@@ -93,7 +98,7 @@ static void allocate(long sizeInBytes) throws IOException {
             MEMORY_ALLOCATE_MAX_RETRIES));
   }
 
-  static void release(long sizeInBytes) {
+  private static void release(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes) {
     if (sizeInBytes <= 0) {
       return;
     }
@@ -107,4 +112,20 @@ public static long getMemoryAllocationCount() {
   public static long getMemoryAllocationFailureCount() {
     return memoryAllocationFailureCount.get();
   }
+
+  static final class MemoryAccount {
+    private final AutoResizingBufferMemoryControl memoryControl;
+
+    private MemoryAccount(AutoResizingBufferMemoryControl memoryControl) {
+      this.memoryControl = memoryControl;
+    }
+
+    void allocate(long sizeInBytes) throws IOException {
+      AutoResizingBufferMemoryManager.allocate(memoryControl, sizeInBytes);
+    }
+
+    void release(long sizeInBytes) {
+      AutoResizingBufferMemoryManager.release(memoryControl, sizeInBytes);
+    }
+  }
 }

iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/AutoResizingBufferTest.java

@@ -152,6 +152,43 @@ public void testSnappyElasticFramedTransportReleasesMemoryWhenConstructorFails()
     }
   }
 
+  @Test
+  public void testBufferReleasesMemoryToTheControlUsedForAllocation() throws Exception {
+    TestMemoryControl firstMemoryControl = new TestMemoryControl();
+    AutoResizingBufferMemoryManager.setMemoryControl(firstMemoryControl);
+    AutoResizingBuffer firstBuffer = new AutoResizingBuffer(100);
+
+    TestMemoryControl secondMemoryControl = new TestMemoryControl();
+    AutoResizingBufferMemoryManager.setMemoryControl(secondMemoryControl);
+    AutoResizingBuffer secondBuffer = new AutoResizingBuffer(200);
+
+    firstBuffer.close();
+    Assert.assertEquals(0, firstMemoryControl.getUsedMemoryInBytes());
+    Assert.assertEquals(200, secondMemoryControl.getUsedMemoryInBytes());
+
+    secondBuffer.close();
+    Assert.assertEquals(0, secondMemoryControl.getUsedMemoryInBytes());
+  }
+
+  @Test
+  public void testBufferCreatedWithoutControlDoesNotReleaseToNewControl() throws Exception {
+    AutoResizingBuffer bufferWithoutControl = new AutoResizingBuffer(100);
+
+    AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+    AutoResizingBuffer bufferWithControl = new AutoResizingBuffer(100);
+
+    bufferWithoutControl.close();
+    Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes());
+
+    bufferWithoutControl.resizeIfNecessary(100);
+    Assert.assertEquals(200, memoryControl.getUsedMemoryInBytes());
+    bufferWithoutControl.close();
+    Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes());
+
+    bufferWithControl.close();
+    Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+  }
+
   private void setLastShrinkTime(AutoResizingBuffer buffer, long lastShrinkTime) throws Exception {
     Field field = AutoResizingBuffer.class.getDeclaredField("lastShrinkTime");
     field.setAccessible(true);

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/ConfigNode.java

@@ -34,7 +34,6 @@
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.exception.IoTDBException;
 import org.apache.iotdb.commons.exception.StartupException;
-import org.apache.iotdb.commons.memory.MemoryConfig;
 import org.apache.iotdb.commons.service.JMXService;
 import org.apache.iotdb.commons.service.RegisterManager;
 import org.apache.iotdb.commons.service.ServiceType;
@@ -143,7 +142,6 @@ protected void start() throws IoTDBException {
       LOGGER.info(ConfigNodeMessages.STARTING_IOTDB, IoTDBConstant.VERSION_WITH_BUILD);
       ConfigNodeStartupCheck checks = new ConfigNodeStartupCheck(IoTDBConstant.CN_ROLE);
       checks.startUpCheck();
-      MemoryConfig.getInstance();
     } catch (StartupException | ConfigurationException | IOException e) {
       LOGGER.error(ConfigNodeMessages.MEET_ERROR_WHEN_DOING_START_CHECKING, e);
       throw new IoTDBException(ConfigNodeMessages.ERROR_STARTING, -1);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/DataNodeMemoryConfig.java

@@ -77,6 +77,9 @@ public class DataNodeMemoryConfig {
   /** The memory manager of on heap */
   private MemoryManager onHeapMemoryManager;
 
+  /** Memory budget for RPC auto-resizing buffers */
+  private long autoResizingBufferMemorySize;
+
   /** Memory manager for the write process */
   private MemoryManager storageEngineMemoryManager;
 
@@ -168,7 +171,7 @@ public void init(TrimProperties properties) {
     long schemaEngineMemorySize = Runtime.getRuntime().maxMemory() / 10;
     long consensusMemorySize = Runtime.getRuntime().maxMemory() / 10;
     long pipeMemorySize = Runtime.getRuntime().maxMemory() / 10;
-    long autoResizingBufferMemorySize = Runtime.getRuntime().maxMemory() / 20;
+    autoResizingBufferMemorySize = Runtime.getRuntime().maxMemory() / 20;
     if (memoryAllocateProportion != null) {
       String[] proportions = memoryAllocateProportion.split(":");
       int proportionSum = 0;
@@ -217,8 +220,6 @@ public void init(TrimProperties properties) {
     consensusMemoryManager =
         onHeapMemoryManager.getOrCreateMemoryManager("Consensus", consensusMemorySize);
     pipeMemoryManager = onHeapMemoryManager.getOrCreateMemoryManager("Pipe", pipeMemorySize);
-    MemoryConfig.getInstance()
-        .setAutoResizingBufferMemoryControl(onHeapMemoryManager, autoResizingBufferMemorySize);
     LOGGER.info(
         DataNodeMiscMessages.MISC_LOG_INITIAL_ALLOCATEMEMORYFORWRITE_B90EC7D9,
         storageEngineMemoryManager.getTotalMemorySizeInBytes());
@@ -260,6 +261,11 @@ public void init(TrimProperties properties) {
             "DirectBuffer", totalDirectBufferMemorySizeLimit);
   }
 
+  public void activateAutoResizingBufferMemoryControl() {
+    MemoryConfig.getInstance()
+        .setAutoResizingBufferMemoryControl(onHeapMemoryManager, autoResizingBufferMemorySize);
+  }
+
   @SuppressWarnings("squid:S3518")
   private void initSchemaMemoryAllocate(
       MemoryManager schemaEngineMemoryManager, TrimProperties properties) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java

@@ -253,6 +253,7 @@ protected void start() {
     logger.info(DataNodeMiscMessages.STARTING_DATANODE);
     boolean isFirstStart;
     try {
+      IoTDBDescriptor.getInstance().getMemoryConfig().activateAutoResizingBufferMemoryControl();
       // Check if this DataNode is start for the first time and do other pre-checks
       isFirstStart = prepareDataNode();
 

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/DataNodeMemoryConfigTest.java

@@ -0,0 +1,40 @@
+/*
+ * 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.conf;
+
+import org.apache.iotdb.commons.memory.MemoryConfig;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class DataNodeMemoryConfigTest {
+
+  @Test
+  public void testRpcMemoryControlIsActivatedOnlyExplicitly() {
+    DataNodeMemoryConfig memoryConfig = IoTDBDescriptor.getInstance().getMemoryConfig();
+
+    Assert.assertEquals(
+        0, MemoryConfig.getInstance().getAutoResizingBufferMemoryTotalSizeInBytes());
+
+    memoryConfig.activateAutoResizingBufferMemoryControl();
+
+    Assert.assertTrue(MemoryConfig.getInstance().getAutoResizingBufferMemoryTotalSizeInBytes() > 0);
+  }
+}

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java

@@ -35,7 +35,6 @@ public class MemoryConfig {
   private boolean isAutoResizingBufferMemoryControlEnabled;
 
   private MemoryConfig() {
-    initAutoResizingBufferMemoryControl();
     MetricService.getInstance().addMetricSet(new AutoResizingBufferMemoryMetrics(this));
   }
 
@@ -53,53 +52,47 @@ private static class MemoryConfigHolder {
     private MemoryConfigHolder() {}
   }
 
-  private void initAutoResizingBufferMemoryControl() {
+  private void initAutoResizingBufferMemoryControl(IMemoryBlock memoryBlock) {
     AutoResizingBufferMemoryManager.setMemoryControl(
         new AutoResizingBufferMemoryControl() {
           @Override
-          public synchronized boolean allocate(long sizeInBytes) {
-            IMemoryBlock memoryBlock = getAutoResizingBufferMemoryBlock();
-            if (memoryBlock == null) {
-              return true;
-            }
+          public boolean allocate(long sizeInBytes) {
             return memoryBlock.allocate(sizeInBytes);
           }
 
           @Override
-          public synchronized void release(long sizeInBytes) {
-            IMemoryBlock memoryBlock = getAutoResizingBufferMemoryBlock();
-            if (memoryBlock != null) {
-              memoryBlock.release(sizeInBytes);
-            }
+          public void release(long sizeInBytes) {
+            memoryBlock.release(sizeInBytes);
           }
         });
   }
 
+  /**
+   * Installs RPC auto-resizing buffer memory control once. Replacing a live control would orphan
+   * the memory blocks retained by existing RPC buffers.
+   */
   public synchronized void setAutoResizingBufferMemoryControl(
       MemoryManager parentMemoryManager, long memorySizeInBytes) {
     if (autoResizingBufferMemoryManagerParent != null) {
-      autoResizingBufferMemoryManagerParent.releaseChildMemoryManager(
-          AUTO_RESIZING_BUFFER_MEMORY_MANAGER_NAME);
+      return;
     }
     autoResizingBufferMemoryManagerParent = parentMemoryManager;
-    autoResizingBufferMemoryBlock = null;
 
     if (memorySizeInBytes <= 0) {
-      isAutoResizingBufferMemoryControlEnabled = false;
       return;
     }
 
     MemoryManager autoResizingBufferMemoryManager =
         parentMemoryManager.getOrCreateMemoryManager(
             AUTO_RESIZING_BUFFER_MEMORY_MANAGER_NAME, memorySizeInBytes, true);
     if (autoResizingBufferMemoryManager == null) {
-      isAutoResizingBufferMemoryControlEnabled = false;
       return;
     }
     autoResizingBufferMemoryBlock =
         autoResizingBufferMemoryManager.exactAllocate(
             AUTO_RESIZING_BUFFER_MEMORY_BLOCK_NAME, memorySizeInBytes, MemoryBlockType.DYNAMIC);
     isAutoResizingBufferMemoryControlEnabled = true;
+    initAutoResizingBufferMemoryControl(autoResizingBufferMemoryBlock);
   }
 
   private synchronized IMemoryBlock getAutoResizingBufferMemoryBlock() {