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


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +128,218 @@ 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; hybrid 
realtime eviction requires offline coverage.
+  /// 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;
+    BigInteger retainedBytes = BigInteger.ZERO;
+    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;
+      }
+      // Avoid overflowing the table total even if individual segment sizes 
are valid longs.
+      retainedBytes = retainedBytes.add(BigInteger.valueOf(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;
+    }
+    BigInteger limit = BigInteger.valueOf(retentionSizeBytes);
+    if (retainedBytes.compareTo(limit) <= 0) {
+      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);
+    List<String> candidateNames = completedSegments.stream()
+        .filter(metadata -> getSizeRetentionTimestamp(metadata) >= 0)
+        .map(SegmentZKMetadata::getSegmentName)
+        .collect(Collectors.toCollection(ArrayList::new));
+    if (realtime) {
+      
candidateNames.removeAll(_pinotHelixResourceManager.getLastLLCCompletedSegments(metadataList));
+      if (_isHybridTableRetentionStrategyEnabled) {
+        TableConfig offlineTableConfig =
+            
_pinotHelixResourceManager.getOfflineTableConfig(TableNameBuilder.extractRawTableName(tableNameWithType));
+        if (offlineTableConfig != null) {
+          long timeBoundaryMs;
+          try {
+            timeBoundaryMs = getHybridTimeBoundaryMs(offlineTableConfig);
+          } catch (Exception e) {
+            LOGGER.warn("Cannot determine offline coverage for table: {}, skip 
size retention", tableNameWithType, e);
+            return false;
+          }
+          Set<String> coveredSegments = completedSegments.stream()
+              .filter(metadata -> metadata.getEndTimeMs() >= 0 && 
metadata.getEndTimeMs() < timeBoundaryMs)
+              
.map(SegmentZKMetadata::getSegmentName).collect(Collectors.toSet());
+          candidateNames.retainAll(coveredSegments);
+        }
+      }
+    } else {
+      // A size typo must not remove the entire OFFLINE table, even if its 
newest segment exceeds the cap alone.
+      completedSegments.stream().filter(metadata -> 
getSizeRetentionTimestamp(metadata) >= 0)
+          
.max(evictionOrder).map(SegmentZKMetadata::getSegmentName).ifPresent(candidateNames::remove);
+    }
+    Set<String> candidates = new HashSet<>(candidateNames);
+    completedSegments.removeIf(metadata -> 
!candidates.contains(metadata.getSegmentName()));
+    completedSegments.sort(evictionOrder);
+
+    List<String> segmentsToDelete = new ArrayList<>();
+    BigInteger projectedRetainedBytes = retainedBytes;
+    for (SegmentZKMetadata metadata : completedSegments) {
+      if (projectedRetainedBytes.compareTo(limit) <= 0) {
+        break;
+      }
+      // Stop at the lineage boundary even when that segment is excluded by 
another protection or from accounting.
+      // This snapshot rule applies regardless of the lineage-exclusive 
deletion flag.
+      if (lineageBarrier != null && evictionOrder.compare(metadata, 
lineageBarrier) >= 0) {
+        LOGGER.info("Stop size retention for table: {} at lineage-owned 
segment: {}", tableNameWithType,
+            lineageBarrier.getSegmentName());
+        break;
+      }
+      long size = completedSegmentSizes.get(metadata.getSegmentName());
+      if (size == 0) {
+        continue;
+      }
+      segmentsToDelete.add(metadata.getSegmentName());
+      projectedRetainedBytes = 
projectedRetainedBytes.subtract(BigInteger.valueOf(size));
+    }
+    if (!segmentsToDelete.isEmpty()) {
+      PinotResourceManagerResponse response;
+      // Coordinate with local lineage writers and reject stale accounting or 
a changed stopping point, even when
+      // lineage-exclusive deletion is disabled. Other controllers can still 
update lineage after this fresh read.
+      synchronized 
(_pinotHelixResourceManager.getLineageUpdaterLock(tableNameWithType)) {
+        SegmentLineage currentLineage =
+            
SegmentLineageAccessHelper.getSegmentLineage(_pinotHelixResourceManager.getPropertyStore(),
+                tableNameWithType);
+        Map<String, LineageEntry> currentEntries =
+            currentLineage == null ? Map.of() : 
currentLineage.getLineageEntries();
+        if (!lineageEntries.equals(currentEntries)) {

Review Comment:
   This check is correct and I verified it rests on `LineageEntry`'s value 
equality (segments, state, timestamp). One property worth being aware of: 
full-map equality aborts the batch on *any* lineage change, including ones that 
can't affect the plan (e.g. a new entry whose members are all newer than the 
barrier, or an entry-timestamp bump). That's conservative and self-heals next 
cycle, so I wouldn't change it — just noting it's deliberate over-rejection, 
and the WARN ("skip size retention deletion") plus gauge=1 makes it observable. 
If merge-heavy tables ever thrash on this, scoping the comparison to entries 
whose members intersect the barrier-or-older region would be the relaxation.



##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/util/ControllerZkHelixUtilsTest.java:
##########
@@ -36,6 +37,12 @@
 
 public class ControllerZkHelixUtilsTest {
 
+  @AfterMethod(alwaysRun = true)
+  public void resetControllerJobLimits() {

Review Comment:
   nit (process): this is a genuine fix — `ControllerJobTypes.init` mutates 
process-wide state and this class was leaking custom limits into later tests in 
the same JVM fork — but it's unrelated to size-based retention (presumably 
found while running the controller suite for this PR). Fine to keep here given 
its size, but worth a line in the PR description so the squash commit explains 
why a ZK-utils test changed in a retention PR.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +128,218 @@ 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; hybrid 
realtime eviction requires offline coverage.
+  /// 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.

Review Comment:
   Worth confirming the operational trade-off is intended: because COMPLETED 
(and REVERTED) members pin the barrier until their lineage entry is removed, a 
table with continuous merge/rollup activity will have its barrier trail the 
merge watermark by the lineage-cleanup window — with defaults, replaced 
segments live 4h (APPEND) / 24h (REFRESH) and entries are pruned after ~1 day. 
Since merge tasks target the oldest data, that's exactly where size retention 
wants to evict, so a continuously-merging table that's over cap may sit at 
`sizeRetentionBlocked=1` for hours after every merge round. It's bounded 
(lineage cleanup runs right before this in `processTable`) and the javadoc 
states the rule, but a sentence in the docs noting the *typical duration* of 
the pinning (tied to 
`replacedSegmentsRetentionPeriod`/`lineageEntryCleanupRetentionPeriod`) would 
help operators interpret the gauge.



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