xiangfu0 commented on code in PR #19995:
URL: https://github.com/apache/pinot/pull/19995#discussion_r4343871352


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java:
##########
@@ -2278,6 +2283,80 @@ private ZkClient getOrBuildMultiWriteZkClient() {
     }
   }
 
+  /// Removes the ZK metadata of the given segments that are IN_PROGRESS and 
not in the IdealState of the table, i.e.
+  /// new consuming segments that were not added to the IdealState, and 
returns the segments whose ZK metadata was
+  /// removed.
+  ///
+  /// The removal is submitted in one ZK `multi()` transaction with 
version-checked deletes of the segment ZK metadata
+  /// read here, and a version-checked write of the unchanged IdealState:
+  /// - An IdealState update that read the IdealState before the transaction 
fails its version check. An updater that
+  ///   reads the segment ZK metadata again when it retries, like the one of 
RealtimeSegmentValidationManager, then no
+  ///   longer finds the removed segments.
+  /// - If the IdealState, or the ZK metadata of one of the segments, is 
updated after it was read here, the transaction
+  ///   fails and is retried.
+  ///
+  /// The data versions cannot tell a segment ZK metadata deleted and 
recreated with the same name since it was read,
+  /// e.g. by an upload under the same name, from the one that was read. As 
the IdealState is outside the property
+  /// store, the transaction uses the multi-write ZkClient directly instead of 
[ZkMultiWriteBuilder].
+  public Set<String> removeInProgressSegmentsZKMetadataNotInIdealState(String 
tableNameWithType,
+      Collection<String> segmentNames) {
+    return 
removeInProgressSegmentsZKMetadataNotInIdealState(tableNameWithType, 
segmentNames, () -> {
+    });
+  }
+
+  /// @param afterRead Runs after reading the IdealState and the segment ZK 
metadata, before removing it
+  @VisibleForTesting
+  Set<String> removeInProgressSegmentsZKMetadataNotInIdealState(String 
tableNameWithType,
+      Collection<String> segmentNames, Runnable afterRead) {
+    ZkClient zkClient = getOrBuildMultiWriteZkClient();
+    String idealStatePath = PropertyPathBuilder.idealState(_helixClusterName, 
tableNameWithType);
+    String propertyStoreRoot = 
PropertyPathBuilder.propertyStore(_helixClusterName);
+    for (int attempt = 1; ; attempt++) {
+      Stat idealStateStat = new Stat();
+      ZNRecord idealState = zkClient.readData(idealStatePath, idealStateStat);

Review Comment:
   Minor/robustness: `zkClient.readData(path, stat)` returns `null` if the 
IdealState node is missing, which would NPE on `idealState.getMapFields()` 
below. Not reachable for a live realtime table mid-commit, so not blocking — 
but a null check that returns an empty set would make the helper safe to call 
defensively.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java:
##########
@@ -2278,6 +2283,80 @@ private ZkClient getOrBuildMultiWriteZkClient() {
     }
   }
 
+  /// Removes the ZK metadata of the given segments that are IN_PROGRESS and 
not in the IdealState of the table, i.e.
+  /// new consuming segments that were not added to the IdealState, and 
returns the segments whose ZK metadata was
+  /// removed.
+  ///
+  /// The removal is submitted in one ZK `multi()` transaction with 
version-checked deletes of the segment ZK metadata
+  /// read here, and a version-checked write of the unchanged IdealState:
+  /// - An IdealState update that read the IdealState before the transaction 
fails its version check. An updater that
+  ///   reads the segment ZK metadata again when it retries, like the one of 
RealtimeSegmentValidationManager, then no
+  ///   longer finds the removed segments.
+  /// - If the IdealState, or the ZK metadata of one of the segments, is 
updated after it was read here, the transaction
+  ///   fails and is retried.
+  ///
+  /// The data versions cannot tell a segment ZK metadata deleted and 
recreated with the same name since it was read,
+  /// e.g. by an upload under the same name, from the one that was read. As 
the IdealState is outside the property
+  /// store, the transaction uses the multi-write ZkClient directly instead of 
[ZkMultiWriteBuilder].
+  public Set<String> removeInProgressSegmentsZKMetadataNotInIdealState(String 
tableNameWithType,
+      Collection<String> segmentNames) {
+    return 
removeInProgressSegmentsZKMetadataNotInIdealState(tableNameWithType, 
segmentNames, () -> {
+    });
+  }
+
+  /// @param afterRead Runs after reading the IdealState and the segment ZK 
metadata, before removing it
+  @VisibleForTesting
+  Set<String> removeInProgressSegmentsZKMetadataNotInIdealState(String 
tableNameWithType,
+      Collection<String> segmentNames, Runnable afterRead) {
+    ZkClient zkClient = getOrBuildMultiWriteZkClient();
+    String idealStatePath = PropertyPathBuilder.idealState(_helixClusterName, 
tableNameWithType);
+    String propertyStoreRoot = 
PropertyPathBuilder.propertyStore(_helixClusterName);
+    for (int attempt = 1; ; attempt++) {
+      Stat idealStateStat = new Stat();
+      ZNRecord idealState = zkClient.readData(idealStatePath, idealStateStat);
+      List<Op> segmentDeletes = new ArrayList<>();
+      Set<String> segmentsToRemove = new HashSet<>();
+      for (String segmentName : segmentNames) {
+        if (idealState.getMapFields().containsKey(segmentName)) {
+          continue;
+        }
+        String segmentPath =
+            propertyStoreRoot + 
ZKMetadataProvider.constructPropertyStorePathForSegment(tableNameWithType, 
segmentName);
+        Stat segmentStat = new Stat();
+        ZNRecord segmentRecord;
+        try {
+          segmentRecord = zkClient.readData(segmentPath, segmentStat);
+        } catch (ZkNoNodeException e) {
+          continue;
+        }
+        if (new SegmentZKMetadata(segmentRecord).getStatus() == 
CommonConstants.Segment.Realtime.Status.IN_PROGRESS) {
+          segmentDeletes.add(Op.delete(segmentPath, segmentStat.getVersion()));
+          segmentsToRemove.add(segmentName);
+        }
+      }
+      if (segmentsToRemove.isEmpty()) {
+        return segmentsToRemove;
+      }
+      afterRead.run();
+      List<Op> ops = new ArrayList<>(segmentDeletes.size() + 1);
+      ops.add(Op.setData(idealStatePath, zkClient.serialize(idealState, 
idealStatePath), idealStateStat.getVersion()));
+      ops.addAll(segmentDeletes);
+      try {
+        zkClient.multi(ops);
+        return segmentsToRemove;
+      } catch (ZkException e) {

Review Comment:
   Minor: this catches `ZkException` and inspects `getCause()`, whereas the 
sibling `ZkMultiWriteBuilder.execute()` unwraps and rethrows the 
`KeeperException` subtype directly. Both work, but the two multi() call sites 
now handle the same failure differently. Worth a brief comment here (or reusing 
the builder's unwrap helper) so the next reader doesn't assume one of them is 
wrong.



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