This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/codex2/table-size-retention in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 496029f48c0f4de362984c902d7aa0cde03391c1 Author: Xiang Fu <[email protected]> AuthorDate: Tue Oct 6 11:07:42 2026 -0700 Add size-based table retention --- .../helix/core/retention/RetentionManager.java | 121 +++++++++ .../helix/core/retention/RetentionManagerTest.java | 288 +++++++++++++++++++++ .../segment/local/utils/TableConfigUtils.java | 12 + .../segment/local/utils/TableConfigUtilsTest.java | 45 ++++ .../SegmentsValidationAndRetentionConfig.java | 19 ++ .../spi/utils/builder/TableConfigBuilder.java | 9 + .../spi/utils/builder/TableConfigBuilderTest.java | 41 +++ 7 files changed, 535 insertions(+) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java index b5ba82d3133..68c75dd048c 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java @@ -21,8 +21,11 @@ package org.apache.pinot.controller.helix.core.retention; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.io.IOException; +import java.math.BigInteger; import java.net.URI; import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -63,6 +66,7 @@ import org.apache.pinot.spi.filesystem.PinotFS; import org.apache.pinot.spi.filesystem.PinotFSFactory; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.CommonConstants.Segment.Realtime.Status; +import org.apache.pinot.spi.utils.DataSizeUtils; import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.apache.pinot.spi.utils.retry.RetryPolicies; @@ -122,6 +126,123 @@ public class RetentionManager extends ControllerPeriodicTask<Void> { // 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 independently of time retention. The limit counts one + /// copy of every active completed segment, including segments protected by lineage or realtime recovery. Unknown + /// sizes or missing metadata disable this pass; consuming segments are excluded from both accounting and eviction. + @VisibleForTesting + protected void manageSizeBasedRetention(TableConfig tableConfig) { + String retentionSize = tableConfig.getValidationConfig().getRetentionSize(); + if (retentionSize == null) { + return; + } + String tableNameWithType = tableConfig.getTableName(); + if (tableConfig.getTableType() == TableType.OFFLINE + && !"APPEND".equalsIgnoreCase(IngestionConfigUtils.getBatchSegmentIngestionType(tableConfig))) { + return; + } + 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; + } + + Set<String> activeSegments = new HashSet<>(_pinotHelixResourceManager.getSegmentsFor(tableNameWithType, false)); + if (activeSegments.isEmpty()) { + return; + } + List<SegmentZKMetadata> metadataList = _pinotHelixResourceManager.getSegmentsZKMetadata(tableNameWithType); + 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 (!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; + } + // 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; + } + BigInteger limit = BigInteger.valueOf(retentionSizeBytes); + if (retainedBytes.compareTo(limit) <= 0) { + return; + } + + 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)); + } + // Protect lineage before selecting victims: protected bytes still count toward the limit. + removeLineageLockedSegments(tableNameWithType, candidateNames); + Set<String> candidates = new HashSet<>(candidateNames); + completedSegments.removeIf(metadata -> !candidates.contains(metadata.getSegmentName())); + completedSegments.sort(Comparator.comparingLong(RetentionManager::getSizeRetentionTimestamp) + .thenComparing(SegmentZKMetadata::getSegmentName)); + + List<String> segmentsToDelete = new ArrayList<>(); + for (SegmentZKMetadata metadata : completedSegments) { + if (retainedBytes.compareTo(limit) <= 0) { + break; + } + long size = completedSegmentSizes.get(metadata.getSegmentName()); + if (size == 0) { + continue; + } + segmentsToDelete.add(metadata.getSegmentName()); + retainedBytes = retainedBytes.subtract(BigInteger.valueOf(size)); + } + if (!segmentsToDelete.isEmpty()) { + LOGGER.info("Deleting {} oldest segments from table: {} for retention size: {}", segmentsToDelete.size(), + tableNameWithType, retentionSize); + // The public deletion path rechecks live lineage in case it changed after candidate selection. + _pinotHelixResourceManager.deleteSegments(tableNameWithType, segmentsToDelete); + } + if (retainedBytes.compareTo(limit) > 0) { + LOGGER.warn("Table: {} remains above retention size: {} due to protected or undated segments", + tableNameWithType, retentionSize); + } + } + + /// Returns the compressed bytes used by size retention, or a negative value when the size is unknown. + /// Subclasses can supply authoritative sizes for logical segments that aggregate multiple physical segments. + protected long getSizeRetentionSegmentSize(String tableNameWithType, SegmentZKMetadata metadata) { + return metadata.getSizeInBytes(); + } + + private static long getSizeRetentionTimestamp(SegmentZKMetadata metadata) { + long endTime = metadata.getEndTimeMs(); + if (endTime >= 0) { + return endTime; + } + long creationTime = metadata.getCreationTime(); + return creationTime >= 0 ? creationTime : metadata.getPushTime(); } @Override diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/retention/RetentionManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/retention/RetentionManagerTest.java index c0023bda878..5a901cbe7b6 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/retention/RetentionManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/retention/RetentionManagerTest.java @@ -37,6 +37,9 @@ import org.apache.helix.model.IdealState; import org.apache.helix.model.InstanceConfig; import org.apache.helix.store.zk.ZkHelixPropertyStore; import org.apache.helix.zookeeper.datamodel.ZNRecord; +import org.apache.pinot.common.lineage.LineageEntry; +import org.apache.pinot.common.lineage.LineageEntryState; +import org.apache.pinot.common.lineage.SegmentLineage; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.ControllerMetrics; import org.apache.pinot.common.utils.LLCSegmentName; @@ -59,8 +62,10 @@ import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Segment.Realtime.Status; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.zookeeper.data.Stat; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -915,4 +920,287 @@ public class RetentionManagerTest { // 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 testSizeRetentionDeletesNewestSegmentWhenItAloneExceedsLimit() { + 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); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("oldest", "newest")); + } + + @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); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager, never()).deleteSegments(anyString(), anyList()); + } + + @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, false)).thenReturn(List.of("oldest", "newest")); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager, never()).deleteSegments(anyString(), anyList()); + } + + @Test + public void testSizeRetentionProtectsBothSidesOfInProgressLineage() { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "50B"); + long now = System.currentTimeMillis(); + List<SegmentZKMetadata> segments = List.of(createSizeRetentionSegment("source", 80, now - 2), + createSizeRetentionSegment("destination", 80, now - 1), createSizeRetentionSegment("unlocked", 60, now)); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + 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); + + // Protected bytes still count, so the unlocked segment must be removed even though the cap is unattainable. + verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("unlocked")); + verify(resourceManager, never()).deleteSegmentsForLineageCleanup(anyString(), anyList()); + } + + @Test + public void testSizeRetentionSkipsRefreshTablesAndUnconfiguredTables() { + for (boolean refreshTable : List.of(true, false)) { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "50B"); + if (refreshTable) { + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setBatchIngestionConfig(new BatchIngestionConfig(null, "REFRESH", null)); + tableConfig.setIngestionConfig(ingestionConfig); + } else { + tableConfig.getValidationConfig().setRetentionSize(null); + } + List<SegmentZKMetadata> segments = + List.of(createSizeRetentionSegment("oversized", 100, System.currentTimeMillis())); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager, never()).deleteSegments(anyString(), anyList()); + } + } + + @Test + public void testSizeRetentionRechecksMetadataAfterTimeRetention() { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "70B"); + tableConfig.getValidationConfig().setRetentionTimeUnit("DAYS"); + tableConfig.getValidationConfig().setRetentionTimeValue("365"); + long now = System.currentTimeMillis(); + List<SegmentZKMetadata> segments = new ArrayList<>(List.of( + createSizeRetentionSegment("expired", 80, now - TimeUnit.DAYS.toMillis(400)), + createSizeRetentionSegment("recent", 60, now - 1), createSizeRetentionSegment("newest", 30, now))); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + when(resourceManager.getSegmentsFor(OFFLINE_TABLE_NAME, false)).thenAnswer( + invocation -> segments.stream().map(SegmentZKMetadata::getSegmentName).toList()); + doAnswer(invocation -> { + List<String> deletedSegments = invocation.getArgument(1); + segments.removeIf(segment -> deletedSegments.contains(segment.getSegmentName())); + return null; + }).when(resourceManager).deleteSegments(eq(OFFLINE_TABLE_NAME), anyList()); + + retentionManager.processTable(OFFLINE_TABLE_NAME); + + var deletionOrder = inOrder(resourceManager); + deletionOrder.verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("expired")); + deletionOrder.verify(resourceManager).getSegmentsZKMetadata(OFFLINE_TABLE_NAME); + deletionOrder.verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("recent")); + assertEquals(segments.stream().map(SegmentZKMetadata::getSegmentName).toList(), List.of("newest")); + } + + @Test + public void testSizeRetentionAccountsForTotalAboveLongMaximum() { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "100B"); + long now = System.currentTimeMillis(); + List<SegmentZKMetadata> segments = List.of(createSizeRetentionSegment("oldest", Long.MAX_VALUE, now - 1), + createSizeRetentionSegment("newest", 1, now)); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("oldest")); + } + + @Test + public void testSizeRetentionSkipsTableWhenActiveMetadataIsMissing() { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "50B"); + List<SegmentZKMetadata> segments = + List.of(createSizeRetentionSegment("oversized", 100, System.currentTimeMillis())); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + when(resourceManager.getSegmentsFor(OFFLINE_TABLE_NAME, false)).thenReturn(List.of("oversized", "missing")); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager, never()).deleteSegments(anyString(), anyList()); + } + + @Test + public void testSizeRetentionKeepsUndatedAndEmptySegments() { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, "50B"); + long now = System.currentTimeMillis(); + SegmentZKMetadata undated = createSizeRetentionSegment("undated", 80, -1); + undated.setCreationTime(-1); + List<SegmentZKMetadata> segments = List.of(createSizeRetentionSegment("empty", 0, now - 1), undated, + createSizeRetentionSegment("dated", 60, now)); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager).deleteSegments(OFFLINE_TABLE_NAME, List.of("dated")); + } + + @Test + public void testSizeRetentionSkipsMalformedStoredConfiguration() { + for (String retentionSize : List.of("invalid", "0B", "-1B", "")) { + TableConfig tableConfig = createSizeRetentionTableConfig(TableType.OFFLINE, retentionSize); + List<SegmentZKMetadata> segments = + List.of(createSizeRetentionSegment("oversized", 100, System.currentTimeMillis())); + PinotHelixResourceManager resourceManager = mock(PinotHelixResourceManager.class); + RetentionManager retentionManager = createSizeRetentionManager(tableConfig, segments, resourceManager); + + retentionManager.manageSizeBasedRetention(tableConfig); + + verify(resourceManager, never()).deleteSegments(anyString(), anyList()); + } + } + + private TableConfig createSizeRetentionTableConfig(TableType tableType, String retentionSize) { + TableConfig tableConfig = new TableConfigBuilder(tableType).setTableName(TEST_TABLE_NAME).build(); + tableConfig.getValidationConfig().setRetentionSize(retentionSize); + return tableConfig; + } + + private SegmentZKMetadata createSizeRetentionSegment(String segmentName, long sizeInBytes, long endTimeMs) { + SegmentZKMetadata metadata = new SegmentZKMetadata(segmentName); + metadata.setSizeInBytes(sizeInBytes); + metadata.setTimeUnit(TimeUnit.MILLISECONDS); + metadata.setEndTime(endTimeMs); + metadata.setCreationTime(System.currentTimeMillis()); + return metadata; + } + + private RetentionManager createSizeRetentionManager(TableConfig tableConfig, List<SegmentZKMetadata> segments, + PinotHelixResourceManager resourceManager) { + String tableName = tableConfig.getTableName(); + when(resourceManager.getTableConfig(tableName)).thenReturn(tableConfig); + when(resourceManager.getSegmentsZKMetadata(tableName)).thenReturn(segments); + when(resourceManager.getSegmentsFor(tableName, false)) + .thenReturn(segments.stream().map(SegmentZKMetadata::getSegmentName).toList()); + when(resourceManager.getPropertyStore()).thenReturn(mock(ZkHelixPropertyStore.class)); + when(resourceManager.getLineageUpdaterLock(tableName)).thenReturn(new Object()); + when(resourceManager.getLastLLCCompletedSegments(anyList())).thenCallRealMethod(); + when(resourceManager.getLastLLCCompletedSegments(tableName)).thenCallRealMethod(); + ControllerConf conf = new ControllerConf(); + conf.setUntrackedSegmentDeletionEnabled(false); + return createRetentionManager(resourceManager, mock(LeadControllerManager.class), conf, + mock(ControllerMetrics.class), mock(BrokerServiceHelper.class)); + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index c77b42930a4..4e76120982b 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -301,6 +301,7 @@ public final class TableConfigUtils { /// Validates retention config. Checks for following things: /// - Valid segmentPushType /// - Valid retentionTimeUnit + /// - Positive retentionSize in bytes, when configured private static void validateRetentionConfig(TableConfig tableConfig) { SegmentsValidationAndRetentionConfig segmentsConfig = tableConfig.getValidationConfig(); String tableName = tableConfig.getTableName(); @@ -329,6 +330,17 @@ public final class TableConfigUtils { } } + String retentionSize = segmentsConfig.getRetentionSize(); + if (retentionSize != null) { + try { + Preconditions.checkArgument(DataSizeUtils.toBytes(retentionSize) > 0, "Retention size must be positive"); + } catch (IllegalArgumentException e) { + throw new IllegalStateException( + String.format("Table: %s, invalid retention size: %s; must be a positive size in bytes", tableName, + retentionSize), e); + } + } + // Untracked segments retention may not be specified. Ignore validation in that case. String untrackedSegmentsRetentionTimeUnitString = segmentsConfig.getUntrackedSegmentsRetentionTimeUnit(); String untrackedSegmentsRetentionTimeValueString = segmentsConfig.getUntrackedSegmentsRetentionTimeValue(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java index ae5af3891de..371b7431096 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java @@ -2564,6 +2564,51 @@ public class TableConfigUtilsTest { } } + @DataProvider + public Object[][] retentionTableTypes() { + return new Object[][]{{TableType.OFFLINE}, {TableType.REALTIME}}; + } + + @Test(dataProvider = "retentionTableTypes") + public void testValidateRetentionSize(TableType tableType) { + Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addDateTime(TIME_COLUMN, DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + TableConfigBuilder builder = new TableConfigBuilder(tableType).setTableName(TABLE_NAME) + .setTimeColumnName(TIME_COLUMN); + if (tableType == TableType.REALTIME) { + builder.setStreamConfigs(getStreamConfigs()); + } + + // Size retention is optional and does not require time retention. + TableConfigUtils.validate(builder.build(), schema); + for (String retentionSize : List.of("1", "1B", "100G", "1.5gB")) { + TableConfigUtils.validate(builder.setRetentionSize(retentionSize).build(), schema); + } + + // The two retention policies can be configured together. + TableConfigUtils.validate(builder.setRetentionTimeUnit("DAYS").setRetentionTimeValue("7").build(), schema); + } + + @Test(dataProvider = "retentionTableTypes") + public void testRejectInvalidRetentionSize(TableType tableType) { + Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addDateTime(TIME_COLUMN, DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + TableConfigBuilder builder = new TableConfigBuilder(tableType).setTableName(TABLE_NAME) + .setTimeColumnName(TIME_COLUMN); + if (tableType == TableType.REALTIME) { + builder.setStreamConfigs(getStreamConfigs()); + } + + for (String retentionSize : List.of("", " ", "0", "0G", "0.5B", "-1", "-1G", "100GBB", "abc")) { + TableConfig tableConfig = builder.setRetentionSize(retentionSize).build(); + IllegalStateException e = + expectThrows(IllegalStateException.class, () -> TableConfigUtils.validate(tableConfig, schema)); + assertTrue(e.getMessage().contains("invalid retention size: " + retentionSize), e.getMessage()); + } + } + @Test public void testValidateDedupConfig() { TableConfig validTableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName(TABLE_NAME) diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/SegmentsValidationAndRetentionConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/SegmentsValidationAndRetentionConfig.java index 9027fd3ec57..2c3173bb9d2 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/SegmentsValidationAndRetentionConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/SegmentsValidationAndRetentionConfig.java @@ -27,6 +27,7 @@ import org.apache.pinot.spi.utils.TimeUtils; public class SegmentsValidationAndRetentionConfig extends BaseJsonConfig { private String _retentionTimeUnit; private String _retentionTimeValue; + private String _retentionSize; private String _deletedSegmentsRetentionPeriod; private String _replacedSegmentsRetentionPeriod; private String _lineageEntryCleanupRetentionPeriod; @@ -88,6 +89,24 @@ public class SegmentsValidationAndRetentionConfig extends BaseJsonConfig { _retentionTimeValue = retentionTimeValue; } + /// Returns the size retention limit, expressed as a positive size understood by + /// [DataSizeUtils#toBytes(String)], for example `"100G"`. + /// + /// The limit applies to the sum of compressed segment archive bytes for one replica of this table. OFFLINE and + /// REALTIME tables enforce their limits independently, including the two types of a hybrid table. When time retention + /// is also configured, a segment can be deleted by either retention criterion. Consuming segments and segments + /// protected by segment lineage are not removed to enforce the limit, so it may be exceeded. This is an asynchronous + /// retention policy, not a limit on server disk usage or an ingestion quota. A null value disables size retention. + public String getRetentionSize() { + return _retentionSize; + } + + /// Sets the compressed segment archive size retention limit described by [#getRetentionSize()]. + /// Set to null to disable size retention. + public void setRetentionSize(String retentionSize) { + _retentionSize = retentionSize; + } + public String getDeletedSegmentsRetentionPeriod() { return _deletedSegmentsRetentionPeriod; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/TableConfigBuilder.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/TableConfigBuilder.java index 8587ea04627..f61f433a4da 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/TableConfigBuilder.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/TableConfigBuilder.java @@ -74,6 +74,7 @@ public class TableConfigBuilder { private String _timeType; private String _retentionTimeUnit; private String _retentionTimeValue; + private String _retentionSize; private String _deletedSegmentsRetentionPeriod = DEFAULT_DELETED_SEGMENTS_RETENTION_PERIOD; private String _replacedSegmentsRetentionPeriod; private String _lineageEntryCleanupRetentionPeriod; @@ -203,6 +204,13 @@ public class TableConfigBuilder { return this; } + /// Sets the compressed segment archive size retention limit described by + /// [SegmentsValidationAndRetentionConfig#getRetentionSize()]. + public TableConfigBuilder setRetentionSize(String retentionSize) { + _retentionSize = retentionSize; + return this; + } + public TableConfigBuilder setDeletedSegmentsRetentionPeriod(String deletedSegmentsRetentionPeriod) { _deletedSegmentsRetentionPeriod = deletedSegmentsRetentionPeriod; return this; @@ -522,6 +530,7 @@ public class TableConfigBuilder { validationConfig.setTimeType(_timeType); validationConfig.setRetentionTimeUnit(_retentionTimeUnit); validationConfig.setRetentionTimeValue(_retentionTimeValue); + validationConfig.setRetentionSize(_retentionSize); validationConfig.setDeletedSegmentsRetentionPeriod(_deletedSegmentsRetentionPeriod); validationConfig.setReplacedSegmentsRetentionPeriod(_replacedSegmentsRetentionPeriod); validationConfig.setLineageEntryCleanupRetentionPeriod(_lineageEntryCleanupRetentionPeriod); diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/builder/TableConfigBuilderTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/builder/TableConfigBuilderTest.java index 2f89899fba5..6a64aa08b7a 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/builder/TableConfigBuilderTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/builder/TableConfigBuilderTest.java @@ -19,11 +19,18 @@ package org.apache.pinot.spi.utils.builder; +import java.io.IOException; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.JsonUtils; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; + /// Tests for the validations in [TableConfigBuilder] public class TableConfigBuilderTest { @@ -31,6 +38,40 @@ public class TableConfigBuilderTest { private static final String TABLE_NAME = "testTable"; private static final String TIME_COLUMN = "timeColumn"; + @DataProvider + public Object[][] tableTypes() { + return new Object[][]{{TableType.OFFLINE}, {TableType.REALTIME}}; + } + + @Test(dataProvider = "tableTypes") + public void testRetentionSizeRoundTrip(TableType tableType) + throws IOException { + TableConfig tableConfig = new TableConfigBuilder(tableType).setTableName(TABLE_NAME) + .setRetentionSize("100G") + .setRetentionTimeUnit("DAYS") + .setRetentionTimeValue("7") + .build(); + assertEquals(tableConfig.getValidationConfig().getRetentionSize(), "100G"); + assertEquals(tableConfig.toJsonNode().get("segmentsConfig").get("retentionSize").asText(), "100G"); + + TableConfig roundTrip = JsonUtils.stringToObject(tableConfig.toJsonString(), TableConfig.class); + assertEquals(roundTrip, tableConfig); + assertEquals(roundTrip.getValidationConfig().getRetentionSize(), "100G"); + assertEquals(roundTrip.getValidationConfig().getRetentionTimeUnit(), "DAYS"); + assertEquals(roundTrip.getValidationConfig().getRetentionTimeValue(), "7"); + assertEquals(new TableConfig(tableConfig).getValidationConfig().getRetentionSize(), "100G"); + } + + @Test(dataProvider = "tableTypes") + public void testRetentionSizeDefaultsToDisabled(TableType tableType) + throws IOException { + TableConfig tableConfig = new TableConfigBuilder(tableType).setTableName(TABLE_NAME).build(); + assertNull(tableConfig.getValidationConfig().getRetentionSize()); + assertFalse(tableConfig.toJsonNode().get("segmentsConfig").has("retentionSize")); + TableConfig roundTrip = JsonUtils.stringToObject(tableConfig.toJsonString(), TableConfig.class); + assertNull(roundTrip.getValidationConfig().getRetentionSize()); + } + @Test public void testValidateSkipSegmentPreprocessFlag() { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
