somandal commented on code in PR #15175:
URL: https://github.com/apache/pinot/pull/15175#discussion_r1999251138


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java:
##########
@@ -138,4 +157,80 @@ private boolean checkIsMinimizeDataMovement(String 
rebalanceJobId, String tableN
       return false;
     }
   }
+
+  private String checkDiskUtilization(String tableNameWithType, Map<String, 
Map<String, String>> currentAssignment,
+      Map<String, Map<String, String>> targetAssignment,
+      TableSizeReader.TableSubTypeSizeDetails tableSubTypeSizeDetails, double 
threshold) {
+    boolean isDiskUtilSafe = true;
+    StringBuilder message = new StringBuilder("UNSAFE. Servers with unsafe 
disk util footprint: ");
+    String sep = "";
+    Map<String, Set<String>> existingServersToSegmentMap = new HashMap<>();
+    Map<String, Set<String>> newServersToSegmentMap = new HashMap<>();
+
+    for (Map.Entry<String, Map<String, String>> entrySet : 
currentAssignment.entrySet()) {
+      for (String segmentKey : entrySet.getValue().keySet()) {
+        existingServersToSegmentMap.computeIfAbsent(segmentKey, k -> new 
HashSet<>()).add(entrySet.getKey());
+      }
+    }
+
+    for (Map.Entry<String, Map<String, String>> entrySet : 
targetAssignment.entrySet()) {
+      for (String segmentKey : entrySet.getValue().keySet()) {
+        newServersToSegmentMap.computeIfAbsent(segmentKey, k -> new 
HashSet<>()).add(entrySet.getKey());
+      }
+    }
+
+    long avgSegmentSize = getAverageSegmentSize(tableSubTypeSizeDetails, 
currentAssignment);
+
+    for (Map.Entry<String, Set<String>> entry : 
newServersToSegmentMap.entrySet()) {
+      String server = entry.getKey();
+      DiskUsageInfo diskUsage = getDiskUsageInfoOfInstance(server);
+
+      if (diskUsage.getTotalSpaceBytes() < 0) {
+        return "Disk usage info not enabled. Try to set 
controller.enable.resource.utilization.check=true";
+      }
+
+      Set<String> segmentSet = entry.getValue();
+
+      Set<String> newSegmentSet = new HashSet<>(segmentSet);
+      Set<String> existingSegmentSet = new HashSet<>();
+      Set<String> intersection = new HashSet<>();
+      if (existingServersToSegmentMap.containsKey(server)) {
+        Set<String> segmentSetForServer = 
existingServersToSegmentMap.get(server);
+        existingSegmentSet.addAll(segmentSetForServer);
+        intersection.addAll(segmentSetForServer);
+        intersection.retainAll(newSegmentSet);
+      }
+      newSegmentSet.removeAll(intersection);
+      Set<String> removedSegmentSet = new HashSet<>(existingSegmentSet);
+      removedSegmentSet.removeAll(intersection);
+
+      long diskUtilizationGain = newSegmentSet.size() * avgSegmentSize;
+      long diskUtilizationLoss = removedSegmentSet.size() * avgSegmentSize;
+
+      long diskUtilizationFootprint = diskUsage.getUsedSpaceBytes() + 
diskUtilizationGain;
+      double diskUtilizationFootprintRate =
+          (double) diskUtilizationFootprint / diskUsage.getTotalSpaceBytes();
+
+      if (diskUtilizationFootprintRate >= threshold) {
+        isDiskUtilSafe = false;
+        message.append(sep)
+            .append(server)
+            .append(String.format(" (%d%%)", (short) 
(diskUtilizationFootprintRate * 100)));
+        sep = ", ";
+      }
+    }
+    return isDiskUtilSafe ? "Within threshold" : message.toString();

Review Comment:
   yes. in case someone changed the threshold config we should know what 
threshold we're checking against



-- 
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: commits-unsubscr...@pinot.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org

Reply via email to