This is an automated email from the ASF dual-hosted git repository.
KKcorps pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new c2e17e36e81 Avoid ZooKeeper reads in segment completion requests for
existing FSMs (#19830)
c2e17e36e81 is described below
commit c2e17e36e812d6f272a25831b315e3ee333a60b9
Author: 9aman <[email protected]>
AuthorDate: Tue Oct 13 10:44:54 2026 +0530
Avoid ZooKeeper reads in segment completion requests for existing FSMs
(#19830)
SegmentCompletionManager read the table config from ZooKeeper for every
segmentConsumed, segmentCommitStart, extendBuildTime and
segmentStoppedConsuming request, only to get the factory that parses the
reported stream offset. segmentCommitStart did not even use the parsed
offset. The read goes through the controller's ZooKeeper session, which
it shares with readers of whole-table segment metadata such as the
retention manager and the segment validation managers, so during a burst
of their reads a request could wait longer than the server's 10 second
request timeout. When a commit start times out, the server retries the
completion from segmentConsumed while the controller still processes the
timed out commit start, so the FSM waits for an upload that never comes,
aborts on the retry, and segment completion has to start over.
The FSM already resolves the factory when it is created, so requests for
an existing FSM now parse the offset with the FSM's factory, and
segmentCommitStart no longer parses it. The FSM lookup also checks the
map before computeIfAbsent(), which can lock the hash bin while another
segment's FSM is created. Failures to read the table config or to parse
the offset now happen inside the request's try block, so they return
FAILED instead of an HTTP error, which the server handles the same way.
---
.../realtime/BlockingSegmentCompletionFSM.java | 9 ++--
.../helix/core/realtime/SegmentCompletionFSM.java | 15 ++++++
.../core/realtime/SegmentCompletionManager.java | 29 +++++++----
.../helix/core/realtime/SegmentCompletionTest.java | 57 ++++++++++++++++++++++
.../tests/SegmentCompletionIntegrationTest.java | 43 +++++++++++++++-
...ureInjectingPinotLLCRealtimeSegmentManager.java | 30 ++++++++++++
6 files changed, 168 insertions(+), 15 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/BlockingSegmentCompletionFSM.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/BlockingSegmentCompletionFSM.java
index 887f37ef641..381bebfdb30 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/BlockingSegmentCompletionFSM.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/BlockingSegmentCompletionFSM.java
@@ -145,11 +145,9 @@ public class BlockingSegmentCompletionFSM implements
SegmentCompletionFSM {
// abort leaving it to realtime segment validation job to fix it.
if (segmentMetadata.getStatus().isCompleted()) {
_state = BlockingSegmentCompletionFSMState.COMMITTED;
- StreamPartitionMsgOffsetFactory factory =
-
_segmentCompletionManager.getStreamPartitionMsgOffsetFactory(_segmentName);
String endOffset = segmentMetadata.getEndOffset();
Preconditions.checkState(endOffset != null, "Failed to find end offset
for segment: %s", segmentName);
- _winningOffset = factory.create(endOffset);
+ _winningOffset = _streamPartitionMsgOffsetFactory.create(endOffset);
_winner = "UNKNOWN";
}
}
@@ -192,6 +190,11 @@ public class BlockingSegmentCompletionFSM implements
SegmentCompletionFSM {
BlockingSegmentCompletionFSMState.COMMITTING) ||
_state.equals(BlockingSegmentCompletionFSMState.COMMITTED);
}
+ @Override
+ public StreamPartitionMsgOffsetFactory getStreamPartitionMsgOffsetFactory() {
+ return _streamPartitionMsgOffsetFactory;
+ }
+
/*
* We just heard from a server that it has reached completion stage, and is
reporting the offset
* that the server is at. Since multiple servers can come in at the same
time for this segment,
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionFSM.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionFSM.java
index e3446f4451a..6c035095801 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionFSM.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionFSM.java
@@ -18,9 +18,11 @@
*/
package org.apache.pinot.controller.helix.core.realtime;
+import javax.annotation.Nullable;
import org.apache.pinot.common.protocols.SegmentCompletionProtocol;
import
org.apache.pinot.controller.helix.core.realtime.segment.CommittingSegmentDescriptor;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
+import org.apache.pinot.spi.stream.StreamPartitionMsgOffsetFactory;
/// Interface for managing the state machine transitions related to segment
completion
@@ -59,6 +61,19 @@ public interface SegmentCompletionFSM {
/// @return `true` if the FSM receives segment created signal, `false`
otherwise.
boolean isImmutableSegmentCreated();
+ /// Returns the factory that parses the stream offsets reported for this
segment. [SegmentCompletionManager] parses
+ /// the offsets of `segmentConsumed`, `extendBuildTime` and
`segmentStoppedConsuming` requests with it, so that a
+ /// request for an existing FSM does not read the table config from
ZooKeeper, where the read can wait behind other
+ /// requests on the controller's ZooKeeper session. It is called by
concurrent requests without synchronizing on the
+ /// FSM, so it must not block, and the returned factory must be thread safe.
+ ///
+ /// @return the offset factory, or `null` if this FSM does not keep one, in
which case the caller resolves it from
+ /// the table config
+ @Nullable
+ default StreamPartitionMsgOffsetFactory getStreamPartitionMsgOffsetFactory()
{
+ return null;
+ }
+
/// Processes the event where a server indicates it has consumed up to a
specified offset.
///
/// This is typically triggered when a server finishes consuming data for a
segment due
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java
index f255eacdf6e..dc54950e0e3 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java
@@ -117,7 +117,22 @@ public class SegmentCompletionManager {
}
private SegmentCompletionFSM lookupOrCreateFsm(LLCSegmentName
llcSegmentName, String msgType) {
- return _fsmMap.computeIfAbsent(llcSegmentName.getSegmentName(), k ->
createFsm(llcSegmentName, msgType));
+ String segmentName = llcSegmentName.getSegmentName();
+ // Look up without locking first. computeIfAbsent() can lock the hash bin
even when the key is present, so it can
+ // wait for the creation of another segment's FSM, which reads ZooKeeper.
+ SegmentCompletionFSM fsm = _fsmMap.get(segmentName);
+ return fsm != null ? fsm : _fsmMap.computeIfAbsent(segmentName, k ->
createFsm(llcSegmentName, msgType));
+ }
+
+ /// Parses an offset reported for the segment with the offset factory of its
FSM, so that a request for an existing
+ /// FSM does not read the table config from ZooKeeper.
+ private StreamPartitionMsgOffset parseOffset(SegmentCompletionFSM fsm,
LLCSegmentName llcSegmentName,
+ String offset) {
+ StreamPartitionMsgOffsetFactory factory =
fsm.getStreamPartitionMsgOffsetFactory();
+ if (factory == null) {
+ factory = getStreamPartitionMsgOffsetFactory(llcSegmentName);
+ }
+ return factory.create(offset);
}
private SegmentCompletionFSM createFsm(LLCSegmentName llcSegmentName, String
msgType) {
@@ -158,13 +173,12 @@ public class SegmentCompletionManager {
}
final String instanceId = reqParams.getInstanceId();
final String stopReason = reqParams.getReason();
- final StreamPartitionMsgOffsetFactory factory =
getStreamPartitionMsgOffsetFactory(segmentName);
- final StreamPartitionMsgOffset offset =
factory.create(reqParams.getStreamPartitionMsgOffset());
SegmentCompletionProtocol.Response response =
SegmentCompletionProtocol.RESP_FAILED;
SegmentCompletionFSM fsm = null;
try {
fsm = lookupOrCreateFsm(segmentName,
SegmentCompletionProtocol.MSG_TYPE_CONSUMED);
+ StreamPartitionMsgOffset offset = parseOffset(fsm, segmentName,
reqParams.getStreamPartitionMsgOffset());
response = fsm.segmentConsumed(instanceId, offset, stopReason);
} catch (Exception e) {
LOGGER.error("Caught exception in segmentConsumed for segment {}",
segmentNameStr, e);
@@ -193,9 +207,6 @@ public class SegmentCompletionManager {
if (!isLeader(tableName) || !_helixManager.isConnected()) {
return SegmentCompletionProtocol.RESP_NOT_LEADER;
}
- final String instanceId = reqParams.getInstanceId();
- final StreamPartitionMsgOffsetFactory factory =
getStreamPartitionMsgOffsetFactory(segmentName);
- final StreamPartitionMsgOffset offset =
factory.create(reqParams.getStreamPartitionMsgOffset());
SegmentCompletionFSM fsm = null;
SegmentCompletionProtocol.Response response =
SegmentCompletionProtocol.RESP_FAILED;
try {
@@ -219,13 +230,12 @@ public class SegmentCompletionManager {
return SegmentCompletionProtocol.RESP_NOT_LEADER;
}
final String instanceId = reqParams.getInstanceId();
- final StreamPartitionMsgOffsetFactory factory =
getStreamPartitionMsgOffsetFactory(segmentName);
- final StreamPartitionMsgOffset offset =
factory.create(reqParams.getStreamPartitionMsgOffset());
final int extTimeSec = reqParams.getExtraTimeSec();
SegmentCompletionFSM fsm = null;
SegmentCompletionProtocol.Response response =
SegmentCompletionProtocol.RESP_FAILED;
try {
fsm = lookupOrCreateFsm(segmentName,
SegmentCompletionProtocol.MSG_TYPE_COMMIT);
+ StreamPartitionMsgOffset offset = parseOffset(fsm, segmentName,
reqParams.getStreamPartitionMsgOffset());
response = fsm.extendBuildTime(instanceId, offset, extTimeSec);
} catch (Exception e) {
LOGGER.error("Caught exception in extendBuildTime for segment {}",
segmentNameStr, e);
@@ -263,13 +273,12 @@ public class SegmentCompletionManager {
return SegmentCompletionProtocol.RESP_NOT_LEADER;
}
final String instanceId = reqParams.getInstanceId();
- final StreamPartitionMsgOffsetFactory factory =
getStreamPartitionMsgOffsetFactory(segmentName);
- final StreamPartitionMsgOffset offset =
factory.create(reqParams.getStreamPartitionMsgOffset());
final String reason = reqParams.getReason();
SegmentCompletionFSM fsm = null;
SegmentCompletionProtocol.Response response =
SegmentCompletionProtocol.RESP_FAILED;
try {
fsm = lookupOrCreateFsm(segmentName,
SegmentCompletionProtocol.MSG_TYPE_STOPPED_CONSUMING);
+ StreamPartitionMsgOffset offset = parseOffset(fsm, segmentName,
reqParams.getStreamPartitionMsgOffset());
response = fsm.stoppedConsuming(instanceId, offset, reason);
} catch (Exception e) {
LOGGER.error("Caught exception in segmentStoppedConsuming for segment
{}", segmentNameStr, e);
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionTest.java
index 95f80ee9a2a..12924be0922 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionTest.java
@@ -39,6 +39,7 @@ import org.apache.pinot.spi.stream.LongMsgOffsetFactory;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffsetFactory;
import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.apache.zookeeper.data.Stat;
import org.testng.Assert;
import org.testng.annotations.BeforeMethod;
@@ -48,6 +49,7 @@ import static
org.apache.pinot.common.protocols.SegmentCompletionProtocol.Contro
import static
org.apache.pinot.common.protocols.SegmentCompletionProtocol.Request;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
public class SegmentCompletionTest {
@@ -203,6 +205,53 @@ public class SegmentCompletionTest {
Assert.assertFalse(_fsmMap.containsKey(_segmentNameStr));
}
+ /// Requests for a segment whose FSM exists must not read the table config,
which is in ZooKeeper, so that a slow
+ /// ZooKeeper cannot delay them past the server's request timeout.
+ @Test
+ public void testRequestsForExistingFsmDoNotReadTableConfig() {
+ SegmentCompletionProtocol.Response response;
+ Request.Params params;
+ // s1 sends offset of 20, gets HOLD at t = 5s. This creates the FSM, which
reads the table config.
+ _segmentCompletionMgr._seconds = 5;
+ params = new
Request.Params().withInstanceId(S_1).withStreamPartitionMsgOffset(_s1Offset.toString())
+ .withSegmentName(_segmentNameStr);
+ response = _segmentCompletionMgr.segmentConsumed(params);
+ assertEquals(response.getStatus(), ControllerResponseStatus.HOLD);
+
+ // Table config reads fail from now on
+ _segmentManager._failTableConfigReads = true;
+
+ // s2 sends offset of 40, gets HOLD
+ _segmentCompletionMgr._seconds += 1;
+ params = new
Request.Params().withInstanceId(S_2).withStreamPartitionMsgOffset(_s2Offset.toString())
+ .withSegmentName(_segmentNameStr);
+ response = _segmentCompletionMgr.segmentConsumed(params);
+ assertEquals(response.getStatus(), ControllerResponseStatus.HOLD);
+ // s3 stops consuming at offset 30
+ _segmentCompletionMgr._seconds += 1;
+ params = new
Request.Params().withInstanceId(S_3).withStreamPartitionMsgOffset(_s3Offset.toString())
+ .withSegmentName(_segmentNameStr).withReason("IAmLazy");
+ response = _segmentCompletionMgr.segmentStoppedConsuming(params);
+ assertEquals(response.getStatus(), ControllerResponseStatus.PROCESSED);
+ // s2 is asked to commit
+ _segmentCompletionMgr._seconds += 1;
+ params = new
Request.Params().withInstanceId(S_2).withStreamPartitionMsgOffset(_s2Offset.toString())
+ .withSegmentName(_segmentNameStr);
+ response = _segmentCompletionMgr.segmentConsumed(params);
+ assertEquals(response.getStatus(), ControllerResponseStatus.COMMIT);
+ // s2 extends its build time, and then starts the commit
+ _segmentCompletionMgr._seconds += 1;
+ params = new
Request.Params().withInstanceId(S_2).withStreamPartitionMsgOffset(_s2Offset.toString())
+ .withSegmentName(_segmentNameStr).withExtraTimeSec(20);
+ response = _segmentCompletionMgr.extendBuildTime(params);
+ assertEquals(response.getStatus(), ControllerResponseStatus.PROCESSED);
+ _segmentCompletionMgr._seconds += 1;
+ params = new
Request.Params().withInstanceId(S_2).withStreamPartitionMsgOffset(_s2Offset.toString())
+ .withSegmentName(_segmentNameStr);
+ response = _segmentCompletionMgr.segmentCommitStart(params);
+ assertEquals(response.getStatus(),
ControllerResponseStatus.COMMIT_CONTINUE);
+ }
+
@Test
public void testStoppedConsumeBeforeHold() {
SegmentCompletionProtocol.Response response;
@@ -1365,6 +1414,8 @@ public class SegmentCompletionTest {
public LLCSegmentName _stoppedSegmentName;
public String _stoppedInstance;
public HelixManager _helixManager = mock(HelixManager.class);
+ // When set, table config reads fail, as they do when ZooKeeper is
unavailable
+ public boolean _failTableConfigReads;
protected MockPinotLLCRealtimeSegmentManager(PinotHelixResourceManager
pinotHelixResourceManager) {
this(pinotHelixResourceManager, new
ControllerMetrics(PinotMetricUtils.getPinotMetricsRegistry()));
@@ -1401,6 +1452,7 @@ public class SegmentCompletionTest {
@Override
public TableConfig getTableConfig(String realtimeTableName) {
+ Preconditions.checkState(!_failTableConfigReads, "Failed to read table
config");
return mock(TableConfig.class);
}
@@ -1412,6 +1464,7 @@ public class SegmentCompletionTest {
}
public static class MockSegmentCompletionManager extends
SegmentCompletionManager {
+ private final PinotLLCRealtimeSegmentManager _realtimeSegmentManager;
private final boolean _isLeader;
public long _seconds;
@@ -1431,11 +1484,15 @@ public class SegmentCompletionTest {
new LeadControllerManager("localhost_1234", helixManager,
controllerMetrics),
SegmentCompletionProtocol.getDefaultMaxSegmentCommitTimeSeconds(),
new SegmentCompletionConfig(new PinotConfiguration()));
+ _realtimeSegmentManager = segmentManager;
_isLeader = isLeader;
}
@Override
protected StreamPartitionMsgOffsetFactory
getStreamPartitionMsgOffsetFactory(LLCSegmentName llcSegmentName) {
+ // Read the table config like the overridden method does, so that tests
can make the read fail
+ _realtimeSegmentManager.getTableConfig(
+
TableNameBuilder.REALTIME.tableNameWithType(llcSegmentName.getTableName()));
return new LongMsgOffsetFactory();
}
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentCompletionIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentCompletionIntegrationTest.java
index d737b7acc8a..e1e5b57d04f 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentCompletionIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentCompletionIntegrationTest.java
@@ -37,8 +37,11 @@ import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.protocols.SegmentCompletionProtocol;
import org.apache.pinot.common.utils.LLCSegmentName;
import org.apache.pinot.common.utils.config.TagNameUtils;
+import org.apache.pinot.controller.BaseControllerStarter;
import
org.apache.pinot.controller.helix.core.PinotHelixSegmentOnlineOfflineStateModelGenerator;
import org.apache.pinot.controller.validation.RealtimeSegmentValidationManager;
+import
org.apache.pinot.integration.tests.realtime.utils.FailureInjectingControllerStarter;
+import
org.apache.pinot.integration.tests.realtime.utils.FailureInjectingPinotLLCRealtimeSegmentManager;
import org.apache.pinot.server.realtime.ControllerLeaderLocator;
import org.apache.pinot.server.realtime.ServerSegmentCompletionProtocolHandler;
import
org.apache.pinot.server.starter.helix.SegmentOnlineOfflineStateModelFactory;
@@ -48,11 +51,12 @@ import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.NetUtils;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.apache.pinot.util.TestUtils;
-import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.testng.Assert.assertEquals;
+
public class SegmentCompletionIntegrationTest extends
BaseClusterIntegrationTest {
private static final int NUM_KAFKA_PARTITIONS = 1;
@@ -66,6 +70,11 @@ public class SegmentCompletionIntegrationTest extends
BaseClusterIntegrationTest
return NUM_KAFKA_PARTITIONS;
}
+ @Override
+ public BaseControllerStarter createControllerStarter() {
+ return new FailureInjectingControllerStarter();
+ }
+
@BeforeClass
public void setUp()
throws Exception {
@@ -141,7 +150,7 @@ public class SegmentCompletionIntegrationTest extends
BaseClusterIntegrationTest
params.withStreamPartitionMsgOffset(new
LongMsgOffset(45688L).toString()).withSegmentName(_currentSegment)
.withReason("RandomReason").withInstanceId(_serverInstance);
SegmentCompletionProtocol.Response response =
protocolHandler.segmentStoppedConsuming(params);
- Assert.assertEquals(response.getStatus(),
SegmentCompletionProtocol.ControllerResponseStatus.PROCESSED);
+ assertEquals(response.getStatus(),
SegmentCompletionProtocol.ControllerResponseStatus.PROCESSED);
// Check if segment get into OFFLINE state
TestUtils.waitForCondition(new Function<Void, Boolean>() {
@@ -187,6 +196,36 @@ public class SegmentCompletionIntegrationTest extends
BaseClusterIntegrationTest
}, 60_000L, "Failed to get a new segment reaching CONSUMING state");
}
+ /// The commit start of a segment whose FSM exists must not wait for a table
config read. The read goes to ZooKeeper,
+ /// where it can wait behind other requests on the controller's ZooKeeper
session for longer than the server's
+ /// request timeout. Runs after [#testStopConsumingAndAutoFix()], on the
consuming segment that test leaves behind.
+ @Test(dependsOnMethods = "testStopConsumingAndAutoFix")
+ public void testCommitStartDoesNotWaitForTableConfigRead() {
+ String realtimeTableName =
TableNameBuilder.REALTIME.tableNameWithType(getTableName());
+ ServerSegmentCompletionProtocolHandler protocolHandler =
+ new ServerSegmentCompletionProtocolHandler(new
ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()),
+ realtimeTableName);
+ SegmentCompletionProtocol.Request.Params params = new
SegmentCompletionProtocol.Request.Params()
+ .withInstanceId(_serverInstance)
+ .withSegmentName(_currentSegment)
+ .withStreamPartitionMsgOffset(new LongMsgOffset(1000L).toString());
+
+ // The server is the only replica, so its first report creates the FSM and
makes the server the committer
+ assertEquals(protocolHandler.segmentConsumed(params).getStatus(),
+ SegmentCompletionProtocol.ControllerResponseStatus.COMMIT);
+
+ // The commit start gets NOT_SENT if it waits for the table config read
longer than the server's request timeout
+ FailureInjectingPinotLLCRealtimeSegmentManager segmentManager =
+ (FailureInjectingPinotLLCRealtimeSegmentManager)
_helixResourceManager.getRealtimeSegmentManager();
+ segmentManager.blockTableConfigReads();
+ try {
+ assertEquals(protocolHandler.segmentCommitStart(params).getStatus(),
+ SegmentCompletionProtocol.ControllerResponseStatus.COMMIT_CONTINUE);
+ } finally {
+ segmentManager.unblockTableConfigReads();
+ }
+ }
+
@AfterClass
public void tearDown()
throws IOException {
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/FailureInjectingPinotLLCRealtimeSegmentManager.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/FailureInjectingPinotLLCRealtimeSegmentManager.java
index 42e175f183e..6b26eab8102 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/FailureInjectingPinotLLCRealtimeSegmentManager.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/FailureInjectingPinotLLCRealtimeSegmentManager.java
@@ -19,21 +19,27 @@
package org.apache.pinot.integration.tests.realtime.utils;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.Uninterruptibles;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.metrics.ControllerMetrics;
import org.apache.pinot.controller.ControllerConf;
import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
import
org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
import org.apache.pinot.controller.helix.core.util.FailureInjectionUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
public class FailureInjectingPinotLLCRealtimeSegmentManager extends
PinotLLCRealtimeSegmentManager {
@VisibleForTesting
private final Map<String, String> _failureConfig;
private final Set<String> _forceExpiredSegments =
ConcurrentHashMap.newKeySet();
+ // Non-null while table config reads are blocked, see blockTableConfigReads()
+ private volatile CountDownLatch _tableConfigReadLatch;
public
FailureInjectingPinotLLCRealtimeSegmentManager(PinotHelixResourceManager
helixResourceManager,
ControllerConf controllerConf, ControllerMetrics controllerMetrics) {
@@ -59,6 +65,30 @@ public class FailureInjectingPinotLLCRealtimeSegmentManager
extends PinotLLCReal
_forceExpiredSegments.clear();
}
+ /// Makes table config reads wait until [#unblockTableConfigReads()] is
called, like reads that queue behind a burst
+ /// of other requests on the controller's ZooKeeper session.
+ public void blockTableConfigReads() {
+ _tableConfigReadLatch = new CountDownLatch(1);
+ }
+
+ public void unblockTableConfigReads() {
+ CountDownLatch latch = _tableConfigReadLatch;
+ if (latch != null) {
+ _tableConfigReadLatch = null;
+ latch.countDown();
+ }
+ }
+
+ @Override
+ public TableConfig getTableConfig(String realtimeTableName) {
+ CountDownLatch latch = _tableConfigReadLatch;
+ if (latch != null) {
+ // Wait longer than the server's request timeout, but not forever in
case a test does not unblock the reads
+ Uninterruptibles.awaitUninterruptibly(latch, 1, TimeUnit.MINUTES);
+ }
+ return super.getTableConfig(realtimeTableName);
+ }
+
@VisibleForTesting
public void enableTestFault(String faultType) {
if (faultType != null) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
