OAK-11155: Allow using a custom segment loading strategy

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/file/AbstractFileStore.java

@@ -37,7 +37,6 @@
 import org.apache.jackrabbit.oak.segment.Segment;
 import org.apache.jackrabbit.oak.segment.Segment.RecordConsumer;
 import org.apache.jackrabbit.oak.segment.SegmentBlob;
-import org.apache.jackrabbit.oak.segment.SegmentBufferMonitor;
 import org.apache.jackrabbit.oak.segment.SegmentCache;
 import org.apache.jackrabbit.oak.segment.SegmentId;
 import org.apache.jackrabbit.oak.segment.SegmentIdFactory;
@@ -50,7 +49,6 @@
 import org.apache.jackrabbit.oak.segment.SegmentWriter;
 import org.apache.jackrabbit.oak.segment.file.tar.EntryRecovery;
 import org.apache.jackrabbit.oak.segment.file.tar.GCGeneration;
-import org.apache.jackrabbit.oak.segment.file.tar.TarFiles;
 import org.apache.jackrabbit.oak.segment.file.tar.TarRecovery;
 import org.apache.jackrabbit.oak.segment.spi.monitor.IOMonitor;
 import org.apache.jackrabbit.oak.segment.spi.monitor.RemoteStoreMonitor;
@@ -124,15 +122,15 @@ public void recoverEntry(UUID uuid, byte[] data, EntryRecovery entryRecovery) th
 
     };
 
-    @NotNull
-    private final SegmentBufferMonitor segmentBufferMonitor;
-
     protected final IOMonitor ioMonitor;
 
     protected final RemoteStoreMonitor remoteStoreMonitor;
 
     protected final int binariesInlineThreshold;
 
+    @NotNull
+    private final SegmentLoader segmentLoader;
+
     AbstractFileStore(final FileStoreBuilder builder) {
         this.directory = builder.getDirectory();
         this.tracker = new SegmentTracker(new SegmentIdFactory() {
@@ -154,8 +152,8 @@ public SegmentId newSegmentId(long msb, long lsb) {
         this.offHeapAccess = builder.getOffHeapAccess();
         this.ioMonitor = builder.getIOMonitor();
         this.remoteStoreMonitor = builder.getRemoteStoreMonitor();
-        this.segmentBufferMonitor = new SegmentBufferMonitor(builder.getStatsProvider());
         this.binariesInlineThreshold = builder.getBinariesInlineThreshold();
+        this.segmentLoader = builder.getSegmentLoader();
     }
 
     static SegmentNotFoundException asSegmentNotFoundException(Exception e, SegmentId id) {
@@ -192,7 +190,12 @@ public SegmentReader getReader() {
     public SegmentIdProvider getSegmentIdProvider() {
         return tracker;
     }
-
+
+    @NotNull
+    public SegmentLoader getSegmentLoader() {
+        return segmentLoader;
+    }
+
     public int getBinariesInlineThreshold() {
         return binariesInlineThreshold;
     }
@@ -292,15 +295,6 @@ static void closeAndLogOnFail(Closeable closeable) {
         }
     }
 
-    Segment readSegmentUncached(TarFiles tarFiles, SegmentId id) {
-        Buffer buffer = tarFiles.readSegment(id.getMostSignificantBits(), id.getLeastSignificantBits());
-        if (buffer == null) {
-            throw new SegmentNotFoundException(id);
-        }
-        segmentBufferMonitor.trackAllocation(buffer);
-        return new Segment(tracker, id, buffer);
-    }
-
     /**
      * Finds all external blob references that are currently accessible
      * in this repository and adds them to the given collector. Useful

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/file/FileStore.java

@@ -509,7 +509,7 @@ public boolean containsSegment(SegmentId id) {
     @NotNull
     public Segment readSegment(final SegmentId id) {
         try (ShutDownCloser ignored = shutDown.keepAlive()) {
-            return segmentCache.getSegment(id, () -> readSegmentUncached(tarFiles, id));
+            return segmentCache.getSegment(id, () -> getSegmentLoader().loadSegment(tarFiles, id, tracker, getWriter()));
         } catch (ExecutionException | UncheckedExecutionException e) {
             if (e.getCause() instanceof RepositoryNotReachableException) {
                 RepositoryNotReachableException re = (RepositoryNotReachableException) e.getCause();

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/file/FileStoreBuilder.java

@@ -35,11 +35,14 @@
 import java.util.HashSet;
 import java.util.Set;
 
+import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.segment.CacheWeights.NodeCacheWeigher;
 import org.apache.jackrabbit.oak.segment.CacheWeights.StringCacheWeigher;
 import org.apache.jackrabbit.oak.segment.CacheWeights.TemplateCacheWeigher;
 import org.apache.jackrabbit.oak.segment.RecordCache;
 import org.apache.jackrabbit.oak.segment.Segment;
+import org.apache.jackrabbit.oak.segment.SegmentBufferMonitor;
+import org.apache.jackrabbit.oak.segment.SegmentNotFoundException;
 import org.apache.jackrabbit.oak.segment.SegmentNotFoundExceptionListener;
 import org.apache.jackrabbit.oak.segment.WriterCacheManager;
 import org.apache.jackrabbit.oak.segment.compaction.SegmentGCOptions;
@@ -143,6 +146,15 @@ public void compactionFailed(@NotNull GCGeneration failedGeneration) {
 
     private boolean built;
 
+    /**
+     * The segment loader to use.
+     *
+     * <p>
+     * The default value is {@code null}, for which a default implementation is used (see {@link #getSegmentLoader()}.
+     */
+    @Nullable
+    private SegmentLoader segmentLoader;
+
     /**
      * Create a new instance of a {@code FileStoreBuilder} for a file store.
      *
@@ -408,6 +420,17 @@ public FileStoreBuilder withBinariesInlineThreshold(int binariesInlineThreshold)
         return this;
     }
 
+    /**
+     * Sets the segment loader to use.
+     * @param segmentLoader a segment loader,
+     *                      or {@code null} to use the default implementation (see {@link #getSegmentLoader()}).
+     * @return this builder.
+     */
+    public FileStoreBuilder withSegmentLoader(@Nullable SegmentLoader segmentLoader) {
+        this.segmentLoader = segmentLoader;
+        return this;
+    }
+
     public Backend buildProcBackend(AbstractFileStore fileStore) throws IOException {
         return new FileStoreProcBackend(fileStore, persistence);
     }
@@ -565,6 +588,29 @@ public WriterCacheManager getCacheManager() {
         return cacheManager;
     }
 
+    /**
+     * Returns the segment loader to use.
+     *
+     * <p>
+     * If not set through {@link #withSegmentLoader(SegmentLoader)}, a default implementation is used, which tracks the
+     * buffers allocated through the {@link StatisticsProvider} given to {@link #withStatisticsProvider(StatisticsProvider)}.
+     */
+    public @NotNull SegmentLoader getSegmentLoader() {
+        if (segmentLoader != null) {
+            return segmentLoader;
+        }
+
+        var segmentBufferMonitor = new SegmentBufferMonitor(statsProvider);
+        return (tarFiles, id, segmentIdProvider, segmentWriter) -> {
+            Buffer buffer = tarFiles.readSegment(id.getMostSignificantBits(), id.getLeastSignificantBits());
+            if (buffer == null) {
+                throw new SegmentNotFoundException(id);
+            }
+            segmentBufferMonitor.trackAllocation(buffer);
+            return new Segment(segmentIdProvider, id, buffer);
+        };
+    }
+
     IOMonitor getIOMonitor() {
         return ioMonitors.isEmpty()
                 ? new IOMonitorAdapter()

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/file/ReadOnlyFileStore.java

@@ -120,12 +120,7 @@ public boolean containsSegment(SegmentId id) {
     @NotNull
     public Segment readSegment(final SegmentId id) {
         try {
-            return segmentCache.getSegment(id, new Callable<Segment>() {
-                @Override
-                public Segment call() throws Exception {
-                    return readSegmentUncached(tarFiles, id);
-                }
-            });
+            return segmentCache.getSegment(id, () -> getSegmentLoader().loadSegment(tarFiles, id, tracker, getWriter()));
         } catch (ExecutionException | UncheckedExecutionException e) {
             throw asSegmentNotFoundException(e, id);
         }

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/file/SegmentLoader.java

@@ -0,0 +1,51 @@
+/*
+ * 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.jackrabbit.oak.segment.file;
+
+import org.apache.jackrabbit.oak.segment.Segment;
+import org.apache.jackrabbit.oak.segment.SegmentId;
+import org.apache.jackrabbit.oak.segment.SegmentIdProvider;
+import org.apache.jackrabbit.oak.segment.SegmentNotFoundException;
+import org.apache.jackrabbit.oak.segment.SegmentWriter;
+import org.apache.jackrabbit.oak.segment.file.tar.TarFiles;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Loads segments that were not previously in the cache from the tar files.
+ */
+@FunctionalInterface
+public interface SegmentLoader {
+    /**
+     * Loads a segment
+     *
+     * @param tarFiles   The tar files
+     * @param id         The ID of the segment
+     * @param idProvider Provides {@link SegmentId} instances
+     * @param writer     The object used to write segments to the tar files (can be used to retrieve a segment that is
+     *                   still being written to, possibly by flushing it).
+     * @return The segment
+     * @throws SegmentNotFoundException if the segment could not be found
+     */
+    @NotNull
+    Segment loadSegment(
+            @NotNull TarFiles tarFiles,
+            @NotNull SegmentId id,
+            @NotNull SegmentIdProvider idProvider,
+            @NotNull SegmentWriter writer
+    ) throws SegmentNotFoundException;
+}