Fix IoTConsensus batch accumulation and config reload

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensus.java

@@ -104,7 +104,7 @@ public class IoTConsensus implements IConsensus {
       new ConcurrentHashMap<>();
   private final IoTConsensusRPCService service;
   private final RegisterManager registerManager = new RegisterManager();
-  private IoTConsensusConfig config;
+  private volatile IoTConsensusConfig config;
 
   /**
    * Optional callback invoked after a new local peer is created via {@link #createLocalPeer}. Used
@@ -537,6 +537,9 @@ public String getRegionDirFromConsensusGroupId(ConsensusGroupId groupId) {
   public void reloadConsensusConfig(ConsensusConfig consensusConfig) {
     config = consensusConfig.getIotConsensusConfig();
 
+    IoTConsensusMemoryManager.getInstance()
+        .updateMaxMemoryRatioForQueue(config.getReplication().getMaxMemoryRatioForQueue());
+
     for (IoTConsensusServerImpl impl : stateMachineMap.values()) {
       impl.reloadConsensusConfig(config);
     }

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java

@@ -147,7 +147,7 @@ public class IoTConsensusServerImpl {
   private final Set<Peer> configuration = ConcurrentHashMap.newKeySet();
   private final AtomicLong searchIndex;
   private final LogDispatcher logDispatcher;
-  private IoTConsensusConfig config;
+  private volatile IoTConsensusConfig config;
   private final ConsensusReqReader consensusReqReader;
   private volatile boolean active;
   private String newSnapshotDirName;
@@ -1350,6 +1350,7 @@ public String getConsensusGroupId() {
   /** This method is used for hot reload of IoTConsensusConfig. */
   public void reloadConsensusConfig(IoTConsensusConfig config) {
     this.config = config;
+    logDispatcher.reloadConfig(config);
   }
 
   /**

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/IoTConsensusMemoryManager.java

@@ -37,7 +37,7 @@ public class IoTConsensusMemoryManager {
   private final AtomicLong syncMemorySizeInByte = new AtomicLong(0);
   private IMemoryBlock memoryBlock =
       new AtomicLongMemoryBlock("Consensus-Default", null, Runtime.getRuntime().maxMemory() / 10);
-  private Double maxMemoryRatioForQueue = 0.6;
+  private volatile double maxMemoryRatioForQueue = 0.6;
 
   private IoTConsensusMemoryManager() {
     MetricService.getInstance().addMetricSet(new IoTConsensusMemoryManagerMetrics(this));
@@ -158,6 +158,10 @@ public void init(IMemoryBlock memoryBlock, double maxMemoryRatioForQueue) {
     this.maxMemoryRatioForQueue = maxMemoryRatioForQueue;
   }
 
+  public void updateMaxMemoryRatioForQueue(double maxMemoryRatioForQueue) {
+    this.maxMemoryRatioForQueue = maxMemoryRatioForQueue;
+  }
+
   @TestOnly
   public void reset() {
     this.memoryBlock.release(this.memoryBlock.getUsedMemoryInBytes());

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/LogDispatcher.java

@@ -183,6 +183,10 @@ public synchronized void checkAndFlushIndex() {
     impl.checkAndUpdateSafeDeletedSearchIndex();
   }
 
+  public synchronized void reloadConfig(IoTConsensusConfig config) {
+    threads.forEach(thread -> thread.reloadConfig(config));
+  }
+
   public void offer(IndexedConsensusRequest request) {
     offer(request, true);
   }
@@ -230,7 +234,7 @@ public class LogDispatcherThread implements Runnable {
 
     private static final long PENDING_REQUEST_TAKING_TIME_OUT_IN_MS = 10_000L;
     private static final long START_INDEX = 1;
-    private final IoTConsensusConfig config;
+    private volatile IoTConsensusConfig config;
     private final Peer peer;
     private final IndexController controller;
     // A sliding window class that manages asynchronous pendingBatches
@@ -289,6 +293,11 @@ public IoTConsensusConfig getConfig() {
       return config;
     }
 
+    private void reloadConfig(IoTConsensusConfig config) {
+      this.config = config;
+      syncStatus.reloadConfig(config);
+    }
+
     public int getPendingEntriesSize() {
       return pendingEntries.size();
     }
@@ -375,11 +384,16 @@ public void run() {
             IndexedConsensusRequest request =
                 pendingEntries.poll(calculateIdlePollTimeoutInMs(), TimeUnit.MILLISECONDS);
             if (request != null) {
+              final IoTConsensusConfig currentConfig = config;
+              final boolean shouldWaitForBatchAccumulation =
+                  pendingEntries.size()
+                          <= currentConfig.getReplication().getMaxLogEntriesNumPerBatch()
+                      && bufferedEntries.isEmpty();
               bufferedEntries.add(request);
               // If write pressure is low, we simply sleep a little to reduce the number of RPC
-              if (pendingEntries.size() <= config.getReplication().getMaxLogEntriesNumPerBatch()
-                  && bufferedEntries.isEmpty()) {
-                Thread.sleep(config.getReplication().getMaxWaitingTimeForAccumulatingBatchInMs());
+              if (shouldWaitForBatchAccumulation) {
+                waitForBatchAccumulation(
+                    currentConfig.getReplication().getMaxWaitingTimeForAccumulatingBatchInMs());
               }
             } else {
               maybeSendIdleWriterSafeTimeBarrier();
@@ -412,6 +426,10 @@ public void run() {
       logger.info(IoTConsensusMessages.DISPATCHER_EXITS, impl.getThisNode(), peer);
     }
 
+    void waitForBatchAccumulation(long waitingTimeInMs) throws InterruptedException {
+      Thread.sleep(waitingTimeInMs);
+    }
+
     public void updateSafelyDeletedSearchIndex() {
       // update safely deleted search index to delete outdated info,
       // indicating that insert nodes whose search index are before this value can be deleted
@@ -428,6 +446,7 @@ public void updateSafelyDeletedSearchIndex() {
 
     public Batch getBatch() {
 
+      final IoTConsensusConfig currentConfig = config;
       long startIndex = syncStatus.getNextSendingIndex();
       long maxIndex;
       synchronized (impl.getIndexObject()) {
@@ -442,7 +461,7 @@ public Batch getBatch() {
         // Use drainTo instead of poll to reduce lock overhead
         pendingEntries.drainTo(
             bufferedEntries,
-            config.getReplication().getMaxLogEntriesNumPerBatch() - bufferedEntries.size());
+            currentConfig.getReplication().getMaxLogEntriesNumPerBatch() - bufferedEntries.size());
       }
       // remove all request that searchIndex < startIndex
       Iterator<IndexedConsensusRequest> iterator = bufferedEntries.iterator();
@@ -456,7 +475,7 @@ public Batch getBatch() {
         }
       }
 
-      Batch batches = new Batch(config);
+      Batch batches = new Batch(currentConfig);
       // This condition will be executed in several scenarios:
       // 1. restart
       // 2. The getBatch() is invoked immediately at the moment the PendingEntries are consumed

iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/logdispatcher/SyncStatus.java

@@ -32,7 +32,7 @@
 public class SyncStatus {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(SyncStatus.class);
-  private final IoTConsensusConfig config;
+  private IoTConsensusConfig config;
   private final IndexController controller;
   private final LinkedList<Batch> pendingBatches = new LinkedList<>();
   private final IoTConsensusMemoryManager iotConsensusMemoryManager =
@@ -43,6 +43,11 @@ public SyncStatus(IndexController controller, IoTConsensusConfig config) {
     this.config = config;
   }
 
+  public synchronized void reloadConfig(IoTConsensusConfig config) {
+    this.config = config;
+    notifyAll();
+  }
+
   /**
    * we may block here if the synchronization pipeline is full.
    *

iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/logdispatcher/IoTConsensusMemoryManagerTest.java

@@ -41,11 +41,14 @@
 public class IoTConsensusMemoryManagerTest {
 
   private IMemoryBlock previousMemoryBlock;
+  private double previousMaxMemoryRatioForQueue;
   private long memoryBlockSize = 16 * 1024L;
 
   @Before
   public void setUp() throws Exception {
     previousMemoryBlock = IoTConsensusMemoryManager.getInstance().getMemoryBlock();
+    previousMaxMemoryRatioForQueue =
+        IoTConsensusMemoryManager.getInstance().getMaxMemoryRatioForQueue();
     IoTConsensusMemoryManager.getInstance()
         .setMemoryBlock(new AtomicLongMemoryBlock("Test", null, memoryBlockSize));
     IoTConsensusMemoryManager.getInstance().reset();
@@ -55,6 +58,8 @@ public void setUp() throws Exception {
   public void tearDown() throws Exception {
     IoTConsensusMemoryManager.getInstance().reset();
     IoTConsensusMemoryManager.getInstance().setMemoryBlock(previousMemoryBlock);
+    IoTConsensusMemoryManager.getInstance()
+        .updateMaxMemoryRatioForQueue(previousMaxMemoryRatioForQueue);
   }
 
   @Test
@@ -121,6 +126,23 @@ public void testClearUnserializedRequest() {
     assertEquals(0L, request.getRetainedMemorySize());
   }
 
+  @Test
+  public void testUpdateMaxMemoryRatioForQueue() {
+    final IndexedConsensusRequest request =
+        new IndexedConsensusRequest(
+            1,
+            Collections.singletonList(
+                new ByteBufferConsensusRequest(ByteBuffer.allocate((int) (memoryBlockSize / 3)))));
+    request.buildSerializedRequests();
+
+    IoTConsensusMemoryManager.getInstance().updateMaxMemoryRatioForQueue(0.25);
+    assertFalse(IoTConsensusMemoryManager.getInstance().reserve(request));
+
+    IoTConsensusMemoryManager.getInstance().updateMaxMemoryRatioForQueue(0.5);
+    assertTrue(IoTConsensusMemoryManager.getInstance().reserve(request));
+    IoTConsensusMemoryManager.getInstance().free(request);
+  }
+
   private void testReserveAndRelease(int numReservation) {
     int allocationSize = 1;
     long allocatedSize = 0;