Add documentation to concurrent multipart upload utility methods

modules/repository-azure/src/main/java/org/elasticsearch/repositories/azure/AzureBlobStore.java

@@ -66,6 +66,7 @@
 import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Nullable;
@@ -74,6 +75,7 @@
 import org.elasticsearch.repositories.azure.AzureRepository.Repository;
 import org.elasticsearch.repositories.blobstore.ChunkedBlobOutputStream;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.threadpool.ThreadPool;
 
 import java.io.FilterInputStream;
 import java.io.IOException;
@@ -101,6 +103,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
@@ -507,7 +510,11 @@ void writeBlobAtomic(
                         return asyncClient.commitBlockList(
                             multiParts.stream().map(MultiPart::blockId).toList(),
                             failIfAlreadyExists == false
-                        ).doOnSuccess(unused -> logger.debug("{}: all {} parts committed", blobName, multiParts.size()));
+                        )
+                            .doOnSuccess(unused -> logger.debug("{}: all {} parts committed", blobName, multiParts.size()))
+                            // Note: non-committed uploaded blocks will be deleted by Azure after a week
+                            // (see https://docs.microsoft.com/en-us/rest/api/storageservices/put-block#remarks)
+                            .doOnError(e -> logger.error(() -> format("%s: failed to commit %d parts", blobName, multiParts.size()), e));
                     })
                     .block();
             }
@@ -562,12 +569,13 @@ private static Mono<String> stageBlock(
             multiPart.blockOffset()
         );
         try {
-            var stream = toSynchronizedInputStream(blobName, provider.apply(multiPart.blockOffset(), multiPart.blockSize()), multiPart);
+            final var stream = provider.apply(multiPart.blockOffset(), multiPart.blockSize());
+            assert stream.markSupported() : "provided input stream must support mark and reset";
             boolean success = false;
             try {
                 var stageBlock = asyncClient.stageBlock(
                     multiPart.blockId(),
-                    toFlux(stream, multiPart.blockSize(), DEFAULT_UPLOAD_BUFFERS_SIZE),
+                    toFlux(wrapInputStream(blobName, stream, multiPart), multiPart.blockSize(), DEFAULT_UPLOAD_BUFFERS_SIZE),
                     multiPart.blockSize()
                 ).doOnSuccess(unused -> {
                     logger.debug(() -> format("%s: part [%s] of size [%s] uploaded", blobName, multiPart.part(), multiPart.blockSize()));
@@ -760,88 +768,105 @@ public synchronized int read() throws IOException {
                                               // we read the input stream (i.e. when it's rate limited)
     }
 
-    private static InputStream toSynchronizedInputStream(String blobName, InputStream delegate, MultiPart multipart) {
-        assert delegate.markSupported() : "An InputStream with mark support was expected";
-        // We need to introduce a read barrier in order to provide visibility for the underlying
-        // input stream state as the input stream can be read from different threads.
-        // TODO See if this is still needed
+    /**
+     * Wraps an {@link InputStream} to assert that it is read only by a single thread at a time and to add log traces.
+     */
+    private static InputStream wrapInputStream(final String blobName, final InputStream delegate, final MultiPart multipart) {
         return new FilterInputStream(delegate) {
 
+            private final AtomicReference<Thread> currentThread = Assertions.ENABLED ? new AtomicReference<>() : null;
             private final boolean isTraceEnabled = logger.isTraceEnabled();
 
             @Override
-            public synchronized int read(byte[] b, int off, int len) throws IOException {
-                var result = super.read(b, off, len);
-                if (isTraceEnabled) {
-                    logger.trace("{} reads {} bytes from {} part {}", Thread.currentThread(), result, blobName, multipart.part());
-                }
-                return result;
-            }
-
-            @Override
-            public synchronized int read() throws IOException {
-                var result = super.read();
-                if (isTraceEnabled) {
-                    logger.trace("{} reads {} byte from {} part {}", Thread.currentThread(), result, blobName, multipart.part());
-                }
-                return result;
-            }
-
-            @Override
-            public synchronized void mark(int readlimit) {
-                if (isTraceEnabled) {
-                    logger.trace("{} marks stream {} part {}", Thread.currentThread(), blobName, multipart.part());
-                }
-                super.mark(readlimit);
-            }
-
-            @Override
-            public synchronized void reset() throws IOException {
-                if (isTraceEnabled) {
-                    logger.trace("{} resets stream {} part {}", Thread.currentThread(), blobName, multipart.part());
+            public int read(byte[] b, int off, int len) throws IOException {
+                assert assertThread(null, Thread.currentThread());
+                assert ThreadPool.assertCurrentThreadPool(AzureRepositoryPlugin.REPOSITORY_THREAD_POOL_NAME);
+                try {
+                    var result = super.read(b, off, len);
+                    if (isTraceEnabled) {
+                        logger.trace("{} reads {} bytes from {} part {}", Thread.currentThread(), result, blobName, multipart.part());
+                    }
+                    return result;
+                } finally {
+                    assert assertThread(Thread.currentThread(), null);
                 }
-                super.reset();
             }
 
             @Override
-            public synchronized void close() throws IOException {
-                if (isTraceEnabled) {
-                    logger.trace("{} closes stream {} part {}", Thread.currentThread(), blobName, multipart.part());
+            public int read() throws IOException {
+                assert assertThread(null, Thread.currentThread());
+                assert ThreadPool.assertCurrentThreadPool(AzureRepositoryPlugin.REPOSITORY_THREAD_POOL_NAME);
+                try {
+                    var result = super.read();
+                    if (isTraceEnabled) {
+                        logger.trace("{} reads {} byte from {} part {}", Thread.currentThread(), result, blobName, multipart.part());
+                    }
+                    return result;
+                } finally {
+                    assert assertThread(Thread.currentThread(), null);
                 }
-                super.close();
             }
 
-            @Override
-            public String toString() {
-                return blobName + " part [" + multipart.part() + "] of size [" + multipart.blockSize() + ']';
+            private boolean assertThread(Thread current, Thread updated) {
+                final Thread witness = currentThread.compareAndExchange(current, updated);
+                assert witness == current
+                    : "Unable to set current thread to ["
+                        + updated
+                        + "]: expected thread ["
+                        + current
+                        + "] to be the thread currently accessing the input stream for reading, but thread "
+                        + witness
+                        + " is already reading "
+                        + blobName
+                        + " part "
+                        + multipart.part();
+                return true;
             }
         };
     }
 
-    private static Flux<ByteBuffer> toFlux(InputStream stream, long length, int chunkSize) {
-        assert stream.markSupported() : "An InputStream with mark support was expected";
-        // We need to mark the InputStream as it's possible that we need to retry for the same chunk
+    /**
+     * Converts an input stream to a Flux of ByteBuffer. This method also checks that the stream has provided the expected number of bytes.
+     *
+     * @param stream            the input stream that needs to be converted
+     * @param length            the expected length in bytes of the input stream
+     * @param byteBufferSize    the size of the ByteBuffer to be created
+     **/
+    private static Flux<ByteBuffer> toFlux(InputStream stream, long length, final int byteBufferSize) {
+        assert stream.markSupported() : "input stream must support mark and reset";
+        // always marks the input stream in case it needs to be retried
         stream.mark(Integer.MAX_VALUE);
+        // defer the creation of the flux until it is subscribed
         return Flux.defer(() -> {
-            // TODO Code in this Flux.defer() can be concurrently executed by multiple threads?
             try {
                 stream.reset();
             } catch (IOException e) {
-                throw new RuntimeException(e);
+                // Flux.defer() catches and propagates the exception
+                throw new UncheckedIOException(e);
             }
+            // the number of bytes read is updated in a thread pool (repository_azure) and later compared to the expected length in another
+            // thread pool (azure_event_loop), so we need this to be atomic.
             final var bytesRead = new AtomicLong(0L);
-            // This flux is subscribed by a downstream operator that finally queues the
-            // buffers into netty output queue. Sadly we are not able to get a signal once
-            // the buffer has been flushed, so we have to allocate those and let the GC to
-            // reclaim them (see MonoSendMany). Additionally, that very same operator requests
-            // 128 elements (that's hardcoded) once it's subscribed (later on, it requests
-            // by 64 elements), that's why we provide 64kb buffers.
 
-            // length is at most 100MB so it's safe to cast back to an integer in this case
-            final int parts = (int) length / chunkSize;
-            final long remaining = length % chunkSize;
-            return Flux.range(0, remaining == 0 ? parts : parts + 1).map(i -> i * chunkSize).concatMap(pos -> Mono.fromCallable(() -> {
-                long count = pos + chunkSize > length ? length - pos : chunkSize;
+            assert length <= ByteSizeValue.ofMb(100L).getBytes() : length;
+            // length is at most 100MB so it's safe to cast back to an integer
+            final int parts = Math.toIntExact(length / byteBufferSize);
+            final long remaining = length % byteBufferSize;
+
+            // This flux is subscribed by a downstream subscriber (reactor.netty.channel.MonoSendMany) that queues the buffers into netty
+            // output queue. Sadly we are not able to get a signal once the buffer has been flushed, so we have to allocate those and let
+            // the GC to reclaim them. Additionally, the MonoSendMany subscriber requests 128 elements from the flux when it subscribes to
+            // it. This 128 value is hardcoded in reactor.netty.channel.MonoSend.MAX_SIZE). After 128 byte buffers have been published by
+            // the flux, the MonoSendMany subscriber requests 64 more byte buffers (see reactor.netty.channel.MonoSend.REFILL_SIZE) and so
+            // on.
+            //
+            // So this flux instantiates 128 ByteBuffer objects of DEFAULT_UPLOAD_BUFFERS_SIZE bytes in heap every time the NettyOutbound in
+            // the Azure's Netty event loop requests byte buffers to write to the network channel. That represents 128 * 64kb = 8 mb per
+            // flux. The creation of the ByteBuffer objects are forked to the repository_azure thread pool, which has a maximum of 15
+            // threads (most of the time, can be less than that for nodes with less than 750mb). It means that max. 15 * 8 = 120mb bytes are
+            // allocated on heap at a time here (omitting the ones already created and pending garbage collection).
+            return Flux.range(0, remaining == 0 ? parts : parts + 1).map(i -> i * byteBufferSize).concatMap(pos -> Mono.fromCallable(() -> {
+                long count = pos + byteBufferSize > length ? length - pos : byteBufferSize;
                 int numOfBytesRead = 0;
                 int offset = 0;
                 int len = (int) count;
@@ -860,16 +885,15 @@ private static Flux<ByteBuffer> toFlux(InputStream stream, long length, int chun
                     );
                 }
                 return ByteBuffer.wrap(buffer);
-            })).doOnComplete(() -> {
+            }), 0 /* no prefetching needed for concatMap as the MonoSendMany subscriber requests 128 elements */).doOnComplete(() -> {
                 if (bytesRead.get() > length) {
                     throw new IllegalStateException(
                         format("Input stream [%s] emitted %d bytes, more than the expected %d bytes.", stream, bytesRead, length)
                     );
                 }
             });
-            // We need to subscribe on a different scheduler to avoid blocking the io threads when we read the input stream
+            // subscribe on a different scheduler to avoid blocking the network io threads when reading bytes from disk
         }).subscribeOn(Schedulers.elastic());
-
     }
 
     /**