[Subscription] Preserve WAL replay progress across pending gaps (#18831)

Author: Caideyipi

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

@@ -2046,7 +2046,7 @@ private MaterializationResult accumulateFromPending(
     } finally {
       // Advance once for the consumed prefix, including when a gap or memory block stops the
       // batch. Per-entry fast-forward can repeatedly inspect the same uncovered WAL file footer.
-      // Gap recovery creates its own iterator at the current cursor before reading WAL.
+      // Gap recovery realigns the same iterator at the current cursor before reading WAL.
       final ProgressWALIterator iterator = subscriptionWALIterator;
       if (nextExpectedSearchIndex.get() > startSearchIndex
           && Objects.nonNull(iterator)
@@ -2077,10 +2077,16 @@ private MaterializationResult fillGapFromWAL(
       final int maxTablets,
       final long maxBatchBytes) {
     pendingWalGapRetryRequested = false;
-    resetSubscriptionWALPosition(fromIndex);
     if (seekGeneration.get() != expectedSeekGeneration || isClosed) {
       return MaterializationResult.STALE;
     }
+    // Preserve the reader position and any memory-blocked request across pending-queue gaps.
+    // Reopening here scans retained WAL prefixes again on every bounded replay round, and drops
+    // the request cached for a memory retry. advanceTo skips covered work without losing uncovered
+    // writer progress, including requests without a local search index.
+    if (Objects.nonNull(subscriptionWALIterator)) {
+      subscriptionWALIterator.advanceTo(fromIndex, this::isWriterProgressCoveredForWalFastForward);
+    }
     // A pending-queue jump can span millions of WAL entries when real-time delivery has overflowed.
     // Keep gap recovery within the normal per-round budget so sparse topic filtering cannot hold
     // the
@@ -2223,6 +2229,9 @@ private MaterializationResult pumpFromSubscriptionWAL(
                 .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_ERROR_READING_SUBSCRIPTION_WAL_A3888AC5,
             this,
             e);
+        // next() may have consumed a request whose materialization failed. Reusing this reader
+        // would lose that request, so recover from the unchanged cursor on the next round.
+        requestSubscriptionWalReset(nextExpectedSearchIndex.get(), expectedSeekGeneration);
         break;
       }
     }

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

@@ -53,6 +53,7 @@
 import org.junit.rules.TemporaryFolder;
 
 import java.io.File;
+import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
@@ -61,6 +62,8 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -332,17 +335,192 @@ public void testWalMemoryBlockRetriesWithoutRecreatingIterator() throws Exceptio
     }
   }
 
+  @Test
+  public void testPendingGapReplayDoesNotRescanPrefixAcrossBoundedRounds() throws Exception {
+    assertPendingGapReplay(false, false);
+  }
+
+  @Test
+  public void testPendingGapReplayKeepsMemoryBlockedRequestAcrossRounds() throws Exception {
+    assertPendingGapReplay(true, false);
+  }
+
+  @Test
+  public void testPendingGapReplayRetriesRequestAfterConversionFailure() throws Exception {
+    assertPendingGapReplay(false, true);
+  }
+
+  private void assertPendingGapReplay(
+      final boolean memoryBackpressure, final boolean materializationFailure) throws Exception {
+    final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+    final CommonConfig config = CommonDescriptor.getInstance().getConfig();
+    final int originalBatchMaxWalEntries = config.getSubscriptionConsensusBatchMaxWalEntries();
+    final int originalBatchMaxTabletCount = config.getSubscriptionConsensusBatchMaxTabletCount();
+    final long originalBatchMaxSize = config.getSubscriptionConsensusBatchMaxSizeInBytes();
+    final int originalBatchMaxDelay = config.getSubscriptionConsensusBatchMaxDelayInMs();
+    final File systemDir = temporaryFolder.newFolder("system-pending-gap-replay");
+    final File walDirectory = temporaryFolder.newFolder("wal-pending-gap-replay");
+    ConsensusPrefetchingQueue queue = null;
+    try {
+      final long oneTabletBytes = createTablet().ramBytesUsed();
+      config.setSubscriptionConsensusBatchMaxWalEntries(2);
+      config.setSubscriptionConsensusBatchMaxTabletCount(1);
+      config.setSubscriptionConsensusBatchMaxSizeInBytes(oneTabletBytes * 16L);
+      config.setSubscriptionConsensusBatchMaxDelayInMs(60_000);
+
+      final int replayedRequests = materializationFailure ? 7 : 6;
+      final int rounds = memoryBackpressure ? replayedRequests : (replayedRequests + 1) / 2;
+      final int requestCount = replayedRequests + 2 * rounds;
+      writeSealedWal(walDirectory, requestCount);
+
+      final WALNode walNode = mock(WALNode.class);
+      when(walNode.getLogDirectory()).thenReturn(walDirectory);
+      when(walNode.getCurrentSearchIndex()).thenReturn((long) requestCount);
+      when(walNode.getCurrentWALFileVersion()).thenReturn(1L);
+      when(walNode.getCurrentWALMetaDataSnapshot()).thenReturn(new WALMetaData());
+
+      final IoTConsensusServerImpl serverImpl = mock(IoTConsensusServerImpl.class);
+      when(serverImpl.getConsensusReqReader()).thenReturn(walNode);
+      when(serverImpl.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker());
+
+      final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class);
+      final AtomicBoolean failureInjected = new AtomicBoolean();
+      when(converter.convert(any()))
+          .thenAnswer(
+              invocation -> {
+                final InsertNode insertNode = (InsertNode) invocation.getArgument(0);
+                if (materializationFailure
+                    && insertNode.getSearchIndex() == 2L
+                    && failureInjected.compareAndSet(false, true)) {
+                  throw new IllegalStateException("transient conversion failure");
+                }
+                return Collections.singletonList(createTablet());
+              });
+      when(converter.getDatabaseName()).thenReturn("db");
+
+      final AtomicLong closedIteratorBodyReads = new AtomicLong();
+      final DataRegionId regionId = new DataRegionId(1);
+      queue =
+          new ConsensusPrefetchingQueue(
+              "consumerGroup",
+              "topic",
+              TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE,
+              regionId,
+              serverImpl,
+              new SubscriptionWalRetentionPolicy(
+                  "topic",
+                  SubscriptionWalRetentionPolicy.UNBOUNDED,
+                  SubscriptionWalRetentionPolicy.UNBOUNDED),
+              converter,
+              newCommitManager(systemDir),
+              new RegionProgress(Collections.emptyMap()),
+              1L,
+              1L,
+              true) {
+            @Override
+            protected ProgressWALIterator createSubscriptionWALIterator(
+                final long startSearchIndex) {
+              return new ProgressWALIterator(walNode, startSearchIndex) {
+                @Override
+                public void close() throws IOException {
+                  closedIteratorBodyReads.addAndGet(walBodyReadCount(this));
+                  super.close();
+                }
+              };
+            }
+          };
+      final TrackingSubscriptionMemoryManager memoryManager =
+          new TrackingSubscriptionMemoryManager(
+              memoryBackpressure
+                  ? oneTabletBytes + Math.max(1L, oneTabletBytes / 2L)
+                  : oneTabletBytes * 16L);
+      queue.setSubscriptionMemoryManager(memoryManager);
+
+      assertNull(queue.poll("consumer"));
+      // Admission failures leave a hole before newer pending requests. Keep a suffix queued so
+      // every bounded round enters fillGapFromWAL, including after ACK releases Tablet memory.
+      for (long index = replayedRequests + 1L; index <= requestCount; index++) {
+        assertTrue(pendingEntries(queue).offer(createRequest(index)));
+      }
+
+      long deliveredLocalSeq = 0L;
+      for (int round = 0; round < rounds; round++) {
+        queue.drivePrefetchOnce();
+        final long expectedLocalSeq =
+            materializationFailure
+                ? round * 2L + 1L
+                : memoryBackpressure ? round + 1L : (round + 1L) * 2L;
+        assertEquals(expectedLocalSeq + 1L, queue.getCurrentReadSearchIndex());
+        assertMemoryBounded(queue, memoryManager);
+        if (memoryBackpressure) {
+          assertTrue(subscriptionWalIterator(queue).hasBufferedNext());
+        }
+        while (deliveredLocalSeq < expectedLocalSeq) {
+          final SubscriptionEvent event = queue.poll("consumer");
+          assertNotNull(event);
+          deliveredLocalSeq++;
+          assertEquals(
+              deliveredLocalSeq, event.getCommitContext().getWriterProgress().getLocalSeq());
+          assertEquals(
+              1000L + deliveredLocalSeq,
+              event.getCommitContext().getWriterProgress().getPhysicalTime());
+          assertTrue(queue.ack("consumer", event.getCommitContext()));
+        }
+        assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+      }
+
+      assertEquals(replayedRequests, deliveredLocalSeq);
+      assertEquals(replayedRequests, queue.getWalPathAcceptedEntries());
+      assertEquals(0L, queue.getPendingPathAcceptedEntries());
+      assertEquals(0L, queue.getWalGapSkippedEntries());
+      assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+      assertEquals(0L, queue.getRetainedTabletBytes());
+      // Grouping reads one request ahead; a blocked request can require one further lookahead.
+      // Count actual body reads, including discarded iterators, rather than timing disk caches.
+      // An exceptional conversion failure needs one reset, rereading its three consumed bodies.
+      assertEquals(
+          replayedRequests + (memoryBackpressure ? 2L : 1L) + (materializationFailure ? 3L : 0L),
+          closedIteratorBodyReads.get() + walBodyReadCount(subscriptionWalIterator(queue)));
+    } finally {
+      if (queue != null) {
+        queue.close();
+      }
+      config.setSubscriptionConsensusBatchMaxWalEntries(originalBatchMaxWalEntries);
+      config.setSubscriptionConsensusBatchMaxTabletCount(originalBatchMaxTabletCount);
+      config.setSubscriptionConsensusBatchMaxSizeInBytes(originalBatchMaxSize);
+      config.setSubscriptionConsensusBatchMaxDelayInMs(originalBatchMaxDelay);
+      IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+    }
+  }
+
+  private static int walBodyReadCount(final ProgressWALIterator iterator) {
+    try {
+      // All requests stay in one historical file, so this offset is the number of body reads.
+      final Field field =
+          ProgressWALIterator.class.getDeclaredField("consumedEntryCountInCurrentFile");
+      field.setAccessible(true);
+      return field.getInt(iterator);
+    } catch (final ReflectiveOperationException e) {
+      throw new AssertionError(e);
+    }
+  }
+
   private static void writeSealedWal(final File walDirectory) throws Exception {
+    writeSealedWal(walDirectory, REQUEST_COUNT);
+  }
+
+  private static void writeSealedWal(final File walDirectory, final int requestCount)
+      throws Exception {
     final File historicalWal =
         new File(
             walDirectory, WALFileUtils.getLogFileName(0L, 0L, WALFileStatus.CONTAINS_SEARCH_INDEX));
     final File liveWal =
         new File(
             walDirectory,
-            WALFileUtils.getLogFileName(1L, REQUEST_COUNT, WALFileStatus.CONTAINS_SEARCH_INDEX));
+            WALFileUtils.getLogFileName(1L, requestCount, WALFileStatus.CONTAINS_SEARCH_INDEX));
 
     try (WALWriter writer = new WALWriter(historicalWal, WALFileVersion.V3)) {
-      for (long searchIndex = 1L; searchIndex <= REQUEST_COUNT; searchIndex++) {
+      for (long searchIndex = 1L; searchIndex <= requestCount; searchIndex++) {
         writeWalEntry(writer, searchIndex);
       }
     }