Add configurable mmap prefetch pages support

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java

@@ -140,6 +140,7 @@ public static ImmutableSegment load(File indexDir, IndexLoadingConfig indexLoadi
         .setSegmentTier(indexLoadingConfig.getSegmentTier())
         .setInstanceTierConfigs(indexLoadingConfig.getInstanceTierConfigs())
         .setSegmentCustomConfigs(zkMetadata != null ? zkMetadata.getCustomMap() : Map.of())
+        .setMaxMmapPrefetchBytes(indexLoadingConfig.getMaxMmapPrefetchBytes())
         .build();
     if (needPreprocess) {
       // Probe with the default (non-tier-aware) loader so this check never physically moves the segment across
@@ -340,6 +341,7 @@ private static void preprocessSegment(File indexDir, String segmentName, long se
         .setSegmentName(segmentName)
         .setSegmentCrc(segmentCrc)
         .setSegmentCustomConfigs(zkMetadata != null ? zkMetadata.getCustomMap() : Map.of())
+        .setMaxMmapPrefetchBytes(indexLoadingConfig.getMaxMmapPrefetchBytes())
         .build();
     SegmentDirectory segmentDirectory =
         SegmentDirectoryLoaderRegistry.getDefaultSegmentDirectoryLoader().load(indexDir.toURI(), segmentLoaderContext);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/IndexLoadingConfig.java

@@ -51,6 +51,7 @@
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.data.OpenStructNaming;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.ReadMode;
 import org.apache.pinot.spi.utils.TimestampIndexUtils;
 
@@ -94,6 +95,7 @@ private static final class ImmutableState {
     private final String _segmentStoreURI;
     @Nullable
     private final String _segmentDirectoryLoader;
+    private final long _maxMmapPrefetchBytes;
     @Nullable
     private final Map<String, Map<String, String>> _instanceTierConfigs;
     private final List<String> _sortedColumns;
@@ -114,6 +116,7 @@ private ImmutableState(@Nullable InstanceDataManagerConfig instanceDataManagerCo
       SegmentVersion segmentVersion = null;
       String segmentStoreURI = null;
       String segmentDirectoryLoader = null;
+      long maxMmapPrefetchBytes = CommonConstants.Server.DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES;
       Map<String, Map<String, String>> instanceTierConfigs = null;
       if (instanceDataManagerConfig != null) {
         ReadMode instanceReadMode = instanceDataManagerConfig.getReadMode();
@@ -133,6 +136,7 @@ private ImmutableState(@Nullable InstanceDataManagerConfig instanceDataManagerCo
         }
         segmentStoreURI = instanceDataManagerConfig.getSegmentStoreUri();
         segmentDirectoryLoader = instanceDataManagerConfig.getSegmentDirectoryLoader();
+        maxMmapPrefetchBytes = instanceDataManagerConfig.getMaxMmapPrefetchBytes();
         Map<String, Map<String, String>> tierConfigs = instanceDataManagerConfig.getTierConfigs();
         instanceTierConfigs = tierConfigs != null ? tierConfigs : Map.of();
       }
@@ -177,6 +181,7 @@ private ImmutableState(@Nullable InstanceDataManagerConfig instanceDataManagerCo
       _realtimeAvgMultiValueCount = realtimeAvgMultiValueCount;
       _segmentStoreURI = segmentStoreURI;
       _segmentDirectoryLoader = segmentDirectoryLoader;
+      _maxMmapPrefetchBytes = maxMmapPrefetchBytes;
       _instanceTierConfigs = instanceTierConfigs;
       _sortedColumns = sortedColumns;
       _columnMinMaxValueGeneratorMode = columnMinMaxValueGeneratorMode;
@@ -378,6 +383,10 @@ public String getSegmentDirectoryLoader() {
         : SegmentDirectoryLoaderRegistry.DEFAULT_SEGMENT_DIRECTORY_LOADER_NAME;
   }
 
+  public long getMaxMmapPrefetchBytes() {
+    return _immutableState._maxMmapPrefetchBytes;
+  }
+
   public String getInstanceId() {
     return _immutableState._instanceId;
   }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectory.java

@@ -53,8 +53,6 @@ public class SegmentLocalFSDirectory extends SegmentDirectory {
 
   // matches most systems
   private static final int PAGE_SIZE_BYTES = 4096;
-  // Prefetch limit...arbitrary but related to common server memory and data size profiles
-  private static final long MAX_MMAP_PREFETCH_PAGES = 100 * 1024 * 1024 * 1024L / PAGE_SIZE_BYTES;
   private static final double PREFETCH_SLOWDOWN_PCT = 0.67;
   private static final AtomicLong PREFETCHED_PAGES = new AtomicLong(0);
 
@@ -63,6 +61,7 @@ public class SegmentLocalFSDirectory extends SegmentDirectory {
   private final ReadMode _readMode;
   private final SegmentDirectoryLoaderContext _segmentDirectoryLoaderContext;
   private final SegmentLock _segmentLock = new SegmentLock();
+  private final long _maxMmapPrefetchPages;
   private SegmentMetadataImpl _segmentMetadata;
   private ColumnIndexDirectory _columnIndexDirectory;
   private StarTreeIndexReader _starTreeIndexReader;
@@ -76,6 +75,7 @@ public SegmentLocalFSDirectory(File indexDir) {
     _segmentDirectory = null;
     _readMode = null;
     _segmentDirectoryLoaderContext = null;
+    _maxMmapPrefetchPages = toPrefetchPages(CommonConstants.Server.DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES);
   }
 
   public SegmentLocalFSDirectory(File indexDir, ReadMode readMode)
@@ -97,6 +97,7 @@ public SegmentLocalFSDirectory(File indexDir, SegmentMetadataImpl metadata, Read
     _segmentMetadata = metadata;
     _readMode = readMode;
     _segmentDirectoryLoaderContext = segmentDirectoryLoaderContext;
+    _maxMmapPrefetchPages = toPrefetchPages(getMaxMmapPrefetchBytes(segmentDirectoryLoaderContext));
 
     try {
       load();
@@ -299,6 +300,17 @@ public void close()
     }
   }
 
+  private static long toPrefetchPages(long maxMmapPrefetchBytes) {
+    Preconditions.checkArgument(maxMmapPrefetchBytes >= 0, "Max mmap prefetch bytes must be non-negative, got: %s",
+        maxMmapPrefetchBytes);
+    return maxMmapPrefetchBytes / PAGE_SIZE_BYTES;
+  }
+
+  private static long getMaxMmapPrefetchBytes(@Nullable SegmentDirectoryLoaderContext segmentDirectoryLoaderContext) {
+    return segmentDirectoryLoaderContext == null ? CommonConstants.Server.DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES
+        : segmentDirectoryLoaderContext.getMaxMmapPrefetchBytes();
+  }
+
   private PinotDataBuffer getIndexForColumn(String column, IndexType<?, ?, ?> type)
       throws IOException {
     PinotDataBuffer buffer;
@@ -331,11 +343,11 @@ private void prefetchMmapData(PinotDataBuffer buffer) {
     // an optimization.
 
     // Prefetch limit and slowdown percentage are arbitrary
-    if (PREFETCHED_PAGES.get() >= MAX_MMAP_PREFETCH_PAGES) {
+    if (PREFETCHED_PAGES.get() >= _maxMmapPrefetchPages) {
       return;
     }
 
-    final long prefetchSlowdownPageLimit = (long) (PREFETCH_SLOWDOWN_PCT * MAX_MMAP_PREFETCH_PAGES);
+    final long prefetchSlowdownPageLimit = (long) (PREFETCH_SLOWDOWN_PCT * _maxMmapPrefetchPages);
     if (PREFETCHED_PAGES.get() >= prefetchSlowdownPageLimit) {
       if (0 < buffer.size()) {
         buffer.getByte(0);

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectoryTest.java

@@ -25,6 +25,7 @@
 import org.apache.pinot.segment.spi.creator.SegmentVersion;
 import org.apache.pinot.segment.spi.index.StandardIndexes;
 import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.segment.spi.loader.SegmentDirectoryLoaderContext;
 import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
 import org.apache.pinot.segment.spi.store.SegmentDirectory;
 import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
@@ -121,6 +122,68 @@ public void testWriteAndReadBackData()
     }
   }
 
+  private static SegmentDirectoryLoaderContext prefetchLoaderContext(long maxMmapPrefetchBytes) {
+    return new SegmentDirectoryLoaderContext.Builder()
+        .setMaxMmapPrefetchBytes(maxMmapPrefetchBytes)
+        .build();
+  }
+
+  /// Prefetching is bounded by a JVM-wide page counter that tests cannot reset, so rather than asserting on how many
+  /// pages got faulted in, these cases pin down the config contract: the byte limit is accepted, zero disables
+  /// prefetching, and reads stay correct either way.
+  @Test
+  public void testPrefetchLimitDisabledStillReadsData()
+      throws Exception {
+    File prefetchDir = new File(SegmentLocalFSDirectoryTest.class.getName() + "-prefetch_disabled");
+    FileUtils.deleteQuietly(prefetchDir);
+    try {
+      FileUtils.copyDirectory(_segmentDirectory.getPath().toFile(), prefetchDir);
+      // 0 bytes disables prefetching entirely
+      try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(prefetchDir, _metadata,
+          ReadMode.mmap, prefetchLoaderContext(0))) {
+        try (SegmentDirectory.Writer writer = segmentDirectory.createWriter()) {
+          PinotDataBuffer buffer = writer.newIndexFor("noPrefetchColumn", StandardIndexes.forward(), 1024);
+          loadData(buffer);
+          writer.save();
+        }
+        try (SegmentDirectory.Reader reader = segmentDirectory.createReader()) {
+          verifyData(reader.getIndexFor("noPrefetchColumn", StandardIndexes.forward()));
+        }
+      }
+    } finally {
+      FileUtils.deleteQuietly(prefetchDir);
+    }
+  }
+
+  @Test
+  public void testSmallPrefetchLimitStillReadsData()
+      throws Exception {
+    File prefetchDir = new File(SegmentLocalFSDirectoryTest.class.getName() + "-prefetch_small");
+    FileUtils.deleteQuietly(prefetchDir);
+    try {
+      FileUtils.copyDirectory(_segmentDirectory.getPath().toFile(), prefetchDir);
+      // 8KB, i.e. a 2 page budget: exercises the slowdown branch that only faults in header pages
+      try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(prefetchDir, _metadata,
+          ReadMode.mmap, prefetchLoaderContext(8 * 1024))) {
+        try (SegmentDirectory.Writer writer = segmentDirectory.createWriter()) {
+          PinotDataBuffer buffer = writer.newIndexFor("smallPrefetchColumn", StandardIndexes.forward(), 1024);
+          loadData(buffer);
+          writer.save();
+        }
+        try (SegmentDirectory.Reader reader = segmentDirectory.createReader()) {
+          verifyData(reader.getIndexFor("smallPrefetchColumn", StandardIndexes.forward()));
+        }
+      }
+    } finally {
+      FileUtils.deleteQuietly(prefetchDir);
+    }
+  }
+
+  @Test(expectedExceptions = IllegalArgumentException.class)
+  public void testNegativePrefetchLimitRejected() {
+    new SegmentLocalFSDirectory(TEST_DIRECTORY, _metadata, ReadMode.mmap, prefetchLoaderContext(-1));
+  }
+
   @Test
   public void testDirectorySize()
       throws Exception {

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/loader/SegmentDirectoryLoaderContext.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.ReadMode;
 
 
@@ -36,10 +37,12 @@ public class SegmentDirectoryLoaderContext {
   private final String _segmentTier;
   private final Map<String, Map<String, String>> _instanceTierConfigs;
   private final Map<String, String> _segmentCustomConfigs;
+  private final long _maxMmapPrefetchBytes;
 
   private SegmentDirectoryLoaderContext(ReadMode readMode, TableConfig tableConfig, Schema schema, String instanceId,
       String tableDataDir, String segmentName, long segmentCrc, String segmentTier,
-      Map<String, Map<String, String>> instanceTierConfigs, Map<String, String> segmentCustomConfigs) {
+      Map<String, Map<String, String>> instanceTierConfigs, Map<String, String> segmentCustomConfigs,
+      long maxMmapPrefetchBytes) {
     _readMode = readMode;
     _tableConfig = tableConfig;
     _schema = schema;
@@ -50,6 +53,7 @@ private SegmentDirectoryLoaderContext(ReadMode readMode, TableConfig tableConfig
     _segmentTier = segmentTier;
     _instanceTierConfigs = instanceTierConfigs;
     _segmentCustomConfigs = segmentCustomConfigs;
+    _maxMmapPrefetchBytes = maxMmapPrefetchBytes;
   }
 
   public ReadMode getReadMode() {
@@ -93,6 +97,12 @@ public Map<String, String> getSegmentCustomConfigs() {
     return _segmentCustomConfigs;
   }
 
+  /// Max amount of mmap'ed segment data to proactively fault into memory on segment load, in bytes. Zero disables
+  /// prefetching. Only applies when [#getReadMode()] is [ReadMode#mmap].
+  public long getMaxMmapPrefetchBytes() {
+    return _maxMmapPrefetchBytes;
+  }
+
   public static class Builder {
     private ReadMode _readMode = ReadMode.DEFAULT_MODE;
     private TableConfig _tableConfig;
@@ -104,6 +114,7 @@ public static class Builder {
     private String _segmentTier;
     private Map<String, Map<String, String>> _instanceTierConfigs;
     private Map<String, String> _segmentCustomConfigs;
+    private long _maxMmapPrefetchBytes = CommonConstants.Server.DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES;
 
     public Builder setReadMode(ReadMode readMode) {
       _readMode = readMode;
@@ -155,9 +166,15 @@ public Builder setSegmentCustomConfigs(Map<String, String> segmentCustomConfigs)
       return this;
     }
 
+    public Builder setMaxMmapPrefetchBytes(long maxMmapPrefetchBytes) {
+      _maxMmapPrefetchBytes = maxMmapPrefetchBytes;
+      return this;
+    }
+
     public SegmentDirectoryLoaderContext build() {
       return new SegmentDirectoryLoaderContext(_readMode, _tableConfig, _schema, _instanceId, _tableDataDir,
-          _segmentName, _segmentCrc, _segmentTier, _instanceTierConfigs, _segmentCustomConfigs);
+          _segmentName, _segmentCrc, _segmentTier, _instanceTierConfigs, _segmentCustomConfigs,
+          _maxMmapPrefetchBytes);
     }
   }
 }

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java

@@ -29,6 +29,7 @@
 import org.apache.pinot.segment.spi.loader.SegmentDirectoryLoaderRegistry;
 import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.DataSizeUtils;
 import org.apache.pinot.spi.utils.ReadMode;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -250,6 +251,11 @@ public boolean isDirectRealtimeOffHeapAllocation() {
   }
 
   @Override
+  public long getMaxMmapPrefetchBytes() {
+    String maxPrefetchSize = _serverConfig.getProperty(MMAP_PREFETCH_MAX_SIZE);
+    return maxPrefetchSize != null ? DataSizeUtils.toBytes(maxPrefetchSize) : DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES;
+  }
+
   public boolean shouldReloadConsumingSegment() {
     return _serverConfig.getProperty(RELOAD_CONSUMING_SEGMENT, DEFAULT_RELOAD_CONSUMING_SEGMENT);
   }

pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java

@@ -59,6 +59,10 @@ public interface InstanceDataManagerConfig {
 
   int getMaxSegmentPreloadThreads();
 
+  /// Max amount of mmap'ed segment data to proactively fault into memory on segment load, in bytes.
+  /// Zero disables prefetching. Only applies when [#getReadMode()] is [ReadMode#mmap].
+  long getMaxMmapPrefetchBytes();
+
   int getMaxParallelSegmentBuilds();
 
   int getMaxParallelSegmentDownloads();

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -1461,6 +1461,14 @@ public static class Server {
     public static final String READ_MODE = "readMode";
     public static final String CONFIG_OF_READ_MODE = INSTANCE_DATA_MANAGER_CONFIG_PREFIX + "." + READ_MODE;
     public static final String DEFAULT_READ_MODE = "mmap";
+    // Upper bound on how much mmap'ed segment data a server proactively faults into memory on load, to reduce the
+    // major page faults (and thus latency variance) seen right after startup. Accepts human readable data sizes
+    // (e.g. '100G'). Set to '0' to disable prefetching entirely. Only applies when readMode is 'mmap'.
+    public static final String MMAP_PREFETCH_MAX_SIZE = "mmap.prefetch.max.size";
+    public static final String CONFIG_OF_MMAP_PREFETCH_MAX_SIZE =
+        INSTANCE_DATA_MANAGER_CONFIG_PREFIX + "." + MMAP_PREFETCH_MAX_SIZE;
+    // 100GB, matching the limit that used to be hard-coded in SegmentLocalFSDirectory
+    public static final long DEFAULT_MMAP_PREFETCH_MAX_SIZE_BYTES = 100L * 1024 * 1024 * 1024;
     public static final String SEGMENT_FORMAT_VERSION = "segment.format.version";
     public static final String CONFIG_OF_SEGMENT_FORMAT_VERSION =
         INSTANCE_DATA_MANAGER_CONFIG_PREFIX + "." + SEGMENT_FORMAT_VERSION;