Fix time segment pruning for remote clusters in multi-cluster routing

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentmetadata/SegmentZkMetadataFetcher.java

@@ -21,13 +21,15 @@
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import org.apache.helix.AccessOption;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.pinot.common.metadata.ZKMetadataProvider;
+import org.apache.pinot.spi.utils.CommonConstants;
 
 
 /**
@@ -82,7 +84,7 @@ public void init(IdealState idealState, ExternalView externalView, Set<String> o
           listener.init(idealState, externalView, segments, znRecords);
         }
         for (int i = 0; i < numSegments; i++) {
-          if (znRecords.get(i) != null) {
+          if (!isConsumingInExternalView(externalView, segments.get(i))) {
             _onlineSegmentsCached.add(segments.get(i));
           }
         }
@@ -98,17 +100,24 @@ public synchronized void onAssignmentChange(IdealState idealState, ExternalView
       List<String> segments = new ArrayList<>();
       List<String> segmentZKMetadataPaths = new ArrayList<>();
       for (String segment : onlineSegments) {
-        if (!_onlineSegmentsCached.contains(segment)) {
-          segments.add(segment);
-          segmentZKMetadataPaths.add(_segmentZKMetadataPathPrefix + segment);
+        if (_onlineSegmentsCached.contains(segment)) {
+          continue;
+        }
+        // Skip segments still in CONSUMING state — they'll be re-evaluated on the next EV change
+        // when they transition to ONLINE (i.e., when they commit).
+        if (isConsumingInExternalView(externalView, segment)) {
+          continue;
         }
+        segments.add(segment);
+        segmentZKMetadataPaths.add(_segmentZKMetadataPathPrefix + segment);
       }
       List<ZNRecord> znRecords = _propertyStore.get(segmentZKMetadataPaths, null, AccessOption.PERSISTENT, false);
       for (SegmentZkMetadataFetchListener listener : _listeners) {
         listener.onAssignmentChange(idealState, externalView, onlineSegments, segments, znRecords);
       }
       int numSegments = segments.size();
       for (int i = 0; i < numSegments; i++) {
+        // All fetched segments are non-consuming (guaranteed by the EV check above), cache them if metadata exists.
         if (znRecords.get(i) != null) {
           _onlineSegmentsCached.add(segments.get(i));
         }
@@ -130,4 +139,17 @@ public synchronized void refreshSegment(String segment) {
       }
     }
   }
+
+  /**
+   * Returns true if the segment is in CONSUMING state on any server in the ExternalView.
+   * Such segments should not be cached in {@code _onlineSegmentsCached} — they will be re-evaluated
+   * on the next ExternalView change, at which point they will have transitioned to ONLINE (committed).
+   */
+  private static boolean isConsumingInExternalView(ExternalView externalView, String segment) {
+    if (externalView == null) {
+      return false;
+    }
+    Map<String, String> stateMap = externalView.getStateMap(segment);
+    return stateMap != null && stateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING);
+  }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java

@@ -107,12 +107,16 @@ private Interval extractIntervalFromSegmentZKMetaZNRecord(String segment, @Nulla
   @Override
   public synchronized void onAssignmentChange(IdealState idealState, ExternalView externalView,
       Set<String> onlineSegments, List<String> pulledSegments, List<ZNRecord> znRecords) {
-    // NOTE: We don't update all the segment ZK metadata for every external view change, but only the new added/removed
-    //       ones. The refreshed segment ZK metadata change won't be picked up.
     for (int idx = 0; idx < pulledSegments.size(); idx++) {
       String segment = pulledSegments.get(idx);
       ZNRecord zNrecord = znRecords.get(idx);
-      _intervalMap.computeIfAbsent(segment, k -> extractIntervalFromSegmentZKMetaZNRecord(k, zNrecord));
+      // Always update segments that have DEFAULT_INTERVAL, which covers two cases:
+      // 1. New segments not yet in the map
+      // 2. Segments that transitioned from CONSUMING (DEFAULT_INTERVAL) to COMMITTED (valid time range)
+      Interval existing = _intervalMap.get(segment);
+      if (existing == null || existing == DEFAULT_INTERVAL) {
+        _intervalMap.put(segment, extractIntervalFromSegmentZKMetaZNRecord(segment, zNrecord));
+      }
     }
     _intervalMap.keySet().retainAll(onlineSegments);
     _intervalTree = new IntervalTree<>(_intervalMap);