Skip to content

Retry consuming-segment stream init on the consumer thread #19370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Vamsi-klu wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Retry consuming-segment stream init on the consumer thread #19370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 85 additions & 24 deletions ...src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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> _partitionGroupConsumer = new AtomicReference<>();
private StreamMetadataProvider _partitionMetadataProvider = null;
private final File _resourceTmpDir;
private final String _tableNameWithType;
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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.
///
Expand Down
129 changes: 129 additions & 0 deletions ...test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -222,6 +333,11 @@ public void setUp() {
SegmentBuildTimeLeaseExtender.initExecutor();
}

@BeforeMethod
public void resetRecordingStreamConsumerFactory() {
RecordingStreamConsumerFactory.reset();
}

@AfterClass
public void tearDown() {
FileUtils.deleteQuietly(TEMP_DIR);
Expand Down Expand Up @@ -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));
Expand All @@ -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);
Expand All @@ -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;
Expand Down
Loading
Loading