You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by kh...@apache.org on 2026/10/13 05:15:00 UTC

(pinot) branch master updated: Avoid ZooKeeper reads in segment completion requests for existing FSMs (#19830)

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 <35...@users.noreply.github.com>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org