Extract more code from frozen to shared blob cache module (#93786)

Extracting effectively all the read path logic to the shared module for
reuse and encapsulate the magic around the buffer permits into a
specific class to make the API nicer.

Author: original-brownbear

x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/common/ByteBufferReference.java

@@ -0,0 +1,70 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.blobcache.common;
+
+import org.elasticsearch.core.Nullable;
+
+import java.nio.ByteBuffer;
+import java.util.concurrent.Semaphore;
+
+/**
+ * Wrapper around a {@link ByteBuffer} that allows making slices of it available for writing to multiple threads concurrently in a safe
+ * manner. Used to populate a read buffer from Lucene concurrently across multiple threads.
+ */
+public final class ByteBufferReference {
+
+    private final Semaphore permits = new Semaphore(Integer.MAX_VALUE);
+
+    private ByteBuffer buffer;
+
+    /**
+     * @param buffer to guard against concurrent manipulations
+     */
+    public ByteBufferReference(ByteBuffer buffer) {
+        this.buffer = buffer;
+    }
+
+    /**
+     * Acquires a permit that must be released via {@link #release()} and returns a {@link ByteBuffer} slice of {@link #buffer} that may be
+     * written to, or {@code null} if {@link #finish(int)} has already been called on this instance and no more writing to the buffer is
+     * allowed.
+     * @param position position relative to the position of {@link #buffer} to start slice at
+     * @param length slice length
+     * @return slice of {@link #buffer} to write to or {@code null} if already closed
+     */
+    @Nullable
+    public ByteBuffer tryAcquire(int position, int length) {
+        if (permits.tryAcquire() == false) {
+            return null;
+        }
+        return buffer.slice(buffer.position() + position, length);
+    }
+
+    /**
+     * Release a permit acquired through {@link #tryAcquire(int, int)}.
+     */
+    public void release() {
+        permits.release();
+    }
+
+    /**
+     * Acquire all permits and then safely advance the position of {@link #buffer} by the given number of bytes.
+     *
+     * @param bytesRead number of bytes to advance the current position of {@link #buffer} by.
+     * @throws Exception on failure
+     */
+    public void finish(int bytesRead) throws Exception {
+        if (buffer != null) {
+            assert bytesRead == 0 || permits.availablePermits() == Integer.MAX_VALUE
+                : "Try to finish [" + bytesRead + "] but only had [" + permits.availablePermits() + "] permits available.";
+            permits.acquire(Integer.MAX_VALUE);
+            buffer.position(buffer.position() + bytesRead); // mark all bytes as accounted for
+            buffer = null;
+        }
+    }
+}

x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBytes.java

@@ -10,6 +10,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.blobcache.BlobCacheUtils;
+import org.elasticsearch.blobcache.common.ByteBufferReference;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.core.AbstractRefCounted;
@@ -164,6 +165,47 @@ private static int positionalWrite(IO fc, long start, ByteBuffer byteBuffer) thr
         return written;
     }
 
+    /**
+     * Read {@code length} bytes from given shared bytes at given {@code channelPos} into {@code byteBufferReference} at given
+     * {@code relativePos}.
+     * @param fc shared bytes channel to read from
+     * @param channelPos position in {@code fc} to read from
+     * @param relativePos position in {@code byteBufferReference}
+     * @param length number of bytes to read
+     * @param byteBufferReference buffer reference
+     * @param cacheFile cache file reference used for exception messages only
+     * @return number of bytes read
+     * @throws IOException on failure
+     */
+    public static int readCacheFile(
+        final IO fc,
+        long channelPos,
+        long relativePos,
+        long length,
+        final ByteBufferReference byteBufferReference,
+        Object cacheFile
+    ) throws IOException {
+        if (length == 0L) {
+            return 0;
+        }
+        final int bytesRead;
+        final ByteBuffer dup = byteBufferReference.tryAcquire(Math.toIntExact(relativePos), Math.toIntExact(length));
+        if (dup != null) {
+            try {
+                bytesRead = fc.read(dup, channelPos);
+                if (bytesRead == -1) {
+                    BlobCacheUtils.throwEOF(channelPos, dup.remaining(), cacheFile);
+                }
+            } finally {
+                byteBufferReference.release();
+            }
+        } else {
+            // return fake response
+            return Math.toIntExact(length);
+        }
+        return bytesRead;
+    }
+
     @Override
     protected void closeInternal() {
         try {

x-pack/plugin/searchable-snapshots/src/main/java/org/elasticsearch/xpack/searchablesnapshots/store/input/FrozenIndexInput.java

@@ -10,7 +10,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.store.IOContext;
-import org.elasticsearch.blobcache.BlobCacheUtils;
+import org.elasticsearch.blobcache.common.ByteBufferReference;
 import org.elasticsearch.blobcache.common.ByteRange;
 import org.elasticsearch.blobcache.shared.SharedBlobCacheService;
 import org.elasticsearch.blobcache.shared.SharedBytes;
@@ -21,10 +21,8 @@
 import org.elasticsearch.xpack.searchablesnapshots.store.IndexInputStats;
 import org.elasticsearch.xpack.searchablesnapshots.store.SearchableSnapshotDirectory;
 
-import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
-import java.util.concurrent.Semaphore;
 
 public class FrozenIndexInput extends MetadataCachingIndexInput {
 
@@ -104,13 +102,10 @@ protected long getDefaultRangeSize() {
     protected void readWithoutBlobCache(ByteBuffer b) throws Exception {
         final long position = getAbsolutePosition();
         final int length = b.remaining();
-        final int originalByteBufPosition = b.position();
-
         // Semaphore that, when all permits are acquired, ensures that async callbacks (such as those used by readCacheFile) are not
         // accessing the byte buffer anymore that was passed to readWithoutBlobCache
         // In particular, it's important to acquire all permits before adapting the ByteBuffer's offset
-        final Semaphore luceneByteBufPermits = new Semaphore(Integer.MAX_VALUE);
-        boolean bufferWriteLocked = false;
+        final ByteBufferReference byteBufferReference = new ByteBufferReference(b);
         logger.trace("readInternal: read [{}-{}] from [{}]", position, position + length, this);
         try {
             final ByteRange startRangeToWrite = computeRange(position);
@@ -122,99 +117,50 @@ protected void readWithoutBlobCache(ByteBuffer b) throws Exception {
                 : "[" + position + "-" + (position + length) + "] vs " + rangeToWrite;
             final ByteRange rangeToRead = ByteRange.of(position, position + length);
 
-            final int bytesRead = cacheFile.populateAndRead(
-                rangeToWrite,
-                rangeToRead,
-                (channel, pos, relativePos, len) -> readCacheFile(
-                    channel,
+            final int bytesRead = cacheFile.populateAndRead(rangeToWrite, rangeToRead, (channel, pos, relativePos, len) -> {
+                logger.trace(
+                    "{}: reading logical {} channel {} pos {} length {} (details: {})",
+                    fileInfo.physicalName(),
+                    rangeToRead.start(),
                     pos,
                     relativePos,
-                    len,
-                    b,
-                    rangeToRead.start(),
-                    luceneByteBufPermits
-                ),
-                (channel, channelPos, relativePos, len, progressUpdater) -> {
-                    final long startTimeNanos = stats.currentTimeNanos();
-                    try (InputStream input = openInputStreamFromBlobStore(rangeToWrite.start() + relativePos, len)) {
-                        assert ThreadPool.assertCurrentThreadPool(SearchableSnapshots.CACHE_FETCH_ASYNC_THREAD_POOL_NAME);
-                        logger.trace(
-                            "{}: writing channel {} pos {} length {} (details: {})",
-                            fileInfo.physicalName(),
-                            channelPos,
-                            relativePos,
-                            len,
-                            cacheFile
-                        );
-                        SharedBytes.copyToCacheFileAligned(
-                            channel,
-                            input,
-                            channelPos,
-                            relativePos,
-                            len,
-                            progressUpdater,
-                            writeBuffer.get().clear(),
-                            cacheFile
-                        );
-                        final long endTimeNanos = stats.currentTimeNanos();
-                        stats.addCachedBytesWritten(len, endTimeNanos - startTimeNanos);
-                    }
-                },
-                SearchableSnapshots.CACHE_FETCH_ASYNC_THREAD_POOL_NAME
-            );
+                    length,
+                    cacheFile
+                );
+                final int read = SharedBytes.readCacheFile(channel, pos, relativePos, len, byteBufferReference, cacheFile);
+                stats.addCachedBytesRead(read);
+                return read;
+            }, (channel, channelPos, relativePos, len, progressUpdater) -> {
+                final long startTimeNanos = stats.currentTimeNanos();
+                try (InputStream input = openInputStreamFromBlobStore(rangeToWrite.start() + relativePos, len)) {
+                    assert ThreadPool.assertCurrentThreadPool(SearchableSnapshots.CACHE_FETCH_ASYNC_THREAD_POOL_NAME);
+                    logger.trace(
+                        "{}: writing channel {} pos {} length {} (details: {})",
+                        fileInfo.physicalName(),
+                        channelPos,
+                        relativePos,
+                        len,
+                        cacheFile
+                    );
+                    SharedBytes.copyToCacheFileAligned(
+                        channel,
+                        input,
+                        channelPos,
+                        relativePos,
+                        len,
+                        progressUpdater,
+                        writeBuffer.get().clear(),
+                        cacheFile
+                    );
+                    final long endTimeNanos = stats.currentTimeNanos();
+                    stats.addCachedBytesWritten(len, endTimeNanos - startTimeNanos);
+                }
+            }, SearchableSnapshots.CACHE_FETCH_ASYNC_THREAD_POOL_NAME);
             assert bytesRead == length : bytesRead + " vs " + length;
-            assert luceneByteBufPermits.availablePermits() == Integer.MAX_VALUE;
-
-            luceneByteBufPermits.acquire(Integer.MAX_VALUE);
-            bufferWriteLocked = true;
-            b.position(originalByteBufPosition + bytesRead); // mark all bytes as accounted for
+            byteBufferReference.finish(bytesRead);
         } finally {
-            if (bufferWriteLocked == false) {
-                luceneByteBufPermits.acquire(Integer.MAX_VALUE);
-            }
-        }
-    }
-
-    private int readCacheFile(
-        final SharedBytes.IO fc,
-        long channelPos,
-        long relativePos,
-        long length,
-        final ByteBuffer buffer,
-        long logicalPos,
-        Semaphore luceneByteBufPermits
-    ) throws IOException {
-        logger.trace(
-            "{}: reading cached {} logical {} channel {} pos {} length {} (details: {})",
-            fileInfo.physicalName(),
-            false,
-            logicalPos,
-            channelPos,
-            relativePos,
-            length,
-            cacheFile
-        );
-        if (length == 0L) {
-            return 0;
-        }
-        final int bytesRead;
-        if (luceneByteBufPermits.tryAcquire()) {
-            try {
-                // create slice that is positioned to read the given values
-                final ByteBuffer dup = buffer.slice(buffer.position() + Math.toIntExact(relativePos), Math.toIntExact(length));
-                bytesRead = fc.read(dup, channelPos);
-                if (bytesRead == -1) {
-                    BlobCacheUtils.throwEOF(channelPos, dup.remaining(), this.cacheFile);
-                }
-            } finally {
-                luceneByteBufPermits.release();
-            }
-        } else {
-            // return fake response
-            return Math.toIntExact(length);
+            byteBufferReference.finish(0);
         }
-        stats.addCachedBytesRead(bytesRead);
-        return bytesRead;
     }
 
     @Override