From 6169f758ad4c9f0ad32ae1f2e9edd96432ab72cc Mon Sep 17 00:00:00 2001 From: Hassan Shaitou <54168508+hash-14@users.noreply.github.com> Date: Tue, 14 Jul 2026 07:05:08 +0000 Subject: [PATCH 1/3] Add lazy segment loading for OFFLINE tables Servers register a metadata-only stub on ONLINE assignment instead of downloading the segment; the first query materializes it from the deep store, and a background sweeper evicts idle segments back to stubs. Server-side only: brokers and clients are unchanged. Config: - Table level: lazyLoadConfig {enabled, idleEvictionSeconds, deleteLocalOnEvict} - Instance level: pinot.server.instance.lazy.load.enabled / lazy.sweep.interval.seconds / lazy.materialize.parallelism New server metrics: LAZY_SEGMENT_COLD_LOADS, LAZY_SEGMENT_EVICTIONS, LAZY_SEGMENT_LOAD_TIME_MS, LAZY_STUBBED_SEGMENT_COUNT --- .../pinot/common/metrics/ServerGauge.java | 2 + .../pinot/common/metrics/ServerMeter.java | 3 + .../pinot/common/metrics/ServerTimer.java | 2 + .../utils/config/TableConfigSerDeUtils.java | 12 + .../data/manager/BaseTableDataManager.java | 252 ++++++++++- .../offline/OfflineTableDataManager.java | 6 + .../OfflineTableDataManagerLazyLoadTest.java | 404 ++++++++++++++++++ .../data/manager/SegmentDataManager.java | 7 + .../segment/local/utils/TableConfigUtils.java | 20 + .../helix/HelixInstanceDataManager.java | 32 ++ .../helix/HelixInstanceDataManagerConfig.java | 24 ++ .../instance/InstanceDataManagerConfig.java | 22 + .../spi/config/table/LazyLoadConfig.java | 65 +++ .../pinot/spi/config/table/TableConfig.java | 20 + 14 files changed, 868 insertions(+), 3 deletions(-) create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/config/table/LazyLoadConfig.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java index 312ace46ee4e..64daca301b18 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java @@ -73,6 +73,8 @@ public enum ServerGauge implements AbstractMetrics.Gauge { // Segment operation metrics - count is the current number of segments undergoing the given operation. // Incremented when the semaphore is acquired and decremented when the semaphore is released SEGMENT_TABLE_DOWNLOAD_COUNT("segmentTableDownloadCount", false), + // Number of lazily stubbed (assigned but not materialized) segments per table + LAZY_STUBBED_SEGMENT_COUNT("segments", false), SEGMENT_DOWNLOAD_COUNT("segmentDownloadCount", true), SEGMENT_ALL_PREPROCESS_COUNT("segmentAllPreprocessCount", true), SEGMENT_STARTREE_PREPROCESS_COUNT("segmentStartreePreprocessCount", true), diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java index 81ff8c3df82a..a2147d45a748 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java @@ -102,6 +102,9 @@ public enum ServerMeter implements AbstractMetrics.Meter { SEGMENT_DOWNLOAD_FAILURES("segments", false), SEGMENT_DOWNLOAD_FROM_REMOTE_FAILURES("segments", false), SEGMENT_DOWNLOAD_FROM_PEERS_FAILURES("segments", false), + // Lazy segment loading + LAZY_SEGMENT_COLD_LOADS("segments", false), + LAZY_SEGMENT_EVICTIONS("segments", false), SEGMENT_BUILD_FAILURE("segments", false), SEGMENT_UPLOAD_FAILURE("segments", false), SEGMENT_UPLOAD_SUCCESS("segments", false), diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerTimer.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerTimer.java index ab5d461d677e..af252a3e517e 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerTimer.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerTimer.java @@ -107,6 +107,8 @@ public enum ServerTimer implements AbstractMetrics.Timer { "Time taken to download a segment from deep store (including untar and move operations)"), SEGMENT_DOWNLOAD_FROM_PEERS_TIME_MS("millis", false, "Time taken to download a segment from peers (including untar and move operations)"), + LAZY_SEGMENT_LOAD_TIME_MS("millis", false, + "Time taken to materialize (download, untar and load) a lazily stubbed segment on first query access"), // Ingestion metrics INGESTION_DELAY_TRACKING_MS("milliseconds", false, "Time taken to run a trackIngestionDelay cycle"), diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtils.java index 436e58f3ae59..b05eb874646c 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtils.java @@ -31,6 +31,7 @@ import org.apache.pinot.spi.config.table.DimensionTableConfig; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.config.table.IndexingConfig; +import org.apache.pinot.spi.config.table.LazyLoadConfig; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.QuotaConfig; import org.apache.pinot.spi.config.table.RoutingConfig; @@ -204,6 +205,13 @@ public static TableConfig fromZNRecord(ZNRecord znRecord) tableSamplerConfigs, isMaterializedView); tableConfig.setDescription(description); tableConfig.setTags(tags); + + // NOTE: lazyLoadConfig is not part of the @JsonCreator constructor, set it explicitly. Any TableConfig + // sub-config missing from this class is silently dropped when the controller persists the config to ZK. + String lazyLoadConfigString = simpleFields.get(TableConfig.LAZY_LOAD_CONFIG_KEY); + if (lazyLoadConfigString != null) { + tableConfig.setLazyLoadConfig(JsonUtils.stringToObject(lazyLoadConfigString, LazyLoadConfig.class)); + } return tableConfig; } @@ -292,6 +300,10 @@ public static ZNRecord toZNRecord(TableConfig tableConfig) if (tags != null && !tags.isEmpty()) { simpleFields.put(TableConfig.TAGS_KEY, JsonUtils.objectToString(tags)); } + LazyLoadConfig lazyLoadConfig = tableConfig.getLazyLoadConfig(); + if (lazyLoadConfig != null) { + simpleFields.put(TableConfig.LAZY_LOAD_CONFIG_KEY, JsonUtils.objectToString(lazyLoadConfig)); + } ZNRecord znRecord = new ZNRecord(tableConfig.getTableName()); znRecord.setSimpleFields(simpleFields); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java index 8bb9393bb4e1..2dcf4f6eb10e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java @@ -22,6 +22,7 @@ import com.google.common.base.Preconditions; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; +import com.google.common.util.concurrent.ThreadFactoryBuilder; import java.io.File; import java.io.IOException; import java.net.URI; @@ -38,7 +39,9 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.Semaphore; +import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -112,6 +115,7 @@ import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig; import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.LazyLoadConfig; import org.apache.pinot.spi.config.table.MultiColumnTextIndexConfig; import org.apache.pinot.spi.config.table.SegmentPartitionConfig; import org.apache.pinot.spi.config.table.StarTreeIndexConfig; @@ -137,6 +141,14 @@ public abstract class BaseTableDataManager implements TableDataManager { protected final ServerMetrics _serverMetrics = ServerMetrics.get(); protected TableUpsertMetadataManager _tableUpsertMetadataManager; + // Lazy segment loading: registry of segments that are assigned ONLINE but not materialized locally. This is a + // parallel map, not a fake SegmentDataManager — all lazy logic is confined to the assignment and acquire touch + // points plus the eviction sweeper. Broker routing is unchanged: the segment looks ONLINE to the cluster. + protected final ConcurrentHashMap _stubbedSegments = new ConcurrentHashMap<>(); + protected boolean _instanceLazyLoadEnabled; + @Nullable + protected ExecutorService _lazyMaterializeExecutor; + protected InstanceDataManagerConfig _instanceDataManagerConfig; protected String _instanceId; protected HelixManager _helixManager; @@ -275,6 +287,20 @@ public void init(InstanceDataManagerConfig instanceDataManagerConfig, HelixManag _logger.info("Async segment refresh is {}!", _enableAsyncSegmentRefresh ? "enabled" : "disabled"); + _instanceLazyLoadEnabled = instanceDataManagerConfig.isLazyLoadEnabled(); + if (_instanceLazyLoadEnabled) { + int materializeParallelism = Math.max(1, instanceDataManagerConfig.getLazyLoadMaterializeParallelism()); + ThreadPoolExecutor lazyMaterializeExecutor = + new ThreadPoolExecutor(materializeParallelism, materializeParallelism, 60, TimeUnit.SECONDS, + new LinkedBlockingQueue<>(), + new ThreadFactoryBuilder().setNameFormat("lazy-materialize-" + _tableNameWithType + "-%d") + .setDaemon(true).build()); + // Threads idle-timeout so the pool costs nothing between cold loads. + lazyMaterializeExecutor.allowCoreThreadTimeOut(true); + _lazyMaterializeExecutor = lazyMaterializeExecutor; + _logger.info("Lazy loading enabled at instance level, materialize parallelism: {}", materializeParallelism); + } + doInit(); _logger.info("Initialized table data manager with data directory: {}", _tableDataDir); @@ -310,6 +336,9 @@ public synchronized void shutDown() { _logger.info("Shutting down table data manager"); _shutDown = true; doShutdown(); + if (_lazyMaterializeExecutor != null) { + _lazyMaterializeExecutor.shutdownNow(); + } _logger.info("Shut down table data manager"); } @@ -319,6 +348,8 @@ public synchronized void shutDown() { * Releases and removes all segments tracked by the table data manager. */ protected void releaseAndRemoveAllSegments() { + _stubbedSegments.clear(); + _serverMetrics.setValueOfTableGauge(_tableNameWithType, ServerGauge.LAZY_STUBBED_SEGMENT_COUNT, 0); List segmentDataManagers; synchronized (_segmentDataManagerMap) { segmentDataManagers = new ArrayList<>(_segmentDataManagerMap.values()); @@ -358,7 +389,8 @@ public Lock getSegmentLock(String segmentName) { @Override public boolean hasSegment(String segmentName) { - return _segmentDataManagerMap.containsKey(segmentName); + // Stubbed segments count as present: they are assigned ONLINE and routable, just not materialized yet. + return _segmentDataManagerMap.containsKey(segmentName) || _stubbedSegments.containsKey(segmentName); } /** @@ -564,6 +596,13 @@ protected void doReplaceSegment(String segmentName) } finally { _segmentReloadSemaphore.release(); } + } else if (_stubbedSegments.containsKey(segmentName)) { + // Segment refreshed while stubbed: update the stub metadata; the next query loads the NEW version (the + // materialize path re-fetches ZK metadata, so this is belt-and-braces consistency). + SegmentZKMetadata zkMetadata = fetchZKMetadata(segmentName); + _stubbedSegments.put(segmentName, zkMetadata); + _logger.info("Segment: {} is stubbed, updated stub metadata on refresh; new version loads on next query", + segmentName); } else { _logger.warn("Failed to find segment: {}, skipping replacing it", segmentName); } @@ -607,11 +646,19 @@ public void offloadSegmentUnsafe(String segmentName) { } protected void doOffloadSegment(String segmentName) { + // Helix OFFLINE/DROP transitions must also clear lazy stubs. + boolean wasStubbed = _stubbedSegments.remove(segmentName) != null; + if (wasStubbed) { + _serverMetrics.setValueOfTableGauge(_tableNameWithType, ServerGauge.LAZY_STUBBED_SEGMENT_COUNT, + _stubbedSegments.size()); + } SegmentDataManager segmentDataManager = unregisterSegment(segmentName); if (segmentDataManager != null) { segmentDataManager.offload(); releaseSegment(segmentDataManager); _logger.info("Offloaded segment: {}", segmentName); + } else if (wasStubbed) { + _logger.info("Offloaded stubbed segment: {} (stub removed, nothing to release)", segmentName); } else { _logger.warn("Failed to find segment: {}, skipping offloading it", segmentName); } @@ -752,9 +799,10 @@ public List acquireSegments(List segmentNames, List< @Override public List acquireSegments(List segmentNames, @Nullable List optionalSegmentNames, List missingSegments) { + materializeStubbedSegmentsInParallel(segmentNames); List segmentDataManagers = new ArrayList<>(); for (String segmentName : segmentNames) { - SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + SegmentDataManager segmentDataManager = getOrMaterializeSegmentDataManager(segmentName); if (segmentDataManager != null && segmentDataManager.increaseReferenceCount()) { segmentDataManagers.add(segmentDataManager); } else { @@ -776,7 +824,7 @@ public List acquireSegments(List segmentNames, @Nullable @Override public SegmentDataManager acquireSegment(String segmentName) { - SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + SegmentDataManager segmentDataManager = getOrMaterializeSegmentDataManager(segmentName); if (segmentDataManager != null && segmentDataManager.increaseReferenceCount()) { return segmentDataManager; } else { @@ -784,6 +832,204 @@ public SegmentDataManager acquireSegment(String segmentName) { } } + /** + * Returns the SegmentDataManager for the given segment, materializing it first if it is currently a lazy stub. + * Returns null if the segment is neither loaded nor stubbed, or if materialization failed (in which case the stub + * stays registered so the next query retries — self-healing). + */ + @Nullable + protected SegmentDataManager getOrMaterializeSegmentDataManager(String segmentName) { + SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + if (segmentDataManager != null) { + return segmentDataManager; + } + if (_stubbedSegments.containsKey(segmentName)) { + return materializeStubbedSegment(segmentName); + } + return null; + } + + /** + * Materializes a stubbed segment: downloads it from the deep store, loads it and removes the stub. Runs under the + * per-segment lock so concurrent cold queries for the same segment dedupe into a single download. On failure the + * stub is kept and null is returned — the segment is reported missing for this query (upstream behavior) and the + * next query retries. + */ + @Nullable + protected SegmentDataManager materializeStubbedSegment(String segmentName) { + Lock segmentLock = getSegmentLock(segmentName); + segmentLock.lock(); + try { + // Re-check under the lock: a concurrent query may have materialized it already. + SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + if (segmentDataManager != null) { + return segmentDataManager; + } + if (!_stubbedSegments.containsKey(segmentName)) { + return null; + } + long startTimeMs = System.currentTimeMillis(); + _logger.info("Materializing stubbed segment: {} on first query access", segmentName); + try { + // Re-fetch ZK metadata so a segment refreshed while stubbed loads its NEW version (CRC check downstream). + SegmentZKMetadata zkMetadata = fetchZKMetadata(segmentName); + IndexLoadingConfig indexLoadingConfig = fetchIndexLoadingConfig(); + indexLoadingConfig.setSegmentTier(zkMetadata.getTier()); + addNewOnlineSegment(zkMetadata, indexLoadingConfig); + _stubbedSegments.remove(segmentName); + long loadTimeMs = System.currentTimeMillis() - startTimeMs; + _serverMetrics.addMeteredTableValue(_tableNameWithType, ServerMeter.LAZY_SEGMENT_COLD_LOADS, 1); + _serverMetrics.addTimedTableValue(_tableNameWithType, ServerTimer.LAZY_SEGMENT_LOAD_TIME_MS, loadTimeMs, + TimeUnit.MILLISECONDS); + _serverMetrics.setValueOfTableGauge(_tableNameWithType, ServerGauge.LAZY_STUBBED_SEGMENT_COUNT, + _stubbedSegments.size()); + _logger.info("Materialized stubbed segment: {} in {}ms", segmentName, loadTimeMs); + return _segmentDataManagerMap.get(segmentName); + } catch (Exception e) { + _logger.error("Failed to materialize stubbed segment: {}, keeping stub so the next query retries", + segmentName, e); + addSegmentError(segmentName, new SegmentErrorInfo(System.currentTimeMillis(), + "Caught exception while materializing stubbed segment", e)); + return null; + } + } finally { + segmentLock.unlock(); + } + } + + /** + * Materializes all stubbed segments needed by a query in parallel on a bounded pool. Failures are swallowed here — + * the per-segment acquire that follows reports them as missing while the stubs stay registered for retry. + */ + protected void materializeStubbedSegmentsInParallel(List segmentNames) { + if (_lazyMaterializeExecutor == null || _stubbedSegments.isEmpty()) { + return; + } + List stubbedSegmentNames = new ArrayList<>(); + for (String segmentName : segmentNames) { + if (_stubbedSegments.containsKey(segmentName) && !_segmentDataManagerMap.containsKey(segmentName)) { + stubbedSegmentNames.add(segmentName); + } + } + if (stubbedSegmentNames.isEmpty()) { + return; + } + if (stubbedSegmentNames.size() == 1) { + materializeStubbedSegment(stubbedSegmentNames.get(0)); + return; + } + _logger.info("Materializing {} stubbed segments in parallel", stubbedSegmentNames.size()); + CompletableFuture.allOf(stubbedSegmentNames.stream() + .map(segmentName -> CompletableFuture.runAsync(() -> materializeStubbedSegment(segmentName), + _lazyMaterializeExecutor)) + .toArray(CompletableFuture[]::new)).join(); + } + + /** + * Registers a metadata-only stub for a segment assigned ONLINE. Zero disk, zero memory — the first query that + * touches the segment materializes it. + */ + protected void registerStubbedSegment(String segmentName, SegmentZKMetadata zkMetadata) { + _stubbedSegments.put(segmentName, zkMetadata); + _recentlyDeletedSegments.invalidate(segmentName); + _serverMetrics.setValueOfTableGauge(_tableNameWithType, ServerGauge.LAZY_STUBBED_SEGMENT_COUNT, + _stubbedSegments.size()); + _logger.info("Lazy loading: registering stub for segment: {} (skipping download)", segmentName); + } + + /** + * Returns the LazyLoadConfig if lazy loading is effectively enabled for this table (table config enabled AND + * instance kill switch on), otherwise null. + */ + @Nullable + protected LazyLoadConfig getEffectiveLazyLoadConfig() { + if (!_instanceLazyLoadEnabled) { + return null; + } + TableConfig tableConfig = _cachedTableConfigAndSchema.getLeft(); + LazyLoadConfig lazyLoadConfig = tableConfig.getLazyLoadConfig(); + return lazyLoadConfig != null && lazyLoadConfig.isEnabled() ? lazyLoadConfig : null; + } + + /** + * Returns true if lazy loading is effectively enabled for this table. + */ + public boolean isLazyLoadEnabled() { + return getEffectiveLazyLoadConfig() != null; + } + + /** + * Evicts segments that have been idle for longer than the table's idleEvictionSeconds back to stubs. Called + * periodically by the instance-level eviction sweeper. Every query acquire resets the idle clock; a segment with a + * query in flight is never evicted. + */ + public void evictIdleLazySegments() { + if (_shutDown) { + return; + } + LazyLoadConfig lazyLoadConfig = getEffectiveLazyLoadConfig(); + if (lazyLoadConfig == null) { + return; + } + long idleEvictionMs = lazyLoadConfig.getIdleEvictionSeconds() * 1000L; + long nowMs = System.currentTimeMillis(); + for (Map.Entry entry : _segmentDataManagerMap.entrySet()) { + SegmentDataManager segmentDataManager = entry.getValue(); + if (!(segmentDataManager instanceof ImmutableSegmentDataManager)) { + continue; + } + if (nowMs - segmentDataManager.getLastAccessTimeMs() < idleEvictionMs) { + continue; + } + evictSegmentToStub(entry.getKey(), lazyLoadConfig); + } + } + + /** + * Evicts a single idle segment back to a stub: offload, release (frees memory when the last reader is done), + * optionally delete the local index dir, re-register the stub. Data is never lost — the deep-store copy remains + * and the segment stays routable. Idle time and in-flight queries are re-checked under the segment lock. + */ + protected void evictSegmentToStub(String segmentName, LazyLoadConfig lazyLoadConfig) { + Lock segmentLock = getSegmentLock(segmentName); + segmentLock.lock(); + try { + SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + if (segmentDataManager == null) { + return; + } + // Re-check under the lock: reference count > 1 means a query is in flight; a fresh access resets the clock. + long idleEvictionMs = lazyLoadConfig.getIdleEvictionSeconds() * 1000L; + if (segmentDataManager.getReferenceCount() > 1 + || System.currentTimeMillis() - segmentDataManager.getLastAccessTimeMs() < idleEvictionMs) { + return; + } + SegmentZKMetadata zkMetadata = fetchZKMetadataNullable(segmentName); + if (zkMetadata == null || zkMetadata.getDownloadUrl() == null + || CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD.equals(zkMetadata.getDownloadUrl())) { + // No authoritative deep-store copy to refetch from — never evict such a segment. + return; + } + _logger.info("Evicting idle lazy segment: {} back to stub (deleteLocalOnEvict: {})", segmentName, + lazyLoadConfig.isDeleteLocalOnEvict()); + TableConfig tableConfig = _cachedTableConfigAndSchema.getLeft(); + File indexDir = getSegmentDataDir(segmentName, zkMetadata.getTier(), tableConfig); + // Register the stub before unregistering the loaded segment so the segment never looks missing. + registerStubbedSegment(segmentName, zkMetadata); + unregisterSegment(segmentName); + _recentlyDeletedSegments.invalidate(segmentName); + segmentDataManager.offload(); + releaseSegment(segmentDataManager); + if (lazyLoadConfig.isDeleteLocalOnEvict()) { + FileUtils.deleteQuietly(indexDir); + } + _serverMetrics.addMeteredTableValue(_tableNameWithType, ServerMeter.LAZY_SEGMENT_EVICTIONS, 1); + _logger.info("Evicted idle segment: {} back to stub", segmentName); + } finally { + segmentLock.unlock(); + } + } + @Override public void releaseSegment(SegmentDataManager segmentDataManager) { if (segmentDataManager.decreaseReferenceCount()) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java index 23105d54a38a..f3f49415b894 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java @@ -77,6 +77,12 @@ protected void doAddOnlineSegment(String segmentName) indexLoadingConfig.setSegmentTier(zkMetadata.getTier()); SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); if (segmentDataManager == null) { + if (isLazyLoadEnabled()) { + // Lazy loading: register a metadata-only stub instead of downloading. The first query that touches the + // segment materializes it from the deep store. + registerStubbedSegment(segmentName, zkMetadata); + return; + } addNewOnlineSegment(zkMetadata, indexLoadingConfig); } else { replaceSegmentIfCrcMismatch(segmentDataManager, zkMetadata, indexLoadingConfig); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java new file mode 100644 index 000000000000..e07937a1cf78 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java @@ -0,0 +1,404 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.data.manager.offline; + +import java.io.DataInputStream; +import java.io.File; +import java.io.FileInputStream; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.commons.io.FileUtils; +import org.apache.helix.HelixManager; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.common.utils.TarCompressionUtils; +import org.apache.pinot.common.utils.fetcher.BaseSegmentFetcher; +import org.apache.pinot.common.utils.fetcher.SegmentFetcherFactory; +import org.apache.pinot.segment.local.data.manager.SegmentDataManager; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.utils.SegmentLocks; +import org.apache.pinot.segment.local.utils.SegmentReloadSemaphore; +import org.apache.pinot.segment.local.utils.ServerReloadJobStatusCache; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; +import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig; +import org.apache.pinot.spi.config.table.LazyLoadConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.util.TestUtils; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Tests for lazy segment loading: stub-on-assignment, materialize-on-first-query, concurrent dedupe, TTL eviction + * back to stub, in-flight query protection, instance kill switch and non-lazy regression. + */ +public class OfflineTableDataManagerLazyLoadTest { + private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "OfflineTableDataManagerLazyLoadTest"); + private static final String RAW_TABLE_NAME = "testLazyTable"; + private static final String OFFLINE_TABLE_NAME = TableNameBuilder.OFFLINE.tableNameWithType(RAW_TABLE_NAME); + private static final File TABLE_DATA_DIR = new File(TEMP_DIR, OFFLINE_TABLE_NAME); + private static final String SEGMENT_NAME = "testLazySegment"; + private static final String STRING_COLUMN = "col1"; + private static final String[] STRING_VALUES = {"A", "D", "E", "B", "C"}; + private static final String LONG_COLUMN = "col2"; + private static final long[] LONG_VALUES = {10000L, 20000L, 50000L, 40000L, 30000L}; + private static final int NUM_ROWS = 5; + private static final long IDLE_EVICTION_SECONDS = 1; + + private static final Schema SCHEMA = + new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME).addSingleValueDimension(STRING_COLUMN, DataType.STRING) + .addMetric(LONG_COLUMN, DataType.LONG).build(); + + // Managers created by the running test — shut down in tearDown so mmap'd segment files can be deleted (Windows + // cannot delete open-mapped files). + private final List _tableDataManagers = new ArrayList<>(); + + @BeforeClass + public void setUp() { + ServerMetrics.register(mock(ServerMetrics.class)); + } + + @BeforeMethod + public void setUpMethod() + throws Exception { + TestUtils.ensureDirectoriesExistAndEmpty(TEMP_DIR); + Map properties = new HashMap<>(); + properties.put(BaseSegmentFetcher.RETRY_COUNT_CONFIG_KEY, 3); + properties.put(BaseSegmentFetcher.RETRY_WAIT_MS_CONFIG_KEY, 100); + properties.put(BaseSegmentFetcher.RETRY_DELAY_SCALE_FACTOR_CONFIG_KEY, 5); + SegmentFetcherFactory.init(new PinotConfiguration(properties)); + } + + @AfterMethod + public void tearDownMethod() + throws Exception { + for (OfflineTableDataManager tableDataManager : _tableDataManagers) { + try { + tableDataManager.shutDown(); + } catch (Exception e) { + // Best effort — deletion below is what we actually need + } + } + _tableDataManagers.clear(); + FileUtils.deleteDirectory(TEMP_DIR); + } + + @Test + public void testStubOnAssignmentAndLoadOnAcquire() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + + // Assignment registers a stub: routable (hasSegment) but nothing loaded and nothing on disk. + tableDataManager.addOnlineSegment(SEGMENT_NAME); + assertTrue(tableDataManager.hasSegment(SEGMENT_NAME)); + assertEquals(tableDataManager.getNumSegments(), 0); + assertFalse(new File(TABLE_DATA_DIR, SEGMENT_NAME).exists()); + verify(tableDataManager, times(0)).addNewOnlineSegment(any(), any()); + + // First acquire (single-segment query path) materializes the segment. + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + assertEquals(segmentDataManager.getSegment().getSegmentMetadata().getTotalDocs(), NUM_ROWS); + assertEquals(tableDataManager.getNumSegments(), 1); + assertTrue(new File(TABLE_DATA_DIR, SEGMENT_NAME).exists()); + verify(tableDataManager, times(1)).addNewOnlineSegment(any(), any()); + tableDataManager.releaseSegment(segmentDataManager); + + // Multi-segment query path (acquireSegments) on a second stubbed segment. + String secondSegmentName = SEGMENT_NAME + "_2"; + SegmentZKMetadata secondZkMetadata = createRawSegment(secondSegmentName); + doReturn(secondZkMetadata).when(tableDataManager).fetchZKMetadata(secondSegmentName); + doReturn(secondZkMetadata).when(tableDataManager).fetchZKMetadataNullable(secondSegmentName); + tableDataManager.addOnlineSegment(secondSegmentName); + assertEquals(tableDataManager.getNumSegments(), 1); + + List missingSegments = new ArrayList<>(); + List segmentDataManagers = + tableDataManager.acquireSegments(List.of(SEGMENT_NAME, secondSegmentName), missingSegments); + assertEquals(segmentDataManagers.size(), 2); + assertTrue(missingSegments.isEmpty()); + assertEquals(tableDataManager.getNumSegments(), 2); + for (SegmentDataManager sdm : segmentDataManagers) { + tableDataManager.releaseSegment(sdm); + } + } + + @Test + public void testConcurrentAcquireDedupesDownload() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + tableDataManager.addOnlineSegment(SEGMENT_NAME); + + int numThreads = 8; + ExecutorService executorService = Executors.newFixedThreadPool(numThreads); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(numThreads); + AtomicInteger successCount = new AtomicInteger(); + for (int i = 0; i < numThreads; i++) { + executorService.submit(() -> { + try { + startLatch.await(); + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + if (segmentDataManager != null) { + successCount.incrementAndGet(); + tableDataManager.releaseSegment(segmentDataManager); + } + } catch (Exception e) { + // Counted as failure via successCount + } finally { + doneLatch.countDown(); + } + }); + } + startLatch.countDown(); + assertTrue(doneLatch.await(60, TimeUnit.SECONDS)); + executorService.shutdownNow(); + + // All 8 concurrent cold queries succeed, but the per-segment lock dedupes them into a single download+load. + assertEquals(successCount.get(), numThreads); + verify(tableDataManager, times(1)).addNewOnlineSegment(any(), any()); + } + + @Test + public void testTtlEvictionAndReloadLoop() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + tableDataManager.addOnlineSegment(SEGMENT_NAME); + + // Materialize. + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + tableDataManager.releaseSegment(segmentDataManager); + assertEquals(tableDataManager.getNumSegments(), 1); + File indexDir = new File(TABLE_DATA_DIR, SEGMENT_NAME); + assertTrue(indexDir.exists()); + + // Idle past the TTL, then sweep: memory freed, local dir deleted, stub re-registered, still routable. + Thread.sleep(IDLE_EVICTION_SECONDS * 1000 + 500); + tableDataManager.evictIdleLazySegments(); + assertEquals(tableDataManager.getNumSegments(), 0); + assertFalse(indexDir.exists()); + assertTrue(tableDataManager.hasSegment(SEGMENT_NAME)); + + // Next query refetches from the deep store and serves correctly. + segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + assertEquals(segmentDataManager.getSegment().getSegmentMetadata().getTotalDocs(), NUM_ROWS); + assertTrue(indexDir.exists()); + tableDataManager.releaseSegment(segmentDataManager); + verify(tableDataManager, times(2)).addNewOnlineSegment(any(), any()); + } + + @Test + public void testNoEvictionWhileQueryInFlight() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + tableDataManager.addOnlineSegment(SEGMENT_NAME); + + // Acquire and hold — simulates a query in flight. + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + + Thread.sleep(IDLE_EVICTION_SECONDS * 1000 + 500); + tableDataManager.evictIdleLazySegments(); + // Reference count > 1 blocks eviction even though the idle TTL has passed. + assertEquals(tableDataManager.getNumSegments(), 1); + assertTrue(new File(TABLE_DATA_DIR, SEGMENT_NAME).exists()); + + tableDataManager.releaseSegment(segmentDataManager); + // After release + idle, the sweep evicts. + Thread.sleep(IDLE_EVICTION_SECONDS * 1000 + 500); + tableDataManager.evictIdleLazySegments(); + assertEquals(tableDataManager.getNumSegments(), 0); + assertTrue(tableDataManager.hasSegment(SEGMENT_NAME)); + } + + @Test + public void testInstanceKillSwitchDisablesLazyLoading() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(false), zkMetadata); + + // Kill switch off: table-level lazy config is ignored, the segment loads eagerly on assignment. + tableDataManager.addOnlineSegment(SEGMENT_NAME); + assertEquals(tableDataManager.getNumSegments(), 1); + assertTrue(new File(TABLE_DATA_DIR, SEGMENT_NAME).exists()); + verify(tableDataManager, times(1)).addNewOnlineSegment(any(), any()); + + // Sweeps are no-ops. + Thread.sleep(IDLE_EVICTION_SECONDS * 1000 + 500); + tableDataManager.evictIdleLazySegments(); + assertEquals(tableDataManager.getNumSegments(), 1); + } + + @Test + public void testNonLazyTableUnchanged() + throws Exception { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + + // Without lazyLoadConfig the table behaves exactly like upstream: eager load on assignment. + tableDataManager.addOnlineSegment(SEGMENT_NAME); + assertEquals(tableDataManager.getNumSegments(), 1); + assertTrue(new File(TABLE_DATA_DIR, SEGMENT_NAME).exists()); + + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + assertEquals(segmentDataManager.getSegment().getSegmentMetadata().getTotalDocs(), NUM_ROWS); + tableDataManager.releaseSegment(segmentDataManager); + + // Sweeps never touch non-lazy tables. + Thread.sleep(IDLE_EVICTION_SECONDS * 1000 + 500); + tableDataManager.evictIdleLazySegments(); + assertEquals(tableDataManager.getNumSegments(), 1); + + // Unknown segments still resolve to null. + assertNull(tableDataManager.acquireSegment("unknownSegment")); + } + + private static TableConfig createLazyTableConfig(long idleEvictionSeconds, boolean deleteLocalOnEvict) { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + tableConfig.setLazyLoadConfig(new LazyLoadConfig(true, idleEvictionSeconds, deleteLocalOnEvict)); + return tableConfig; + } + + private static InstanceDataManagerConfig createLazyInstanceConfig(boolean lazyLoadEnabled) { + InstanceDataManagerConfig config = mock(InstanceDataManagerConfig.class); + when(config.getInstanceDataDir()).thenReturn(TEMP_DIR.getAbsolutePath()); + when(config.shouldCheckCRCOnSegmentLoad()).thenReturn(true); + when(config.isLazyLoadEnabled()).thenReturn(lazyLoadEnabled); + when(config.getLazyLoadMaterializeParallelism()).thenReturn(4); + return config; + } + + /** + * Creates a spied OfflineTableDataManager with the ZK-touching methods stubbed out, following the + * BaseTableDataManagerTest pattern: a real segment is built and tarred locally, and a file:// download URL acts + * as the deep store. + */ + private OfflineTableDataManager createLazyTableDataManager(TableConfig tableConfig, + InstanceDataManagerConfig instanceDataManagerConfig, SegmentZKMetadata zkMetadata) { + OfflineTableDataManager tableDataManager = spy(new OfflineTableDataManager()); + tableDataManager.init(instanceDataManagerConfig, mock(HelixManager.class), new SegmentLocks(), tableConfig, + SCHEMA, new SegmentReloadSemaphore(1), Executors.newSingleThreadExecutor(), null, null, null, false, + mock(ServerReloadJobStatusCache.class)); + doReturn(zkMetadata).when(tableDataManager).fetchZKMetadata(anyString()); + doReturn(zkMetadata).when(tableDataManager).fetchZKMetadataNullable(anyString()); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(tableConfig, SCHEMA); + indexLoadingConfig.setTableDataDir(TABLE_DATA_DIR.getAbsolutePath()); + doReturn(indexLoadingConfig).when(tableDataManager).fetchIndexLoadingConfig(); + _tableDataManagers.add(tableDataManager); + return tableDataManager; + } + + /** + * Builds a real segment, tars it into TEMP_DIR and deletes the index dir — the file:// URL acts as the deep + * store, so the only way for the server to serve the segment is to "download" it. + */ + private static SegmentZKMetadata createRawSegment(String segmentName) + throws Exception { + File indexDir = createSegment(segmentName); + File rawSegmentFile = new File(TEMP_DIR, segmentName + TarCompressionUtils.TAR_COMPRESSED_FILE_EXTENSION); + TarCompressionUtils.createCompressedTarFile(indexDir, rawSegmentFile); + long crc = getCrc(indexDir); + SegmentZKMetadata zkMetadata = new SegmentZKMetadata(segmentName); + // NOTE: toURI() keeps this portable — a raw "file://" + absolute path is an invalid URI on Windows. + zkMetadata.setDownloadUrl(rawSegmentFile.toURI().toString()); + zkMetadata.setCrc(crc); + FileUtils.deleteQuietly(indexDir); + return zkMetadata; + } + + private static File createSegment(String segmentName) + throws Exception { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, SCHEMA); + config.setOutDir(TABLE_DATA_DIR.getAbsolutePath()); + config.setSegmentName(segmentName); + List rows = new ArrayList<>(NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + GenericRow row = new GenericRow(); + row.putValue(STRING_COLUMN, STRING_VALUES[i]); + row.putValue(LONG_COLUMN, LONG_VALUES[i]); + rows.add(row); + } + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(rows)); + driver.build(); + return new File(TABLE_DATA_DIR, segmentName); + } + + private static long getCrc(File indexDir) + throws Exception { + File creationMetaFile = SegmentDirectoryPaths.findCreationMetaFile(indexDir); + assertNotNull(creationMetaFile); + try (DataInputStream in = new DataInputStream(new FileInputStream(creationMetaFile))) { + return in.readLong(); + } + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java index eb2026e63f8f..f791d1ec2ea7 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java @@ -33,11 +33,17 @@ public abstract class SegmentDataManager { private final AtomicBoolean _offloaded = new AtomicBoolean(); private final AtomicBoolean _destroyed = new AtomicBoolean(); private int _referenceCount = 1; + // Idle clock for lazy-loading TTL eviction. Stamped on every successful query acquire. + private volatile long _lastAccessTimeMs = _loadTimeMs; public long getLoadTimeMs() { return _loadTimeMs; } + public long getLastAccessTimeMs() { + return _lastAccessTimeMs; + } + public synchronized int getReferenceCount() { return _referenceCount; } @@ -52,6 +58,7 @@ public synchronized boolean increaseReferenceCount() { return false; } else { _referenceCount++; + _lastAccessTimeMs = System.currentTimeMillis(); return true; } } 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 5dd51a79e317..6cc19b19685b 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 @@ -63,6 +63,7 @@ import org.apache.pinot.spi.config.table.FieldConfig.EncodingType; import org.apache.pinot.spi.config.table.HashFunction; import org.apache.pinot.spi.config.table.IndexingConfig; +import org.apache.pinot.spi.config.table.LazyLoadConfig; import org.apache.pinot.spi.config.table.MultiColumnTextIndexConfig; import org.apache.pinot.spi.config.table.QuotaConfig; import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig; @@ -214,12 +215,31 @@ private static void validateEffectiveTableConfig(TableConfig tableConfig, Schema validateTaskConfig(tableConfig); validateMaterializedViewInvariants(tableConfig); + validateLazyLoadConfig(tableConfig); if (_enforcePoolBasedAssignment) { validateInstancePoolsAndReplicaGroups(tableConfig); } } + /** + * Lazy segment loading is only supported on plain OFFLINE tables: REALTIME, upsert, dedup and dimension tables + * rely on segments being materialized locally at assignment time. + */ + static void validateLazyLoadConfig(TableConfig tableConfig) { + LazyLoadConfig lazyLoadConfig = tableConfig.getLazyLoadConfig(); + if (lazyLoadConfig == null || !lazyLoadConfig.isEnabled()) { + return; + } + Preconditions.checkState(tableConfig.getTableType() == TableType.OFFLINE, + "lazyLoadConfig can only be enabled on OFFLINE tables"); + Preconditions.checkState(!tableConfig.isUpsertEnabled(), "lazyLoadConfig cannot be enabled with upsert"); + Preconditions.checkState(!tableConfig.isDedupEnabled(), "lazyLoadConfig cannot be enabled with dedup"); + Preconditions.checkState(!tableConfig.isDimTable(), "lazyLoadConfig cannot be enabled on dimension tables"); + Preconditions.checkState(lazyLoadConfig.getIdleEvictionSeconds() > 0, + "lazyLoadConfig.idleEvictionSeconds must be positive"); + } + /** * Validates the table config is using instance pool and replica group configuration. * @param tableConfig Table config to validate diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java index 4c2e7aaa7bd9..827ae4f09caf 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java @@ -32,6 +32,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; @@ -122,6 +123,10 @@ public class HelixInstanceDataManager implements InstanceDataManager { @Nullable private ExecutorService _segmentPreloadExecutor; + // Lazy loading: single scheduled daemon thread evicting idle lazy segments across all tables. + @Nullable + private ScheduledExecutorService _lazyEvictionSweeper; + @Override public void setSupplierOfIsServerReadyToConsumeData(BooleanSupplier isServerReadyToConsumeData) { _isServerReadyToConsumeData = isServerReadyToConsumeData; @@ -187,6 +192,30 @@ public synchronized void init(PinotConfiguration config, HelixManager helixManag _errorCache = CacheBuilder.newBuilder().maximumSize(_instanceDataManagerConfig.getErrorCacheSize()).build(); _recentlyDeletedTables = CacheBuilder.newBuilder() .expireAfterWrite(_instanceDataManagerConfig.getDeletedTablesCacheTtlMinutes(), TimeUnit.MINUTES).build(); + + // Start the lazy-loading eviction sweeper only when the instance kill switch is on. Tables without + // lazyLoadConfig enabled are no-ops inside evictIdleLazySegments(). + if (_instanceDataManagerConfig.isLazyLoadEnabled()) { + int sweepIntervalSeconds = Math.max(1, _instanceDataManagerConfig.getLazyLoadSweepIntervalSeconds()); + _lazyEvictionSweeper = Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryBuilder().setNameFormat("lazy-eviction-sweeper").setDaemon(true).build()); + _lazyEvictionSweeper.scheduleWithFixedDelay(this::sweepIdleLazySegments, sweepIntervalSeconds, + sweepIntervalSeconds, TimeUnit.SECONDS); + LOGGER.info("Started lazy-loading eviction sweeper with sweep interval: {}s", sweepIntervalSeconds); + } + } + + private void sweepIdleLazySegments() { + for (TableDataManager tableDataManager : _tableDataManagerMap.values()) { + try { + if (tableDataManager instanceof BaseTableDataManager) { + ((BaseTableDataManager) tableDataManager).evictIdleLazySegments(); + } + } catch (Throwable t) { + LOGGER.error("Caught exception while evicting idle lazy segments for table: {}", + tableDataManager.getTableName(), t); + } + } } ServerIngestionOomProtectionManager.ServerThrottleState getServerIngestionOomProtectionThrottleState() { @@ -280,6 +309,9 @@ public synchronized void start() { @Override public synchronized void shutDown() { + if (_lazyEvictionSweeper != null) { + _lazyEvictionSweeper.shutdownNow(); + } _segmentReloadRefreshExecutor.shutdownNow(); if (_segmentPreloadExecutor != null) { _segmentPreloadExecutor.shutdownNow(); diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java index b62763b462c2..8134d742f247 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java @@ -123,6 +123,15 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig public static final String UPLOAD_SEGMENT_TO_DEEP_STORE = "segment.upload.to.deep.store"; public static final boolean DEFAULT_UPLOAD_SEGMENT_TO_DEEP_STORE = false; + // Lazy segment loading. Keys resolve under the "pinot.server.instance." prefix, e.g. + // pinot.server.instance.lazy.load.enabled + private static final String LAZY_LOAD_ENABLED = "lazy.load.enabled"; + private static final boolean DEFAULT_LAZY_LOAD_ENABLED = true; + private static final String LAZY_SWEEP_INTERVAL_SECONDS = "lazy.sweep.interval.seconds"; + private static final int DEFAULT_LAZY_SWEEP_INTERVAL_SECONDS = 60; + private static final String LAZY_MATERIALIZE_PARALLELISM = "lazy.materialize.parallelism"; + private static final int DEFAULT_LAZY_MATERIALIZE_PARALLELISM = 4; + private final static String[] REQUIRED_KEYS = {INSTANCE_ID}; private static final long DEFAULT_ERROR_CACHE_SIZE = 100L; private static final int DEFAULT_DELETED_TABLES_CACHE_TTL_MINUTES = 60; @@ -359,4 +368,19 @@ public boolean isDimensionTablePreloadDisabled() { return _serverConfig.getProperty(DISABLE_DIMENSION_TABLE_PRELOAD, DEFAULT_DISABLE_DIMENSION_TABLE_PRELOAD); } + + @Override + public boolean isLazyLoadEnabled() { + return _serverConfig.getProperty(LAZY_LOAD_ENABLED, DEFAULT_LAZY_LOAD_ENABLED); + } + + @Override + public int getLazyLoadSweepIntervalSeconds() { + return _serverConfig.getProperty(LAZY_SWEEP_INTERVAL_SECONDS, DEFAULT_LAZY_SWEEP_INTERVAL_SECONDS); + } + + @Override + public int getLazyLoadMaterializeParallelism() { + return _serverConfig.getProperty(LAZY_MATERIALIZE_PARALLELISM, DEFAULT_LAZY_MATERIALIZE_PARALLELISM); + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java index dfe37da0ea32..0f2fa12e2415 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java @@ -92,4 +92,26 @@ public interface InstanceDataManagerConfig { boolean shouldCheckCRCOnSegmentLoad(); boolean isDimensionTablePreloadDisabled(); + + /** + * Instance-level kill switch for lazy segment loading. When false, tables with lazyLoadConfig enabled fall back + * to the regular eager loading path on this instance. Default methods keep other implementors compiling. + */ + default boolean isLazyLoadEnabled() { + return true; + } + + /** + * Interval in seconds between lazy-loading idle-eviction sweeps. + */ + default int getLazyLoadSweepIntervalSeconds() { + return 60; + } + + /** + * Maximum number of stubbed segments a single query materializes in parallel. + */ + default int getLazyLoadMaterializeParallelism() { + return 4; + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/LazyLoadConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/LazyLoadConfig.java new file mode 100644 index 000000000000..52762b82b63c --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/LazyLoadConfig.java @@ -0,0 +1,65 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.spi.config.table; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import javax.annotation.Nullable; +import org.apache.pinot.spi.config.BaseJsonConfig; + + +/** + * Table-level config for lazy segment loading (OFFLINE tables only). + * + *

When enabled, ONLINE segment assignments register a metadata-only stub on the server instead of downloading + * the segment from the deep store. The first query that touches a stubbed segment downloads, untars and loads it, + * after which it serves at native speed. A background sweeper evicts segments that have been idle for longer than + * {@code idleEvictionSeconds} back to stubs, optionally deleting the local index directory to free disk. The + * segment stays routable throughout — the authoritative copy remains in the deep store. + */ +public class LazyLoadConfig extends BaseJsonConfig { + public static final long DEFAULT_IDLE_EVICTION_SECONDS = 3600; + + private final boolean _enabled; + private final long _idleEvictionSeconds; + private final boolean _deleteLocalOnEvict; + + @JsonCreator + public LazyLoadConfig(@JsonProperty("enabled") boolean enabled, + @JsonProperty("idleEvictionSeconds") @Nullable Long idleEvictionSeconds, + @JsonProperty("deleteLocalOnEvict") @Nullable Boolean deleteLocalOnEvict) { + _enabled = enabled; + _idleEvictionSeconds = idleEvictionSeconds != null ? idleEvictionSeconds : DEFAULT_IDLE_EVICTION_SECONDS; + // Default true: eviction frees disk as well as memory. Set to false for two-tier mode where the local index + // directory is kept for a fast re-warm with zero deep-store traffic. + _deleteLocalOnEvict = deleteLocalOnEvict == null || deleteLocalOnEvict; + } + + public boolean isEnabled() { + return _enabled; + } + + public long getIdleEvictionSeconds() { + return _idleEvictionSeconds; + } + + public boolean isDeleteLocalOnEvict() { + return _deleteLocalOnEvict; + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java index c38e0af84a4a..3598624a79f7 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java @@ -69,6 +69,7 @@ public class TableConfig extends BaseJsonConfig { public static final String TABLE_SAMPLERS_KEY = "tableSamplers"; public static final String DESCRIPTION_KEY = "description"; public static final String TAGS_KEY = "tags"; + public static final String LAZY_LOAD_CONFIG_KEY = "lazyLoadConfig"; // Double underscore is reserved for real-time segment name delimiter public static final String TABLE_NAME_FORBIDDEN_SUBSTRING = "__"; @@ -139,6 +140,9 @@ public class TableConfig extends BaseJsonConfig { @JsonPropertyDescription(value = "Configs for table samplers") private List _tableSamplers; + @JsonPropertyDescription(value = "Config for lazy segment loading (OFFLINE tables only)") + private LazyLoadConfig _lazyLoadConfig; + /// Legacy constructor preserved for binary backward-compatibility on the public SPI surface. /// Callers compiled against the pre-`isMaterializedView` signature still link against this entry /// point; it forwards to the canonical constructor with `materializedView=false`. Prefer the @@ -251,6 +255,7 @@ public TableConfig(TableConfig tableConfig) { _tableSamplers = sanitizeAndValidateTableSamplers(tableConfig.getTableSamplers()); _description = tableConfig.getDescription(); _tags = tableConfig.getTags(); + _lazyLoadConfig = tableConfig.getLazyLoadConfig(); } @JsonProperty(TABLE_NAME_KEY) @@ -567,6 +572,21 @@ public void setSegmentAssignmentConfigMap(Map s _segmentAssignmentConfigMap = segmentAssignmentConfigMap; } + @JsonProperty(LAZY_LOAD_CONFIG_KEY) + @Nullable + public LazyLoadConfig getLazyLoadConfig() { + return _lazyLoadConfig; + } + + public void setLazyLoadConfig(LazyLoadConfig lazyLoadConfig) { + _lazyLoadConfig = lazyLoadConfig; + } + + @JsonIgnore + public boolean isLazyLoadEnabled() { + return _lazyLoadConfig != null && _lazyLoadConfig.isEnabled(); + } + @JsonIgnore public int getReplication() { if (_tableType == TableType.REALTIME) { From db53a741b39e957749c051286436430b4c6b3c26 Mon Sep 17 00:00:00 2001 From: Hassan Shaitou <54168508+hash-14@users.noreply.github.com> Date: Tue, 14 Jul 2026 07:05:54 +0000 Subject: [PATCH 2/3] Default instance-level lazy loading to off New behavior must be opt-in: operators enable it per instance with pinot.server.instance.lazy.load.enabled=true, and per table with lazyLoadConfig.enabled=true. --- .../server/starter/helix/HelixInstanceDataManagerConfig.java | 2 +- .../pinot/spi/config/instance/InstanceDataManagerConfig.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java index 8134d742f247..af1295d7d133 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java @@ -126,7 +126,7 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig // Lazy segment loading. Keys resolve under the "pinot.server.instance." prefix, e.g. // pinot.server.instance.lazy.load.enabled private static final String LAZY_LOAD_ENABLED = "lazy.load.enabled"; - private static final boolean DEFAULT_LAZY_LOAD_ENABLED = true; + private static final boolean DEFAULT_LAZY_LOAD_ENABLED = false; private static final String LAZY_SWEEP_INTERVAL_SECONDS = "lazy.sweep.interval.seconds"; private static final int DEFAULT_LAZY_SWEEP_INTERVAL_SECONDS = 60; private static final String LAZY_MATERIALIZE_PARALLELISM = "lazy.materialize.parallelism"; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java index 0f2fa12e2415..bb556ef85ae8 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java @@ -98,7 +98,7 @@ public interface InstanceDataManagerConfig { * to the regular eager loading path on this instance. Default methods keep other implementors compiling. */ default boolean isLazyLoadEnabled() { - return true; + return false; } /** From 50b61961093da422985b80316ec3aab17128f8cb Mon Sep 17 00:00:00 2001 From: Hassan Shaitou <54168508+hash-14@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:25:00 +0000 Subject: [PATCH 3/3] Harden lazy loading query paths and stub registration Fixes issues found in pre-submission review: - Retry acquire when it races the eviction sweeper (both the map-miss and refcount-failure variants), instead of surfacing a false missing-segment error for a servable segment. - Trigger materialization of stubbed optional segments asynchronously: previously they were silently skipped forever (dropping rows with no error flag); now they warm up for subsequent queries without ever blocking the current one. - Bound the multi-segment materialization wait with a configurable timeout (pinot.server.instance.lazy.materialize.timeout.seconds, default 60); still-stubbed segments are reported missing for that query while downloads continue in the background. Handle interrupts and executor shutdown instead of parking query workers on an untimed, uninterruptible join. A lone cold segment keeps the first-query-pays synchronous path, immune to pool backlog. - Dedup in-flight materializations across concurrent queries so retries share one download instead of stacking duplicate tasks that pin pool workers on the segment lock. - Only register a stub when the deep store holds an authoritative copy (otherwise the OFFLINE->ONLINE transition attempts the load so permanent failures surface as Helix ERROR states) and no local index directory exists in the target-tier or default directory (so restarts warm up from disk instead of resetting hot segments to stubs). Clear stale stub entries when a segment takes the eager path. The deep-store-copy predicate is shared with eviction so the two can never diverge. - Stamp the idle clock on release as well as acquire, so a query running longer than the eviction TTL does not leave its segment instantly evictable. --- .../data/manager/BaseTableDataManager.java | 225 ++++++++++++++++-- .../offline/OfflineTableDataManager.java | 9 +- .../OfflineTableDataManagerLazyLoadTest.java | 25 +- .../data/manager/SegmentDataManager.java | 4 +- .../helix/HelixInstanceDataManagerConfig.java | 7 + .../instance/InstanceDataManagerConfig.java | 9 + 6 files changed, 250 insertions(+), 29 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java index 2dcf4f6eb10e..0012704a1f6b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java @@ -30,6 +30,7 @@ import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -37,12 +38,15 @@ import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.Semaphore; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; @@ -145,7 +149,10 @@ public abstract class BaseTableDataManager implements TableDataManager { // parallel map, not a fake SegmentDataManager — all lazy logic is confined to the assignment and acquire touch // points plus the eviction sweeper. Broker routing is unchanged: the segment looks ONLINE to the cluster. protected final ConcurrentHashMap _stubbedSegments = new ConcurrentHashMap<>(); + // In-flight materializations, keyed by segment name: dedups task submission across concurrent cold queries. + private final ConcurrentHashMap> _materializeFutures = new ConcurrentHashMap<>(); protected boolean _instanceLazyLoadEnabled; + protected long _lazyMaterializeTimeoutSeconds; @Nullable protected ExecutorService _lazyMaterializeExecutor; @@ -289,6 +296,7 @@ public void init(InstanceDataManagerConfig instanceDataManagerConfig, HelixManag _instanceLazyLoadEnabled = instanceDataManagerConfig.isLazyLoadEnabled(); if (_instanceLazyLoadEnabled) { + _lazyMaterializeTimeoutSeconds = Math.max(1, instanceDataManagerConfig.getLazyLoadMaterializeTimeoutSeconds()); int materializeParallelism = Math.max(1, instanceDataManagerConfig.getLazyLoadMaterializeParallelism()); ThreadPoolExecutor lazyMaterializeExecutor = new ThreadPoolExecutor(materializeParallelism, materializeParallelism, 60, TimeUnit.SECONDS, @@ -799,11 +807,19 @@ public List acquireSegments(List segmentNames, List< @Override public List acquireSegments(List segmentNames, @Nullable List optionalSegmentNames, List missingSegments) { - materializeStubbedSegmentsInParallel(segmentNames); + // Materialize required stubs in parallel, bounded by the materialize timeout, and trigger fire-and-forget + // downloads for stubbed optional segments (best-effort for this query, warm for the next). Segments returned + // must not be materialized inline below — they either already had their inline attempt or may still be + // downloading in the background after a timeout. + Set noInlineMaterializeSegments = materializeStubbedSegments(segmentNames, optionalSegmentNames); List segmentDataManagers = new ArrayList<>(); for (String segmentName : segmentNames) { - SegmentDataManager segmentDataManager = getOrMaterializeSegmentDataManager(segmentName); - if (segmentDataManager != null && segmentDataManager.increaseReferenceCount()) { + // Inline materialization here covers the sweeper-eviction race (a segment loaded when this query started but + // stubbed now) and a single bounded retry for a pool materialization that failed fast — never a blocking + // wait behind a timed-out download. + SegmentDataManager segmentDataManager = + acquireLoadedSegment(segmentName, !noInlineMaterializeSegments.contains(segmentName)); + if (segmentDataManager != null) { segmentDataManagers.add(segmentDataManager); } else { missingSegments.add(segmentName); @@ -811,9 +827,9 @@ public List acquireSegments(List segmentNames, } if (optionalSegmentNames != null) { for (String segmentName : optionalSegmentNames) { - SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + SegmentDataManager segmentDataManager = acquireLoadedSegment(segmentName, false); // Optional segments are not counted to missing segments that are reported back in query exception. - if (segmentDataManager != null && segmentDataManager.increaseReferenceCount()) { + if (segmentDataManager != null) { segmentDataManagers.add(segmentDataManager); } } @@ -824,11 +840,56 @@ public List acquireSegments(List segmentNames, @Nullable @Override public SegmentDataManager acquireSegment(String segmentName) { - SegmentDataManager segmentDataManager = getOrMaterializeSegmentDataManager(segmentName); - if (segmentDataManager != null && segmentDataManager.increaseReferenceCount()) { - return segmentDataManager; - } else { - return null; + while (true) { + SegmentDataManager segmentDataManager = getOrMaterializeSegmentDataManager(segmentName); + if (segmentDataManager == null) { + return null; + } + if (segmentDataManager.increaseReferenceCount()) { + return segmentDataManager; + } + if (!_stubbedSegments.containsKey(segmentName)) { + // Genuinely destroyed (dropped/replaced) — matches upstream behavior for a failed acquire. + return null; + } + // Raced with evictSegmentToStub: the sweeper destroyed the manager but the segment is stubbed and fully + // servable. Loop to re-materialize. Terminates because a freshly materialized segment is not idle and thus + // cannot be immediately re-evicted. + } + } + + /** + * Acquires a segment for a multi-segment query. When {@code materializeInline} is true, a stubbed segment is + * materialized on the calling thread; this covers the rare race where the eviction sweeper evicts the segment + * around the lookup / reference-count increment (the segment is stubbed and fully servable, so returning null + * would surface a false missing-segment error — the re-load is served from the kept local copy when + * {@code deleteLocalOnEvict=false}, and is a cold re-download otherwise, the accepted price for never returning + * a false miss) and a single bounded retry after a pool materialization failed fast. When false (optional + * segments, or required segments whose download may still be running after a timed-out wait), a stub is reported + * as null without blocking. Returns null if the segment is stubbed and not materialized here, or genuinely gone. + * The loop terminates because a freshly materialized segment is not idle and thus cannot be immediately + * re-evicted. + */ + @Nullable + private SegmentDataManager acquireLoadedSegment(String segmentName, boolean materializeInline) { + while (true) { + SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); + if (segmentDataManager == null) { + if (!materializeInline || !_stubbedSegments.containsKey(segmentName)) { + return null; + } + segmentDataManager = materializeStubbedSegment(segmentName); + if (segmentDataManager == null) { + return null; + } + } + if (segmentDataManager.increaseReferenceCount()) { + return segmentDataManager; + } + if (!materializeInline || !_stubbedSegments.containsKey(segmentName)) { + return null; + } + // Eviction race: loop to re-materialize the stub. } } @@ -898,31 +959,146 @@ protected SegmentDataManager materializeStubbedSegment(String segmentName) { } /** - * Materializes all stubbed segments needed by a query in parallel on a bounded pool. Failures are swallowed here — - * the per-segment acquire that follows reports them as missing while the stubs stay registered for retry. + * Materializes the stubbed segments among the required names in parallel on the bounded pool, waiting at most the + * configured materialize timeout, and triggers fire-and-forget materialization of stubbed optional segments — + * strictly after the required submissions, so best-effort prefetch never queues ahead of the segments this query + * needs. Failures and timeouts are swallowed here: still-stubbed segments are reported missing for this query + * while their stubs stay registered (and any in-flight downloads keep running) so subsequent queries succeed. + * + * @return the segment names the caller must NOT materialize inline: segments whose single inline attempt already + * ran (lone-stub fast path) or whose downloads may still be running (timed-out / interrupted / rejected + * wait). An empty result means every required stub either materialized or failed fast, so one bounded + * inline retry is safe. */ - protected void materializeStubbedSegmentsInParallel(List segmentNames) { - if (_lazyMaterializeExecutor == null || _stubbedSegments.isEmpty()) { - return; + protected Set materializeStubbedSegments(List segmentNames, + @Nullable List optionalSegmentNames) { + ExecutorService lazyMaterializeExecutor = _lazyMaterializeExecutor; + if (lazyMaterializeExecutor == null || _stubbedSegments.isEmpty()) { + return Set.of(); } - List stubbedSegmentNames = new ArrayList<>(); + Set stubbedSegmentNames = new HashSet<>(); for (String segmentName : segmentNames) { if (_stubbedSegments.containsKey(segmentName) && !_segmentDataManagerMap.containsKey(segmentName)) { stubbedSegmentNames.add(segmentName); } } if (stubbedSegmentNames.isEmpty()) { - return; + materializeStubbedSegmentsAsync(optionalSegmentNames); + return stubbedSegmentNames; } if (stubbedSegmentNames.size() == 1) { - materializeStubbedSegment(stubbedSegmentNames.get(0)); - return; + // Fast path: a lone cold segment downloads on the query thread (first-query-pays), immune to pool backlog. + // The inline attempt already happened, so the caller must not retry. + materializeStubbedSegment(stubbedSegmentNames.iterator().next()); + materializeStubbedSegmentsAsync(optionalSegmentNames); + return stubbedSegmentNames; } _logger.info("Materializing {} stubbed segments in parallel", stubbedSegmentNames.size()); - CompletableFuture.allOf(stubbedSegmentNames.stream() - .map(segmentName -> CompletableFuture.runAsync(() -> materializeStubbedSegment(segmentName), - _lazyMaterializeExecutor)) - .toArray(CompletableFuture[]::new)).join(); + try { + List> futures = new ArrayList<>(stubbedSegmentNames.size()); + for (String segmentName : stubbedSegmentNames) { + futures.add(submitMaterialization(segmentName, lazyMaterializeExecutor)); + } + // Optional prefetch is submitted only after every required download is enqueued. + materializeStubbedSegmentsAsync(optionalSegmentNames); + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])) + .get(_lazyMaterializeTimeoutSeconds, TimeUnit.SECONDS); + // The wait completed: anything still stubbed failed fast, and one bounded inline retry is allowed. + return Set.of(); + } catch (TimeoutException e) { + _logger.warn("Timed out after {}s waiting for {} stubbed segments to materialize; segments still stubbed are" + + " reported missing for this query while materialization continues in the background", + _lazyMaterializeTimeoutSeconds, stubbedSegmentNames.size()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + _logger.warn("Interrupted while waiting for {} stubbed segments to materialize", stubbedSegmentNames.size()); + } catch (ExecutionException e) { + // materializeStubbedSegment handles its own failures (keeps the stub); this guards unexpected errors. + _logger.error("Caught exception while materializing stubbed segments", e); + } catch (RejectedExecutionException e) { + // Executor concurrently shut down (table shutdown): remaining stubs count as missing for this query. + _logger.warn("Lazy materialize executor is shut down, skipping materialization of {} stubbed segments", + stubbedSegmentNames.size()); + } + return stubbedSegmentNames; + } + + /** + * Triggers materialization of the stubbed segments among the given names without waiting. Used for optional + * segments: best-effort for the current query, warm for the next one. + */ + protected void materializeStubbedSegmentsAsync(@Nullable List segmentNames) { + ExecutorService lazyMaterializeExecutor = _lazyMaterializeExecutor; + if (segmentNames == null || segmentNames.isEmpty() || lazyMaterializeExecutor == null + || _stubbedSegments.isEmpty()) { + return; + } + for (String segmentName : segmentNames) { + if (_stubbedSegments.containsKey(segmentName) && !_segmentDataManagerMap.containsKey(segmentName)) { + try { + submitMaterialization(segmentName, lazyMaterializeExecutor); + } catch (RejectedExecutionException e) { + // Executor concurrently shut down (table shutdown): nothing to trigger. + return; + } + } + } + } + + /** + * Submits a materialization task for the segment, deduplicating against in-flight submissions: concurrent queries + * cold-touching the same segment share one future instead of stacking duplicate tasks that would each pin a pool + * worker on the segment lock. The cleanup callback is attached OUTSIDE computeIfAbsent (whose mapping function + * must not modify the map) and removes by (key, value) so a completed future can never shadow a newer one. + */ + private CompletableFuture submitMaterialization(String segmentName, ExecutorService lazyMaterializeExecutor) { + CompletableFuture future = _materializeFutures.computeIfAbsent(segmentName, + name -> CompletableFuture.runAsync(() -> materializeStubbedSegment(name), lazyMaterializeExecutor)); + future.whenComplete((v, t) -> _materializeFutures.remove(segmentName, future)); + return future; + } + + /** + * Returns true if the segment may be registered as a lazy stub instead of eagerly loaded. Two conditions guard + * this: (1) the deep store must hold an authoritative copy to materialize from later — otherwise the + * OFFLINE->ONLINE transition must attempt the load now so a permanent failure surfaces as a Helix ERROR state + * instead of being silently deferred to query time; (2) no local index directory may exist — a restarted server + * warms up from its valid local copy instead of resetting hot segments to stubs. + */ + protected boolean canRegisterStub(String segmentName, SegmentZKMetadata zkMetadata) { + if (!hasAuthoritativeDeepStoreCopy(zkMetadata)) { + return false; + } + TableConfig tableConfig = _cachedTableConfigAndSchema.getLeft(); + File tierDataDir = getSegmentDataDir(segmentName, zkMetadata.getTier(), tableConfig); + if (tierDataDir.exists()) { + return false; + } + // A copy may also sit in the default-tier directory when the segment's target tier changed while the server + // was down; the eager path relocates/reuses it, a stub would orphan it and re-download from the deep store. + File defaultDataDir = getSegmentDataDir(segmentName); + return defaultDataDir.equals(tierDataDir) || !defaultDataDir.exists(); + } + + /** + * Returns true if the segment's ZK metadata points at an authoritative deep-store copy that a stub can later be + * materialized from (peer-download segments have no such copy). Shared by stub registration and eviction so the + * two predicates can never diverge. + */ + protected static boolean hasAuthoritativeDeepStoreCopy(SegmentZKMetadata zkMetadata) { + String downloadUrl = zkMetadata.getDownloadUrl(); + return downloadUrl != null && !CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD.equals(downloadUrl); + } + + /** + * Removes a stale stub entry, if present, when a segment takes the eager loading path (e.g. re-sent ONLINE after + * an eviction that kept the local copy), so the stub registry and its gauge never over-report. + */ + protected void unregisterStubbedSegment(String segmentName) { + if (_stubbedSegments.remove(segmentName) != null) { + _serverMetrics.setValueOfTableGauge(_tableNameWithType, ServerGauge.LAZY_STUBBED_SEGMENT_COUNT, + _stubbedSegments.size()); + } } /** @@ -1005,8 +1181,7 @@ protected void evictSegmentToStub(String segmentName, LazyLoadConfig lazyLoadCon return; } SegmentZKMetadata zkMetadata = fetchZKMetadataNullable(segmentName); - if (zkMetadata == null || zkMetadata.getDownloadUrl() == null - || CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD.equals(zkMetadata.getDownloadUrl())) { + if (zkMetadata == null || !hasAuthoritativeDeepStoreCopy(zkMetadata)) { // No authoritative deep-store copy to refetch from — never evict such a segment. return; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java index f3f49415b894..b5231f3c902b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManager.java @@ -77,13 +77,18 @@ protected void doAddOnlineSegment(String segmentName) indexLoadingConfig.setSegmentTier(zkMetadata.getTier()); SegmentDataManager segmentDataManager = _segmentDataManagerMap.get(segmentName); if (segmentDataManager == null) { - if (isLazyLoadEnabled()) { + if (isLazyLoadEnabled() && canRegisterStub(segmentName, zkMetadata)) { // Lazy loading: register a metadata-only stub instead of downloading. The first query that touches the - // segment materializes it from the deep store. + // segment materializes it from the deep store. Segments without an authoritative deep-store copy or with an + // existing local copy load eagerly (see canRegisterStub). registerStubbedSegment(segmentName, zkMetadata); return; } addNewOnlineSegment(zkMetadata, indexLoadingConfig); + // Clear any stale stub (e.g. re-sent ONLINE after an eviction that kept the local copy) only AFTER the eager + // load succeeds: while loading, concurrent queries can still self-heal through the stub, and if the load + // throws, the stub keeps the segment servable from the deep store. + unregisterStubbedSegment(segmentName); } else { replaceSegmentIfCrcMismatch(segmentDataManager, zkMetadata, indexLoadingConfig); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java index e07937a1cf78..4f0bb0cb160c 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/OfflineTableDataManagerLazyLoadTest.java @@ -79,7 +79,8 @@ /** * Tests for lazy segment loading: stub-on-assignment, materialize-on-first-query, concurrent dedupe, TTL eviction - * back to stub, in-flight query protection, instance kill switch and non-lazy regression. + * back to stub, in-flight query protection, eager load when a local copy exists (restart warm-up), instance kill + * switch and non-lazy regression. */ public class OfflineTableDataManagerLazyLoadTest { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "OfflineTableDataManagerLazyLoadTest"); @@ -274,6 +275,27 @@ public void testNoEvictionWhileQueryInFlight() assertTrue(tableDataManager.hasSegment(SEGMENT_NAME)); } + @Test + public void testLocalCopySkipsStubOnAssignment() + throws Exception { + TableConfig tableConfig = createLazyTableConfig(IDLE_EVICTION_SECONDS, true); + SegmentZKMetadata zkMetadata = createRawSegment(SEGMENT_NAME); + OfflineTableDataManager tableDataManager = + createLazyTableDataManager(tableConfig, createLazyInstanceConfig(true), zkMetadata); + + // Recreate the local index dir to simulate a server restart with a valid local copy: the assignment must load + // eagerly from disk instead of resetting the segment to a stub (which would force a deep-store re-download). + createSegment(SEGMENT_NAME); + tableDataManager.addOnlineSegment(SEGMENT_NAME); + assertEquals(tableDataManager.getNumSegments(), 1); + verify(tableDataManager, times(1)).addNewOnlineSegment(any(), any()); + + SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(SEGMENT_NAME); + assertNotNull(segmentDataManager); + assertEquals(segmentDataManager.getSegment().getSegmentMetadata().getTotalDocs(), NUM_ROWS); + tableDataManager.releaseSegment(segmentDataManager); + } + @Test public void testInstanceKillSwitchDisablesLazyLoading() throws Exception { @@ -333,6 +355,7 @@ private static InstanceDataManagerConfig createLazyInstanceConfig(boolean lazyLo when(config.shouldCheckCRCOnSegmentLoad()).thenReturn(true); when(config.isLazyLoadEnabled()).thenReturn(lazyLoadEnabled); when(config.getLazyLoadMaterializeParallelism()).thenReturn(4); + when(config.getLazyLoadMaterializeTimeoutSeconds()).thenReturn(300L); return config; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java index f791d1ec2ea7..fc2550f0f7e2 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/SegmentDataManager.java @@ -33,7 +33,8 @@ public abstract class SegmentDataManager { private final AtomicBoolean _offloaded = new AtomicBoolean(); private final AtomicBoolean _destroyed = new AtomicBoolean(); private int _referenceCount = 1; - // Idle clock for lazy-loading TTL eviction. Stamped on every successful query acquire. + // Idle clock for lazy-loading TTL eviction. Stamped on every successful acquire AND release, so a query running + // longer than the eviction TTL does not leave its segment instantly evictable the moment it finishes. private volatile long _lastAccessTimeMs = _loadTimeMs; public long getLoadTimeMs() { @@ -69,6 +70,7 @@ public synchronized boolean increaseReferenceCount() { * @return Whether the segment can be destroyed (i.e. reference count is 0) */ public synchronized boolean decreaseReferenceCount() { + _lastAccessTimeMs = System.currentTimeMillis(); if (_referenceCount <= 1) { _referenceCount = 0; return true; diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java index af1295d7d133..1cb04351b189 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java @@ -131,6 +131,8 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig private static final int DEFAULT_LAZY_SWEEP_INTERVAL_SECONDS = 60; private static final String LAZY_MATERIALIZE_PARALLELISM = "lazy.materialize.parallelism"; private static final int DEFAULT_LAZY_MATERIALIZE_PARALLELISM = 4; + private static final String LAZY_MATERIALIZE_TIMEOUT_SECONDS = "lazy.materialize.timeout.seconds"; + private static final long DEFAULT_LAZY_MATERIALIZE_TIMEOUT_SECONDS = 60; private final static String[] REQUIRED_KEYS = {INSTANCE_ID}; private static final long DEFAULT_ERROR_CACHE_SIZE = 100L; @@ -383,4 +385,9 @@ public int getLazyLoadSweepIntervalSeconds() { public int getLazyLoadMaterializeParallelism() { return _serverConfig.getProperty(LAZY_MATERIALIZE_PARALLELISM, DEFAULT_LAZY_MATERIALIZE_PARALLELISM); } + + @Override + public long getLazyLoadMaterializeTimeoutSeconds() { + return _serverConfig.getProperty(LAZY_MATERIALIZE_TIMEOUT_SECONDS, DEFAULT_LAZY_MATERIALIZE_TIMEOUT_SECONDS); + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java index bb556ef85ae8..19c4c17164a1 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java @@ -114,4 +114,13 @@ default int getLazyLoadSweepIntervalSeconds() { default int getLazyLoadMaterializeParallelism() { return 4; } + + /** + * Maximum time in seconds a multi-segment query waits for its stubbed segments to materialize before reporting + * the still-stubbed ones as missing (their downloads continue in the background for subsequent queries). Brokers + * typically give up well before long downloads finish, so waiting longer mostly pins server query workers. + */ + default long getLazyLoadMaterializeTimeoutSeconds() { + return 60; + } }