GH-3719: Fix vectored read allocation limits and fallback safety

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadBufferAllocator.java

@@ -0,0 +1,110 @@
+/*
+ * 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.parquet.hadoop;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import org.apache.parquet.bytes.ByteBufferAllocator;
+import org.apache.parquet.bytes.ByteBufferReleaser;
+import org.apache.parquet.util.AutoCloseables;
+
+/**
+ * Owns the original allocations for one vectored read. Filesystems may return slices of these buffers and allocate
+ * additional buffers for checksums, so the range results do not identify all the allocations that need releasing.
+ *
+ * <p>The caller must wait for submission and all reads to finish before transferring or closing this owner. Stopping
+ * allocations, cancelling a result future, or closing the stream does not establish that the filesystem has stopped
+ * using a previously allocated buffer.
+ */
+final class VectoredReadBufferAllocator implements ByteBufferAllocator, AutoCloseable {
+  private final ByteBufferAllocator allocator;
+  private final Map<ByteBuffer, Boolean> buffers = new IdentityHashMap<>();
+  private volatile boolean acceptingAllocations = true;
+
+  VectoredReadBufferAllocator(ByteBufferAllocator allocator) {
+    this.allocator = Objects.requireNonNull(allocator, "allocator");
+  }
+
+  @Override
+  public synchronized ByteBuffer allocate(int size) {
+    if (!acceptingAllocations) {
+      throw new IllegalStateException("Vectored read is no longer accepting buffer allocations");
+    }
+    ByteBuffer buffer = Objects.requireNonNull(allocator.allocate(size), "allocated buffer");
+    // stopAllocating() may run while the delegate is allocating. The already accepted allocation still belongs
+    // to this owner and must remain available to the filesystem until the caller establishes completion.
+    buffers.put(buffer, Boolean.TRUE);
+    return buffer;
+  }
+
+  @Override
+  public synchronized void release(ByteBuffer buffer) {
+    if (buffers.remove(buffer) == null) {
+      throw new IllegalArgumentException("Buffer is not owned by this vectored read");
+    }
+    allocator.release(buffer);
+  }
+
+  @Override
+  public boolean isDirect() {
+    return allocator.isDirect();
+  }
+
+  /** Rejects new allocations without waiting for a potentially blocked delegate allocation or releasing buffers. */
+  void stopAllocating() {
+    acceptingAllocations = false;
+  }
+
+  /**
+   * Transfers original buffers to a releaser associated with the delegate allocator. This is only valid after
+   * successful submission and completion of every range. The owner must not have been aborted or transferred.
+   */
+  synchronized void transferTo(ByteBufferReleaser releaser) {
+    Objects.requireNonNull(releaser, "releaser");
+    if (!acceptingAllocations) {
+      throw new IllegalStateException("Vectored read buffer ownership has already been stopped or transferred");
+    }
+    acceptingAllocations = false;
+    Iterator<ByteBuffer> iterator = buffers.keySet().iterator();
+    while (iterator.hasNext()) {
+      releaser.releaseLater(iterator.next());
+      iterator.remove();
+    }
+  }
+
+  /**
+   * Releases the still-owned originals exactly once. The caller must establish that no backend work can use them;
+   * this method itself does not wait for asynchronous IO. Buffers already transferred to a row group are unaffected.
+   */
+  @Override
+  public synchronized void close() {
+    acceptingAllocations = false;
+    List<AutoCloseable> releases = new ArrayList<>(buffers.size());
+    for (ByteBuffer buffer : buffers.keySet()) {
+      releases.add(() -> allocator.release(buffer));
+    }
+    buffers.clear();
+    AutoCloseables.uncheckedClose(releases);
+  }
+}

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadOperation.java

@@ -0,0 +1,179 @@
+/*
+ * 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.parquet.hadoop;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import org.apache.parquet.bytes.ByteBufferAllocator;
+import org.apache.parquet.bytes.ByteBufferReleaser;
+import org.apache.parquet.hadoop.util.wrapped.io.FutureIO;
+import org.apache.parquet.io.ParquetFileRange;
+import org.apache.parquet.io.SeekableInputStream;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Owns one vectored submission and its allocations until the reader can take ownership.
+ * The executor must have a single worker: queued failure cleanup must not overtake a
+ * submission which is still running after its Future has been cancelled.
+ */
+final class VectoredReadOperation {
+  private static final Logger LOG = LoggerFactory.getLogger(VectoredReadOperation.class);
+
+  private final SeekableInputStream stream;
+  private final List<ParquetFileRange> ranges;
+  private final VectoredReadBufferAllocator allocator;
+  private final ExecutorService executor;
+  private final long timeoutNanos;
+  private final long readStart = System.nanoTime();
+  private Future<Void> submission;
+  private volatile boolean submissionSucceeded;
+  private boolean aborted;
+  private boolean releaseRegistered;
+
+  VectoredReadOperation(
+      SeekableInputStream stream,
+      List<ParquetFileRange> ranges,
+      ByteBufferAllocator allocator,
+      ExecutorService executor,
+      long timeout,
+      TimeUnit unit) {
+    this.stream = stream;
+    this.ranges = ranges;
+    this.allocator = new VectoredReadBufferAllocator(allocator);
+    this.executor = executor;
+    this.timeoutNanos = unit.toNanos(timeout);
+  }
+
+  void awaitSubmission() throws IOException, TimeoutException {
+    submission = executor.submit(() -> {
+      stream.readVectored(ranges, allocator);
+      submissionSucceeded = true;
+      return null;
+    });
+    FutureIO.awaitFuture(submission, remainingNanos(), TimeUnit.NANOSECONDS);
+  }
+
+  boolean submissionSucceeded() {
+    return submissionSucceeded;
+  }
+
+  long remainingNanos() {
+    return Math.max(timeoutNanos - (System.nanoTime() - readStart), 0L);
+  }
+
+  void transferTo(ByteBufferReleaser releaser) {
+    if (!submissionSucceeded || aborted) {
+      throw new IllegalStateException("Cannot transfer buffers from an unsuccessful vectored read");
+    }
+    for (ParquetFileRange range : ranges) {
+      CompletableFuture<ByteBuffer> future = range.getDataReadFuture();
+      if (future == null || !future.isDone() || future.isCompletedExceptionally()) {
+        throw new IllegalStateException("Cannot transfer buffers before all vectored reads succeed");
+      }
+    }
+    allocator.transferTo(releaser);
+  }
+
+  /**
+   * Stop the caller's wait without treating interruption as proof that backend IO stopped.
+   * Once aborted, the reader must not reuse the stream or this executor.
+   */
+  void abort(Throwable failure) {
+    if (aborted) {
+      return;
+    }
+    aborted = true;
+    allocator.stopAllocating();
+    if (submission != null) {
+      submission.cancel(true);
+    }
+    if (submissionSucceeded) {
+      // All futures are published. If the failed read and its siblings already
+      // finished, reclaim their buffers before the caller closes its allocator.
+      releaseWhenReadsFinish(failure);
+    }
+    try {
+      // Future.cancel(true) may return while the callable is still running. A task on
+      // the same single worker cannot close the stream until that callable has exited.
+      executor.execute(() -> {
+        releaseWhenReadsFinish(failure);
+        try {
+          stream.close();
+        } catch (IOException | RuntimeException closeFailure) {
+          if (failure != closeFailure) {
+            failure.addSuppressed(closeFailure);
+          }
+          LOG.warn("Failed to close a stream after a vectored read failure", closeFailure);
+        }
+      });
+    } catch (RejectedExecutionException cleanupFailure) {
+      // The reader owns this executor and must not shut it down before queuing cleanup.
+      // Do not mask the read error or recycle buffers whose IO lifetime is now unknown.
+      failure.addSuppressed(cleanupFailure);
+      LOG.warn("Could not schedule cleanup after a vectored read failure", cleanupFailure);
+    } finally {
+      // shutdownNow would discard the queued cleanup or interrupt its close operation.
+      executor.shutdown();
+    }
+  }
+
+  private void releaseWhenReadsFinish(Throwable failure) {
+    if (releaseRegistered) {
+      return;
+    }
+    releaseRegistered = true;
+    CompletableFuture<?>[] futures = new CompletableFuture<?>[ranges.size()];
+    for (int i = 0; i < ranges.size(); i++) {
+      CompletableFuture<ByteBuffer> future = ranges.get(i).getDataReadFuture();
+      if (future == null) {
+        // A backend may have started IO without publishing its result. Neither closing
+        // its stream nor cancelling a future establishes that pooled memory is reusable.
+        LOG.debug("Retaining allocations after a vectored submission with unpublished reads");
+        return;
+      }
+      futures[i] = future;
+    }
+    // A failed submission may also publish a future for work never submitted. Such a
+    // future need not complete: retain its ownership without blocking a cleanup worker.
+    CompletableFuture.allOf(futures).whenComplete((ignored, readFailure) -> {
+      for (CompletableFuture<?> future : futures) {
+        if (future.isCancelled()) {
+          LOG.debug("Retaining allocations after a cancelled vectored result with unknown IO lifetime");
+          return;
+        }
+      }
+      try {
+        allocator.close();
+      } catch (RuntimeException releaseFailure) {
+        if (failure != releaseFailure) {
+          failure.addSuppressed(releaseFailure);
+        }
+        LOG.warn("Failed to release buffers after a vectored read failure", releaseFailure);
+      }
+    });
+  }
+}

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java

@@ -124,7 +124,8 @@ private VectorIoBridge() {
    * @param stream input stream to query.
    * @param allocator allocator to be used.
    *
-   * @return true if the stream declares the capability is available.
+   * @return true if the runtime API and allocator are supported; an individual request
+   * may still be rejected.
    */
   public boolean readVectoredAvailable(final FSDataInputStream stream, final ByteBufferAllocator allocator) {
     return available() && !allocator.isDirect();
@@ -155,13 +156,15 @@ private void checkAvailable() {
    * The default iterates through the ranges to read each synchronously, but
    * the intent is that FSDataInputStream subclasses can make more efficient
    * readers.
-   * The {@link ParquetFileRange} parameters all have their
-   * data read futures set to the range reads of the associated
-   * operations; callers must await these to complete.
+   * After a successful return, the {@link ParquetFileRange} parameters all have
+   * their data read futures set to the associated operations; callers can await
+   * these futures for the read results.
    * <p>
-   * As a result of the call, each range will have FileRange.setData(CompletableFuture)
-   * called with a future that when complete will have a ByteBuffer with the
-   * data from the file's range.
+   * If submission throws, futures already assigned by Hadoop are still published.
+   * A null future does not prove that no work started. A non-null future may have
+   * been created for a read that was never submitted and may never complete after
+   * a submission failure. Callers must handle that failure without waiting for
+   * every such future.
    * <p>
    *   The position returned by getPos() after readVectored() is undefined.
    * </p>
@@ -198,15 +201,36 @@ public void readVectoredRanges(
     // Setting the parquet range as a reference.
     List<FileRangeBridge.WrappedFileRange> fileRanges =
         sorted.stream().map(rangeBridge::toFileRange).collect(Collectors.toList());
-    readWrappedRanges(stream, fileRanges, allocator::allocate);
-
-    // copy back the completable futures from the scheduled
-    // vector reads to the ParquetFileRange entries passed in.
-    fileRanges.forEach(fileRange -> {
-      // toFileRange() sets up this back reference
-      ParquetFileRange parquetFileRange = (ParquetFileRange) fileRange.getReference();
-      parquetFileRange.setDataReadFuture(fileRange.getData());
-    });
+    Throwable submissionFailure = null;
+    try {
+      readWrappedRanges(stream, fileRanges, allocator::allocate);
+    } catch (IOException | RuntimeException | Error failure) {
+      submissionFailure = failure;
+      throw failure;
+    } finally {
+      publishReadFutures(fileRanges, submissionFailure);
+    }
+  }
+
+  /**
+   * Publish available futures even after partial submission, without replacing
+   * the submission failure if a future cannot be retrieved.
+   */
+  static void publishReadFutures(List<FileRangeBridge.WrappedFileRange> fileRanges, Throwable submissionFailure) {
+    for (FileRangeBridge.WrappedFileRange fileRange : fileRanges) {
+      try {
+        // toFileRange() sets up this back reference.
+        ParquetFileRange parquetFileRange = (ParquetFileRange) fileRange.getReference();
+        parquetFileRange.setDataReadFuture(fileRange.getData());
+      } catch (RuntimeException | Error publicationFailure) {
+        if (submissionFailure == null) {
+          throw publicationFailure;
+        }
+        if (submissionFailure != publicationFailure) {
+          submissionFailure.addSuppressed(publicationFailure);
+        }
+      }
+    }
   }
 
   /**