From ed1761ba9b3f3a0254843b2eeb38541b35a89d35 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Wed, 26 Aug 2026 06:24:29 +0000 Subject: [PATCH] Retry consuming-segment stream init on the consumer thread First create used the no-policy Kafka path on the Helix constructor (~10s then OFFLINE). #17062 already retries recreate on the consumer thread. Apply the same policy to first create so a transient Kafka/DNS failure does not leave a partition under-replicated until natural flush. Co-authored-by: Cursor --- .../realtime/RealtimeSegmentDataManager.java | 109 +++++++++++---- .../RealtimeSegmentDataManagerTest.java | 129 ++++++++++++++++++ .../RecordingStreamConsumerFactory.java | 109 +++++++++++++++ 3 files changed, 323 insertions(+), 24 deletions(-) create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RecordingStreamConsumerFactory.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java index 71a61b92d5bc..c789ef51f32d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java @@ -238,8 +238,11 @@ public void deleteSegmentFile() { private static final int MSG_COUNT_THRESHOLD_FOR_LOG = 100000; private static final int BUILD_TIME_LEASE_SECONDS = 30; private static final int MAX_CONSECUTIVE_ERROR_COUNT = 5; - // 8 min max timeout for the retry policy - private static final RetryPolicy CONSUMER_RECREATE_RETRY_POLICY = + // 8 min max timeout for the retry policy. Used for both mid-consume recreate (#17062) and first + // consumer creation on the consumer thread, so a retryable Kafka/DNS init failure does not mark + // the replica OFFLINE after the short Helix-thread path (~10s). + @VisibleForTesting + static final RetryPolicy CONSUMER_RECREATE_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(10, 1000L, 2.0f); // Interrupt consumer thread every 10 seconds in case it doesn't stop, e.g. interrupt flag getting cleared somehow @@ -317,7 +320,7 @@ public void deleteSegmentFile() { private final PartitionGroupConsumptionStatus _partitionGroupConsumptionStatus; final String _clientId; private final TransformPipeline _transformPipeline; - private PartitionGroupConsumer _partitionGroupConsumer = null; + private final AtomicReference _partitionGroupConsumer = new AtomicReference<>(); private StreamMetadataProvider _partitionMetadataProvider = null; private final File _resourceTmpDir; private final String _tableNameWithType; @@ -512,7 +515,8 @@ protected boolean consumeLoop() // Update _currentOffset upon return from this method MessageBatch messageBatch; try { - messageBatch = _partitionGroupConsumer.fetchMessages(_currentOffset, _streamConfig.getFetchTimeoutMillis()); + messageBatch = + _partitionGroupConsumer.get().fetchMessages(_currentOffset, _streamConfig.getFetchTimeoutMillis()); //track realtime rows fetched on a table level. This included valid + invalid rows _serverMetrics.addMeteredTableValue(_clientId, ServerMeter.REALTIME_ROWS_FETCHED, messageBatch.getUnfilteredMessageCount()); @@ -820,6 +824,22 @@ public void run() { _segmentLogger.info("Starting consumption on segment: {}, maxRowCount: {}, maxEndTime: {}.", _llcSegmentName, _segmentMaxRowCount, new DateTime(_consumeEndTime, DateTimeZone.UTC)); + // Create the stream consumer here (consumer thread) with CONSUMER_RECREATE_RETRY_POLICY. Doing this in + // the Helix constructor would either block the state-transition thread for the long retry window or keep + // the short Kafka 5x2s path and mark the replica OFFLINE. See + // https://github.com/apache/pinot/issues/11314 and https://github.com/apache/pinot/pull/17062. + if (_shouldStop) { + return; + } + makeStreamConsumer("Starting"); + if (_shouldStop) { + // Leave a successfully created consumer in place so CONSUMING -> ONLINE can catch up. + return; + } + if (_partitionGroupConsumer.get() == null) { + throw new IllegalStateException("Stream consumer was not created for " + _clientId); + } + // TODO: // When reaching here, the current consuming segment has already acquired the consumer semaphore, but there is // no guarantee that the previous consuming segment is already persisted (replaced with immutable segment). It @@ -957,7 +977,7 @@ public void run() { } case COMMIT: { _state = State.COMMITTING; - _currentOffset = _partitionGroupConsumer.checkpoint(_currentOffset); + _currentOffset = _partitionGroupConsumer.get().checkpoint(_currentOffset); // Lock the segment to avoid multiple threads touching the same segment. Lock segmentLock = _realtimeTableDataManager.getSegmentLock(_segmentNameStr); // NOTE: We need to lock interruptibly because the lock might already be held by the Helix thread for the @@ -1454,11 +1474,12 @@ private void closeStreamConsumer() { } } + /// Closes the current consumer in place and leaves the reference set. COMMIT/CATCH_UP may still + /// call [PartitionGroupConsumer#checkpoint] on that closed instance (the default is identity). private void closePartitionGroupConsumer() { - try { - _partitionGroupConsumer.close(); - } catch (Exception e) { - _segmentLogger.warn("Could not close stream consumer", e); + PartitionGroupConsumer consumer = _partitionGroupConsumer.get(); + if (consumer != null) { + closeQuietly(consumer); } } @@ -1680,16 +1701,18 @@ public void goOnlineFromConsuming(SegmentZKMetadata segmentZKMetadata) } } else { boolean success = false; - // Since online helix transition for a segment can arrive before segment's consumer acquires the - // semaphore, check _consumerSemaphoreAcquired before catching up. - // This is to avoid consuming in parallel to another segment's consumer. - if (_consumerSemaphoreAcquired.get()) { + // ONLINE can arrive before the consumer thread has acquired the semaphore or finished first + // create. Catch up only when both are true; otherwise download. The semaphore used to imply a + // live consumer because the Helix constructor created it. It no longer does. + if (_consumerSemaphoreAcquired.get() && _partitionGroupConsumer.get() != null) { _segmentLogger.info("Attempting to catch up from offset {} to {} ", _currentOffset, endOffset); success = catchupToFinalOffset(endOffset, TimeUnit.MILLISECONDS.convert(MAX_TIME_FOR_CONSUMING_TO_ONLINE_IN_SECONDS, TimeUnit.SECONDS)); } else { - _segmentLogger.warn("Consumer semaphore was not acquired, Skipping catch up from offset {} to {} ", - _currentOffset, endOffset); + _segmentLogger.warn( + "Skipping catch up from offset {} to {} (consumerSemaphoreAcquired={}, consumerPresent={})", + _currentOffset, endOffset, _consumerSemaphoreAcquired.get(), + _partitionGroupConsumer.get() != null); } if (success) { @@ -1978,7 +2001,8 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf // Initialize stopOnDecodeError configuration with proper validation _stopOnDecodeError = parseStopOnDecodeErrorConfig(_streamConfig); - makeStreamConsumer("Starting"); + // Stream consumer creation is deferred to PartitionConsumer.run() so a retryable init failure + // uses CONSUMER_RECREATE_RETRY_POLICY on the consumer thread instead of blocking Helix. createPartitionMetadataProvider("Starting"); setPartitionParameters(realtimeSegmentConfigBuilder, indexingConfig.getSegmentPartitionConfig()); _realtimeSegment = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), serverMetrics); @@ -2152,16 +2176,33 @@ private void setPartitionParameters(RealtimeSegmentConfig.Builder realtimeSegmen } } - /// Creates a new stream consumer + /// Creates a new stream consumer using [CONSUMER_RECREATE_RETRY_POLICY]. + /// + /// Called from the consumer thread at the start of [PartitionConsumer#run], not from the Helix + /// constructor, so a retryable stream/init error does not block the state-transition thread. private void makeStreamConsumer(String reason) { - if (_partitionGroupConsumer != null) { + if (_streamConsumerClosed.get()) { + return; + } + if (_partitionGroupConsumer.get() != null) { closePartitionGroupConsumer(); } _segmentLogger.info("Creating new stream consumer for topic partition {} , reason: {}", _clientId, reason); try { - _partitionGroupConsumer = - _streamConsumerFactory.createPartitionGroupConsumer(_clientId, _partitionGroupConsumptionStatus); - _partitionGroupConsumer.start(_currentOffset); + PartitionGroupConsumer consumer = + _streamConsumerFactory.createPartitionGroupConsumer(_clientId, _partitionGroupConsumptionStatus, + CONSUMER_RECREATE_RETRY_POLICY); + consumer.start(_currentOffset); + // Offload may have closed a still-null field while create was in flight. Do not close on + // _shouldStop: stop() is also used for CONSUMING -> ONLINE catchup. + if (_streamConsumerClosed.get()) { + closeQuietly(consumer); + return; + } + _partitionGroupConsumer.set(consumer); + if (_streamConsumerClosed.get()) { + closePartitionGroupConsumer(); + } } catch (Exception e) { _segmentLogger.error("Faced exception while trying to create stream consumer for topic partition {} reason {}", _clientId, reason, e); @@ -2170,17 +2211,26 @@ private void makeStreamConsumer(String reason) { } } + private void closeQuietly(PartitionGroupConsumer consumer) { + try { + consumer.close(); + } catch (Exception e) { + _segmentLogger.warn("Could not close stream consumer", e); + } + } + /// Checkpoints existing consumer before creating a new consumer instance /// Assumes there is a valid instance of [PartitionGroupConsumer] private void recreateStreamConsumer(String reason) { _segmentLogger.info("Recreating stream consumer for topic partition {}, reason: {}", _clientId, reason); - _currentOffset = _partitionGroupConsumer.checkpoint(_currentOffset); + _currentOffset = _partitionGroupConsumer.get().checkpoint(_currentOffset); closePartitionGroupConsumer(); try { - _partitionGroupConsumer = + PartitionGroupConsumer consumer = _streamConsumerFactory.createPartitionGroupConsumer(_clientId, _partitionGroupConsumptionStatus, CONSUMER_RECREATE_RETRY_POLICY); - _partitionGroupConsumer.start(_currentOffset); + consumer.start(_currentOffset); + _partitionGroupConsumer.set(consumer); } catch (Exception e) { _segmentLogger.error("Faced exception while trying to recreate stream consumer for topic partition {}", _clientId, e); @@ -2285,6 +2335,17 @@ AtomicBoolean getConsumerSemaphoreAcquired() { return _consumerSemaphoreAcquired; } + @VisibleForTesting + @Nullable + PartitionGroupConsumer getPartitionGroupConsumer() { + return _partitionGroupConsumer.get(); + } + + @VisibleForTesting + void setPartitionGroupConsumer(@Nullable PartitionGroupConsumer partitionGroupConsumer) { + _partitionGroupConsumer.set(partitionGroupConsumer); + } + /// Parses the stopOnDecodeError configuration with proper validation and type safety. /// Implements the suggested improvement from code review to add input validation. /// diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java index 9ef007b2e491..e808b5cfd31c 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java @@ -67,6 +67,7 @@ import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.LongMsgOffsetFactory; +import org.apache.pinot.spi.stream.PartitionGroupConsumer; import org.apache.pinot.spi.stream.PermanentConsumerException; import org.apache.pinot.spi.stream.StreamConfigProperties; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; @@ -75,6 +76,7 @@ import org.testng.Assert; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import static org.mockito.ArgumentMatchers.any; @@ -181,6 +183,115 @@ public void testPostStopConsumedMsgDoesNotCheckRegisteredSegmentManager() } } + @Test + public void testHelixConstructorDoesNotCreateStreamConsumer() + throws Exception { + TableConfig tableConfig = Fixtures.createTableConfig(RecordingStreamConsumerFactory.class.getName(), + FakeStreamMessageDecoder.class.getName()); + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager(false, new TimeSupplier(), null, + null, tableConfig)) { + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITHOUT_POLICY_COUNT.get(), 0); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITH_POLICY_COUNT.get(), 0); + Assert.assertFalse(segmentDataManager._postConsumeStoppedCalled); + } + } + + @Test + public void testPartitionConsumerInitUsesRecreateRetryPolicy() + throws Exception { + TableConfig tableConfig = Fixtures.createTableConfig(RecordingStreamConsumerFactory.class.getName(), + FakeStreamMessageDecoder.class.getName()); + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager(false, new TimeSupplier(), null, + null, tableConfig)) { + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + + RealtimeSegmentDataManager.PartitionConsumer consumer = segmentDataManager.createPartitionConsumer(); + LongMsgOffset endOffset = new LongMsgOffset(START_OFFSET_VALUE + 500); + segmentDataManager._consumeOffsets.add(endOffset); + segmentDataManager._responses.add(new SegmentCompletionProtocol.Response( + new SegmentCompletionProtocol.Response.Params().withStatus( + SegmentCompletionProtocol.ControllerResponseStatus.HOLD) + .withStreamPartitionMsgOffset(endOffset.toString()))); + + consumer.run(); + + Assert.assertSame(RecordingStreamConsumerFactory.LAST_RETRY_POLICY.get(), + RealtimeSegmentDataManager.CONSUMER_RECREATE_RETRY_POLICY); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITH_POLICY_COUNT.get(), 1); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITHOUT_POLICY_COUNT.get(), 0); + Assert.assertNotNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertFalse(segmentDataManager._postConsumeStoppedCalled); + } + } + + @Test + public void testConsumerInitFailureAfterRetryExhaustionPostsStopConsumed() + throws Exception { + RecordingStreamConsumerFactory.FAIL_CREATE.set(true); + TableConfig tableConfig = Fixtures.createTableConfig(RecordingStreamConsumerFactory.class.getName(), + FakeStreamMessageDecoder.class.getName()); + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager(false, new TimeSupplier(), null, + null, tableConfig)) { + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertFalse(segmentDataManager._postConsumeStoppedCalled); + + segmentDataManager.createPartitionConsumer().run(); + + Assert.assertSame(RecordingStreamConsumerFactory.LAST_RETRY_POLICY.get(), + RealtimeSegmentDataManager.CONSUMER_RECREATE_RETRY_POLICY); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITH_POLICY_COUNT.get(), 1); + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITHOUT_POLICY_COUNT.get(), 0); + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertTrue(segmentDataManager._postConsumeStoppedCalled); + Assert.assertEquals(segmentDataManager._state.get(segmentDataManager), RealtimeSegmentDataManager.State.ERROR); + } + } + + @Test + public void testOffloadSucceedsWhenStreamConsumerWasNeverCreated() + throws Exception { + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + } + } + + @Test + public void testStopDuringConsumerInitKeepsConsumerForCatchup() + throws Exception { + TableConfig tableConfig = Fixtures.createTableConfig(RecordingStreamConsumerFactory.class.getName(), + FakeStreamMessageDecoder.class.getName()); + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager(false, new TimeSupplier(), null, + null, tableConfig)) { + RecordingStreamConsumerFactory.BEFORE_CREATE_WITH_POLICY.set(segmentDataManager::stop); + + segmentDataManager.createPartitionConsumer().run(); + + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITH_POLICY_COUNT.get(), 1); + Assert.assertNotNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertEquals(RecordingStreamConsumerFactory.CLOSE_COUNT.get(), 0); + Assert.assertFalse(segmentDataManager._postConsumeStoppedCalled); + } + } + + @Test + public void testOffloadDuringConsumerInitClosesCreatedConsumer() + throws Exception { + TableConfig tableConfig = Fixtures.createTableConfig(RecordingStreamConsumerFactory.class.getName(), + FakeStreamMessageDecoder.class.getName()); + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager(false, new TimeSupplier(), null, + null, tableConfig)) { + RecordingStreamConsumerFactory.BEFORE_CREATE_WITH_POLICY.set(segmentDataManager::offload); + + segmentDataManager.createPartitionConsumer().run(); + + Assert.assertEquals(RecordingStreamConsumerFactory.CREATE_WITH_POLICY_COUNT.get(), 1); + Assert.assertNull(segmentDataManager.getPartitionGroupConsumer()); + Assert.assertEquals(RecordingStreamConsumerFactory.CLOSE_COUNT.get(), 1); + Assert.assertFalse(segmentDataManager._postConsumeStoppedCalled); + } + } + private FakeRealtimeSegmentDataManager createFakeSegmentManager(boolean noUpsert, TimeSupplier timeSupplier, @Nullable String maxRows, @Nullable String maxDuration, @Nullable TableConfig tableConfig) throws Exception { @@ -222,6 +333,11 @@ public void setUp() { SegmentBuildTimeLeaseExtender.initExecutor(); } + @BeforeMethod + public void resetRecordingStreamConsumerFactory() { + RecordingStreamConsumerFactory.reset(); + } + @AfterClass public void tearDown() { FileUtils.deleteQuietly(TEMP_DIR); @@ -611,6 +727,7 @@ public void testOnlineTransitionAfterStop() // If catching up, but we did not get to the final offset, then download and replace try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { segmentDataManager.getConsumerSemaphoreAcquired().set(true); + segmentDataManager.setPartitionGroupConsumer(mock(PartitionGroupConsumer.class)); segmentDataManager._stopWaitTimeMs = 0; segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CATCHING_UP); segmentDataManager._consumeOffsets.add(new LongMsgOffset(finalOffsetValue - 1)); @@ -622,6 +739,7 @@ public void testOnlineTransitionAfterStop() // But then if we get to the exact offset, we get to build and replace, not download try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { segmentDataManager.getConsumerSemaphoreAcquired().set(true); + segmentDataManager.setPartitionGroupConsumer(mock(PartitionGroupConsumer.class)); segmentDataManager._stopWaitTimeMs = 0; segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CATCHING_UP); segmentDataManager._consumeOffsets.add(finalOffset); @@ -630,6 +748,17 @@ public void testOnlineTransitionAfterStop() Assert.assertTrue(segmentDataManager._buildAndReplaceCalled); } + // Semaphore acquired but first create never finished: skip catchup and download. + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { + segmentDataManager.getConsumerSemaphoreAcquired().set(true); + segmentDataManager._stopWaitTimeMs = 0; + segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CATCHING_UP); + segmentDataManager._consumeOffsets.add(finalOffset); + segmentDataManager.goOnlineFromConsuming(metadata); + Assert.assertTrue(segmentDataManager._downloadAndReplaceCalled); + Assert.assertFalse(segmentDataManager._buildAndReplaceCalled); + } + // But then if we get to the exact offset, we download the segment because consumer semaphore was never acquired. try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { segmentDataManager._stopWaitTimeMs = 0; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RecordingStreamConsumerFactory.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RecordingStreamConsumerFactory.java new file mode 100644 index 000000000000..130302c84381 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RecordingStreamConsumerFactory.java @@ -0,0 +1,109 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.data.manager.realtime; + +import java.io.IOException; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.pinot.core.realtime.impl.fakestream.FakeStreamConsumerFactory; +import org.apache.pinot.spi.stream.MessageBatch; +import org.apache.pinot.spi.stream.PartitionGroupConsumer; +import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus; +import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; +import org.apache.pinot.spi.utils.retry.RetryPolicy; + +/// Test [org.apache.pinot.spi.stream.StreamConsumerFactory] that records which create overload was used +/// and can fail create to simulate exhausted stream-consumer init. +public class RecordingStreamConsumerFactory extends FakeStreamConsumerFactory { + static final AtomicInteger CREATE_WITHOUT_POLICY_COUNT = new AtomicInteger(); + static final AtomicInteger CREATE_WITH_POLICY_COUNT = new AtomicInteger(); + static final AtomicReference LAST_RETRY_POLICY = new AtomicReference<>(); + static final AtomicBoolean FAIL_CREATE = new AtomicBoolean(); + static final AtomicReference BEFORE_CREATE_WITH_POLICY = new AtomicReference<>(); + static final AtomicInteger CLOSE_COUNT = new AtomicInteger(); + + static void reset() { + CREATE_WITHOUT_POLICY_COUNT.set(0); + CREATE_WITH_POLICY_COUNT.set(0); + LAST_RETRY_POLICY.set(null); + FAIL_CREATE.set(false); + BEFORE_CREATE_WITH_POLICY.set(null); + CLOSE_COUNT.set(0); + } + + @Override + public PartitionGroupConsumer createPartitionGroupConsumer(String clientId, + PartitionGroupConsumptionStatus partitionGroupConsumptionStatus) { + CREATE_WITHOUT_POLICY_COUNT.incrementAndGet(); + if (FAIL_CREATE.get()) { + throw new RuntimeException("stream consumer create failed"); + } + return super.createPartitionGroupConsumer(clientId, partitionGroupConsumptionStatus); + } + + @Override + public PartitionGroupConsumer createPartitionGroupConsumer(String clientId, + PartitionGroupConsumptionStatus partitionGroupConsumptionStatus, RetryPolicy retryPolicy) { + CREATE_WITH_POLICY_COUNT.incrementAndGet(); + LAST_RETRY_POLICY.set(retryPolicy); + Runnable beforeCreate = BEFORE_CREATE_WITH_POLICY.get(); + if (beforeCreate != null) { + beforeCreate.run(); + } + if (FAIL_CREATE.get()) { + throw new RuntimeException("stream consumer create exhausted"); + } + PartitionGroupConsumer delegate = + super.createPartitionGroupConsumer(clientId, partitionGroupConsumptionStatus); + return new CloseCountingPartitionGroupConsumer(delegate); + } + + private static final class CloseCountingPartitionGroupConsumer implements PartitionGroupConsumer { + private final PartitionGroupConsumer _delegate; + + CloseCountingPartitionGroupConsumer(PartitionGroupConsumer delegate) { + _delegate = delegate; + } + + @Override + public void start(StreamPartitionMsgOffset startOffset) { + _delegate.start(startOffset); + } + + @Override + public MessageBatch fetchMessages(StreamPartitionMsgOffset startOffset, int timeoutMs) + throws TimeoutException { + return _delegate.fetchMessages(startOffset, timeoutMs); + } + + @Override + public StreamPartitionMsgOffset checkpoint(StreamPartitionMsgOffset lastOffset) { + return _delegate.checkpoint(lastOffset); + } + + @Override + public void close() + throws IOException { + CLOSE_COUNT.incrementAndGet(); + _delegate.close(); + } + } +}