This is an automated email from the ASF dual-hosted git repository.

jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git


The following commit(s) were added to refs/heads/master by this push:
     new 07dc28102b8 Fix subscription commit recovery after RPC timeouts 
(#19023)
07dc28102b8 is described below

commit 07dc28102b8840d82a7b7d7e47530f0660b4005f
Author: Caideyipi <[email protected]>
AuthorDate: Tue Oct 27 15:44:35 2026 +0800

    Fix subscription commit recovery after RPC timeouts (#19023)
    
    * Retry explicit subscription commits after connection failures
    
    * Preserve completed subscription commit batches during retries
    
    * Preserve newer subscription checkpoints across commit retries
    
    * Retry transient subscription poll cleanup NACKs
    
    * Retry only unconfirmed subscription acknowledgements
    
    * Preserve subscription ACK progress across concurrent commits and server 
errors
---
 .../base/AbstractSubscriptionConsumer.java         | 182 ++++-
 .../base/AbstractSubscriptionProvider.java         |  40 +-
 .../base/SubscriptionConsumerLifecycleTest.java    |   2 +-
 .../SubscriptionConsumerMultiProviderPollTest.java | 887 ++++++++++++++++++++-
 4 files changed, 1079 insertions(+), 32 deletions(-)

diff --git 
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
 
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
index 840fe5ac6ea..f6cbf193fd3 100644
--- 
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
+++ 
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
@@ -110,6 +110,9 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
   private static final long SLEEP_MS = 100L;
   private static final long SLEEP_DELTA_MS = 50L;
   private static final long TIMER_DELTA_MS = 250L;
+  private static final int MAX_COMMIT_RETRIES = 3;
+  private static final long COMMIT_RETRY_INTERVAL_MS =
+      ConsumerConstant.ENDPOINTS_SYNC_INTERVAL_MS_MIN_VALUE;
   private static final AtomicLong LAST_CONSUMER_INSTANCE_EPOCH = new 
AtomicLong();
 
   private final String username;
@@ -1060,9 +1063,7 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
         inFlightFilesCommitContextSet.remove(commitContext);
       }
       // construct temporary message to nack
-      nack(
-          Collections.singletonList(
-              new SubscriptionMessage(commitContext, file.getAbsolutePath(), 
null)));
+      nackAfterPollFailure(new SubscriptionMessage(commitContext, 
file.getAbsolutePath(), null), e);
       throw new SubscriptionRuntimeNonCriticalException(e.getMessage(), e);
     }
   }
@@ -1271,13 +1272,25 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
         throw (SubscriptionConsumerFencedException) e;
       }
       // construct temporary message to nack
-      nack(
-          Collections.singletonList(
-              new SubscriptionMessage(response.getCommitContext(), 
Collections.emptyMap())));
+      nackAfterPollFailure(
+          new SubscriptionMessage(response.getCommitContext(), 
Collections.emptyMap()), e);
       throw new SubscriptionRuntimeNonCriticalException(e.getMessage(), e);
     }
   }
 
+  private void nackAfterPollFailure(final SubscriptionMessage message, final 
Exception pollFailure)
+      throws SubscriptionException {
+    try {
+      nack(Collections.singletonList(message));
+    } catch (final SubscriptionConnectionException nackFailure) {
+      // Recovery was scheduled by commitInternal. A NACK replay can increment 
the server's poison
+      // message counter again. Preserve the poll failure and let the server 
recycle the event.
+      if (nackFailure != pollFailure) {
+        pollFailure.addSuppressed(nackFailure);
+      }
+    }
+  }
+
   private Optional<SubscriptionMessage> pollTabletsInternal(
       final SubscriptionPollResponse initialResponse, final PollTimer timer)
       throws SubscriptionException {
@@ -1541,15 +1554,9 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
         dataNodeIdToSubscriptionCommitContexts.entrySet()) {
       final List<SubscriptionCommitContext> groupedCommitContexts = 
entry.getValue();
       requestedCount += groupedCommitContexts.size();
-      final AbstractSubscriptionProvider.CommitResult commitResp =
-          commitInternal(entry.getKey(), groupedCommitContexts, nack);
       final List<SubscriptionCommitContext> acceptedCommitContexts =
-          commitResp.getAcceptedCommitContexts();
+          commitWithRetry(entry.getKey(), groupedCommitContexts, nack);
       acceptedCount += acceptedCommitContexts.size();
-      if (!nack) {
-        overlayCommittedPositions(commitResp.getCommittedProgressByTopic());
-        onAckedCommitContexts(acceptedCommitContexts);
-      }
       if (acceptedCommitContexts.size() != groupedCommitContexts.size()) {
         final List<SubscriptionCommitContext> failedInGroup =
             new ArrayList<>(groupedCommitContexts);
@@ -1570,6 +1577,88 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
     }
   }
 
+  private List<SubscriptionCommitContext> commitWithRetry(
+      final int dataNodeId,
+      final List<SubscriptionCommitContext> commitContexts,
+      final boolean nack)
+      throws SubscriptionException {
+    final List<SubscriptionCommitContext> remainingCommitContexts = new 
ArrayList<>(commitContexts);
+    final List<SubscriptionCommitContext> acceptedCommitContexts = new 
ArrayList<>();
+    for (int retries = 0; ; retries++) {
+      checkIfFenced();
+      AbstractSubscriptionProvider.CommitResult result =
+          AbstractSubscriptionProvider.CommitResult.empty();
+      SubscriptionConnectionException connectionFailure = null;
+      SubscriptionException failure = null;
+      try {
+        result = commitInternal(dataNodeId, remainingCommitContexts, nack);
+      } catch (final SubscriptionException e) {
+        failure = getCommitFailure(e);
+        result = getCompletedCommitResult(e);
+        if (failure instanceof SubscriptionConnectionException) {
+          connectionFailure = (SubscriptionConnectionException) failure;
+        }
+      }
+
+      final List<SubscriptionCommitContext> acceptedInAttempt = 
result.getAcceptedCommitContexts();
+      acceptedCommitContexts.addAll(acceptedInAttempt);
+      if (!nack) {
+        // Apply each response once, before waiting. Reapplying older progress 
at retry completion
+        // could overwrite a concurrent commit; accepted contexts must also 
leave the buffer now.
+        overlayCommittedPositions(result.getCommittedProgressByTopic());
+        onAckedCommitContexts(acceptedInAttempt);
+      }
+      if (failure != null && connectionFailure == null) {
+        throw failure;
+      }
+      remainingCommitContexts.removeAll(new HashSet<>(acceptedInAttempt));
+      if (nack
+          || remainingCommitContexts.isEmpty()
+          || retries >= MAX_COMMIT_RETRIES
+          || isClosed()
+          || Thread.currentThread().isInterrupted()) {
+        if (connectionFailure != null) {
+          throw connectionFailure;
+        }
+        return acceptedCommitContexts;
+      }
+      checkIfFenced();
+      // A partial response has no permanent-rejection reason. Retry only its 
rejected contexts
+      // and any batches whose response was lost, using the same budget as 
connection recovery.
+      // commitInternal has released the provider read lock before we wait for 
recovery.
+      try {
+        Thread.sleep(COMMIT_RETRY_INTERVAL_MS);
+      } catch (final InterruptedException interruptedException) {
+        Thread.currentThread().interrupt();
+        if (connectionFailure != null) {
+          connectionFailure.addSuppressed(interruptedException);
+          throw connectionFailure;
+        }
+        return acceptedCommitContexts;
+      }
+      checkIfFenced();
+      if (isClosed()) {
+        if (connectionFailure != null) {
+          throw connectionFailure;
+        }
+        return acceptedCommitContexts;
+      }
+    }
+  }
+
+  private SubscriptionException getCommitFailure(final SubscriptionException 
e) {
+    return e instanceof AbstractSubscriptionProvider.CommitException
+        ? ((AbstractSubscriptionProvider.CommitException) e).getFailure()
+        : e;
+  }
+
+  private AbstractSubscriptionProvider.CommitResult getCompletedCommitResult(
+      final SubscriptionException e) {
+    return e instanceof AbstractSubscriptionProvider.CommitException
+        ? ((AbstractSubscriptionProvider.CommitException) 
e).getCompletedResult()
+        : AbstractSubscriptionProvider.CommitResult.empty();
+  }
+
   /** Invoked after the server accepts commit contexts from an acknowledgement 
request. */
   protected void onAckedCommitContexts(
       final Collection<SubscriptionCommitContext> acceptedCommitContexts) {
@@ -1615,10 +1704,26 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
             commitInternal(entry.getKey(), groupedCommitContexts, false);
         overlayCommittedPositions(commitResp.getCommittedProgressByTopic());
         removableCommitContexts.addAll(commitResp.getAcceptedCommitContexts());
-      } catch (final SubscriptionConnectionException e) {
+      } catch (final SubscriptionException e) {
+        final SubscriptionException failure = getCommitFailure(e);
+        final AbstractSubscriptionProvider.CommitResult completedResult =
+            getCompletedCommitResult(e);
+        
overlayCommittedPositions(completedResult.getCommittedProgressByTopic());
+        final Set<SubscriptionCommitContext> acceptedBeforeFailure =
+            new HashSet<>(completedResult.getAcceptedCommitContexts());
+        removableCommitContexts.addAll(acceptedBeforeFailure);
+        if (!(failure instanceof SubscriptionConnectionException)) {
+          // The caller cannot receive the removable set when an error 
propagates. Remove confirmed
+          // contexts from every completed group before leaving the 
automatic/close-time ACK path.
+          onAckedCommitContexts(removableCommitContexts);
+          throw failure;
+        }
         int stagedCount = 0;
         int retainedCount = 0;
         for (final SubscriptionCommitContext commitContext : 
groupedCommitContexts) {
+          if (acceptedBeforeFailure.contains(commitContext)) {
+            continue;
+          }
           if (isConsensusCommitContext(commitContext)) {
             stagePendingRedirectAck(commitContext);
             removableCommitContexts.add(commitContext);
@@ -1642,7 +1747,7 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
               this,
               retainedCount,
               entry.getKey(),
-              e);
+              failure);
         }
       }
     }
@@ -1681,6 +1786,7 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
       final List<SubscriptionCommitContext> subscriptionCommitContexts,
       final boolean nack)
       throws SubscriptionException {
+    checkIfFenced();
     providers.acquireReadLock();
     try {
       final AbstractSubscriptionProvider provider = 
providers.getProvider(dataNodeId);
@@ -1699,11 +1805,13 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
       }
       try {
         return provider.commit(subscriptionCommitContexts, nack);
-      } catch (final SubscriptionConsumerFencedException e) {
-        fence(e);
-        throw e;
-      } catch (final SubscriptionConnectionException e) {
-        providers.syncIfUnavailable(this);
+      } catch (final SubscriptionException e) {
+        final SubscriptionException failure = getCommitFailure(e);
+        if (failure instanceof SubscriptionConsumerFencedException) {
+          fence((SubscriptionConsumerFencedException) failure);
+        } else if (failure instanceof SubscriptionConnectionException) {
+          providers.syncIfUnavailable(this);
+        }
         throw e;
       }
     } finally {
@@ -2095,15 +2203,11 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
       }
 
       final List<SubscriptionCommitContext> contextsToRedirect = new 
ArrayList<>(pendingContexts);
+      AbstractSubscriptionProvider.CommitResult commitResp;
       try {
-        final AbstractSubscriptionProvider.CommitResult commitResp =
-            commitInternal(entry.getValue(), contextsToRedirect, false);
-        overlayCommittedPositions(commitResp.getCommittedProgressByTopic());
-        
commitResp.getAcceptedCommitContexts().forEach(pendingContexts::remove);
-        if (pendingContexts.isEmpty()) {
-          pendingRedirectAcksByTopicRegion.remove(entry.getKey(), 
pendingContexts);
-        }
+        commitResp = commitInternal(entry.getValue(), contextsToRedirect, 
false);
       } catch (final SubscriptionException e) {
+        commitResp = getCompletedCommitResult(e);
         LOGGER.warn(
             SubscriptionMessages
                 
.LOG_ARG_FAILED_REDIRECT_ARG_PENDING_CONSENSUS_ACK_S_ARG_VIA_E6A10AC5,
@@ -2113,6 +2217,11 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
             entry.getValue(),
             e);
       }
+      overlayCommittedPositions(commitResp.getCommittedProgressByTopic());
+      commitResp.getAcceptedCommitContexts().forEach(pendingContexts::remove);
+      if (pendingContexts.isEmpty()) {
+        pendingRedirectAcksByTopicRegion.remove(entry.getKey(), 
pendingContexts);
+      }
     }
   }
 
@@ -2167,7 +2276,24 @@ abstract class AbstractSubscriptionConsumer implements 
AutoCloseable {
     if (Objects.isNull(progressByTopic) || progressByTopic.isEmpty()) {
       return;
     }
-    progressByTopic.forEach(this::overlayCommittedPositions);
+    // ACK responses from different providers can arrive out of order, 
including completed batches
+    // carried by a failure. Preserve the newest position for each writer. 
Explicit seek operations
+    // use the topic-scoped overlay above and can still move a checkpoint 
backward.
+    for (final Entry<String, TopicProgress> entry : 
progressByTopic.entrySet()) {
+      if (Objects.isNull(entry.getValue())) {
+        continue;
+      }
+      for (final RegionProgress regionProgress : 
entry.getValue().getRegionProgress().values()) {
+        if (Objects.nonNull(regionProgress)) {
+          regionProgress
+              .getWriterPositions()
+              .forEach(
+                  (writerId, writerProgress) ->
+                      mergeTopicProgress(
+                          committedPositionsByTopic, entry.getKey(), writerId, 
writerProgress));
+        }
+      }
+    }
   }
 
   private void overlayTopicProgress(
diff --git 
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
 
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
index 56f73249103..125005763dc 100644
--- 
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
+++ 
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
@@ -804,20 +804,32 @@ public abstract class AbstractSubscriptionProvider {
 
     final List<SubscriptionCommitContext> acceptedCommitContexts = new 
ArrayList<>();
     final Map<String, TopicProgress> committedProgressByTopic = new 
LinkedHashMap<>();
+    int completedCommitContextCount = 0;
     rpcLock.lock();
     try {
       for (final List<SubscriptionCommitContext> batch : batches) {
         final CommitResult result = commitOnce(batch, nack);
+        completedCommitContextCount += batch.size();
         acceptedCommitContexts.addAll(result.getAcceptedCommitContexts());
         mergeTopicProgress(committedProgressByTopic, 
result.getCommittedProgressByTopic());
       }
+    } catch (final SubscriptionException e) {
+      if (nack || completedCommitContextCount == 0) {
+        throw e;
+      }
+      // Preserve every completed response even if a later batch fails with a 
server error. The
+      // consumer applies these results before propagating the original 
failure or retrying.
+      throw new CommitException(
+          e,
+          completedCommitContextCount,
+          new CommitResult(acceptedCommitContexts, committedProgressByTopic));
     } finally {
       rpcLock.unlock();
     }
     return new CommitResult(acceptedCommitContexts, committedProgressByTopic);
   }
 
-  private CommitResult commitOnce(
+  CommitResult commitOnce(
       final List<SubscriptionCommitContext> subscriptionCommitContexts, final 
boolean nack)
       throws SubscriptionException {
     final PipeSubscribeCommitReq req;
@@ -898,6 +910,32 @@ public abstract class AbstractSubscriptionProvider {
     TPipeSubscribeResp send(TPipeSubscribeReq req) throws TException;
   }
 
+  static final class CommitException extends SubscriptionException {
+    private final int completedCommitContextCount;
+    private final CommitResult completedResult;
+
+    private CommitException(
+        final SubscriptionException cause,
+        final int completedCommitContextCount,
+        final CommitResult completedResult) {
+      super(cause.getMessage(), cause);
+      this.completedCommitContextCount = completedCommitContextCount;
+      this.completedResult = completedResult;
+    }
+
+    int getCompletedCommitContextCount() {
+      return completedCommitContextCount;
+    }
+
+    CommitResult getCompletedResult() {
+      return completedResult;
+    }
+
+    SubscriptionException getFailure() {
+      return (SubscriptionException) getCause();
+    }
+  }
+
   static final class CommitResult {
     private static final CommitResult EMPTY =
         new CommitResult(Collections.emptyList(), Collections.emptyMap());
diff --git 
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
 
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
index 51f35c5d820..7a05318acf8 100644
--- 
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
+++ 
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
@@ -436,7 +436,7 @@ public class SubscriptionConsumerLifecycleTest {
       Assert.assertEquals(1, uncommittedCommitContexts.size());
       Assert.assertEquals(
           Collections.singleton(rejectedCommitContext), 
uncommittedCommitContexts.get(0L));
-      Assert.assertEquals(1, consumer.commitRequestCount);
+      Assert.assertEquals(4, consumer.commitRequestCount);
     } finally {
       consumer.close();
     }
diff --git 
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
 
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
index 3ad5935b337..49f5292694f 100644
--- 
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
+++ 
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
@@ -24,33 +24,48 @@ import org.apache.iotdb.rpc.subscription.config.TopicConfig;
 import 
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
 import 
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
 import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
+import 
org.apache.iotdb.rpc.subscription.exception.SubscriptionOwnerFencedException;
+import 
org.apache.iotdb.rpc.subscription.exception.SubscriptionRuntimeCriticalException;
+import 
org.apache.iotdb.rpc.subscription.exception.SubscriptionRuntimeNonCriticalException;
+import org.apache.iotdb.rpc.subscription.payload.poll.FileInitPayload;
+import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress;
 import 
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
 import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse;
 import 
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType;
 import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload;
 import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress;
+import org.apache.iotdb.rpc.subscription.payload.poll.WriterId;
+import org.apache.iotdb.rpc.subscription.payload.poll.WriterProgress;
 import 
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHeartbeatResp;
 import org.apache.iotdb.session.AbstractSessionBuilder;
 import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder;
 import org.apache.iotdb.session.subscription.payload.SubscriptionMessage;
+import org.apache.iotdb.session.subscription.util.PollTimer;
 
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.write.record.Tablet;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.net.SocketTimeoutException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.SortedMap;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
@@ -59,6 +74,8 @@ import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 
 public class SubscriptionConsumerMultiProviderPollTest {
 
@@ -69,6 +86,8 @@ public class SubscriptionConsumerMultiProviderPollTest {
   private static final String CONSUMER_GROUP_ID = "test_consumer_group";
   private static final long LONG_INTERVAL_MS = 86_400_000L;
 
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
   @Test
   public void testPollTriesNextProviderBeforeBackingOff() throws 
SubscriptionException {
     final TestPullConsumer consumer = new TestPullConsumer(2, 0);
@@ -367,6 +386,820 @@ public class SubscriptionConsumerMultiProviderPollTest {
     }
   }
 
+  @Test(timeout = 20_000)
+  public void 
testSyncCommitRetriesFailedProviderWithoutRepeatingAcceptedGroups() throws 
Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0, false, false, 
3, true);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider healthy = 
consumer.createdProviders.get(0);
+      final TestSubscriptionProvider failed = consumer.createdProviders.get(1);
+      failed.failContinuation = true;
+      final SubscriptionMessage first =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
1L), 1L);
+      final SubscriptionMessage second =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(2, 0, TOPIC, CONSUMER_GROUP_ID, 
2L), 2L);
+      final Field field =
+          
AbstractSubscriptionPullConsumer.class.getDeclaredField("uncommittedCommitContexts");
+      field.setAccessible(true);
+      @SuppressWarnings("unchecked")
+      final SortedMap<Long, Set<SubscriptionCommitContext>> uncommitted =
+          (SortedMap<Long, Set<SubscriptionCommitContext>>) 
field.get(consumer);
+      uncommitted.put(0L, ConcurrentHashMap.newKeySet());
+      uncommitted.get(0L).add(first.getCommitContext());
+      uncommitted.get(0L).add(second.getCommitContext());
+
+      consumer.commitSync(Arrays.asList(first, second));
+
+      Assert.assertTrue(uncommitted.isEmpty());
+      Assert.assertEquals(1, healthy.commitCount.get());
+      Assert.assertEquals(1, failed.commitCount.get());
+      Assert.assertEquals(0, failed.closeCount.get());
+      Assert.assertEquals(1, failed.sessionCloseCount.get());
+      final TestSubscriptionProvider replacement =
+          (TestSubscriptionProvider) getProviders(consumer).getProvider(2);
+      Assert.assertNotSame(failed, replacement);
+      Assert.assertEquals(1, replacement.commitCount.get());
+      Assert.assertEquals(
+          Collections.singletonList(second.getCommitContext()), 
replacement.committed);
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test(timeout = 20_000)
+  public void testSyncCommitRecoversAlreadyUnavailableProvider() throws 
Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider failed = consumer.createdProviders.get(0);
+      failed.setUnavailable();
+
+      consumer.commitSync(
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
1L), 1L));
+
+      Assert.assertEquals(0, failed.commitCount.get());
+      final TestSubscriptionProvider replacement =
+          (TestSubscriptionProvider) getProviders(consumer).getProvider(1);
+      Assert.assertNotSame(failed, replacement);
+      Assert.assertEquals(1, replacement.commitCount.get());
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test(timeout = 25_000)
+  public void testSyncCommitConnectionRetriesAreBounded() throws Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+    final SubscriptionConnectionException failure =
+        new SubscriptionConnectionException("read timed out", new 
SocketTimeoutException());
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      provider.commitAction =
+          () -> {
+            throw failure;
+          };
+
+      Assert.assertSame(
+          failure,
+          Assert.assertThrows(
+              SubscriptionConnectionException.class, () -> 
consumer.commitSync(commitMessage())));
+
+      Assert.assertEquals(4, provider.commitCount.get());
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test
+  public void testSyncCommitDoesNotRetryServerErrors() throws Exception {
+    for (final SubscriptionException failure :
+        new SubscriptionException[] {
+          new SubscriptionRuntimeNonCriticalException("commit rejected"),
+          new SubscriptionRuntimeCriticalException("commit failed"),
+          new SubscriptionConsumerFencedException("consumer fenced")
+        }) {
+      final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+      try {
+        consumer.open();
+        final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+        provider.commitAction =
+            () -> {
+              throw failure;
+            };
+
+        Assert.assertSame(
+            failure,
+            Assert.assertThrows(
+                SubscriptionException.class, () -> 
consumer.commitSync(commitMessage())));
+        Assert.assertEquals(1, provider.commitCount.get());
+        Assert.assertEquals(
+            failure instanceof SubscriptionConsumerFencedException, 
consumer.isFenced());
+      } finally {
+        consumer.close();
+      }
+    }
+  }
+
+  @Test(timeout = 20_000)
+  public void testSyncCommitRetriesOnlyRejectedContexts() throws Exception {
+    assertPartialCommitRetry(false);
+  }
+
+  @Test(timeout = 25_000)
+  public void 
testSyncCommitRetainsAcceptedContextsAfterPartialRetryExhaustion() throws 
Exception {
+    assertPartialCommitRetry(true);
+  }
+
+  private void assertPartialCommitRetry(final boolean exhaustRetries) throws 
Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0, false, false, 
3, true);
+    try {
+      consumer.open();
+      final List<SubscriptionMessage> messages = new ArrayList<>();
+      final List<SubscriptionCommitContext> contexts = new ArrayList<>();
+      for (int id = 1; id <= 10; id++) {
+        final SubscriptionCommitContext context =
+            new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, id);
+        contexts.add(context);
+        messages.add(new SubscriptionMessage(context, (long) id));
+      }
+      // Accepted contexts need not form a prefix of the request.
+      final List<SubscriptionCommitContext> rejected =
+          Arrays.asList(contexts.get(2), contexts.get(6));
+      final List<SubscriptionCommitContext> accepted = new 
ArrayList<>(contexts);
+      accepted.removeAll(rejected);
+      final SortedMap<Long, Set<SubscriptionCommitContext>> uncommitted =
+          getUncommittedCommitContexts(consumer);
+      uncommitted.put(0L, ConcurrentHashMap.newKeySet());
+      uncommitted.get(0L).addAll(contexts);
+      final List<List<SubscriptionCommitContext>> requests = new ArrayList<>();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      provider.commitResponse =
+          requested -> {
+            requests.add(new ArrayList<>(requested));
+            if (requests.size() == 1) {
+              return new AbstractSubscriptionProvider.CommitResult(
+                  accepted, Collections.singletonMap(TOPIC, 
batchProgress(contexts.get(0))));
+            }
+            Assert.assertEquals(2, uncommitted.get(0L).size());
+            Assert.assertTrue(uncommitted.get(0L).containsAll(rejected));
+            if (!exhaustRetries && requests.size() == 2) {
+              throw new SubscriptionConnectionException("retry read timed 
out");
+            }
+            return new AbstractSubscriptionProvider.CommitResult(
+                exhaustRetries ? Collections.emptyList() : new 
ArrayList<>(requested),
+                Collections.emptyMap());
+          };
+
+      if (exhaustRetries) {
+        Assert.assertThrows(
+            SubscriptionRuntimeNonCriticalException.class, () -> 
consumer.commitSync(messages));
+        Assert.assertEquals(2, uncommitted.get(0L).size());
+        Assert.assertTrue(uncommitted.get(0L).containsAll(rejected));
+      } else {
+        consumer.commitSync(messages);
+        Assert.assertTrue(uncommitted.isEmpty());
+      }
+      Assert.assertEquals(exhaustRetries ? 4 : 3, requests.size());
+      Assert.assertEquals(contexts, requests.get(0));
+      for (int i = 1; i < requests.size(); i++) {
+        Assert.assertEquals(rejected, requests.get(i));
+      }
+      Assert.assertEquals(batchProgress(contexts.get(0)), 
consumer.committedPositions(TOPIC));
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test(timeout = 10_000)
+  public void testSyncCommitRetryPreservesInterruption() throws Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+    final ExecutorService executor = Executors.newSingleThreadExecutor();
+    final CountDownLatch commitFailed = new CountDownLatch(1);
+    final AtomicReference<Thread> commitThread = new AtomicReference<>();
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      provider.commitAction =
+          () -> {
+            commitFailed.countDown();
+            throw new SubscriptionConnectionException("read timed out");
+          };
+      final Future<Boolean> interrupted =
+          executor.submit(
+              () -> {
+                commitThread.set(Thread.currentThread());
+                Assert.assertThrows(
+                    SubscriptionConnectionException.class,
+                    () -> consumer.commitSync(commitMessage()));
+                return Thread.currentThread().isInterrupted();
+              });
+      Assert.assertTrue(commitFailed.await(5, TimeUnit.SECONDS));
+      commitThread.get().interrupt();
+
+      Assert.assertTrue(interrupted.get(1, TimeUnit.SECONDS));
+      Assert.assertEquals(1, provider.commitCount.get());
+    } finally {
+      consumer.close();
+      executor.shutdownNow();
+      executor.awaitTermination(5, TimeUnit.SECONDS);
+    }
+  }
+
+  @Test(timeout = 10_000)
+  public void testSyncCommitStopsRetryingAfterCloseOrFencing() throws 
Exception {
+    for (final boolean closeConsumer : new boolean[] {true, false}) {
+      final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+      try {
+        consumer.open();
+        final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+        final SubscriptionConsumerFencedException fenced =
+            new SubscriptionConsumerFencedException("consumer fenced");
+        provider.commitAction =
+            () -> {
+              if (closeConsumer) {
+                consumer.close();
+              } else {
+                consumer.fence(fenced);
+              }
+              throw new SubscriptionConnectionException("read timed out");
+            };
+
+        final SubscriptionException failure =
+            Assert.assertThrows(
+                SubscriptionException.class, () -> 
consumer.commitSync(commitMessage()));
+
+        if (closeConsumer) {
+          Assert.assertTrue(failure instanceof 
SubscriptionConnectionException);
+        } else {
+          Assert.assertSame(fenced, failure);
+        }
+        Assert.assertEquals(1, provider.commitCount.get());
+      } finally {
+        consumer.close();
+      }
+    }
+  }
+
+  private static SubscriptionMessage commitMessage() {
+    return new SubscriptionMessage(
+        new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 1L), 1L);
+  }
+
+  @Test
+  public void testPollCleanupConnectionFailureDoesNotMakePollCritical() throws 
Exception {
+    for (final boolean filePoll : new boolean[] {false, true}) {
+      final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+      final SubscriptionRuntimeNonCriticalException pollFailure =
+          new SubscriptionRuntimeNonCriticalException("continuation poll 
failed");
+      final SubscriptionConnectionException nackFailure =
+          new SubscriptionConnectionException("NACK read timed out");
+      try {
+        consumer.open();
+        final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+        provider.continuationAction =
+            () -> {
+              throw pollFailure;
+            };
+        provider.commitAction =
+            () -> {
+              throw nackFailure;
+            };
+
+        final Throwable failure = invokeFailingContinuationPoll(consumer, 
filePoll);
+
+        Assert.assertTrue(failure instanceof 
SubscriptionRuntimeNonCriticalException);
+        Assert.assertSame(pollFailure, failure.getCause());
+        Assert.assertArrayEquals(new Throwable[] {nackFailure}, 
pollFailure.getSuppressed());
+        Assert.assertEquals(1, provider.commitCount.get());
+        Assert.assertFalse(consumer.isFenced());
+      } finally {
+        consumer.close();
+      }
+    }
+  }
+
+  @Test
+  public void testPollCleanupStillPropagatesFencingAndServerErrors() throws 
Exception {
+    for (final boolean filePoll : new boolean[] {false, true}) {
+      for (final SubscriptionException nackFailure :
+          new SubscriptionException[] {
+            new SubscriptionConsumerFencedException("consumer fenced"),
+            new SubscriptionRuntimeCriticalException("NACK failed"),
+            new SubscriptionRuntimeNonCriticalException("NACK rejected")
+          }) {
+        final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+        try {
+          consumer.open();
+          final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+          provider.continuationAction =
+              () -> {
+                throw new 
SubscriptionRuntimeNonCriticalException("continuation poll failed");
+              };
+          provider.commitAction =
+              () -> {
+                throw nackFailure;
+              };
+
+          Assert.assertSame(nackFailure, 
invokeFailingContinuationPoll(consumer, filePoll));
+          Assert.assertEquals(1, provider.commitCount.get());
+          Assert.assertEquals(
+              nackFailure instanceof SubscriptionConsumerFencedException, 
consumer.isFenced());
+        } finally {
+          consumer.close();
+        }
+      }
+    }
+  }
+
+  private Throwable invokeFailingContinuationPoll(
+      final TestPullConsumer consumer, final boolean filePoll) throws 
Exception {
+    final Field directory = 
AbstractSubscriptionConsumer.class.getDeclaredField("fileSaveDir");
+    directory.setAccessible(true);
+    directory.set(consumer, temporaryFolder.newFolder().getAbsolutePath());
+    final SubscriptionCommitContext context = 
commitMessage().getCommitContext();
+    final SubscriptionPollResponse response =
+        filePoll
+            ? new SubscriptionPollResponse(
+                SubscriptionPollResponseType.FILE_INIT.getType(),
+                new FileInitPayload("test.tsfile"),
+                context)
+            : new SubscriptionPollResponse(
+                SubscriptionPollResponseType.TABLETS.getType(),
+                new TabletsPayload(Collections.emptyMap(), 1),
+                context);
+    final Method poll =
+        AbstractSubscriptionConsumer.class.getDeclaredMethod(
+            filePoll ? "pollFile" : "pollTablets", 
SubscriptionPollResponse.class, PollTimer.class);
+    poll.setAccessible(true);
+    return Assert.assertThrows(
+            InvocationTargetException.class,
+            () ->
+                poll.invoke(consumer, response, new 
PollTimer(System.currentTimeMillis(), 1_000L)))
+        .getCause();
+  }
+
+  @Test(timeout = 20_000)
+  public void testSyncCommitRetriesOnlyUnconfirmedBatches() throws Exception {
+    assertCommitBatchRetry(false);
+  }
+
+  @Test(timeout = 25_000)
+  public void testSyncCommitRetainsAcceptedBatchesWhenRetriesAreExhausted() 
throws Exception {
+    assertCommitBatchRetry(true);
+  }
+
+  @Test(timeout = 20_000)
+  public void testAsyncCommitRetryDoesNotOverwriteNewerCommittedProgress() 
throws Exception {
+    final TestPullConsumer consumer =
+        new TestPullConsumer(1, 0) {
+          @Override
+          protected void ack(final Iterable<SubscriptionMessage> messages)
+              throws SubscriptionException {
+            commitSync(messages);
+          }
+        };
+    final CountDownLatch retryStarted = new CountDownLatch(1);
+    final CountDownLatch newerCommitCompleted = new CountDownLatch(1);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      final SubscriptionMessage first = commitMessage();
+      final SubscriptionMessage second =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
2L), 2L);
+      enableCommitBatching(provider, first.getCommitContext());
+      provider.commitAction =
+          () -> {
+            if (provider.commitBatches.size() == 3) {
+              retryStarted.countDown();
+              await(newerCommitCompleted);
+            }
+          };
+
+      final CompletableFuture<Void> retry = 
consumer.commitAsync(Arrays.asList(first, second));
+      Assert.assertTrue(retryStarted.await(10, TimeUnit.SECONDS));
+      Assert.assertEquals(
+          batchProgress(first.getCommitContext()), 
consumer.committedPositions(TOPIC));
+
+      final TopicProgress newerProgress =
+          new TopicProgress(
+              Collections.singletonMap(
+                  "region1",
+                  new RegionProgress(
+                      Collections.singletonMap(
+                          new WriterId("region1", 1), new WriterProgress(1L, 
2L)))));
+      final TestSubscriptionProvider otherProvider = 
consumer.createdProviders.get(1);
+      otherProvider.committedProgressByTopic = Collections.singletonMap(TOPIC, 
newerProgress);
+      consumer.commitSync(
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(2, 0, TOPIC, CONSUMER_GROUP_ID, 
3L), 3L));
+      Assert.assertEquals(newerProgress, consumer.committedPositions(TOPIC));
+
+      newerCommitCompleted.countDown();
+      retry.get(5, TimeUnit.SECONDS);
+
+      Assert.assertEquals(
+          newerProgress.getRegionProgress().get("region1"),
+          
consumer.committedPositions(TOPIC).getRegionProgress().get("region1"));
+      Assert.assertEquals(2, 
consumer.committedPositions(TOPIC).getRegionProgress().size());
+    } finally {
+      newerCommitCompleted.countDown();
+      consumer.close();
+    }
+  }
+
+  private void assertCommitBatchRetry(final boolean exhaustRetries) throws 
Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0, false, false, 
3, true);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      final SubscriptionMessage first = commitMessage();
+      final SubscriptionMessage second =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
2L), 2L);
+      enableCommitBatching(provider, first.getCommitContext());
+      provider.failAllRemainingBatches = exhaustRetries;
+      final SortedMap<Long, Set<SubscriptionCommitContext>> uncommitted =
+          getUncommittedCommitContexts(consumer);
+      uncommitted.put(0L, ConcurrentHashMap.newKeySet());
+      uncommitted.get(0L).add(first.getCommitContext());
+      uncommitted.get(0L).add(second.getCommitContext());
+
+      if (exhaustRetries) {
+        Assert.assertThrows(
+            SubscriptionConnectionException.class,
+            () -> consumer.commitSync(Arrays.asList(first, second)));
+        Assert.assertEquals(Collections.singleton(second.getCommitContext()), 
uncommitted.get(0L));
+      } else {
+        consumer.commitSync(Arrays.asList(first, second));
+        Assert.assertTrue(uncommitted.isEmpty());
+      }
+
+      Assert.assertEquals(
+          
batchProgress(first.getCommitContext()).getRegionProgress().get("region1"),
+          
consumer.committedPositions(TOPIC).getRegionProgress().get("region1"));
+      Assert.assertEquals(
+          exhaustRetries ? 1 : 2, 
consumer.committedPositions(TOPIC).getRegionProgress().size());
+
+      Assert.assertEquals(exhaustRetries ? 5 : 3, 
provider.commitBatches.size());
+      Assert.assertEquals(
+          Collections.singletonList(first.getCommitContext()), 
provider.commitBatches.get(0));
+      for (int i = 1; i < provider.commitBatches.size(); i++) {
+        Assert.assertEquals(
+            Collections.singletonList(second.getCommitContext()), 
provider.commitBatches.get(i));
+      }
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test
+  public void testAutoCommitPreservesConfirmedBatchesAfterConnectionFailure() 
throws Exception {
+    for (final boolean consensus : new boolean[] {false, true}) {
+      final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+      try {
+        consumer.open();
+        final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+        final SubscriptionCommitContext first = batchCommitContext(1L, 
consensus);
+        final SubscriptionCommitContext second = batchCommitContext(2L, 
consensus);
+        enableCommitBatching(provider, first);
+
+        final Set<SubscriptionCommitContext> removable =
+            consumer.ackCommitContextsWithPartialProgress(Arrays.asList(first, 
second));
+
+        Assert.assertTrue(removable.contains(first));
+        Assert.assertEquals(consensus, removable.contains(second));
+        Assert.assertEquals(consensus ? 2 : 1, removable.size());
+        Assert.assertEquals(batchProgress(first), 
consumer.committedPositions(TOPIC));
+        final Map<String, Set<SubscriptionCommitContext>> pending =
+            getPendingRedirectAcks(consumer);
+        if (consensus) {
+          Assert.assertEquals(1, pending.size());
+          Assert.assertEquals(
+              Collections.singleton(second), pending.get(TOPIC + '\u0001' + 
"region1"));
+        } else {
+          Assert.assertTrue(pending.isEmpty());
+          Assert.assertEquals(
+              Collections.singleton(second),
+              
consumer.ackCommitContextsWithPartialProgress(Collections.singletonList(second)));
+          Assert.assertEquals(Collections.singletonList(second), 
provider.commitBatches.get(2));
+        }
+      } finally {
+        consumer.close();
+      }
+    }
+  }
+
+  @Test
+  public void testRedirectAckRetriesOnlyUnconfirmedBatches() throws Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      final SubscriptionCommitContext first = batchCommitContext(1L, true);
+      final SubscriptionCommitContext second = batchCommitContext(2L, true);
+      enableCommitBatching(provider, first);
+      final Map<String, Set<SubscriptionCommitContext>> pending = 
getPendingRedirectAcks(consumer);
+      final Set<SubscriptionCommitContext> contexts = 
ConcurrentHashMap.newKeySet();
+      contexts.add(first);
+      contexts.add(second);
+      pending.put(TOPIC + '\u0001' + "region1", contexts);
+      final Method flush =
+          AbstractSubscriptionConsumer.class.getDeclaredMethod(
+              "flushPendingRedirectAcks", List.class);
+      flush.setAccessible(true);
+      final List<SubscriptionMessage> currentMessages =
+          Collections.singletonList(new SubscriptionMessage(first, 1L));
+
+      flush.invoke(consumer, currentMessages);
+
+      final SubscriptionCommitContext accepted = 
provider.commitBatches.get(0).get(0);
+      final SubscriptionCommitContext unconfirmed = 
provider.commitBatches.get(1).get(0);
+      Assert.assertEquals(Collections.singleton(unconfirmed), contexts);
+      Assert.assertEquals(batchProgress(accepted), 
consumer.committedPositions(TOPIC));
+
+      flush.invoke(consumer, currentMessages);
+
+      Assert.assertTrue(pending.isEmpty());
+      Assert.assertEquals(3, provider.commitBatches.size());
+      Assert.assertEquals(Collections.singletonList(unconfirmed), 
provider.commitBatches.get(2));
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test
+  public void testSyncCommitPreservesConfirmedBatchesAfterServerErrors() 
throws Exception {
+    for (final SubscriptionException failure : commitServerFailures()) {
+      assertConfirmedBatchesAfterServerError(failure, "sync");
+    }
+  }
+
+  @Test
+  public void testAsyncCommitPreservesConfirmedBatchesAfterServerErrors() 
throws Exception {
+    for (final SubscriptionException failure : commitServerFailures()) {
+      assertConfirmedBatchesAfterServerError(failure, "async");
+    }
+  }
+
+  @Test
+  public void testAutoCommitPreservesConfirmedBatchesAfterServerErrors() 
throws Exception {
+    for (final SubscriptionException failure : commitServerFailures()) {
+      assertConfirmedBatchesAfterServerError(failure, "auto");
+    }
+  }
+
+  @Test
+  public void testCloseCommitPreservesConfirmedBatchesAfterServerErrors() 
throws Exception {
+    for (final SubscriptionException failure : commitServerFailures()) {
+      assertConfirmedBatchesAfterServerError(failure, "close");
+    }
+  }
+
+  @Test
+  public void testRedirectAckPreservesConfirmedBatchesAfterServerErrors() 
throws Exception {
+    for (final SubscriptionException failure : commitServerFailures()) {
+      assertConfirmedBatchesAfterServerError(failure, "redirect");
+    }
+  }
+
+  private static List<SubscriptionException> commitServerFailures() {
+    return Arrays.asList(
+        new SubscriptionRuntimeNonCriticalException("commit rejected"),
+        new SubscriptionRuntimeCriticalException("commit failed"),
+        new SubscriptionOwnerFencedException("owner fenced"),
+        new SubscriptionConsumerFencedException("consumer fenced"));
+  }
+
+  private void assertConfirmedBatchesAfterServerError(
+      final SubscriptionException failure, final String operation) throws 
Exception {
+    final TestPullConsumer consumer =
+        new TestPullConsumer(1, 0, false, false, 3, true) {
+          @Override
+          protected void ack(final Iterable<SubscriptionMessage> messages)
+              throws SubscriptionException {
+            commitSync(messages);
+          }
+        };
+    try {
+      consumer.open();
+      final SubscriptionCommitContext first = batchCommitContext(1L, false);
+      final SubscriptionCommitContext second = batchCommitContext(2L, false);
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      enableCommitBatching(provider, first);
+      provider.commitAction =
+          () -> {
+            if (provider.commitBatches.size() > 1) {
+              throw failure;
+            }
+          };
+      final SortedMap<Long, Set<SubscriptionCommitContext>> uncommitted =
+          getUncommittedCommitContexts(consumer);
+      final Set<SubscriptionCommitContext> contexts = 
ConcurrentHashMap.newKeySet();
+      contexts.addAll(Arrays.asList(first, second));
+      uncommitted.put(0L, contexts);
+      final List<SubscriptionMessage> messages =
+          Arrays.asList(new SubscriptionMessage(first, 1L), new 
SubscriptionMessage(second, 2L));
+
+      switch (operation) {
+        case "sync":
+          Assert.assertSame(
+              failure,
+              Assert.assertThrows(
+                  SubscriptionException.class, () -> 
consumer.commitSync(messages)));
+          break;
+        case "async":
+          Assert.assertSame(
+              failure,
+              Assert.assertThrows(
+                      java.util.concurrent.ExecutionException.class,
+                      () -> consumer.commitAsync(messages).get(5, 
TimeUnit.SECONDS))
+                  .getCause());
+          break;
+        case "auto":
+          final Class<?> worker =
+              Class.forName(AbstractSubscriptionPullConsumer.class.getName() + 
"$AutoCommitWorker");
+          final java.lang.reflect.Constructor<?> constructor =
+              
worker.getDeclaredConstructor(AbstractSubscriptionPullConsumer.class);
+          constructor.setAccessible(true);
+          ((Runnable) constructor.newInstance(consumer)).run();
+          break;
+        case "close":
+          consumer.close();
+          break;
+        case "redirect":
+          getPendingRedirectAcks(consumer).put(TOPIC + '\u0001' + "region1", 
contexts);
+          final Method flush =
+              AbstractSubscriptionConsumer.class.getDeclaredMethod(
+                  "flushPendingRedirectAcks", List.class);
+          flush.setAccessible(true);
+          flush.invoke(
+              consumer,
+              Collections.singletonList(new 
SubscriptionMessage(batchCommitContext(1L, true), 1L)));
+          break;
+        default:
+          throw new AssertionError(operation);
+      }
+
+      final SubscriptionCommitContext accepted = 
provider.commitBatches.get(0).get(0);
+      final SubscriptionCommitContext unconfirmed = 
provider.commitBatches.get(1).get(0);
+      Assert.assertEquals(Collections.singleton(unconfirmed), contexts);
+      final Field positions =
+          
AbstractSubscriptionConsumer.class.getDeclaredField("committedPositionsByTopic");
+      positions.setAccessible(true);
+      Assert.assertEquals(
+          batchProgress(accepted), ((Map<?, ?>) 
positions.get(consumer)).get(TOPIC));
+      Assert.assertEquals(2, provider.commitBatches.size());
+      Assert.assertEquals(
+          failure instanceof SubscriptionConsumerFencedException, 
consumer.isFenced());
+    } finally {
+      consumer.close();
+    }
+  }
+
+  @Test(timeout = 20_000)
+  public void testCommitFailureDoesNotOverwriteNewerConcurrentProgress() 
throws Exception {
+    final TestPullConsumer consumer =
+        new TestPullConsumer(1, 0) {
+          @Override
+          protected void ack(final Iterable<SubscriptionMessage> messages)
+              throws SubscriptionException {
+            commitSync(messages);
+          }
+        };
+    final CountDownLatch firstBatchCompleted = new CountDownLatch(1);
+    final CountDownLatch newerCommitCompleted = new CountDownLatch(1);
+    try {
+      consumer.open();
+      final SubscriptionMessage first = commitMessage();
+      final SubscriptionMessage second =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
2L), 2L);
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      enableCommitBatching(provider, first.getCommitContext());
+      provider.commitAction =
+          () -> {
+            if (provider.commitBatches.size() == 2) {
+              firstBatchCompleted.countDown();
+              await(newerCommitCompleted);
+            }
+          };
+      final CompletableFuture<Void> older = 
consumer.commitAsync(Arrays.asList(first, second));
+      Assert.assertTrue(firstBatchCompleted.await(5, TimeUnit.SECONDS));
+
+      final Map<WriterId, WriterProgress> writers = new HashMap<>();
+      writers.put(new WriterId("region1", 1), new WriterProgress(1L, 2L));
+      writers.put(new WriterId("region1", 2), new WriterProgress(1L, 3L));
+      final TopicProgress newerProgress =
+          new TopicProgress(Collections.singletonMap("region1", new 
RegionProgress(writers)));
+      final TestSubscriptionProvider otherProvider = 
consumer.createdProviders.get(1);
+      otherProvider.committedProgressByTopic = Collections.singletonMap(TOPIC, 
newerProgress);
+      consumer.commitSync(
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(2, 0, TOPIC, CONSUMER_GROUP_ID, 
3L), 3L));
+
+      newerCommitCompleted.countDown();
+      older.get(10, TimeUnit.SECONDS);
+
+      Assert.assertEquals(
+          newerProgress.getRegionProgress().get("region1"),
+          
consumer.committedPositions(TOPIC).getRegionProgress().get("region1"));
+      Assert.assertEquals(2, 
consumer.committedPositions(TOPIC).getRegionProgress().size());
+    } finally {
+      newerCommitCompleted.countDown();
+      consumer.close();
+    }
+  }
+
+  private static SubscriptionCommitContext batchCommitContext(
+      final long commitId, final boolean consensus) {
+    return consensus
+        ? new SubscriptionCommitContext(
+            1,
+            0,
+            TOPIC,
+            CONSUMER_GROUP_ID,
+            1L,
+            new WriterId("region1", 1),
+            new WriterProgress(1L, commitId))
+        : new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
commitId);
+  }
+
+  @Test(timeout = 20_000)
+  public void testSyncCommitRetriesRejectedBatchAfterLaterConnectionFailure() 
throws Exception {
+    final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+    try {
+      consumer.open();
+      final TestSubscriptionProvider provider = 
consumer.createdProviders.get(0);
+      final SubscriptionMessage first = commitMessage();
+      final SubscriptionMessage second =
+          new SubscriptionMessage(
+              new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 
2L), 2L);
+      enableCommitBatching(provider, first.getCommitContext());
+      provider.rejectFirstBatch = true;
+
+      consumer.commitSync(Arrays.asList(first, second));
+
+      Assert.assertEquals(4, provider.commitBatches.size());
+      Assert.assertEquals(
+          Arrays.asList(
+              Collections.singletonList(first.getCommitContext()),
+              Collections.singletonList(second.getCommitContext()),
+              Collections.singletonList(first.getCommitContext()),
+              Collections.singletonList(second.getCommitContext())),
+          provider.commitBatches);
+    } finally {
+      consumer.close();
+    }
+  }
+
+  private static void enableCommitBatching(
+      final TestSubscriptionProvider provider, final SubscriptionCommitContext 
context)
+      throws Exception {
+    final Field field = 
AbstractSubscriptionProvider.class.getDeclaredField("requestBodySizeLimit");
+    field.setAccessible(true);
+    // Reserve exactly enough room for one context per commit RPC, including 
the response headroom.
+    field.setInt(
+        provider,
+        3
+            * (Integer.BYTES
+                + Byte.BYTES
+                + SubscriptionCommitContext.serialize(context).remaining()));
+    provider.batchCommits = true;
+  }
+
+  private static TopicProgress batchProgress(final SubscriptionCommitContext 
context) {
+    final String regionId = "region" + context.getCommitId();
+    return new TopicProgress(
+        Collections.singletonMap(
+            regionId,
+            new RegionProgress(
+                Collections.singletonMap(
+                    new WriterId(regionId, 1), new WriterProgress(1L, 
context.getCommitId())))));
+  }
+
+  @SuppressWarnings("unchecked")
+  private static SortedMap<Long, Set<SubscriptionCommitContext>> 
getUncommittedCommitContexts(
+      final TestPullConsumer consumer) throws Exception {
+    final Field field =
+        
AbstractSubscriptionPullConsumer.class.getDeclaredField("uncommittedCommitContexts");
+    field.setAccessible(true);
+    return (SortedMap<Long, Set<SubscriptionCommitContext>>) 
field.get(consumer);
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Map<String, Set<SubscriptionCommitContext>> 
getPendingRedirectAcks(
+      final TestPullConsumer consumer) throws Exception {
+    final Field field =
+        
AbstractSubscriptionConsumer.class.getDeclaredField("pendingRedirectAcksByTopicRegion");
+    field.setAccessible(true);
+    return (Map<String, Set<SubscriptionCommitContext>>) field.get(consumer);
+  }
+
   private static AbstractSubscriptionProviders getProviders(
       final AbstractSubscriptionConsumer consumer) throws Exception {
     final Field field = 
AbstractSubscriptionConsumer.class.getDeclaredField("providers");
@@ -412,6 +1245,22 @@ public class SubscriptionConsumerMultiProviderPollTest {
         final boolean emptyTabletBatch,
         final boolean failFirstPoll,
         final int providerCount) {
+      this(
+          dataProviderId,
+          emptyPollsBeforeData,
+          emptyTabletBatch,
+          failFirstPoll,
+          providerCount,
+          false);
+    }
+
+    private TestPullConsumer(
+        final int dataProviderId,
+        final int emptyPollsBeforeData,
+        final boolean emptyTabletBatch,
+        final boolean failFirstPoll,
+        final int providerCount,
+        final boolean autoCommit) {
       super(
           new AbstractSubscriptionPullConsumerBuilder()
               .host(HOST)
@@ -420,7 +1269,8 @@ public class SubscriptionConsumerMultiProviderPollTest {
               .consumerGroupId(CONSUMER_GROUP_ID)
               .heartbeatIntervalMs(LONG_INTERVAL_MS)
               .endpointsSyncIntervalMs(LONG_INTERVAL_MS)
-              .autoCommit(false));
+              .autoCommit(autoCommit)
+              .autoCommitIntervalMs(LONG_INTERVAL_MS));
       this.dataProviderId = dataProviderId;
       this.emptyPollsBeforeData = emptyPollsBeforeData;
       this.emptyTabletBatch = emptyTabletBatch;
@@ -513,6 +1363,15 @@ public class SubscriptionConsumerMultiProviderPollTest {
     private boolean emptyHeartbeatTopics;
     private Runnable handshakeAction = () -> {};
     private boolean failContinuation;
+    private Runnable continuationAction = () -> {};
+    private Runnable commitAction = () -> {};
+    private Function<List<SubscriptionCommitContext>, CommitResult> 
commitResponse;
+    private List<SubscriptionCommitContext> committed = 
Collections.emptyList();
+    private Map<String, TopicProgress> committedProgressByTopic = 
Collections.emptyMap();
+    private boolean batchCommits;
+    private boolean failAllRemainingBatches;
+    private boolean rejectFirstBatch;
+    private final List<List<SubscriptionCommitContext>> commitBatches = new 
ArrayList<>();
 
     private final int dataNodeId;
     private final Map<Integer, Integer> pollCounts;
@@ -663,6 +1522,7 @@ public class SubscriptionConsumerMultiProviderPollTest {
     List<SubscriptionPollResponse> pollTablets(
         final SubscriptionCommitContext commitContext, final int offset, final 
long timeoutMs)
         throws SubscriptionException {
+      continuationAction.run();
       failContinuationIfRequested();
       if (emptyTabletBatch && dataNodeId == dataProviderId && offset == 1) {
         return Collections.singletonList(
@@ -679,6 +1539,7 @@ public class SubscriptionConsumerMultiProviderPollTest {
         final SubscriptionCommitContext commitContext,
         final long writingOffset,
         final long timeoutMs) {
+      continuationAction.run();
       failContinuationIfRequested();
       return Collections.emptyList();
     }
@@ -686,8 +1547,30 @@ public class SubscriptionConsumerMultiProviderPollTest {
     @Override
     CommitResult commit(final List<SubscriptionCommitContext> contexts, final 
boolean nack) {
       commitCount.incrementAndGet();
+      if (batchCommits) {
+        return super.commit(contexts, nack);
+      }
+      commitAction.run();
       failContinuationIfRequested();
-      return new CommitResult(contexts, Collections.emptyMap());
+      if (commitResponse != null) {
+        return commitResponse.apply(contexts);
+      }
+      committed = new ArrayList<>(contexts);
+      return new CommitResult(contexts, committedProgressByTopic);
+    }
+
+    @Override
+    CommitResult commitOnce(final List<SubscriptionCommitContext> contexts, 
final boolean nack) {
+      commitBatches.add(new ArrayList<>(contexts));
+      commitAction.run();
+      if (commitBatches.size() == 2 || (failAllRemainingBatches && 
commitBatches.size() > 1)) {
+        throw new SubscriptionConnectionException("read timed out", new 
SocketTimeoutException());
+      }
+      if (rejectFirstBatch && commitBatches.size() == 1) {
+        return CommitResult.empty();
+      }
+      return new CommitResult(
+          contexts, Collections.singletonMap(TOPIC, 
batchProgress(contexts.get(0))));
     }
 
     private void failContinuationIfRequested() {

Reply via email to