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


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +128,173 @@ 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, and an in-progress replacement 
postpones the entire pass. 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.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+      return;
+    }
+    // Leave the gauge blocked on unsafe early returns or exceptions, and 
clear it only after a successful pass.
+    _controllerMetrics.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 1);
+    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;
+    }
+
+    SegmentLineage segmentLineage =
+        
SegmentLineageAccessHelper.getSegmentLineage(_pinotHelixResourceManager.getPropertyStore(),
 tableNameWithType);
+    if (segmentLineage != null && 
segmentLineage.getLineageEntries().values().stream()
+        .anyMatch(entry -> entry.getState() == LineageEntryState.IN_PROGRESS)) 
{
+      LOGGER.warn("In-progress segment lineage for table: {}, skip size 
retention", tableNameWithType);

Review Comment:
   Two small things on this check:
   
   1. With the default `controller.lineage.exclusive.delete.enabled=true` this 
is airtight (the `deleteSegments` batch rejection plus your response check 
close the race). With the flag **off**, though, there's a residual TOCTOU: a 
`startReplaceSegments` landing between this read and `deleteSegments` isn't 
caught anywhere, and size retention targets exactly the oldest segments that 
merge/rollup replaces, so the collision odds are higher than for time 
retention. Fine to accept as the documented legacy-mode trade-off, but worth a 
sentence in the `getRetentionSize` javadoc (or holding `getLineageUpdaterLock` 
across check+delete like `manageSegmentLineageCleanupForTable` does).
   
   2. nit: a single pass now reads the lineage znode up to four times — here, 
inside `getSegmentsFor(..., true)` → `excludeReplacedSegments`, in 
`removeLineageLockedSegments`, and the recheck inside `deleteSegments`. The 
last one must stay fresh, but the first three could share one read if you fetch 
the `SegmentLineage` once and pass it down.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +128,173 @@ 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, and an in-progress replacement 
postpones the entire pass. 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.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);

Review Comment:
   This branch runs for **every** table the controller leads on every retention 
pass, and `setOrUpdateTableGauge` registers the gauge on first call — so every 
table in the cluster gets a permanent `sizeRetentionBlocked.<table>` series 
pinned at 0, even though (typically) almost none configure `retentionSize`. In 
a cluster with thousands of tables that's a lot of metric cardinality for a 
niche feature.
   
   Suggest calling `_controllerMetrics.removeTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED)` here instead — it's safe when the 
gauge was never registered, and it still correctly clears a stale `1` when an 
operator removes `retentionSize` (the transition 
`testSizeRetentionGaugeResetsWhenRetentionIsDisabled` covers). Would also want 
the `ControllerGauge` doc updated from "0 when the policy is disabled" to 
"absent when the policy is disabled".



##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/retention/RetentionManagerTest.java:
##########
@@ -915,4 +923,577 @@ public void 
testRetentionWithInvalidEndTimeAndCreationTimeFallback() {
     // already asserts the correct segments via TestNG assertions
     verify(pinotHelixResourceManager, 
times(1)).deleteSegments(eq(OFFLINE_TABLE_NAME), anyList());
   }
+
+  @Test
+  public void testSizeOnlyRetentionDeletesOldestSegmentsUntilWithinLimit() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "100B");
+    long now = System.currentTimeMillis();
+    List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("newest", 60, now),
+        createSizeRetentionSegment("oldest", 60, now - 2),
+        createSizeRetentionSegment("middle", 60, now - 1));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager);
+
+    // Exercise the complete table pass: a time-retention configuration is 
deliberately absent.
+    retentionManager.processTable(OFFLINE_TABLE_NAME);
+
+    verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, 
List.of("oldest", "middle"));
+  }
+
+  @Test
+  public void testSizeRetentionKeepsSegmentsAtOrBelowLimit() {
+    for (String limit : List.of("100B", "101B")) {
+      TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, limit);
+      tableConfig.getValidationConfig().setReplication("3");
+      long now = System.currentTimeMillis();
+      List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("oldest", 40, now - 1),
+          createSizeRetentionSegment("newest", 60, now));
+      PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+      RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager);
+
+      retentionManager.manageSizeBasedRetention(tableConfig);
+
+      // Stored archive sizes count once, regardless of table replication.
+      verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+    }
+  }
+
+  @Test
+  public void 
testSizeRetentionKeepsNewestOfflineSegmentWhenItAloneExceedsLimit() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "100B");
+    long now = System.currentTimeMillis();
+    List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("newest", 120, now),
+        createSizeRetentionSegment("oldest", 40, now - 1));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    ControllerMetrics metrics = mock(ControllerMetrics.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager, metrics);
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, 
List.of("oldest"));
+    verify(metrics, atLeastOnce()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 1);
+    verify(metrics, never()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+  }
+
+  @Test
+  public void testSizeRetentionNeverDeletesOnlyOfflineSegment() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "1B");
+    List<SegmentZKMetadata> segments =
+        List.of(createSizeRetentionSegment("only", 100, 
System.currentTimeMillis()));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    ControllerMetrics metrics = mock(ControllerMetrics.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager, metrics);
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+    verify(metrics, atLeastOnce()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 1);
+    verify(metrics, never()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+  }
+
+  @Test
+  public void testSizeRetentionOrderingWithCreationAndPushTimeFallbacks() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "20B");
+    long now = System.currentTimeMillis();
+    SegmentZKMetadata creationFallback = 
createSizeRetentionSegment("creationFallback", 10, -1);
+    creationFallback.setCreationTime(now - 1);
+    SegmentZKMetadata pushFallback = 
createSizeRetentionSegment("pushFallback", 10, -1);
+    pushFallback.setCreationTime(-1);
+    pushFallback.setPushTime(now - 4);
+    List<SegmentZKMetadata> segments = List.of(creationFallback,
+        createSizeRetentionSegment("b", 10, now - 3), pushFallback,
+        createSizeRetentionSegment("oldest", 10, now - 5), 
createSizeRetentionSegment("a", 10, now - 3));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager);
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, 
List.of("oldest", "pushFallback", "a"));
+  }
+
+  @Test
+  public void 
testSizeRetentionProtectsIncompleteAndLastCompletedRealtimeSegments() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.REALTIME, "50B");
+    long now = System.currentTimeMillis();
+    SegmentZKMetadata oldest = createSizeRetentionSegment(
+        new LLCSegmentName(TEST_TABLE_NAME, 0, 4, now - 3).getSegmentName(), 
20, now - 3);
+    oldest.setStatus(Status.DONE);
+    SegmentZKMetadata lastCompleted = createSizeRetentionSegment(
+        new LLCSegmentName(TEST_TABLE_NAME, 0, 5, now - 1).getSegmentName(), 
40, now - 1);
+    lastCompleted.setStatus(Status.DONE);
+    SegmentZKMetadata consuming = createSizeRetentionSegment("consuming", -1, 
now - 5);
+    consuming.setStatus(Status.IN_PROGRESS);
+    SegmentZKMetadata committing = createSizeRetentionSegment("committing", 
999, now - 4);
+    committing.setStatus(Status.COMMITTING);
+    SegmentZKMetadata uploaded = createSizeRetentionSegment("uploaded", 60, 
now - 2);
+    List<SegmentZKMetadata> segments = List.of(lastCompleted, committing, 
uploaded, consuming, oldest);
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager);
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager).deleteSegments(REALTIME_TABLE_NAME, 
List.of(oldest.getSegmentName(), "uploaded"));
+  }
+
+  @Test
+  public void testSizeRetentionSkipsTableWhenCompletedSegmentSizeIsUnknown() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "50B");
+    long now = System.currentTimeMillis();
+    List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("oldest", 60, now - 2),
+        createSizeRetentionSegment("unknown", -1, now - 1), 
createSizeRetentionSegment("newest", 60, now));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    ControllerMetrics metrics = mock(ControllerMetrics.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager, metrics);
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+    verify(metrics, atLeastOnce()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 1);
+    verify(metrics, never()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+  }
+
+  @Test
+  public void testSizeRetentionIgnoresMetadataOutsideIdealState() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "100B");
+    long now = System.currentTimeMillis();
+    List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("oldest", 40, now - 2),
+        createSizeRetentionSegment("orphan", -1, now - 1), 
createSizeRetentionSegment("newest", 40, now));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager);
+    when(resourceManager.getSegmentsFor(OFFLINE_TABLE_NAME, 
true)).thenReturn(List.of("oldest", "newest"));
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+  }
+
+  @Test
+  public void 
testSizeRetentionSkipsInProgressLineageEvenWhenExclusiveDeleteIsDisabled() {
+    for (boolean exclusiveDelete : List.of(true, false)) {
+      TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "50B");
+      long now = System.currentTimeMillis();
+      List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("olderLive", 60, now - 3),
+          createSizeRetentionSegment("source", 80, now - 2),
+          createSizeRetentionSegment("destination", 80, now - 1), 
createSizeRetentionSegment("newestLive", 60, now));
+      PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+      ControllerMetrics metrics = mock(ControllerMetrics.class);
+      ControllerConf conf = new ControllerConf();
+      conf.setProperty(ControllerConf.LINEAGE_EXCLUSIVE_DELETE_ENABLED, 
exclusiveDelete);
+      RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager, metrics,
+          conf, mock(BrokerServiceHelper.class));
+      SegmentLineage lineage = new SegmentLineage(OFFLINE_TABLE_NAME);
+      lineage.addLineageEntry("replacement",
+          new LineageEntry(List.of("source"), List.of("destination"), 
LineageEntryState.IN_PROGRESS, now));
+      when(resourceManager.getPropertyStore().get(anyString(), 
any(Stat.class), eq(AccessOption.PERSISTENT)))
+          .thenReturn(lineage.toZNRecord());
+
+      retentionManager.manageSizeBasedRetention(tableConfig);
+
+      // Transient duplicate bytes must not trigger deletion of unrelated live 
data.
+      verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+      verify(resourceManager, 
never()).deleteSegmentsForLineageCleanup(anyString(), anyList());
+      verify(resourceManager, never()).getSegmentsFor(anyString(), 
anyBoolean());
+      verify(metrics, atLeastOnce()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME,
+          ControllerGauge.SIZE_RETENTION_BLOCKED, 1);
+      verify(metrics, never()).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+    }
+  }
+
+  @Test
+  public void 
testSizeRetentionExcludesCompletedReplacedSegmentsFromAccounting() {
+    TableConfig tableConfig = 
createSizeRetentionTableConfig(TableType.OFFLINE, "100B");
+    long now = System.currentTimeMillis();
+    List<SegmentZKMetadata> segments = 
List.of(createSizeRetentionSegment("source", -1, now - 2),
+        createSizeRetentionSegment("destination", 40, now - 1), 
createSizeRetentionSegment("newest", 40, now));
+    PinotHelixResourceManager resourceManager = 
mock(PinotHelixResourceManager.class);
+    ControllerMetrics metrics = mock(ControllerMetrics.class);
+    RetentionManager retentionManager = 
createSizeRetentionManager(tableConfig, segments, resourceManager, metrics);
+    SegmentLineage lineage = new SegmentLineage(OFFLINE_TABLE_NAME);
+    lineage.addLineageEntry("replacement",
+        new LineageEntry(List.of("source"), List.of("destination"), 
LineageEntryState.COMPLETED, now));
+    when(resourceManager.getPropertyStore().get(anyString(), any(Stat.class), 
eq(AccessOption.PERSISTENT)))
+        .thenReturn(lineage.toZNRecord());
+    when(resourceManager.getSegmentsFor(OFFLINE_TABLE_NAME, 
true)).thenReturn(List.of("destination", "newest"));
+
+    retentionManager.manageSizeBasedRetention(tableConfig);
+
+    verify(resourceManager).getSegmentsFor(OFFLINE_TABLE_NAME, true);
+    verify(resourceManager, never()).getSegmentsFor(OFFLINE_TABLE_NAME, false);
+    verify(resourceManager, never()).deleteSegments(anyString(), anyList());
+    verify(metrics).setOrUpdateTableGauge(OFFLINE_TABLE_NAME, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);

Review Comment:
   Coverage on the new behaviors is genuinely strong (each one has a test that 
asserts the opposite of the v1 behavior, and I verified the failure-response 
branch actually executes in the run log). One gap worth closing while you're 
here: the success-path gauge assertions (`verify(..., 0)` here and in the 
retry/disabled/recovery tests) don't pin the *final* gauge value — a 1-write 
legitimately precedes the 0, so `never(..., 1)` can't be added and no `InOrder` 
is used. A regression that writes 1 after the success-path 0 would pass all of 
these. Either Mockito `InOrder` asserting the 0 is last, or a real 
`ControllerMetrics` with `MetricValueUtils.getTableGaugeValue(...)` (the 
`SegmentStatusCheckerTest` pattern) would close it — the latter would also let 
the leadership test assert the gauge is truly gone after `nonLeaderCleanup`.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java:
##########
@@ -122,6 +128,173 @@ 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, and an in-progress replacement 
postpones the entire pass. 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.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 0);
+      return;
+    }
+    // Leave the gauge blocked on unsafe early returns or exceptions, and 
clear it only after a successful pass.
+    _controllerMetrics.setOrUpdateTableGauge(tableNameWithType, 
ControllerGauge.SIZE_RETENTION_BLOCKED, 1);

Review Comment:
   Minor: because the gauge is set to 1 at pass entry and only cleared at the 
end, a metrics reporter that scrapes mid-pass sees `sizeRetentionBlocked=1` on 
a perfectly healthy table — and the window isn't tiny (several ZK reads, plus a 
broker RPC for hybrid tables). That contradicts the documented semantics and 
can trip zero-tolerance alerts.
   
   The leave-1-on-exception intent is sound and you can keep it without the 
flap: track a local `boolean blocked = true`, flip it to false on the success 
paths, and write the gauge exactly once in a `finally` block. Exceptions still 
leave 1; healthy passes never publish a transient 1. (That would also let the 
success-path tests pin the *final* value — right now they verify a 0-write 
happened but can't assert no 1-write follows it, so a regression that writes 1 
after the 0 would pass; alternatively use a real `ControllerMetrics` + 
`MetricValueUtils.getTableGaugeValue` as `SegmentStatusCheckerTest` does.)



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