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


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/tablestate/TableStateUtils.java:
##########
@@ -82,30 +91,162 @@ 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. Resolve 
partitions from names or segment metadata,
+  /// matching upsert and dedup loading; unresolved segments 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<>();
+    Function<String, Integer> partitionIdFromMetadata =
+        segmentName -> getPartitionIdFromMetadata(helixManager, 
tableNameWithType, segmentName);
+    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, partitionIdFromMetadata) && 
!isLocallyReady.test(segmentName)) {
+        requiredSegments.add(segmentName);
+      }
+    }
+    return areSegmentsLoaded(helixManager, tableNameWithType, 
requiredSegments);
+  }
+
+  /// Reads the table assignment, this live instance's current state, and 
partition metadata for custom-named ONLINE
+  /// segments once for evaluation by multiple consumers. Subsequent Helix 
updates cannot mutate the 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(), Map.of());
+    }
+    String instanceName = helixManager.getInstanceName();
+    Map<String, String> idealStates = new HashMap<>();
+    Map<String, Integer> metadataPartitionIds = new HashMap<>();
+    for (Map.Entry<String, Map<String, String>> entry : 
idealState.getRecord().getMapFields().entrySet()) {
+      String segmentName = entry.getKey();
+      String state = entry.getValue().get(instanceName);
+      if (state != null) {
+        idealStates.put(segmentName, state);
+        if (SegmentStateModel.ONLINE.equals(state) && 
SegmentUtils.getPartitionIdFromSegmentName(segmentName) == null) {
+          Integer partitionId = getPartitionIdFromMetadata(helixManager, 
tableNameWithType, segmentName);

Review Comment:
   Addressed in c624a6630e. Snapshot creation now skips partition metadata when 
a custom segment is already ONLINE in CurrentState or locally initialized. 
Fresh confirmation uses the same snapshot factory and evaluator, so both paths 
apply this rule and the separate readiness implementation is removed.
   
   The snapshot retains the assignment but never retains local readiness as a 
cached result. Evaluation rechecks the current manager; if initialization 
disappeared after a lookup was skipped, the absent partition ID remains unknown 
and conservatively blocks until refresh. Cached evaluation performs no 
ZooKeeper reads.
   
   Added regressions asserting zero segment-metadata lookups for loaded/locally 
initialized custom history in both initial and cached-positive checks, 
including manager replacement. Formatting, checkstyle, license checks, and 
production/test compilation passed. No local tests were run; CI will execute 
the regressions. Fresh checks now consistently fetch assignment/current state 
even when there is no required history; runtime performance remains 
unbenchmarked.



-- 
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