Optimize subscription WAL acknowledgement path

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/io/WALMetaData.java

@@ -83,12 +83,17 @@ public WALMetaData() {
   }
 
   public WALMetaData(long firstSearchIndex, List<Integer> buffersSize, Set<Long> memTablesId) {
+    this(firstSearchIndex, buffersSize, memTablesId, 0);
+  }
+
+  private WALMetaData(
+      long firstSearchIndex, List<Integer> buffersSize, Set<Long> memTablesId, int entriesNum) {
     this.firstSearchIndex = firstSearchIndex;
     this.buffersSize = buffersSize;
     this.memTablesId = memTablesId;
-    this.physicalTimes = new ArrayList<>();
-    this.nodeIds = new ArrayList<>();
-    this.localSeqs = new ArrayList<>();
+    this.physicalTimes = new ArrayList<>(entriesNum);
+    this.nodeIds = new ArrayList<>(entriesNum);
+    this.localSeqs = new ArrayList<>(entriesNum);
   }
 
   /** Adds an entry without explicit writer progress metadata. */
@@ -249,7 +254,12 @@ public static WALMetaData deserialize(ByteBuffer buffer, WALFileVersion version)
         memTablesId.add(buffer.getLong());
       }
     }
-    WALMetaData result = new WALMetaData(firstSearchIndex, buffersSize, memTablesId);
+    WALMetaData result =
+        new WALMetaData(
+            firstSearchIndex,
+            buffersSize,
+            memTablesId,
+            version == WALFileVersion.V3 ? entriesNum : 0);
     // V3 extension: file-level timestamp range + per-entry writer progress metadata
     if (version == WALFileVersion.V3 && buffer.hasRemaining()) {
       result.minDataTs = buffer.getLong();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java

@@ -200,6 +200,10 @@
   private final Map<SubscriptionCommitContext, SubscriptionEvent> outstandingEventsByCommitContext =
       new ConcurrentHashMap<>();
 
+  /** Events whose commit is being persisted outside the in-flight map callback. */
+  private final Set<SubscriptionEvent> acknowledgementInProgressEvents =
+      ConcurrentHashMap.newKeySet();
+
   /**
    * Compact batch boundaries and source positions; no Tablet or serialized response is retained.
    */
@@ -1580,6 +1584,10 @@
         }
 
         synchronized (event) {
+          if (acknowledgementInProgressEvents.contains(event)) {
+            prefetchingQueue.add(event);
+            continue;
+          }
           // ACK may have removed a recycled event while it was being claimed from the queue.
           if (event.isCommitted()
               || (event.getCommitContext().hasWriterProgress()
@@ -3549,6 +3557,22 @@
     requestPrefetch();
   }
 
+  /** Releases a requeued event's materialized payload; it can be restored from its WAL record. */
+  private void detachPayloadForReplay(final SubscriptionEvent event) {
+    synchronized (event) {
+      if (Objects.isNull(deliveryBatches.get(event))
+          || event.isCommitted()
+          || outstandingEventsByCommitContext.get(event.getCommitContext()) != event
+          || !event.detachCurrentPayload(event.getCurrentResponse())) {
+        return;
+      }
+      final Long retainedBytes = retainedBytesByEvent.remove(event);
+      if (Objects.nonNull(retainedBytes)) {
+        releaseTabletMemory(retainedBytes);
+      }
+    }
+  }
+
   private boolean restorePayloadForDelivery(final SubscriptionEvent event) {
     if (!event.isPayloadDetached()) {
       return true;
@@ -4095,6 +4119,7 @@
     final WriterProgress commitWriterProgress = extractCommitWriterProgress(commitContext);
     final AtomicBoolean found = new AtomicBoolean(false);
     final AtomicBoolean acked = new AtomicBoolean(false);
+    final SubscriptionEvent[] eventToCommit = new SubscriptionEvent[1];
     inFlightEvents.compute(
         new InFlightEventKey(consumerId, commitContext),
         (key, ev) -> {
@@ -4115,29 +4140,40 @@
             return null;
           }
 
-          final boolean committed = commitWriterProgress(commitWriterId, commitWriterProgress);
-          if (!committed) {
-            if (!silent) {
-              LOGGER.warn(
-                  DataNodePipeMessages
-                      .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_ADVANCE_COMMIT_FRONTIER_56E606C0,
-                  this,
-                  commitContext);
-            }
-            return ev;
-          }
-
-          ev.ack();
-          ev.recordCommittedTimestamp();
-          acked.set(true);
-          cleanUpEvent(ev, false);
-          return null;
+          eventToCommit[0] = ev;
+          acknowledgementInProgressEvents.add(ev);
+          return ev;
         });
-    if (acked.get()) {
-      // Retention refresh can read metadata from every newly relevant WAL file. Keep that I/O out
-      // of ConcurrentHashMap.compute so recycling and requeueing other in-flight events can
-      // proceed while an ACK advances committed progress.
-      refreshCommittedWalRetentionBoundAndNotify();
+    final SubscriptionEvent event = eventToCommit[0];
+    if (event != null) {
+      try {
+        final boolean committed = commitWriterProgress(commitWriterId, commitWriterProgress);
+        if (committed) {
+          final InFlightEventKey key = new InFlightEventKey(consumerId, commitContext);
+          inFlightEvents.compute(
+              key,
+              (ignored, inFlightEvent) -> {
+                synchronized (event) {
+                  event.ack();
+                  event.recordCommittedTimestamp();
+                  acked.set(true);
+                  cleanUpEvent(event, false);
+                }
+                return inFlightEvent == event ? null : inFlightEvent;
+              });
+          inFlightEvents.entrySet().removeIf(entry -> entry.getValue() == event);
+          prefetchingQueue.remove(event);
+          refreshCommittedWalRetentionBoundAndNotify();
+        } else if (!silent) {
+          LOGGER.warn(
+              DataNodePipeMessages
+                  .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_ADVANCE_COMMIT_FRONTIER_56E606C0,
+              this,
+              commitContext);
+        }
+      } finally {
+        acknowledgementInProgressEvents.remove(event);
+      }
     }
     return !found.get()
         ? InFlightAckResult.MISSING
@@ -4296,6 +4332,9 @@
                 cleanUpEvent(ev, false);
                 return null;
               }
+              if (acknowledgementInProgressEvents.contains(ev)) {
+                detachPayloadForReplay(ev);
+              }
               prefetchingQueue.add(ev);
               LOGGER.debug(
                   DataNodePipeMessages
@@ -6270,7 +6309,7 @@
 
   private static final class WriterState {
     private long effectiveSafePt = 0L;
     private boolean closed = false;
   }
 
   private static final class PreparedEntry implements WriterBufferedEntry {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java

@@ -2686,6 +2686,7 @@ ProgressWALReader openProgressWALReader(final File walFile) throws IOException {
         // A slow retention read must not keep the event-map bucket locked by ACK and stall
         // a shared prefetch worker trying to recycle that event.
         executor.submit(() -> recycle.invoke(activeQueue)).get(5, TimeUnit.SECONDS);
+        assertTrue(event.isCommitted());
         assertEquals(0L, activeQueue.getSubscriptionUncommittedEventCount());
         assertEquals(0L, activeQueue.getRetainedTabletBytes());
         assertFalse(ack.isDone());