[Subscription] Prevent event loss on poll payload overflow

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java

@@ -125,12 +125,23 @@ public List<SubscriptionEvent> poll(
       }
       final List<SubscriptionEvent> events =
           broker.poll(consumerId, topicNames, remainingBytes, progressByTopic);
-      allEvents.addAll(events);
       for (final SubscriptionEvent event : events) {
         try {
-          remainingBytes -= event.getCurrentResponseSize();
+          final long currentSize = event.getCurrentResponseSize();
+          // Each broker preserves the existing handling for its first oversized event. If another
+          // broker already used part of this response, put the event back so it can be retried with
+          // the full budget on the next poll.
+          if (!allEvents.isEmpty()
+              && currentSize > remainingBytes
+              && broker.requeue(consumerId, event.getCommitContext())) {
+            remainingBytes = 0;
+            break;
+          }
+          allEvents.add(event);
+          remainingBytes -= currentSize;
         } catch (final IOException ignored) {
           // best effort
+          allEvents.add(event);
         }
       }
     }
@@ -204,6 +215,18 @@ public List<SubscriptionCommitContext> commit(
     return allSuccessful;
   }
 
+  public boolean requeue(
+      final ConsumerConfig consumerConfig, final SubscriptionCommitContext commitContext) {
+    final String consumerGroupId = consumerConfig.getConsumerGroupId();
+    final String consumerId = consumerConfig.getConsumerId();
+    for (final ISubscriptionBroker broker : getBrokers(consumerGroupId)) {
+      if (broker.acceptsCommitContext(commitContext) && broker.requeue(consumerId, commitContext)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   public int refreshInFlightEventLeases(
       final ConsumerConfig consumerConfig,
       final List<SubscriptionCommitContext> processorBufferedCommitContexts) {

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

@@ -124,6 +124,7 @@ public List<SubscriptionEvent> poll(
     final List<SubscriptionEvent> eventsToPoll = new ArrayList<>();
     final List<SubscriptionEvent> eventsToNack = new ArrayList<>();
     long totalSize = 0;
+    boolean responseFull = false;
 
     for (final String topicName : topicNames) {
       final List<ConsensusPrefetchingQueue> queues =
@@ -169,14 +170,23 @@ public List<SubscriptionEvent> poll(
           continue;
         }
 
+        // Preserve the existing handling for a single oversized event. Once this response already
+        // contains data, defer an event that does not fit instead of returning an oversized batch.
+        if (totalSize > 0
+            && currentSize > maxBytes - totalSize
+            && consensusQueue.requeue(consumerId, event.getCommitContext())) {
+          responseFull = true;
+          break;
+        }
+
         eventsToPoll.add(event);
         totalSize += currentSize;
 
         if (totalSize >= maxBytes) {
           break;
         }
       }
-      if (totalSize >= maxBytes) {
+      if (responseFull || totalSize >= maxBytes) {
         break;
       }
     }
@@ -280,6 +290,17 @@ public List<SubscriptionCommitContext> commit(
     return successfulCommitContexts;
   }
 
+  @Override
+  public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) {
+    final List<ConsensusPrefetchingQueue> queues =
+        topicNameToConsensusPrefetchingQueues.get(commitContext.getTopicName());
+    if (Objects.isNull(queues) || queues.isEmpty()) {
+      return false;
+    }
+    final ConsensusPrefetchingQueue queue = getQueueForCommitContext(queues, commitContext);
+    return Objects.nonNull(queue) && queue.requeue(consumerId, commitContext);
+  }
+
   @Override
   public int refreshInFlightEventLeases(
       final String consumerId, final List<SubscriptionCommitContext> commitContexts) {

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

@@ -48,6 +48,9 @@ List<SubscriptionEvent> pollTablets(
   List<SubscriptionCommitContext> commit(
       String consumerId, List<SubscriptionCommitContext> commitContexts, boolean nack);
 
+  /** Returns an in-flight event to its prefetching queue without incrementing its nack count. */
+  boolean requeue(String consumerId, SubscriptionCommitContext commitContext);
+
   default List<SubscriptionCommitContext> selectAcceptedCommitContexts(
       final List<SubscriptionCommitContext> commitContexts) {
     if (Objects.isNull(commitContexts) || commitContexts.isEmpty()) {

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

@@ -166,6 +166,14 @@ public List<SubscriptionEvent> poll(
         continue;
       }
 
+      // Preserve the existing handling for a single oversized event. Once this response already
+      // contains data, defer an event that does not fit instead of returning an oversized batch.
+      if (totalSize > 0
+          && currentSize > maxBytes - totalSize
+          && prefetchingQueue.requeue(consumerId, event.getCommitContext())) {
+        break;
+      }
+
       // Add the event to the poll list
       eventsToPoll.add(event);
 
@@ -175,8 +183,8 @@ public List<SubscriptionEvent> poll(
       // Update the total size
       totalSize += currentSize;
 
-      // If adding this event exceeds the maxBytes (pessimistic estimation), break the loop
-      if (totalSize + currentSize > maxBytes) {
+      // If the response has reached maxBytes, stop polling more events.
+      if (totalSize >= maxBytes) {
         break;
       }
     }
@@ -375,6 +383,15 @@ public List<SubscriptionCommitContext> commit(
     return successfulCommitContexts;
   }
 
+  @Override
+  public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) {
+    final SubscriptionPrefetchingQueue prefetchingQueue =
+        topicNameToPrefetchingQueue.get(commitContext.getTopicName());
+    return Objects.nonNull(prefetchingQueue)
+        && !prefetchingQueue.isClosed()
+        && prefetchingQueue.requeue(consumerId, commitContext);
+  }
+
   @Override
   public int refreshInFlightEventLeases(
       final String consumerId, final List<SubscriptionCommitContext> commitContexts) {

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

@@ -817,6 +817,39 @@ private boolean refreshInFlightEventLeaseInternal(
     return refreshed.get();
   }
 
+  /**
+   * Returns an event to the prefetching queue without modifying its response or nack count.
+   *
+   * <p>This is used when the server polled the event but cannot fit it in the current response.
+   */
+  public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) {
+    acquireReadLock();
+    try {
+      if (isClosed()) {
+        return false;
+      }
+      final AtomicBoolean requeued = new AtomicBoolean(false);
+      inFlightEvents.compute(
+          new Pair<>(consumerId, commitContext),
+          (key, ev) -> {
+            if (Objects.isNull(ev)) {
+              return null;
+            }
+            if (ev.isCommitted()) {
+              ev.cleanUp(false);
+              return null;
+            }
+            ev.resetLastPolledTimestamp();
+            prefetchEvent(ev);
+            requeued.set(true);
+            return null;
+          });
+      return requeued.get();
+    } finally {
+      releaseReadLock();
+    }
+  }
+
   /**
    * @return {@code true} if ack successfully
    */

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

@@ -2579,6 +2579,44 @@ private boolean refreshInFlightEventLeaseInternal(
     return refreshed.get();
   }
 
+  /**
+   * Returns an event to the prefetching queue without modifying its response or nack count.
+   *
+   * <p>This is used when the server polled the event but cannot fit it in the current response.
+   */
+  public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) {
+    acquireReadLock();
+    try {
+      if (isClosed || closeRequested || pendingSeekRequest != null || !isActive) {
+        return false;
+      }
+      if (Objects.isNull(commitContext)
+          || !commitContext.hasWriterProgress()
+          || isCommitContextOutdated(commitContext)) {
+        return false;
+      }
+      final AtomicBoolean requeued = new AtomicBoolean(false);
+      inFlightEvents.compute(
+          new InFlightEventKey(consumerId, commitContext),
+          (key, ev) -> {
+            if (Objects.isNull(ev)) {
+              return null;
+            }
+            if (ev.isCommitted()) {
+              cleanUpEvent(ev, false);
+              return null;
+            }
+            ev.resetLastPolledTimestamp();
+            prefetchingQueue.add(ev);
+            requeued.set(true);
+            return null;
+          });
+      return requeued.get();
+    } finally {
+      releaseReadLock();
+    }
+  }
+
   private boolean canAcceptCommitContext(
       final SubscriptionCommitContext commitContext, final String action, final boolean silent) {
     if (isClosed || closeRequested || pendingSeekRequest != null) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java

@@ -275,6 +275,11 @@ public void nack() {
     }
   }
 
+  /** Makes this event pollable again without treating local response-size control as a nack. */
+  public void resetLastPolledTimestamp() {
+    lastPolledTimestamp.set(INVALID_TIMESTAMP);
+  }
+
   /** Returns the current nack count for this event. */
   public long getNackCount() {
     return nackCount.get();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java

@@ -798,10 +798,18 @@ private TPipeSubscribeResp handlePipeSubscribePollInternal(final PipeSubscribePo
                           req.getRequest(),
                           e);
                     }
-                    // nack
+                    // A response-size overflow caused by events already added to this response is
+                    // local batching backpressure, not a consumer rejection. Requeue it without
+                    // increasing the poison-message nack counter.
                     if (!isOutdated) {
-                      SubscriptionAgent.broker()
-                          .commit(consumerConfig, Collections.singletonList(commitContext), true);
+                      final boolean requeued =
+                          e instanceof SubscriptionPayloadExceedException
+                              && totalSize.get() > 0
+                              && SubscriptionAgent.broker().requeue(consumerConfig, commitContext);
+                      if (!requeued) {
+                        SubscriptionAgent.broker()
+                            .commit(consumerConfig, Collections.singletonList(commitContext), true);
+                      }
                     }
                     return null;
                   }

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

@@ -0,0 +1,93 @@
+/*
+ * 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.subscription.broker;
+
+import org.apache.iotdb.commons.consensus.DataRegionId;
+import org.apache.iotdb.db.subscription.broker.consensus.ConsensusPrefetchingQueue;
+import org.apache.iotdb.db.subscription.event.SubscriptionEvent;
+import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
+
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ConsensusSubscriptionBrokerPayloadLimitTest {
+
+  private static final String CONSUMER_GROUP_ID = "consumerGroup";
+  private static final String CONSUMER_ID = "consumer";
+  private static final String TOPIC_NAME = "topic";
+
+  @Test
+  public void testPollRequeuesEventThatWouldExceedPayloadLimit() throws Exception {
+    final ConsensusSubscriptionBroker broker = new ConsensusSubscriptionBroker(CONSUMER_GROUP_ID);
+    final ConsensusPrefetchingQueue firstQueue = mock(ConsensusPrefetchingQueue.class);
+    final ConsensusPrefetchingQueue secondQueue = mock(ConsensusPrefetchingQueue.class);
+    final SubscriptionEvent firstEvent = mock(SubscriptionEvent.class);
+    final SubscriptionEvent secondEvent = mock(SubscriptionEvent.class);
+    final SubscriptionCommitContext firstCommitContext = newCommitContext(1, 1);
+    final SubscriptionCommitContext secondCommitContext = newCommitContext(2, 2);
+
+    when(firstQueue.getConsensusGroupId()).thenReturn(new DataRegionId(1));
+    when(secondQueue.getConsensusGroupId()).thenReturn(new DataRegionId(2));
+    when(firstQueue.poll(CONSUMER_ID, null)).thenReturn(firstEvent);
+    when(secondQueue.poll(CONSUMER_ID, null)).thenReturn(secondEvent);
+    when(firstEvent.getCurrentResponseSize()).thenReturn(40);
+    when(secondEvent.getCurrentResponseSize()).thenReturn(30);
+    when(firstEvent.getCommitContext()).thenReturn(firstCommitContext);
+    when(secondEvent.getCommitContext()).thenReturn(secondCommitContext);
+    when(secondQueue.requeue(CONSUMER_ID, secondCommitContext)).thenReturn(true);
+    bindQueues(broker, Arrays.asList(firstQueue, secondQueue));
+
+    final List<SubscriptionEvent> events =
+        broker.poll(CONSUMER_ID, Collections.singleton(TOPIC_NAME), 60L);
+
+    assertEquals(1, events.size());
+    assertSame(firstEvent, events.get(0));
+    verify(secondQueue).requeue(CONSUMER_ID, secondCommitContext);
+  }
+
+  private static SubscriptionCommitContext newCommitContext(
+      final int regionId, final int commitId) {
+    return new SubscriptionCommitContext(
+        1, 1, TOPIC_NAME, CONSUMER_GROUP_ID, commitId, "DataRegion[" + regionId + "]", 0L);
+  }
+
+  @SuppressWarnings("unchecked")
+  private static void bindQueues(
+      final ConsensusSubscriptionBroker broker, final List<ConsensusPrefetchingQueue> queues)
+      throws Exception {
+    final Field field =
+        ConsensusSubscriptionBroker.class.getDeclaredField("topicNameToConsensusPrefetchingQueues");
+    field.setAccessible(true);
+    final Map<String, List<ConsensusPrefetchingQueue>> queuesByTopic =
+        (Map<String, List<ConsensusPrefetchingQueue>>) field.get(broker);
+    queuesByTopic.put(TOPIC_NAME, queues);
+  }
+}