From 79929e13d29453745e608146840df6f18a93194a Mon Sep 17 00:00:00 2001
From: Ziqi Wang <waziqi@yelp.com>
Date: Wed, 15 Jan 2025 11:39:15 -0800
Subject: [PATCH] enable preload if default FSDirectory is resolved as
 MMapDirectory

---
 .../server/luceneserver/DirectoryFactory.java | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/DirectoryFactory.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/DirectoryFactory.java
index 4bf6cd19b..0160049d1 100644
--- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/DirectoryFactory.java
+++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/DirectoryFactory.java
@@ -110,7 +110,25 @@ public static DirectoryFactory get(final String dirImpl) {
       return new DirectoryFactory() {
         @Override
         public Directory open(Path path, IndexPreloadConfig preloadConfig) throws IOException {
-          return FSDirectory.open(path);
+          Directory directory = FSDirectory.open(path);
+          if (directory instanceof MMapDirectory){
+            MMapDirectory mMapDirectory = (MMapDirectory) directory;
+            // no preloading
+            if (!preloadConfig.getShouldPreload()) {
+              return mMapDirectory;
+            }
+
+            // all preloading
+            if (preloadConfig.getPreloadAll()) {
+              mMapDirectory.setPreload(true);
+              return mMapDirectory;
+            }
+
+            // some preloading
+            mMapDirectory.setPreload(true);
+            return new SplitMMapDirectory(FSLockFactory.getDefault(), mMapDirectory, preloadConfig.getExtensions());
+          }
+          return directory;
         }
       };
     } else if (dirImpl.equals("MMapDirectory")) {