From 7a1f65fce94e2a8a99a87059469f0917993aa1d9 Mon Sep 17 00:00:00 2001 From: Shounak kulkarni Date: Thu, 6 Aug 2026 11:32:04 +0530 Subject: [PATCH 1/3] Move realtime stream offset fetch out of the ideal-state update lock RealtimeSegmentValidationManager's ensureAllPartitionsConsuming fetched the stream offsets inside the Helix ideal-state update lambda, so on a table with many partitions the offset I/O (which can take minutes) was held under the per-table ideal-state lock and re-run on every ZK CAS retry, stalling concurrent segment commits. Pre-fetch the offsets from a read-only snapshot of the ideal state, outside the lock (preFetchOffsets), and pass them into the package-private ensureAllPartitionsConsuming, whose updater lambda now performs only in-memory ideal-state mutation. Lock hold-time is proportional to the mutation, and the offset fetch runs once regardless of CAS retries. The smallest-offset stream fetch is gated (anyPartitionNeedsSmallestOffset) so it runs only on a reset or when a partition actually needs a new CONSUMING segment. A partition that starts needing repair after the lock-free snapshot is deferred to the next validation run rather than being repaired with substituted start offsets. --- .../PinotLLCRealtimeSegmentManager.java | 152 ++++++++++++++---- .../PinotLLCRealtimeSegmentManagerTest.java | 138 +++++++++++++++- 2 files changed, 260 insertions(+), 30 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index b48a41f74093..ed2b68e70fe4 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -1448,28 +1448,34 @@ public void ensureAllPartitionsConsuming(TableConfig tableConfig, List { assert idealState != null; boolean isTableEnabled = idealState.isEnabled(); boolean isTablePaused = isTablePaused(idealState); - boolean offsetsHaveToChange = offsetCriteria != null; if (isTableEnabled && !isTablePaused) { - List currentPartitionGroupConsumptionStatusList = - offsetsHaveToChange ? List.of() - // offsets from metadata are not valid anymore; fetch for all partitions - : getPartitionGroupConsumptionStatusList(idealState, streamConfigs); - // FIXME: Right now, we assume topics are sharing same offset criteria - OffsetCriteria originalOffsetCriteria = streamConfigs.get(0).getOffsetCriteria(); - // Read the smallest offset when a new partition is detected - streamConfigs.stream() - .forEach(streamConfig -> streamConfig.setOffsetCriteria( - offsetsHaveToChange ? offsetCriteria : OffsetCriteria.SMALLEST_OFFSET_CRITERIA)); - List streamMetadataList = - getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); - streamConfigs.stream().forEach(streamConfig -> streamConfig.setOffsetCriteria(originalOffsetCriteria)); - return ensureAllPartitionsConsuming(tableConfig, streamConfigs, idealState, streamMetadataList, - offsetCriteria); + return ensureAllPartitionsConsuming(tableConfig, streamConfigs, idealState, + preFetchedOffsets._streamMetadataList, offsetCriteria, + preFetchedOffsets._partitionIdToSmallestOffset); } else { LOGGER.info("Skipping LLC segments validation for table: {}, isTableEnabled: {}, isTablePaused: {}", realtimeTableName, isTableEnabled, isTablePaused); @@ -1481,6 +1487,78 @@ public void ensureAllPartitionsConsuming(TableConfig tableConfig, List streamConfigs, String realtimeTableName, + IdealState snapshotIdealState, OffsetCriteria offsetCriteria) { + boolean offsetsHaveToChange = offsetCriteria != null; + List currentPartitionGroupConsumptionStatusList = + offsetsHaveToChange ? List.of() + // offsets from metadata are not valid anymore; fetch for all partitions + : getPartitionGroupConsumptionStatusList(snapshotIdealState, streamConfigs); + // FIXME: Right now, we assume topics are sharing same offset criteria + OffsetCriteria originalOffsetCriteria = streamConfigs.get(0).getOffsetCriteria(); + // For the periodic run, compute start offsets with SMALLEST so a newly detected partition starts from the + // beginning; for a reset, use the requested criteria. Restored in the finally below. + streamConfigs.forEach(streamConfig -> streamConfig.setOffsetCriteria( + offsetsHaveToChange ? offsetCriteria : OffsetCriteria.SMALLEST_OFFSET_CRITERIA)); + try { + List streamMetadataList = + getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, snapshotIdealState); + Map partitionIdToSmallestOffset = null; + if (offsetCriteria == null || !offsetCriteria.equals(OffsetCriteria.SMALLEST_OFFSET_CRITERIA)) { + Map latestSegmentZKMetadataMap = getLatestSegmentZKMetadataMap(realtimeTableName); + if (offsetsHaveToChange || anyPartitionNeedsSmallestOffset(snapshotIdealState, latestSegmentZKMetadataMap)) { + partitionIdToSmallestOffset = + fetchPartitionGroupIdToSmallestOffset(streamConfigs, snapshotIdealState, latestSegmentZKMetadataMap); + } + } + return new PreFetchedOffsets(streamMetadataList, partitionIdToSmallestOffset); + } finally { + streamConfigs.forEach(streamConfig -> streamConfig.setOffsetCriteria(originalOffsetCriteria)); + } + } + + /// Returns true if at least one partition's latest segment is present in the ideal state but has no replica in the + /// CONSUMING state, i.e. it may need a new CONSUMING segment created (the only repair path that consults the + /// smallest stream offset). Mirrors the trigger in [#ensureAllPartitionsConsuming]. + private boolean anyPartitionNeedsSmallestOffset(IdealState idealState, + Map latestSegmentZKMetadataMap) { + Map> instanceStatesMap = idealState.getRecord().getMapFields(); + for (SegmentZKMetadata latestSegmentZKMetadata : latestSegmentZKMetadataMap.values()) { + Map instanceStateMap = instanceStatesMap.get(latestSegmentZKMetadata.getSegmentName()); + if (instanceStateMap != null && !instanceStateMap.containsValue(SegmentStateModel.CONSUMING)) { + return true; + } + } + return false; + } + + /// Holder for the stream state pre-fetched by [#preFetchOffsets] outside the ideal-state update lock. + /// `_partitionIdToSmallestOffset` is null when the smallest offsets were not fetched (see [#preFetchOffsets]). + @VisibleForTesting + static class PreFetchedOffsets { + final List _streamMetadataList; + @Nullable + final Map _partitionIdToSmallestOffset; + + PreFetchedOffsets(List streamMetadataList, + @Nullable Map partitionIdToSmallestOffset) { + _streamMetadataList = streamMetadataList; + _partitionIdToSmallestOffset = partitionIdToSmallestOffset; + } + } + /// Updates ideal state after completion of a realtime segment @VisibleForTesting IdealState updateIdealStateOnSegmentCompletion(String realtimeTableName, String committingSegmentName, @@ -1718,7 +1796,8 @@ private boolean isAllInstancesInState(Map instanceStateMap, Stri */ @VisibleForTesting IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List streamConfigs, - IdealState idealState, List streamMetadataList, OffsetCriteria offsetCriteria) { + IdealState idealState, List streamMetadataList, OffsetCriteria offsetCriteria, + @Nullable Map preFetchedPartitionIdToSmallestOffset) { String realtimeTableName = tableConfig.getTableName(); InstancePartitions instancePartitions = getConsumingInstancePartitions(tableConfig); @@ -1745,10 +1824,22 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List partitionIdToSmallestOffset = null; - if (offsetCriteria != null && offsetCriteria.equals(OffsetCriteria.SMALLEST_OFFSET_CRITERIA)) { + // Map from partition id to the smallest stream offset, pre-fetched outside the ideal-state lock (see + // preFetchOffsets). Three cases: + // - non-null: the fetched map. A partition absent from it has reached end of life. + // - null with SMALLEST offset criteria: the start offsets computed above already are the smallest offsets, so + // reuse them (they were fetched with SMALLEST for every partition). + // - null otherwise: the lock-free snapshot gate saw no partition needing a new CONSUMING segment, so the + // smallest offsets were not fetched. Start offsets are NOT the stream-smallest in this case, so they must + // not be substituted; a partition that turns out to need a new segment now (it started needing repair after + // the snapshot) is deferred to the next validation run below. + Map partitionIdToSmallestOffset; + if (preFetchedPartitionIdToSmallestOffset != null) { + partitionIdToSmallestOffset = preFetchedPartitionIdToSmallestOffset; + } else if (offsetCriteria != null && offsetCriteria.equals(OffsetCriteria.SMALLEST_OFFSET_CRITERIA)) { partitionIdToSmallestOffset = partitionIdToStartOffset; + } else { + partitionIdToSmallestOffset = null; } // Walk over all partitions that we have metadata for, and repair any partitions necessary. @@ -1861,16 +1952,21 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List> instanceStatesMap = segmentManager._idealState.getRecord().getMapFields(); + + // Both partition 0 and partition 3 lose all replicas of their CONSUMING segment, so both need a new CONSUMING + // segment created from the smallest offset. + turnNewConsumingSegmentOffline(instanceStatesMap, + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, CURRENT_TIME_MS).getSegmentName()); + turnNewConsumingSegmentOffline(instanceStatesMap, + new LLCSegmentName(RAW_TABLE_NAME, 3, 0, CURRENT_TIME_MS).getSegmentName()); + segmentManager._exceededMaxSegmentCompletionTime = true; + + List streamMetadataList = + segmentManager.getNewStreamMetadataList(segmentManager._streamConfigs, List.of(), mock(IdealState.class)); + // Pre-fetched smallest offsets: partition 0 gets an offset ahead of its committed segment (so the new segment + // must start from it, proving the pre-fetched map is used); partition 3 is intentionally omitted. + LongMsgOffset partition0SmallestOffset = new LongMsgOffset(PARTITION_OFFSET.getOffset() + 10_000); + Map preFetchedSmallestOffset = new HashMap<>(); + preFetchedSmallestOffset.put(0, partition0SmallestOffset); + preFetchedSmallestOffset.put(1, PARTITION_OFFSET); + preFetchedSmallestOffset.put(2, PARTITION_OFFSET); + + segmentManager.ensureAllPartitionsConsuming(segmentManager._tableConfig, segmentManager._streamConfigs, + segmentManager._idealState, streamMetadataList, null, preFetchedSmallestOffset); + + // Partition 0 (present in the pre-fetched map) gets a fresh CONSUMING segment whose start offset comes from the + // pre-fetched smallest offset; partition 3 (absent from the populated map -> treated as end of life) is skipped + // and no exception is thrown. + String partition0NewSegment = new LLCSegmentName(RAW_TABLE_NAME, 0, 1, CURRENT_TIME_MS).getSegmentName(); + assertTrue(instanceStatesMap.containsKey(partition0NewSegment)); + assertEquals(segmentManager._segmentZKMetadataMap.get(partition0NewSegment).getStartOffset(), + partition0SmallestOffset.toString()); + assertFalse(instanceStatesMap.containsKey( + new LLCSegmentName(RAW_TABLE_NAME, 3, 1, CURRENT_TIME_MS).getSegmentName())); + } + + /// Phase-2: when the smallest offsets were not pre-fetched (null map, periodic/non-SMALLEST criteria) but a + /// partition needs a new CONSUMING segment (it started needing repair after the lock-free snapshot), the repair is + /// deferred to the next run rather than substituting start offsets. Healthy partitions are untouched and no + /// exception is thrown. + @Test + public void testEnsureAllPartitionsConsumingDefersRepairWhenSmallestOffsetsNotPreFetched() { + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(); + setUpNewTable(segmentManager, 2, 5, 4); + Map> instanceStatesMap = segmentManager._idealState.getRecord().getMapFields(); + + // Partition 0 loses all replicas of its CONSUMING segment, so it needs a new one. + turnNewConsumingSegmentOffline(instanceStatesMap, + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, CURRENT_TIME_MS).getSegmentName()); + segmentManager._exceededMaxSegmentCompletionTime = true; + + List streamMetadataList = + segmentManager.getNewStreamMetadataList(segmentManager._streamConfigs, List.of(), mock(IdealState.class)); + + // Null pre-fetched map with a periodic (null) offset criteria -> the defer branch. + segmentManager.ensureAllPartitionsConsuming(segmentManager._tableConfig, segmentManager._streamConfigs, + segmentManager._idealState, streamMetadataList, null, null); + + // Partition 0's repair was deferred: no new CONSUMING segment created for it. + assertFalse(instanceStatesMap.containsKey( + new LLCSegmentName(RAW_TABLE_NAME, 0, 1, CURRENT_TIME_MS).getSegmentName())); + // Healthy partition 1 still has its original CONSUMING segment, untouched. + assertTrue(instanceStatesMap.containsKey( + new LLCSegmentName(RAW_TABLE_NAME, 1, 0, CURRENT_TIME_MS).getSegmentName())); + } + + /// Phase-2: the smallest-offset stream fetch is only performed when it can be used. On a healthy table (every + /// partition has a CONSUMING segment) [PinotLLCRealtimeSegmentManager#preFetchOffsets] returns a null + /// smallest-offset map; once a partition loses all CONSUMING replicas it is fetched. + @Test + public void testPreFetchOffsetsSkipsSmallestOffsetFetchForHealthyTable() { + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(); + setUpNewTable(segmentManager, 2, 5, 4); + + PinotLLCRealtimeSegmentManager.PreFetchedOffsets healthy = segmentManager.preFetchOffsets( + segmentManager._streamConfigs, REALTIME_TABLE_NAME, segmentManager._idealState, null); + assertNull(healthy._partitionIdToSmallestOffset); + assertNotNull(healthy._streamMetadataList); + + // Turn all replicas of partition 0's CONSUMING segment OFFLINE - now the smallest offset is needed. + turnNewConsumingSegmentOffline(segmentManager._idealState.getRecord().getMapFields(), + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, CURRENT_TIME_MS).getSegmentName()); + PinotLLCRealtimeSegmentManager.PreFetchedOffsets needsRepair = segmentManager.preFetchOffsets( + segmentManager._streamConfigs, REALTIME_TABLE_NAME, segmentManager._idealState, null); + assertNotNull(needsRepair._partitionIdToSmallestOffset); + } + + /// Phase-2: preFetchOffsets temporarily overrides the shared streamConfigs offset criteria; it must restore the + /// original even when the stream fetch throws, so subsequent segment creation is not corrupted. + @Test + public void testPreFetchOffsetsRestoresOffsetCriteriaOnFailure() { + // Toggled on only after table setup, so the simulated failure occurs during preFetchOffsets and not during + // setUpNewTable (which also builds stream metadata). + boolean[] failStreamFetch = {false}; + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager() { + @Override + List getNewStreamMetadataList(List streamConfigs, + List currentPartitionGroupConsumptionStatusList, IdealState idealState) { + if (failStreamFetch[0]) { + throw new RuntimeException("simulated stream failure"); + } + return super.getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); + } + }; + setUpNewTable(segmentManager, 2, 5, 4); + OffsetCriteria originalOffsetCriteria = segmentManager._streamConfigs.get(0).getOffsetCriteria(); + // Guard the test's discriminating power: the reset criteria below must differ from the original, otherwise a + // missing restore would be indistinguishable from the mutation. + assertNotEquals(originalOffsetCriteria, OffsetCriteria.LARGEST_OFFSET_CRITERIA); + failStreamFetch[0] = true; + + // Pass a reset criteria (LARGEST) so preFetchOffsets mutates the shared streamConfigs to LARGEST before the + // fetch throws. If the finally-restore were removed, the criteria would remain LARGEST and the assertion below + // would fail. + try { + segmentManager.preFetchOffsets(segmentManager._streamConfigs, REALTIME_TABLE_NAME, segmentManager._idealState, + OffsetCriteria.LARGEST_OFFSET_CRITERIA); + fail("Expected the simulated stream failure to propagate"); + } catch (RuntimeException e) { + // Expected + } + assertEquals(segmentManager._streamConfigs.get(0).getOffsetCriteria(), originalOffsetCriteria); + } + /// Removes the new CONSUMING segment and sets the latest committed (ONLINE) segment to CONSUMING if exists in the /// ideal state. private void removeNewConsumingSegment(Map> instanceStatesMap, String consumingSegment, @@ -2494,8 +2625,11 @@ public void setUpNewTable() { } public void ensureAllPartitionsConsuming() { - ensureAllPartitionsConsuming(_tableConfig, _streamConfigs, _idealState, - getNewStreamMetadataList(_streamConfigs, List.of(), mock(IdealState.class)), null); + // Mirror the production flow: pre-fetch offsets (gated) outside the ideal-state update, then pass them into the + // package-private repair method. + PreFetchedOffsets preFetchedOffsets = preFetchOffsets(_streamConfigs, REALTIME_TABLE_NAME, _idealState, null); + ensureAllPartitionsConsuming(_tableConfig, _streamConfigs, _idealState, preFetchedOffsets._streamMetadataList, + null, preFetchedOffsets._partitionIdToSmallestOffset); } @Override From 07dc518ffe68b154bd591f968df1aa96402b8418 Mon Sep 17 00:00:00 2001 From: Shounak kulkarni Date: Mon, 17 Aug 2026 13:11:43 +0530 Subject: [PATCH 2/3] Derive smallest-offset gate from ideal state; dedupe latest-LLC-segment lookup Addresses review feedback on the offset pre-fetch: anyPartitionNeedsSmallestOffset previously built a full latest-segment ZK metadata map just to decide whether the smallest-offset fetch was needed, adding a third O(partitions) PropertyStore scan on every healthy validation cycle (on top of getPartitionGroupConsumptionStatusList and the updater's own read). The gate now derives the latest LLC segment per partition from the snapshot ideal state's segment names alone (no ZK reads) and only builds the latest-segment ZK metadata map when the smallest-offset fetch is actually required. Extract the repeated "latest LLC segment per partition" logic into getLatestLLCSegmentPerPartition(Collection) and reuse it in getLatestSegmentZKMetadataMap, getPartitionGroupConsumptionStatusList, and the gate. --- .../PinotLLCRealtimeSegmentManager.java | 83 ++++++++----------- 1 file changed, 35 insertions(+), 48 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index ed2b68e70fe4..9c04cd2299db 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -263,25 +263,10 @@ public List getPartitionGroupConsumptionStatusL List streamConfigs) { List partitionGroupConsumptionStatusList = new ArrayList<>(); - // From all segment names in the ideal state, find unique partition group ids and their latest segment - Map partitionGroupIdToLatestSegment = new HashMap<>(); - for (String segment : idealState.getRecord().getMapFields().keySet()) { - // With Pinot upsert table allowing uploads of segments, the segment name of an upsert table segment may not - // conform to LLCSegment format. We can skip such segments because they are NOT the consuming segments. - LLCSegmentName llcSegmentName = LLCSegmentName.of(segment); - if (llcSegmentName == null) { - continue; - } - int partitionGroupId = llcSegmentName.getPartitionGroupId(); - partitionGroupIdToLatestSegment.compute(partitionGroupId, (k, latestSegment) -> { - if (latestSegment == null) { - return llcSegmentName; - } else { - return latestSegment.getSequenceNumber() > llcSegmentName.getSequenceNumber() ? latestSegment - : llcSegmentName; - } - }); - } + // From all segment names in the ideal state, find unique partition group ids and their latest segment. + // Non-LLC segments (e.g. uploaded upsert segments) are skipped because they are NOT consuming segments. + Map partitionGroupIdToLatestSegment = + getLatestLLCSegmentPerPartition(idealState.getRecord().getMapFields().keySet()); // Create a {@link PartitionGroupConsumptionStatus} for each latest segment String tableNameWithType = streamConfigs.get(0).getTableNameWithType(); @@ -1386,32 +1371,31 @@ public void reduceSegmentSizeAndReset(LLCSegmentName llcSegmentName, int prevNum /// /// @param realtimeTableName Realtime table name /// @return Map from partition group id to the latest LLC realtime segment ZK metadata - private Map getLatestSegmentZKMetadataMap(String realtimeTableName) { - List segments = getLLCSegments(realtimeTableName); - - Map latestLLCSegmentNameMap = new HashMap<>(); - for (String segmentName : segments) { - LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - latestLLCSegmentNameMap.compute(llcSegmentName.getPartitionGroupId(), (partitionId, latestLLCSegmentName) -> { - if (latestLLCSegmentName == null) { - return llcSegmentName; - } else { - if (llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { - return llcSegmentName; - } else { - return latestLLCSegmentName; - } - } - }); + /// Returns the latest (highest sequence number) LLC segment per partition group, parsed from the given segment + /// names. Non-LLC segment names (e.g. uploaded upsert segments) are ignored. + private static Map getLatestLLCSegmentPerPartition(Collection segmentNames) { + Map partitionGroupIdToLatestSegment = new HashMap<>(); + for (String segmentName : segmentNames) { + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); + if (llcSegmentName == null) { + continue; + } + partitionGroupIdToLatestSegment.merge(llcSegmentName.getPartitionGroupId(), llcSegmentName, + (existing, candidate) -> candidate.getSequenceNumber() > existing.getSequenceNumber() ? candidate + : existing); } + return partitionGroupIdToLatestSegment; + } + private Map getLatestSegmentZKMetadataMap(String realtimeTableName) { + Map latestLLCSegmentNameMap = + getLatestLLCSegmentPerPartition(getLLCSegments(realtimeTableName)); Map latestSegmentZKMetadataMap = new HashMap<>(); for (Map.Entry entry : latestLLCSegmentNameMap.entrySet()) { SegmentZKMetadata latestSegmentZKMetadata = getSegmentZKMetadata(realtimeTableName, entry.getValue().getSegmentName()); latestSegmentZKMetadataMap.put(entry.getKey(), latestSegmentZKMetadata); } - return latestSegmentZKMetadataMap; } @@ -1517,10 +1501,11 @@ PreFetchedOffsets preFetchOffsets(List streamConfigs, String realt getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, snapshotIdealState); Map partitionIdToSmallestOffset = null; if (offsetCriteria == null || !offsetCriteria.equals(OffsetCriteria.SMALLEST_OFFSET_CRITERIA)) { - Map latestSegmentZKMetadataMap = getLatestSegmentZKMetadataMap(realtimeTableName); - if (offsetsHaveToChange || anyPartitionNeedsSmallestOffset(snapshotIdealState, latestSegmentZKMetadataMap)) { - partitionIdToSmallestOffset = - fetchPartitionGroupIdToSmallestOffset(streamConfigs, snapshotIdealState, latestSegmentZKMetadataMap); + // Decide whether the smallest-offset stream fetch is needed from the snapshot ideal state alone (no ZK + // metadata reads); only build the latest-segment ZK metadata map when the fetch is actually required. + if (offsetsHaveToChange || anyPartitionNeedsSmallestOffset(snapshotIdealState)) { + partitionIdToSmallestOffset = fetchPartitionGroupIdToSmallestOffset(streamConfigs, snapshotIdealState, + getLatestSegmentZKMetadataMap(realtimeTableName)); } } return new PreFetchedOffsets(streamMetadataList, partitionIdToSmallestOffset); @@ -1529,14 +1514,16 @@ PreFetchedOffsets preFetchOffsets(List streamConfigs, String realt } } - /// Returns true if at least one partition's latest segment is present in the ideal state but has no replica in the - /// CONSUMING state, i.e. it may need a new CONSUMING segment created (the only repair path that consults the - /// smallest stream offset). Mirrors the trigger in [#ensureAllPartitionsConsuming]. - private boolean anyPartitionNeedsSmallestOffset(IdealState idealState, - Map latestSegmentZKMetadataMap) { + /// Returns true if at least one partition's latest LLC segment in the ideal state has no replica in the CONSUMING + /// state, i.e. it may need a new CONSUMING segment created (the only repair path that consults the smallest stream + /// offset). Derived from the snapshot ideal state alone (no ZK metadata reads), picking the latest segment per + /// partition the same way [#getPartitionGroupConsumptionStatusList] does. + private boolean anyPartitionNeedsSmallestOffset(IdealState idealState) { Map> instanceStatesMap = idealState.getRecord().getMapFields(); - for (SegmentZKMetadata latestSegmentZKMetadata : latestSegmentZKMetadataMap.values()) { - Map instanceStateMap = instanceStatesMap.get(latestSegmentZKMetadata.getSegmentName()); + Map partitionGroupIdToLatestSegment = + getLatestLLCSegmentPerPartition(instanceStatesMap.keySet()); + for (LLCSegmentName latestSegment : partitionGroupIdToLatestSegment.values()) { + Map instanceStateMap = instanceStatesMap.get(latestSegment.getSegmentName()); if (instanceStateMap != null && !instanceStateMap.containsValue(SegmentStateModel.CONSUMING)) { return true; } From d53b1f4ff274a4d92071448dff0cbf4d0eb8d75f Mon Sep 17 00:00:00 2001 From: Shounak kulkarni Date: Mon, 17 Aug 2026 13:28:11 +0530 Subject: [PATCH 3/3] Test offset pre-fetch happens once, before and outside the ideal-state update The existing test helper called preFetchOffsets and the package-private repair sequentially, so no test covered the public ensureAllPartitionsConsuming path's core guarantee. Add a test that drives the public method with a controlled static HelixHelper: it supplies the snapshot ideal state and applies the updater twice to simulate a ZK CAS conflict, then asserts the stream offset fetch happened before HelixHelper.updateIdealState (outside the updater) and was not repeated on retry. --- .../PinotLLCRealtimeSegmentManagerTest.java | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java index 8b759c73720f..1d7d3adcd0e0 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java @@ -40,6 +40,7 @@ import java.util.TreeSet; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; import javax.annotation.Nullable; @@ -66,6 +67,7 @@ import org.apache.pinot.common.utils.FileUploadDownloadClient; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.URIUtils; +import org.apache.pinot.common.utils.helix.HelixHelper; import org.apache.pinot.controller.ControllerConf; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignment; @@ -1186,6 +1188,51 @@ List getNewStreamMetadataList(List streamConfigs, assertEquals(segmentManager._streamConfigs.get(0).getOffsetCriteria(), originalOffsetCriteria); } + /// Phase-2 core regression, exercised through the public ensureAllPartitionsConsuming path (not the package-private + /// helper): the stream offset fetches must happen BEFORE HelixHelper.updateIdealState (outside the updater), and a + /// CAS retry that re-applies the updater must NOT repeat them. A controlled static HelixHelper supplies the + /// snapshot ideal state and drives the updater twice (simulating a ZK version conflict). + @Test + public void testEnsureAllPartitionsConsumingFetchesOffsetsOnceOutsideIdealStateUpdate() { + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(); + setUpNewTable(segmentManager, 2, 5, 4); + // Make partition 0 need a new CONSUMING segment so the smallest-offset fetch is also triggered. + turnNewConsumingSegmentOffline(segmentManager._idealState.getRecord().getMapFields(), + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, CURRENT_TIME_MS).getSegmentName()); + segmentManager._exceededMaxSegmentCompletionTime = true; + + // Fetch counts sampled inside the (mocked) updateIdealState: when the updater is first entered, and again after + // it has been applied twice. + int[] fetchCountWhenUpdaterEntered = {-1}; + int[] fetchCountAfterUpdaterRetries = {-1}; + + try (MockedStatic helixHelperMock = mockStatic(HelixHelper.class)) { + helixHelperMock.when(() -> HelixHelper.getTableIdealState(any(), eq(REALTIME_TABLE_NAME))) + .thenReturn(segmentManager._idealState); + helixHelperMock.when( + () -> HelixHelper.updateIdealState(any(), eq(REALTIME_TABLE_NAME), any(), any(), anyBoolean())) + .thenAnswer(invocation -> { + // Everything fetched so far happened before the updater ran, i.e. outside the ideal-state lock. + fetchCountWhenUpdaterEntered[0] = segmentManager._getNewStreamMetadataListCallCount; + Function updater = invocation.getArgument(2); + // Simulate a ZK CAS conflict by applying the updater more than once. + IdealState result = updater.apply(segmentManager._idealState); + updater.apply(segmentManager._idealState); + fetchCountAfterUpdaterRetries[0] = segmentManager._getNewStreamMetadataListCallCount; + return result; + }); + + segmentManager.ensureAllPartitionsConsuming(segmentManager._tableConfig, segmentManager._streamConfigs, null); + } + + // Stream offsets were fetched before HelixHelper.updateIdealState was invoked (outside the updater). + assertTrue(fetchCountWhenUpdaterEntered[0] > 0, + "Expected stream offsets to be fetched before updateIdealState"); + // Re-applying the updater (CAS retry) did not trigger any additional stream fetch. + assertEquals(fetchCountAfterUpdaterRetries[0], fetchCountWhenUpdaterEntered[0], + "Updater retries must not repeat the stream offset fetch"); + } + /// Removes the new CONSUMING segment and sets the latest committed (ONLINE) segment to CONSUMING if exists in the /// ideal state. private void removeNewConsumingSegment(Map> instanceStatesMap, String consumingSegment, @@ -2546,6 +2593,9 @@ private static class FakePinotLLCRealtimeSegmentManager extends PinotLLCRealtime int _numPartitions; List _streamMetadataList = null; boolean _exceededMaxSegmentCompletionTime = false; + // Counts every stream-metadata (offset) fetch; both getNewStreamMetadataList overloads funnel through the 3-arg + // one below, so this captures the streamMetadataList fetch and the fetchPartitionGroupIdToSmallestOffset fetch. + int _getNewStreamMetadataListCallCount = 0; FileUploadDownloadClient _mockedFileUploadDownloadClient; PinotHelixResourceManager _mockResourceManager; @@ -2716,6 +2766,7 @@ Set getPartitionIds(StreamConfig streamConfig) { @Override List getNewStreamMetadataList(List streamConfigs, List currentPartitionGroupConsumptionStatusList, IdealState idealState) { + _getNewStreamMetadataListCallCount++; if (_streamMetadataList != null) { return _streamMetadataList; } else {