Add POSIX madvise support and a closeable MappableByteBuffer API to improve mmap-based blob cache access with explicit prefetching

libs/native/src/main/java/org/elasticsearch/nativeaccess/AbstractNativeAccess.java

@@ -15,6 +15,10 @@
 import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
 
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+
 abstract class AbstractNativeAccess implements NativeAccess {
 
     protected static final Logger logger = LogManager.getLogger(NativeAccess.class);
@@ -57,6 +61,12 @@ public CloseableByteBuffer newConfinedBuffer(int len) {
         return javaLib.newConfinedBuffer(len);
     }
 
+    @Override
+    public CloseableMappedByteBuffer map(FileChannel fileChannel, MapMode mode, long position, long size) throws IOException {
+        assert fileChannel != null && position >= 0 && size > 0;
+        return javaLib.map(fileChannel, mode, position, size);
+    }
+
     @Override
     public boolean isMemoryLocked() {
         return isMemoryLocked;

libs/native/src/main/java/org/elasticsearch/nativeaccess/CloseableMappedByteBuffer.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.nativeaccess;
+
+/** A closeable buffer backed by a mapped file. */
+public interface CloseableMappedByteBuffer extends CloseableByteBuffer {
+
+    /**
+     * Returns a slice of this buffer. Closing a slice does not close it's parent.
+     */
+    CloseableMappedByteBuffer slice(long index, long length);
+
+    /**
+     * Prefetches the given offset and length.
+     */
+    void prefetch(long offset, long length);
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/NativeAccess.java

@@ -9,6 +9,9 @@
 
 package org.elasticsearch.nativeaccess;
 
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
 import java.nio.file.Path;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -103,6 +106,13 @@ default WindowsFunctions getWindowsFunctions() {
      */
     CloseableByteBuffer newConfinedBuffer(int len);
 
+    /**
+     * Creates a new {@link CloseableMappedByteBuffer} using a shared arena. The buffer can be used
+     * across multiple threads, and should be closed.
+     * @return the buffer
+     */
+    CloseableMappedByteBuffer map(FileChannel fileChannel, MapMode mode, long position, long size) throws IOException;
+
     /**
      * Possible stats for execution filtering.
      */

libs/native/src/main/java/org/elasticsearch/nativeaccess/NoopNativeAccess.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
+import java.nio.channels.FileChannel;
 import java.nio.file.Path;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -89,6 +90,12 @@ public CloseableByteBuffer newConfinedBuffer(int len) {
         return null;
     }
 
+    @Override
+    public CloseableMappedByteBuffer map(FileChannel fileChannel, FileChannel.MapMode mode, long position, long size) {
+        logger.warn("cannot map because native access is not available");
+        return null;
+    }
+
     @Override
     public Optional<VectorSimilarityFunctions> getVectorSimilarityFunctions() {
         logger.warn("cannot get vector distance because native access is not available");

libs/native/src/main/java/org/elasticsearch/nativeaccess/PosixNativeAccess.java

@@ -10,10 +10,13 @@
 package org.elasticsearch.nativeaccess;
 
 import org.elasticsearch.core.SuppressForbidden;
+import org.elasticsearch.nativeaccess.jdk.PosixCloseableMappedByteBuffer;
 import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 import org.elasticsearch.nativeaccess.lib.PosixCLibrary;
 import org.elasticsearch.nativeaccess.lib.VectorLibrary;
 
+import java.io.IOException;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Optional;
@@ -191,6 +194,11 @@ public Optional<VectorSimilarityFunctions> getVectorSimilarityFunctions() {
         return Optional.ofNullable(vectorDistance);
     }
 
+    @Override
+    public CloseableMappedByteBuffer map(FileChannel fileChannel, FileChannel.MapMode mode, long position, long size) throws IOException {
+        return PosixCloseableMappedByteBuffer.ofShared(fileChannel, mode, position, size);
+    }
+
     String rlimitToString(long value) {
         if (value == constants.RLIMIT_INFINITY()) {
             return "unlimited";

libs/native/src/main/java/org/elasticsearch/nativeaccess/WindowsNativeAccess.java

@@ -9,10 +9,13 @@
 
 package org.elasticsearch.nativeaccess;
 
+import org.elasticsearch.nativeaccess.jdk.JdkCloseableMappedByteBuffer;
 import org.elasticsearch.nativeaccess.lib.Kernel32Library;
 import org.elasticsearch.nativeaccess.lib.Kernel32Library.Handle;
 import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 
+import java.io.IOException;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Optional;
@@ -171,4 +174,9 @@ public WindowsFunctions getWindowsFunctions() {
     public Optional<VectorSimilarityFunctions> getVectorSimilarityFunctions() {
         return Optional.empty(); // not supported yet
     }
+
+    @Override
+    public CloseableMappedByteBuffer map(FileChannel fileChannel, FileChannel.MapMode mode, long position, long size) throws IOException {
+        return JdkCloseableMappedByteBuffer.ofShared(fileChannel, mode, position, size);
+    }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkCloseableMappedByteBuffer.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import org.elasticsearch.nativeaccess.CloseableMappedByteBuffer;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+import java.util.Objects;
+
+public class JdkCloseableMappedByteBuffer implements CloseableMappedByteBuffer {
+
+    private final Arena arena;
+    protected final MemorySegment segment;
+    private final MappedByteBuffer bufferView;
+
+    public static JdkCloseableMappedByteBuffer ofShared(FileChannel fileChannel, MapMode mode, long position, long size)
+        throws IOException {
+        var arena = Arena.ofShared();
+        var seg = fileChannel.map(mode, position, size, arena);
+        return new JdkCloseableMappedByteBuffer(seg, arena);
+    }
+
+    protected JdkCloseableMappedByteBuffer(MemorySegment seg, Arena arena) {
+        this.arena = arena;
+        this.segment = seg;
+        this.bufferView = (MappedByteBuffer) seg.asByteBuffer();
+    }
+
+    @Override
+    public MappedByteBuffer buffer() {
+        return bufferView;
+    }
+
+    @Override
+    public void close() {
+        if (arena != null) {
+            arena.close();
+        }
+    }
+
+    @Override
+    public CloseableMappedByteBuffer slice(long index, long length) {
+        var slice = segment.asSlice(index, length);
+        return new JdkCloseableMappedByteBuffer(slice, null); // closing a slice does not close the parent.
+    }
+
+    @Override
+    public void prefetch(long offset, long length) {
+        Objects.checkFromIndexSize(offset, length, segment.byteSize());
+        // no explicit action, override in subclass if needed.
+    }
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkJavaLibrary.java

@@ -10,8 +10,12 @@
 package org.elasticsearch.nativeaccess.jdk;
 
 import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+import org.elasticsearch.nativeaccess.CloseableMappedByteBuffer;
 import org.elasticsearch.nativeaccess.lib.JavaLibrary;
 
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+
 class JdkJavaLibrary implements JavaLibrary {
 
     @Override
@@ -23,4 +27,9 @@ public CloseableByteBuffer newSharedBuffer(int len) {
     public CloseableByteBuffer newConfinedBuffer(int len) {
         return JdkCloseableByteBuffer.ofConfined(len);
     }
+
+    @Override
+    public CloseableMappedByteBuffer map(FileChannel fileChannel, FileChannel.MapMode mode, long position, long size) throws IOException {
+        return JdkCloseableMappedByteBuffer.ofShared(fileChannel, mode, position, size);
+    }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkPosixCLibrary.java

@@ -23,6 +23,8 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
+import java.lang.ref.Reference;
+import java.util.Objects;
 
 import static java.lang.foreign.MemoryLayout.PathElement.groupElement;
 import static java.lang.foreign.ValueLayout.ADDRESS;
@@ -37,6 +39,8 @@ class JdkPosixCLibrary implements PosixCLibrary {
 
     private static final Logger logger = LogManager.getLogger(JdkPosixCLibrary.class);
 
+    private static final int PAGE_SIZE;
+
     // errno can change between system calls, so we capture it
     private static final StructLayout CAPTURE_ERRNO_LAYOUT = Linker.Option.captureStateLayout();
     static final Linker.Option CAPTURE_ERRNO_OPTION = Linker.Option.captureCallState("errno");
@@ -53,6 +57,10 @@ class JdkPosixCLibrary implements PosixCLibrary {
         FunctionDescriptor.of(JAVA_INT, JAVA_INT, ADDRESS)
     );
     private static final MethodHandle mlockall$mh = downcallHandleWithErrno("mlockall", FunctionDescriptor.of(JAVA_INT, JAVA_INT));
+    private static final MethodHandle madvise$mh = downcallHandleWithErrno(
+        "madvise",
+        FunctionDescriptor.of(JAVA_INT, JAVA_LONG, JAVA_LONG, JAVA_INT)
+    );
     private static final MethodHandle fcntl$mh = downcallHandle(
         "fcntl",
         FunctionDescriptor.of(JAVA_INT, JAVA_INT, JAVA_INT, ADDRESS),
@@ -92,6 +100,12 @@ class JdkPosixCLibrary implements PosixCLibrary {
             );
         }
         fstat$mh = fstat;
+
+        try {
+            PAGE_SIZE = (int) downcallHandle("getpagesize", FunctionDescriptor.of(JAVA_INT)).invokeExact();
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
     }
     private static final MethodHandle socket$mh = downcallHandleWithErrno(
         "socket",
@@ -182,6 +196,33 @@ public int mlockall(int flags) {
         }
     }
 
+    @Override
+    public int madvise(MemorySegment segment, long offset, long length, int advice) {
+        if (segment.isNative() == false) {
+            throw new IllegalArgumentException("unexpected non-native segment: " + segment);
+        }
+        Objects.checkFromIndexSize(offset, length, segment.byteSize());
+        long base = segment.address() + offset;
+        try {
+            return (int) madvise$mh.invokeExact(errnoState, base, length, advice);
+        } catch (Throwable t) {
+            throw madviseError(t, segment);
+        } finally {
+            // protects the segment from being potentially being GC'ed during out downcall
+            Reference.reachabilityFence(segment);
+        }
+    }
+
+    static Error madviseError(Throwable t, MemorySegment segment) {
+        String msg = "madvise failed: segment=" + segment + ", scope=" + segment.scope() + ", isAlive=" + segment.scope().isAlive();
+        return new AssertionError(msg, t);
+    }
+
+    @Override
+    public int getPageSize() {
+        return PAGE_SIZE;
+    }
+
     @Override
     public int fcntl(int fd, int cmd, FStore fst) {
         assert fst instanceof JdkFStore;

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/PosixCloseableMappedByteBuffer.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
+import org.elasticsearch.nativeaccess.lib.PosixCLibrary;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+import java.util.Objects;
+
+public class PosixCloseableMappedByteBuffer extends JdkCloseableMappedByteBuffer {
+
+    static final PosixCLibrary LIB = NativeLibraryProvider.instance().getLibrary(PosixCLibrary.class);
+    static final int PAGE_SIZE = LIB.getPageSize();
+
+    public static PosixCloseableMappedByteBuffer ofShared(FileChannel fileChannel, MapMode mode, long position, long size)
+        throws IOException {
+        var arena = Arena.ofShared();
+        var seg = fileChannel.map(mode, position, size, arena);
+        return new PosixCloseableMappedByteBuffer(seg, arena);
+    }
+
+    protected PosixCloseableMappedByteBuffer(MemorySegment seg, Arena arena) {
+        super(seg, arena);
+    }
+
+    @Override
+    public void prefetch(long offset, long length) {
+        Objects.checkFromIndexSize(offset, length, segment.byteSize());
+        // Align offset with the page size, this is required for madvise.
+        // Compute the offset of the current position in the OS's page.
+        final long offsetInPage = (segment.address() + offset) % PAGE_SIZE;
+        offset -= offsetInPage;
+        length += offsetInPage;
+        if (offset < 0) {
+            // start of the page is before the start of this segment, ignore the first page.
+            offset += PAGE_SIZE;
+            length -= PAGE_SIZE;
+            if (length <= 0) {
+                // This segment has no data beyond the first page.
+                return;
+            }
+        }
+        int ret = LIB.madvise(segment, offset, length, PosixCLibrary.POSIX_MADV_WILLNEED);
+        if (ret != 0) {
+            int errno = LIB.errno();
+            throw new RuntimeException("madvise failed with (error=" + errno + "): " + LIB.strerror(errno));
+        }
+    }
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/JavaLibrary.java

@@ -10,9 +10,17 @@
 package org.elasticsearch.nativeaccess.lib;
 
 import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+import org.elasticsearch.nativeaccess.CloseableMappedByteBuffer;
+
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
 
 public non-sealed interface JavaLibrary extends NativeLibrary {
     CloseableByteBuffer newSharedBuffer(int len);
 
     CloseableByteBuffer newConfinedBuffer(int len);
+
+    CloseableMappedByteBuffer map(FileChannel fileChannel, MapMode mode, long position, long size) throws IOException;
+
 }