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 840fe5ac6eaf..f6cbf193fd35 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 @@ private Optional pollFile( 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 @@ private Optional pollTablets( 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 pollTabletsInternal( final SubscriptionPollResponse initialResponse, final PollTimer timer) throws SubscriptionException { @@ -1541,15 +1554,9 @@ private void commit(final Iterable commitContexts, fi dataNodeIdToSubscriptionCommitContexts.entrySet()) { final List groupedCommitContexts = entry.getValue(); requestedCount += groupedCommitContexts.size(); - final AbstractSubscriptionProvider.CommitResult commitResp = - commitInternal(entry.getKey(), groupedCommitContexts, nack); final List 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 failedInGroup = new ArrayList<>(groupedCommitContexts); @@ -1570,6 +1577,88 @@ private void commit(final Iterable commitContexts, fi } } + private List commitWithRetry( + final int dataNodeId, + final List commitContexts, + final boolean nack) + throws SubscriptionException { + final List remainingCommitContexts = new ArrayList<>(commitContexts); + final List 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 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 acceptedCommitContexts) { @@ -1615,10 +1704,26 @@ protected Set ackCommitContextsWithPartialProgress( 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 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 @@ protected Set ackCommitContextsWithPartialProgress( this, retainedCount, entry.getKey(), - e); + failure); } } } @@ -1681,6 +1786,7 @@ private AbstractSubscriptionProvider.CommitResult commitInternal( final List subscriptionCommitContexts, final boolean nack) throws SubscriptionException { + checkIfFenced(); providers.acquireReadLock(); try { final AbstractSubscriptionProvider provider = providers.getProvider(dataNodeId); @@ -1699,11 +1805,13 @@ private AbstractSubscriptionProvider.CommitResult commitInternal( } 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 @@ private void flushPendingRedirectAcks(final List currentMes } final List 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 @@ private void flushPendingRedirectAcks(final List currentMes entry.getValue(), e); } + overlayCommittedPositions(commitResp.getCommittedProgressByTopic()); + commitResp.getAcceptedCommitContexts().forEach(pendingContexts::remove); + if (pendingContexts.isEmpty()) { + pendingRedirectAcksByTopicRegion.remove(entry.getKey(), pendingContexts); + } } } @@ -2167,7 +2276,24 @@ private void overlayCommittedPositions(final Map progress 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 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 56f732491031..125005763dc5 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 @@ CommitResult commit( final List acceptedCommitContexts = new ArrayList<>(); final Map committedProgressByTopic = new LinkedHashMap<>(); + int completedCommitContextCount = 0; rpcLock.lock(); try { for (final List 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 subscriptionCommitContexts, final boolean nack) throws SubscriptionException { final PipeSubscribeCommitReq req; @@ -898,6 +910,32 @@ interface PipeSubscribeRequestSender { 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 51f35c5d8205..7a05318acf80 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 void testPartialCommitKeepsRejectedContextsInAutoCommitBuffer() throws Ex 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 3ad5935b3373..49f5292694f2 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.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.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 void testContinuationAndCommitFailuresSchedulePartialRecovery() throws Ex } } + @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> uncommitted = + (SortedMap>) 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 messages = new ArrayList<>(); + final List 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 rejected = + Arrays.asList(contexts.get(2), contexts.get(6)); + final List accepted = new ArrayList<>(contexts); + accepted.removeAll(rejected); + final SortedMap> uncommitted = + getUncommittedCommitContexts(consumer); + uncommitted.put(0L, ConcurrentHashMap.newKeySet()); + uncommitted.get(0L).addAll(contexts); + final List> 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 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 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 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 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> 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 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> 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> pending = getPendingRedirectAcks(consumer); + final Set 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 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 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 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> uncommitted = + getUncommittedCommitContexts(consumer); + final Set contexts = ConcurrentHashMap.newKeySet(); + contexts.addAll(Arrays.asList(first, second)); + uncommitted.put(0L, contexts); + final List 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 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 older = consumer.commitAsync(Arrays.asList(first, second)); + Assert.assertTrue(firstBatchCompleted.await(5, TimeUnit.SECONDS)); + + final Map 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> getUncommittedCommitContexts( + final TestPullConsumer consumer) throws Exception { + final Field field = + AbstractSubscriptionPullConsumer.class.getDeclaredField("uncommittedCommitContexts"); + field.setAccessible(true); + return (SortedMap>) field.get(consumer); + } + + @SuppressWarnings("unchecked") + private static Map> getPendingRedirectAcks( + final TestPullConsumer consumer) throws Exception { + final Field field = + AbstractSubscriptionConsumer.class.getDeclaredField("pendingRedirectAcksByTopicRegion"); + field.setAccessible(true); + return (Map>) field.get(consumer); + } + private static AbstractSubscriptionProviders getProviders( final AbstractSubscriptionConsumer consumer) throws Exception { final Field field = AbstractSubscriptionConsumer.class.getDeclaredField("providers"); @@ -412,6 +1245,22 @@ private TestPullConsumer( 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 @@ private TestPullConsumer( .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 @@ private static class TestSubscriptionProvider extends AbstractSubscriptionProvid private boolean emptyHeartbeatTopics; private Runnable handshakeAction = () -> {}; private boolean failContinuation; + private Runnable continuationAction = () -> {}; + private Runnable commitAction = () -> {}; + private Function, CommitResult> commitResponse; + private List committed = Collections.emptyList(); + private Map committedProgressByTopic = Collections.emptyMap(); + private boolean batchCommits; + private boolean failAllRemainingBatches; + private boolean rejectFirstBatch; + private final List> commitBatches = new ArrayList<>(); private final int dataNodeId; private final Map pollCounts; @@ -663,6 +1522,7 @@ List poll( List 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 @@ List pollFile( final SubscriptionCommitContext commitContext, final long writingOffset, final long timeoutMs) { + continuationAction.run(); failContinuationIfRequested(); return Collections.emptyList(); } @@ -686,8 +1547,30 @@ List pollFile( @Override CommitResult commit(final List 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 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() {