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


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +127,197 @@ protected void processTable(String tableNameWithType) {
 
     // Delete segments based on segment lineage and clean up segment lineage 
metadata.
     manageSegmentLineageCleanupForTable(tableConfig);
+
+    // Re-read active segments after time and lineage cleanup so 
already-deleted bytes do not cause extra eviction.
+    manageSizeBasedRetention(tableConfig);
+  }
+
+  /// Enforces the optional compressed-segment size limit on live data 
independently of time retention. Replaced
+  /// segments are excluded from accounting. Eviction proceeds oldest first 
and stops before the first active segment
+  /// listed on either side of any lineage entry, including replacement copies 
excluded from accounting. The newest
+  /// OFFLINE segment and realtime recovery segments are preserved. Each typed 
table uses its own size limit and
+  /// eviction order without consulting the OFFLINE time boundary.
+  /// Unknown sizes or missing metadata disable this pass; consuming segments 
are excluded from accounting and eviction.
+  @VisibleForTesting
+  protected void manageSizeBasedRetention(TableConfig tableConfig) {
+    String tableNameWithType = tableConfig.getTableName();
+    String retentionSize = 
tableConfig.getValidationConfig().getRetentionSize();
+    if (retentionSize == null || (tableConfig.getTableType() == 
TableType.OFFLINE
+        && 
!"APPEND".equalsIgnoreCase(IngestionConfigUtils.getBatchSegmentIngestionType(tableConfig))))
 {
+      _controllerMetrics.removeTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED);
+      return;
+    }
+    boolean blocked = true;
+    try {
+      blocked = !enforceSizeBasedRetention(tableConfig, retentionSize);
+    } finally {
+      // Publish only the completed pass's result, including a blocked result 
if an exception interrupted it.
+      _controllerMetrics.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED,
+          blocked ? 1 : 0);
+    }
+  }
+
+  /// Returns whether the table is within its size cap after this pass; unsafe 
skips and failed evictions return false.
+  private boolean enforceSizeBasedRetention(TableConfig tableConfig, String 
retentionSize) {
+    String tableNameWithType = tableConfig.getTableName();
+    long retentionSizeBytes;
+    try {
+      retentionSizeBytes = DataSizeUtils.toBytes(retentionSize);
+      Preconditions.checkArgument(retentionSizeBytes > 0, "Retention size must 
be positive");
+    } catch (Exception e) {
+      LOGGER.warn("Invalid retention size: {} for table: {}, skip size 
retention", retentionSize, tableNameWithType);
+      return false;
+    }
+
+    // Read IdealState before lineage, matching the existing live-segment 
filtering order. A replacement starting
+    // after this snapshot cannot add both its source and destination to the 
bytes counted by this pass.
+    Set<String> activeSegments = new 
HashSet<>(_pinotHelixResourceManager.getSegmentsFor(tableNameWithType, false));
+    SegmentLineage segmentLineage =
+        
SegmentLineageAccessHelper.getSegmentLineage(_pinotHelixResourceManager.getPropertyStore(),
 tableNameWithType);
+    Map<String, LineageEntry> lineageEntries = segmentLineage == null ? 
Map.of() : segmentLineage.getLineageEntries();
+    Set<String> lineageOwnedSegments = new HashSet<>();
+    // Size retention protects both sides of every retained entry, regardless 
of state, age, or the deletion flag.
+    lineageEntries.values().forEach(entry -> {
+      lineageOwnedSegments.addAll(entry.getSegmentsFrom());
+      lineageOwnedSegments.addAll(entry.getSegmentsTo());
+    });
+    lineageOwnedSegments.retainAll(activeSegments);
+    // Replaced segments awaiting lineage cleanup are no longer queryable and 
must not trigger live-data eviction.
+    SegmentLineageUtils.filterSegmentsBasedOnLineageInPlace(activeSegments, 
segmentLineage);
+    if (activeSegments.isEmpty()) {
+      return true;
+    }
+    List<SegmentZKMetadata> metadataList = 
_pinotHelixResourceManager.getSegmentsZKMetadata(tableNameWithType);
+    Map<String, SegmentZKMetadata> lineageOwnedMetadata = new HashMap<>();
+    Set<String> segmentsWithMetadata = new HashSet<>();
+    List<SegmentZKMetadata> completedSegments = new ArrayList<>();
+    Map<String, Long> completedSegmentSizes = new HashMap<>();
+    boolean realtime = tableConfig.getTableType() == TableType.REALTIME;
+    long retainedBytes = 0L;
+    for (SegmentZKMetadata metadata : metadataList) {
+      String segmentName = metadata.getSegmentName();
+      if (lineageOwnedSegments.contains(segmentName)) {
+        lineageOwnedMetadata.put(segmentName, metadata);
+      }
+      if (!activeSegments.contains(segmentName)) {
+        continue;
+      }
+      segmentsWithMetadata.add(segmentName);
+      if (realtime && !metadata.getStatus().isCompleted()) {
+        continue;
+      }
+      long size = getSizeRetentionSegmentSize(tableNameWithType, metadata);
+      if (size < 0) {
+        LOGGER.warn("Unknown size for segment: {} in table: {}, skip size 
retention", segmentName, tableNameWithType);
+        return false;
+      }
+      // Plain addition assumes realistic table sizes; corrupt totals above 
Long.MAX_VALUE may wrap and report healthy.
+      retainedBytes += size;
+      completedSegments.add(metadata);
+      completedSegmentSizes.put(segmentName, size);
+    }
+    if (!segmentsWithMetadata.containsAll(activeSegments)) {
+      LOGGER.warn("Missing active segment metadata for table: {}, skip size 
retention", tableNameWithType);
+      return false;
+    }
+    if (retainedBytes <= retentionSizeBytes) {
+      return true;
+    }
+
+    Comparator<SegmentZKMetadata> evictionOrder = 
Comparator.comparingLong(RetentionManager::getSizeRetentionTimestamp)
+        .thenComparing(SegmentZKMetadata::getSegmentName);
+    // Shadow copies do not count toward the cap, but still stop eviction at 
their position in the ordering.
+    // Without metadata or a timestamp for any active lineage-owned segment, 
its position cannot be determined safely.
+    if (lineageOwnedMetadata.size() != lineageOwnedSegments.size() || 
lineageOwnedMetadata.values().stream()
+        .anyMatch(metadata -> getSizeRetentionTimestamp(metadata) < 0)) {
+      LOGGER.warn("Cannot order lineage-owned segments for table: {}, skip 
size retention", tableNameWithType);
+      return false;
+    }
+    SegmentZKMetadata lineageBarrier = 
lineageOwnedMetadata.values().stream().min(evictionOrder).orElse(null);
+    completedSegments.removeIf(metadata -> getSizeRetentionTimestamp(metadata) 
< 0);
+    if (realtime) {
+      Set<String> lastCompletedSegments =
+          new 
HashSet<>(_pinotHelixResourceManager.getLastLLCCompletedSegments(metadataList));
+      completedSegments.removeIf(metadata -> 
lastCompletedSegments.contains(metadata.getSegmentName()));
+    } else {
+      // A size typo must not remove the entire OFFLINE table, even if its 
newest segment exceeds the cap alone.
+      
completedSegments.stream().max(evictionOrder).ifPresent(completedSegments::remove);

Review Comment:
   Fixed in b53e33f3cf43e7c45663a71bbd0dc132c88957fc. The selected newest 
metadata object is now removed from the eviction candidates by identity using 
`removeIf(metadata -> metadata == newest)`, avoiding `SegmentZKMetadata.equals` 
and its map/JSON work. Timestamp/name selection and the newest-segment floor 
stay the same. All 78 targeted tests passed, including the newest-only and 
timestamp/name fallback cases, along with scoped formatting, Checkstyle, 
license checks, and the added-line compiler-warning check.



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