Fix IoTConsensus batch accumulation by byte size

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

@@ -63,6 +63,10 @@ public void addTLogEntry(TLogEntry entry) {
   }
 
   public boolean canAccumulate() {
+    return canAccumulate(config, logEntries.size(), memorySize);
+  }
+
+  static boolean canAccumulate(IoTConsensusConfig config, int logEntriesSize, long memorySize) {
     // When reading entries from the WAL, the memory size is calculated based on the serialized
     // size, which can be significantly smaller than the actual size.
     // Thus, we add a multiplier to sender's memory size to estimate the receiver's memory cost.
@@ -71,7 +75,7 @@ public boolean canAccumulate() {
     long senderMemSize = LogDispatcher.getSenderMemSizeSum().get();
     double multiplier = senderMemSize > 0 ? (double) receiverMemSize / senderMemSize : 1.0;
     multiplier = Math.max(multiplier, 1.0);
-    return logEntries.size() < config.getReplication().getMaxLogEntriesNumPerBatch()
+    return logEntriesSize < config.getReplication().getMaxLogEntriesNumPerBatch()
         && ((long) (memorySize * multiplier)) < config.getReplication().getMaxSizePerBatch();
   }
 

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

@@ -432,12 +432,16 @@ void waitForBatchAccumulation(long waitingTimeInMs) throws InterruptedException
       }
 
       final long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(waitingTimeInMs);
-      final int maxLogEntriesNumPerBatch = config.getReplication().getMaxLogEntriesNumPerBatch();
-
-      // Keep collecting while the batch is below its entry limit. A plain sleep makes the
-      // dispatcher wait for the full accumulation interval even when the batch becomes full
-      // immediately, which unnecessarily throttles IoTConsensus under sustained write load.
-      while (bufferedEntries.size() < maxLogEntriesNumPerBatch) {
+      final IoTConsensusConfig currentConfig = config;
+      int accumulatedEntries = bufferedEntries.size();
+      long accumulatedMemorySize =
+          bufferedEntries.stream().mapToLong(IndexedConsensusRequest::getMemorySize).sum();
+
+      // Keep collecting while the batch is below both its entry and memory limits. A plain sleep,
+      // or checking only the entry limit, makes the dispatcher wait for the full accumulation
+      // interval after a batch has already reached its memory limit. This unnecessarily throttles
+      // IoTConsensus when each request contains a large tablet.
+      while (Batch.canAccumulate(currentConfig, accumulatedEntries, accumulatedMemorySize)) {
         final long remainingNanos = deadlineNanos - System.nanoTime();
         if (remainingNanos <= 0) {
           return;
@@ -449,6 +453,8 @@ void waitForBatchAccumulation(long waitingTimeInMs) throws InterruptedException
           return;
         }
         bufferedEntries.add(request);
+        accumulatedEntries++;
+        accumulatedMemorySize += request.getMemorySize();
       }
     }
 

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

@@ -177,6 +177,72 @@ public void sendBatchAsync(Batch sentBatch, DispatchLogHandler handler) {
     }
   }
 
+  @Test
+  public void testBatchAccumulationStopsWhenMemoryLimitIsReached() throws Exception {
+    final Peer localPeer = createPeer(1, 6697);
+    final Peer remotePeer = createPeer(2, 6698);
+    final IoTConsensusConfig config =
+        IoTConsensusConfig.newBuilder()
+            .setReplication(
+                IoTConsensusConfig.Replication.newBuilder()
+                    .setMaxLogEntriesNumPerBatch(1024)
+                    .setMaxSizePerBatch(1)
+                    .setMaxWaitingTimeForAccumulatingBatchInMs(10_000)
+                    .build())
+            .build();
+    final ScheduledExecutorService backgroundTaskService =
+        Executors.newSingleThreadScheduledExecutor();
+    final ExecutorService executorService = Executors.newSingleThreadExecutor();
+    LogDispatcher.LogDispatcherThread dispatcherThread = null;
+    Future<?> dispatcherFuture = null;
+    try {
+      final IoTConsensusServerImpl server =
+          createServer(
+              localPeer, Collections.singletonList(localPeer), config, backgroundTaskService);
+      final CountDownLatch batchSent = new CountDownLatch(1);
+      final AtomicInteger getBatchInvocations = new AtomicInteger();
+      dispatcherThread =
+          server.getLogDispatcher().new LogDispatcherThread(remotePeer, config, 0) {
+            @Override
+            public Batch getBatch() {
+              return getBatchInvocations.getAndIncrement() == 0
+                  ? new Batch(config)
+                  : createBatch(config, 1);
+            }
+
+            @Override
+            public void sendBatchAsync(Batch sentBatch, DispatchLogHandler handler) {
+              assertEquals(1, getPendingEntriesSize());
+              assertEquals(1, getBufferRequestSize());
+              batchSent.countDown();
+              Thread.currentThread().interrupt();
+            }
+          };
+      final IndexedConsensusRequest firstRequest =
+          new IndexedConsensusRequest(1, Collections.singletonList(new TestEntry(1, localPeer)));
+      firstRequest.buildSerializedRequests();
+      final IndexedConsensusRequest secondRequest =
+          new IndexedConsensusRequest(2, Collections.singletonList(new TestEntry(2, localPeer)));
+      secondRequest.buildSerializedRequests();
+      assertTrue(dispatcherThread.offer(firstRequest));
+      assertTrue(dispatcherThread.offer(secondRequest));
+
+      dispatcherFuture = executorService.submit(dispatcherThread);
+      assertTrue(batchSent.await(2, TimeUnit.SECONDS));
+      dispatcherFuture.get(2, TimeUnit.SECONDS);
+    } finally {
+      if (dispatcherFuture != null) {
+        dispatcherFuture.cancel(true);
+      }
+      executorService.shutdownNow();
+      executorService.awaitTermination(5, TimeUnit.SECONDS);
+      if (dispatcherThread != null) {
+        dispatcherThread.stop();
+      }
+      backgroundTaskService.shutdownNow();
+    }
+  }
+
   @Test
   public void testReloadConfigUpdatesExistingDispatcherPipeline() throws Exception {
     final Peer localPeer = createPeer(1, 6677);