Skip to content

Commit a23c610

Browse files
noob-se7ennoob-se7en
authored
Fixes all race conditions from stream consumer (#17089)
* Fixes all race condition from stream consumer * Adds test for tableDataManager cache * Fixes bug related to fetching partition offset * Fixes test and addresses comments * Reduces code duplication
1 parent 8006230 commit a23c610

17 files changed

Lines changed: 501 additions & 114 deletions

File tree

‎pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java‎

Lines changed: 27 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,6 @@
9494
import org.apache.pinot.spi.stream.OffsetCriteria;
9595
import org.apache.pinot.spi.stream.PartitionGroupConsumer;
9696
import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
97-
import org.apache.pinot.spi.stream.PartitionLagState;
9897
import org.apache.pinot.spi.stream.PermanentConsumerException;
9998
import org.apache.pinot.spi.stream.StreamConfig;
10099
import org.apache.pinot.spi.stream.StreamConsumerFactory;
@@ -1082,21 +1081,12 @@ public long getLastConsumedTimestamp() {
10821081
/**
10831082
* Returns the {@link ConsumerPartitionState} for the partition group.
10841083
*/
1085-
public Map<String, ConsumerPartitionState> getConsumerPartitionState() {
1084+
public Map<String, ConsumerPartitionState> getConsumerPartitionState(
1085+
@Nullable StreamPartitionMsgOffset latestMsgOffset) {
10861086
String partitionGroupId = String.valueOf(_partitionGroupId);
1087-
return Collections.singletonMap(partitionGroupId, new ConsumerPartitionState(partitionGroupId, getCurrentOffset(),
1088-
getLastConsumedTimestamp(), fetchLatestStreamOffset(5_000), _lastRowMetadata));
1089-
}
1090-
1091-
/**
1092-
* Returns the {@link PartitionLagState} for the partition group.
1093-
*/
1094-
public Map<String, PartitionLagState> getPartitionToLagState(
1095-
Map<String, ConsumerPartitionState> consumerPartitionStateMap) {
1096-
if (_partitionMetadataProvider == null) {
1097-
createPartitionMetadataProvider("Get Partition Lag State");
1098-
}
1099-
return _partitionMetadataProvider.getCurrentPartitionLagState(consumerPartitionStateMap);
1087+
return Collections.singletonMap(partitionGroupId,
1088+
new ConsumerPartitionState(partitionGroupId, getCurrentOffset(), getLastConsumedTimestamp(), latestMsgOffset,
1089+
_lastRowMetadata));
11001090
}
11011091

11021092
/**
@@ -1119,6 +1109,18 @@ public StreamPartitionMsgOffset getCurrentOffset() {
11191109
return _currentOffset;
11201110
}
11211111

1112+
public String getTableStreamName() {
1113+
return _tableStreamName;
1114+
}
1115+
1116+
public int getStreamPartitionId() {
1117+
return _streamPartitionId;
1118+
}
1119+
1120+
public StreamConsumerFactory getStreamConsumerFactory() {
1121+
return _streamConsumerFactory;
1122+
}
1123+
11221124
@Nullable
11231125
public StreamPartitionMsgOffset getLatestStreamOffsetAtStartupTime() {
11241126
return _latestStreamOffsetAtStartupTime;
@@ -1436,6 +1438,10 @@ protected void postStopConsumedMsg(String reason) {
14361438
} while (!_shouldStop);
14371439
}
14381440

1441+
public StreamMetadataProvider getPartitionMetadataProvider() {
1442+
return _partitionMetadataProvider;
1443+
}
1444+
14391445
protected SegmentCompletionProtocol.Response postSegmentConsumedMsg() {
14401446
// Post segmentConsumed to current leader.
14411447
// Retry maybe once if leader is not found.
@@ -1905,41 +1911,21 @@ public long getTimeSinceEventLastConsumedMs() {
19051911
}
19061912

19071913
@Nullable
1908-
public StreamPartitionMsgOffset fetchLatestStreamOffset(long maxWaitTimeMs, boolean useDebugLog) {
1909-
return fetchStreamOffset(OffsetCriteria.LARGEST_OFFSET_CRITERIA, maxWaitTimeMs, useDebugLog);
1910-
}
1911-
1912-
@Nullable
1913-
public StreamPartitionMsgOffset fetchLatestStreamOffset(long maxWaitTimeMs) {
1914-
return fetchLatestStreamOffset(maxWaitTimeMs, false);
1915-
}
1916-
1917-
@Nullable
1918-
public StreamPartitionMsgOffset fetchEarliestStreamOffset(long maxWaitTimeMs, boolean useDebugLog) {
1919-
return fetchStreamOffset(OffsetCriteria.SMALLEST_OFFSET_CRITERIA, maxWaitTimeMs, useDebugLog);
1914+
private StreamPartitionMsgOffset fetchLatestStreamOffset(long maxWaitTimeMs) {
1915+
return fetchStreamOffset(OffsetCriteria.LARGEST_OFFSET_CRITERIA, maxWaitTimeMs);
19201916
}
19211917

19221918
@Nullable
1923-
public StreamPartitionMsgOffset fetchEarliestStreamOffset(long maxWaitTimeMs) {
1924-
return fetchEarliestStreamOffset(maxWaitTimeMs, false);
1925-
}
1926-
1927-
@Nullable
1928-
private StreamPartitionMsgOffset fetchStreamOffset(OffsetCriteria offsetCriteria, long maxWaitTimeMs,
1929-
boolean useDebugLog) {
1919+
private StreamPartitionMsgOffset fetchStreamOffset(OffsetCriteria offsetCriteria, long maxWaitTimeMs) {
19301920
if (_partitionMetadataProvider == null) {
19311921
createPartitionMetadataProvider("Fetch latest stream offset");
19321922
}
19331923
try {
19341924
return _partitionMetadataProvider.fetchStreamPartitionOffset(offsetCriteria, maxWaitTimeMs);
19351925
} catch (Exception e) {
1936-
String logMessage = "Cannot fetch stream offset with criteria " + offsetCriteria + " for clientId " + _clientId
1937-
+ " and partitionGroupId " + _partitionGroupId + " with maxWaitTime " + maxWaitTimeMs;
1938-
if (!useDebugLog) {
1939-
_segmentLogger.warn(logMessage, e);
1940-
} else {
1941-
_segmentLogger.debug(logMessage, e);
1942-
}
1926+
_segmentLogger.warn(
1927+
"Cannot fetch stream offset with criteria {} for clientId {} and partitionGroupId {} with maxWaitTime {}",
1928+
offsetCriteria, _clientId, _partitionGroupId, maxWaitTimeMs, e);
19431929
}
19441930
return null;
19451931
}

‎pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentMetadataUtils.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.pinot.core.data.manager.realtime;
20+
21+
import java.util.Collections;
22+
import java.util.Map;
23+
import javax.annotation.Nullable;
24+
import org.apache.pinot.spi.stream.StreamMetadataProvider;
25+
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
26+
import org.slf4j.Logger;
27+
import org.slf4j.LoggerFactory;
28+
29+
30+
/**
31+
* Utility class for realtime segment metadata operations.
32+
*/
33+
public class RealtimeSegmentMetadataUtils {
34+
private static final Logger LOGGER = LoggerFactory.getLogger(RealtimeSegmentMetadataUtils.class);
35+
private static final long STREAM_METADATA_FETCH_TIMEOUT_MS = 5000;
36+
37+
private RealtimeSegmentMetadataUtils() {
38+
}
39+
40+
/**
41+
* Fetches the latest stream offset for a segment's partition using the provided metadata provider.
42+
* This encapsulates the common pattern of getting the partition ID from the segment and fetching
43+
* the offset from the metadata provider.
44+
*
45+
* @param realtimeSegmentDataManager The segment data manager to get partition ID from
46+
* @param streamMetadataProvider The stream metadata provider to use for fetching
47+
* @return The latest stream offset for the partition, or null if not available
48+
* @throws RuntimeException if fetching fails
49+
*/
50+
@Nullable
51+
public static StreamPartitionMsgOffset fetchLatestStreamOffset(RealtimeSegmentDataManager realtimeSegmentDataManager,
52+
StreamMetadataProvider streamMetadataProvider) {
53+
try {
54+
int partitionId = realtimeSegmentDataManager.getStreamPartitionId();
55+
Map<Integer, StreamPartitionMsgOffset> partitionMsgOffsetMap =
56+
streamMetadataProvider.fetchLatestStreamOffset(Collections.singleton(partitionId),
57+
STREAM_METADATA_FETCH_TIMEOUT_MS);
58+
return partitionMsgOffsetMap.get(partitionId);
59+
} catch (Exception e) {
60+
String segmentName = realtimeSegmentDataManager.getSegmentName();
61+
LOGGER.error("Failed to fetch latest stream offset for segment: {}", segmentName, e);
62+
throw new RuntimeException("Failed to fetch latest stream offset for segment: " + segmentName, e);
63+
}
64+
}
65+
}

‎pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,13 @@
2020

2121
import com.google.common.annotations.VisibleForTesting;
2222
import com.google.common.base.Preconditions;
23+
import com.google.common.cache.Cache;
24+
import com.google.common.cache.CacheBuilder;
25+
import com.google.common.cache.RemovalNotification;
2326
import java.io.File;
2427
import java.io.IOException;
2528
import java.net.URI;
29+
import java.time.Duration;
2630
import java.util.ArrayList;
2731
import java.util.Collections;
2832
import java.util.HashSet;
@@ -31,6 +35,7 @@
3135
import java.util.Set;
3236
import java.util.UUID;
3337
import java.util.concurrent.ConcurrentHashMap;
38+
import java.util.concurrent.ExecutionException;
3439
import java.util.concurrent.Semaphore;
3540
import java.util.concurrent.TimeUnit;
3641
import java.util.concurrent.TimeoutException;
@@ -82,7 +87,9 @@
8287
import org.apache.pinot.spi.data.FieldSpec;
8388
import org.apache.pinot.spi.data.Schema;
8489
import org.apache.pinot.spi.stream.StreamConfigProperties;
90+
import org.apache.pinot.spi.stream.StreamConsumerFactory;
8591
import org.apache.pinot.spi.stream.StreamMessageMetadata;
92+
import org.apache.pinot.spi.stream.StreamMetadataProvider;
8693
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
8794
import org.apache.pinot.spi.utils.CommonConstants;
8895
import org.apache.pinot.spi.utils.CommonConstants.Segment.Realtime.Status;
@@ -133,6 +140,9 @@ public class RealtimeTableDataManager extends BaseTableDataManager {
133140
public static final long SLEEP_INTERVAL_MS = 30000; // 30 seconds sleep interval
134141
@Deprecated
135142
private static final String SEGMENT_DOWNLOAD_TIMEOUT_MINUTES = "segmentDownloadTimeoutMinutes";
143+
private static final Duration STREAM_METADATA_PROVIDER_CACHE_TTL = Duration.ofMinutes(10);
144+
145+
protected Cache<String, StreamMetadataProvider> _streamMetadataProviderCache;
136146

137147
private final BooleanSupplier _isServerReadyToServeQueries;
138148

@@ -216,6 +226,7 @@ protected void doInit() {
216226
}
217227

218228
_enforceConsumptionInOrder = isEnforceConsumptionInOrder();
229+
_streamMetadataProviderCache = getStreamMetadataProviderCache();
219230

220231
// For dedup and partial-upsert, need to wait for all segments loaded before starting consuming data
221232
if (isDedupEnabled() || isPartialUpsertEnabled()) {
@@ -248,6 +259,23 @@ public boolean getAsBoolean() {
248259
}
249260
}
250261

262+
@VisibleForTesting
263+
protected Cache<String, StreamMetadataProvider> getStreamMetadataProviderCache() {
264+
return CacheBuilder.newBuilder()
265+
.expireAfterAccess(STREAM_METADATA_PROVIDER_CACHE_TTL)
266+
.removalListener((RemovalNotification<String, StreamMetadataProvider> notification) -> {
267+
StreamMetadataProvider provider = notification.getValue();
268+
if (provider != null) {
269+
try {
270+
provider.close();
271+
} catch (Exception e) {
272+
LOGGER.warn("Failed to close StreamMetadataProvider for key {}", notification.getKey(), e);
273+
}
274+
}
275+
})
276+
.build();
277+
}
278+
251279
@Override
252280
protected void doStart() {
253281
}
@@ -355,6 +383,21 @@ public List<SegmentContext> getSegmentContexts(List<IndexSegment> selectedSegmen
355383
return segmentContexts;
356384
}
357385

386+
/**
387+
* Returns thread safe StreamMetadataProvider which is shared across different callers.
388+
*/
389+
public StreamMetadataProvider getStreamMetadataProvider(RealtimeSegmentDataManager realtimeSegmentDataManager) {
390+
String tableStreamName = realtimeSegmentDataManager.getTableStreamName();
391+
StreamConsumerFactory streamConsumerFactory = realtimeSegmentDataManager.getStreamConsumerFactory();
392+
try {
393+
return _streamMetadataProviderCache.get(tableStreamName,
394+
() -> streamConsumerFactory.createStreamMetadataProvider(tableStreamName, true));
395+
} catch (ExecutionException e) {
396+
LOGGER.error("Failed to get stream metadata provider for tableStream: {}", tableStreamName);
397+
throw new RuntimeException(e);
398+
}
399+
}
400+
358401
/**
359402
* Returns all partitionGroupIds for the partitions hosted by this server for current table.
360403
* @apiNote this involves Zookeeper read and should not be used frequently due to efficiency concerns.

‎pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManagerTest.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,27 @@
1818
*/
1919
package org.apache.pinot.core.data.manager.realtime;
2020

21+
import com.google.common.cache.Cache;
22+
import com.google.common.cache.CacheBuilder;
23+
import com.google.common.cache.RemovalNotification;
24+
import java.time.Duration;
25+
import java.util.concurrent.Semaphore;
2126
import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
27+
import org.apache.pinot.core.realtime.impl.fakestream.FakeStreamConfigUtils;
2228
import org.apache.pinot.spi.config.table.TableConfig;
2329
import org.apache.pinot.spi.config.table.TableType;
2430
import org.apache.pinot.spi.data.DateTimeFieldSpec;
2531
import org.apache.pinot.spi.data.FieldSpec;
2632
import org.apache.pinot.spi.data.Schema;
33+
import org.apache.pinot.spi.stream.StreamConfig;
34+
import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
35+
import org.apache.pinot.spi.stream.StreamMetadataProvider;
2736
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
37+
import org.apache.pinot.util.TestUtils;
2838
import org.joda.time.DateTimeZone;
2939
import org.joda.time.format.DateTimeFormat;
40+
import org.mockito.Mockito;
41+
import org.testng.Assert;
3042
import org.testng.annotations.Test;
3143

3244
import static org.mockito.Mockito.mock;
@@ -60,4 +72,58 @@ public void testSetDefaultTimeValueIfInvalid() {
6072
assertEquals(timeFieldSpec.getDefaultNullValue(),
6173
Integer.parseInt(DateTimeFormat.forPattern("yyyyMMdd").withZone(DateTimeZone.UTC).print(currentTimeMs)));
6274
}
75+
76+
@Test
77+
public void testStreamMetadataProviderCache() {
78+
class FakeRealtimeTableDataManager extends RealtimeTableDataManager {
79+
private boolean _cacheEntryRemovalNotified;
80+
public FakeRealtimeTableDataManager(Semaphore segmentBuildSemaphore) {
81+
super(segmentBuildSemaphore);
82+
_streamMetadataProviderCache = getStreamMetadataProviderCache();
83+
}
84+
public void updateCache() {
85+
_streamMetadataProviderCache = CacheBuilder.newBuilder()
86+
.expireAfterAccess(Duration.ofMillis(1))
87+
.removalListener((RemovalNotification<String, StreamMetadataProvider> notification) -> {
88+
StreamMetadataProvider provider = notification.getValue();
89+
if (provider != null) {
90+
try {
91+
provider.close();
92+
_cacheEntryRemovalNotified = true;
93+
} catch (Exception e) {
94+
LOGGER.warn("Failed to close StreamMetadataProvider for key {}", notification.getKey(), e);
95+
}
96+
}
97+
})
98+
.build();
99+
}
100+
public Cache<String, StreamMetadataProvider> getCache() {
101+
return _streamMetadataProviderCache;
102+
}
103+
}
104+
FakeRealtimeTableDataManager fakeRealtimeTableDataManager = new FakeRealtimeTableDataManager(null);
105+
106+
RealtimeSegmentDataManager mockRealtimeSegmentDataManager = Mockito.mock(RealtimeSegmentDataManager.class);
107+
when(mockRealtimeSegmentDataManager.getTableStreamName()).thenReturn("testTable-testTopic");
108+
StreamConfig streamConfig = FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs();
109+
when(mockRealtimeSegmentDataManager.getStreamConsumerFactory()).thenReturn(
110+
StreamConsumerFactoryProvider.create(streamConfig));
111+
112+
StreamMetadataProvider streamMetadataProvider =
113+
fakeRealtimeTableDataManager.getStreamMetadataProvider(mockRealtimeSegmentDataManager);
114+
Assert.assertEquals(fakeRealtimeTableDataManager.getStreamMetadataProvider(mockRealtimeSegmentDataManager),
115+
streamMetadataProvider);
116+
117+
fakeRealtimeTableDataManager.updateCache();
118+
Assert.assertEquals(fakeRealtimeTableDataManager.getCache().size(), 0);
119+
120+
StreamMetadataProvider streamMetadataProvider1 =
121+
fakeRealtimeTableDataManager.getStreamMetadataProvider(mockRealtimeSegmentDataManager);
122+
Assert.assertNotEquals(streamMetadataProvider, streamMetadataProvider1);
123+
124+
TestUtils.waitForCondition(
125+
aVoid -> !(fakeRealtimeTableDataManager.getStreamMetadataProvider(mockRealtimeSegmentDataManager)
126+
.equals(streamMetadataProvider1)), 5, 2000, "streamMetadataProvider returned from cache must be new.");
127+
Assert.assertTrue(fakeRealtimeTableDataManager._cacheEntryRemovalNotified);
128+
}
63129
}

‎pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaConsumerFactory.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,15 @@ public StreamMetadataProvider createStreamMetadataProvider(String clientId) {
4545
return new KafkaStreamMetadataProvider(clientId, _streamConfig);
4646
}
4747

48+
@Override
49+
public StreamMetadataProvider createStreamMetadataProvider(String clientId, boolean concurrentAccessExpected) {
50+
if (concurrentAccessExpected) {
51+
return new SynchronizedKafkaStreamMetadataProvider(clientId, _streamConfig);
52+
} else {
53+
return createStreamMetadataProvider(clientId);
54+
}
55+
}
56+
4857
@Override
4958
public PartitionGroupConsumer createPartitionGroupConsumer(String clientId,
5059
PartitionGroupConsumptionStatus partitionGroupConsumptionStatus) {

0 commit comments

Comments
 (0)