Jackie-Jiang commented on code in PR #20106:
URL: https://github.com/apache/pinot/pull/20106#discussion_r4558475553


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/tablestate/TableStateUtils.java:
##########
@@ -82,30 +87,130 @@ public static List<String> 
getSegmentsInGivenStateForThisInstance(HelixManager h
   public static boolean isAllSegmentsLoaded(HelixManager helixManager, String 
tableNameWithType) {
     List<String> onlineSegments =
         getSegmentsInGivenStateForThisInstance(helixManager, 
tableNameWithType, SegmentStateModel.ONLINE);
-    if (onlineSegments.isEmpty()) {
-      LOGGER.info("No ONLINE segment found for table: {}", tableNameWithType);
-      return true;
-    }
+    return areSegmentsLoaded(helixManager, tableNameWithType, onlineSegments);
+  }
 
-    // Check if ideal state and current state matches for all segments 
assigned to the current instance
+  /// Checks the current assignment before a consumer starts. Earlier ONLINE 
segments of the same partition must be
+  /// loaded, as must uploaded segments from that partition. Segments whose 
partition cannot be inferred from their
+  /// name are conservatively included. The consumer and its successors are 
excluded to avoid waiting on themselves.
+  ///
+  /// @param isLocallyReady identifies segments whose metadata is already 
ready locally despite a pending ONLINE
+  ///                      transition, such as a predecessor that finished 
consumption and released its permit
+  public static boolean isReadyToConsume(HelixManager helixManager, String 
tableNameWithType, LLCSegmentName consumer,
+      Predicate<String> isLocallyReady) {
     HelixDataAccessor dataAccessor = helixManager.getHelixDataAccessor();
-    PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();
-    String instanceName = helixManager.getInstanceName();
-    LiveInstance liveInstance = 
dataAccessor.getProperty(keyBuilder.liveInstance(instanceName));
-    if (liveInstance == null) {
-      LOGGER.warn("Failed to find live instance for instance: {}", 
instanceName);
+    IdealState idealState = 
dataAccessor.getProperty(dataAccessor.keyBuilder().idealStates(tableNameWithType));
+    if (idealState == null) {
+      LOGGER.warn("Failed to find ideal state for table: {}", 
tableNameWithType);
       return false;
     }
-    String sessionId = liveInstance.getEphemeralOwner();
-    CurrentState currentState =
-        dataAccessor.getProperty(keyBuilder.currentState(instanceName, 
sessionId, tableNameWithType));
-    if (currentState == null) {
-      LOGGER.warn("Failed to find current state for instance: {}, sessionId: 
{}, table: {}", instanceName, sessionId,
-          tableNameWithType);
+    String instanceName = helixManager.getInstanceName();
+    Map<String, String> consumerAssignment = 
idealState.getInstanceStateMap(consumer.getSegmentName());
+    String consumerState = consumerAssignment != null ? 
consumerAssignment.get(instanceName) : null;
+    if (!SegmentStateModel.CONSUMING.equals(consumerState) && 
!SegmentStateModel.ONLINE.equals(consumerState)) {
       return false;
     }
+
+    List<String> requiredSegments = new ArrayList<>();
+    for (Map.Entry<String, Map<String, String>> entry : 
idealState.getRecord().getMapFields().entrySet()) {
+      if 
(!SegmentStateModel.ONLINE.equals(entry.getValue().get(instanceName))) {
+        continue;
+      }
+      String segmentName = entry.getKey();
+      if (isRequiredHistory(segmentName, consumer) && 
!isLocallyReady.test(segmentName)) {
+        requiredSegments.add(segmentName);
+      }
+    }
+    return areSegmentsLoaded(helixManager, tableNameWithType, 
requiredSegments);
+  }
+
+  /// Reads the table assignment and this live instance's current state once 
for evaluation by multiple consumers.
+  /// The records are copied so subsequent Helix updates cannot mutate the 
returned snapshot. A cached snapshot must
+  /// only reject readiness; callers must perform a fresh check before 
authorizing consumption from a cached result.
+  public static ReadinessSnapshot getReadinessSnapshot(HelixManager 
helixManager, String tableNameWithType) {
+    HelixDataAccessor dataAccessor = helixManager.getHelixDataAccessor();
+    IdealState idealState = 
dataAccessor.getProperty(dataAccessor.keyBuilder().idealStates(tableNameWithType));
+    if (idealState == null) {
+      LOGGER.warn("Failed to find ideal state for table: {}", 
tableNameWithType);
+      return new ReadinessSnapshot(Map.of(), Map.of());
+    }
+    String instanceName = helixManager.getInstanceName();
+    Map<String, String> idealStates = new HashMap<>();
+    for (Map.Entry<String, Map<String, String>> entry : 
idealState.getRecord().getMapFields().entrySet()) {
+      String state = entry.getValue().get(instanceName);
+      if (state != null) {
+        idealStates.put(entry.getKey(), state);
+      }
+    }
+    return new ReadinessSnapshot(idealStates, getCurrentStates(helixManager, 
tableNameWithType));
+  }
+
+  /// An immutable view of this instance's assignment and current state, safe 
to share between consumer threads.
+  /// Evaluation performs no ZooKeeper reads and reevaluates the supplied 
local-readiness predicate for each required
+  /// segment. The predicate must itself be safe for its caller's threads; its 
results are never retained here.
+  public static final class ReadinessSnapshot {
+    private final Map<String, String> _idealStates;
+    private final Map<String, String> _currentStates;
+
+    private ReadinessSnapshot(Map<String, String> idealStates, Map<String, 
String> currentStates) {
+      _idealStates = Collections.unmodifiableMap(new HashMap<>(idealStates));
+      _currentStates = Collections.unmodifiableMap(new 
HashMap<>(currentStates));
+    }
+
+    public boolean isReadyToConsume(LLCSegmentName consumer, Predicate<String> 
isLocallyReady) {
+      return isReadyToConsume(consumer, isLocallyReady, false);
+    }
+
+    /// Optionally logs the first unavailable history segment. Enable logging 
for a freshly fetched snapshot, keeping
+    /// repeated rejection checks against a cached snapshot quiet.
+    public boolean isReadyToConsume(LLCSegmentName consumer, Predicate<String> 
isLocallyReady,
+        boolean logUnavailableSegments) {
+      String consumerState = _idealStates.get(consumer.getSegmentName());
+      if (!SegmentStateModel.CONSUMING.equals(consumerState) && 
!SegmentStateModel.ONLINE.equals(consumerState)) {
+        return false;
+      }
+      for (Map.Entry<String, String> entry : _idealStates.entrySet()) {
+        String segmentName = entry.getKey();
+        if (SegmentStateModel.ONLINE.equals(entry.getValue()) && 
isRequiredHistory(segmentName, consumer)
+            && !isLocallyReady.test(segmentName)) {
+          String actualState = _currentStates.get(segmentName);
+          if (!SegmentStateModel.ONLINE.equals(actualState)) {
+            if (logUnavailableSegments) {
+              if (SegmentStateModel.ERROR.equals(actualState)) {
+                LOGGER.error("Consumer: {} is waiting for history segment: {} 
in ERROR state, expected: {}",
+                    consumer.getSegmentName(), segmentName, 
SegmentStateModel.ONLINE);
+              } else {
+                LOGGER.info("Consumer: {} is waiting for history segment: {} 
in state: {}, expected: {}",
+                    consumer.getSegmentName(), segmentName, actualState, 
SegmentStateModel.ONLINE);
+              }
+            }
+            return false;
+          }
+        }
+      }
+      return true;
+    }
+  }
+
+  private static boolean isRequiredHistory(String segmentName, LLCSegmentName 
consumer) {
+    LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
+    if (llcSegmentName != null) {
+      return llcSegmentName.getPartitionGroupId() == 
consumer.getPartitionGroupId()
+          && llcSegmentName.getSequenceNumber() < consumer.getSequenceNumber();
+    }
+    Integer partitionId = 
SegmentUtils.getPartitionIdFromSegmentName(segmentName);
+    return partitionId == null || partitionId == 
consumer.getPartitionGroupId();

Review Comment:
   Addressed in c393bb6d55. Both snapshot creation and fresh confirmation now 
fall back to segment metadata through 
`SegmentUtils.getSegmentPartitionId(metadata, null)`, matching upsert/dedup 
loading. Name-derived IDs retain precedence; missing or ambiguous metadata 
remains conservatively required.
   
   The snapshot stores immutable resolved IDs and performs no ZooKeeper reads 
during evaluation. Fresh confirmation re-reads metadata, so a cached partition 
ID cannot hide newly required history after an upload changes the metadata. 
Missing records remain unknown, and ZooKeeper failures follow the existing 
retry path.
   
   Added coverage for custom segments in the same/other partition (including 
ERROR), missing/ambiguous metadata, snapshot immutability, fresh confirmation 
after metadata changes, and lookup-timeout recovery. Formatting, checkstyle, 
license checks, and production/test compilation passed. No tests were run 
locally; the new regressions await CI.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to