ParallelBlockInputStream.java

/** Copyright (c) 2026 Leonard Ehrenfried. <mail@leonard.io>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Lesser General Public License as
   published by the Free Software Foundation, either version 3 of the
   License, or (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.

*/

package crosby.binary.file;

import java.io.Closeable;
import java.io.DataInputStream;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayDeque;
import java.util.Queue;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * A drop-in, parallel replacement for {@link BlockInputStream}.
 *
 * Reading a fileblock's header and raw bytes must stay strictly sequential,
 * since it consumes an {@link InputStream}, but decompressing and parsing a
 * block's contents (the CPU-heavy step, dominated by zlib inflate for large
 * files) is independent per block. This class overlaps that work across a
 * thread pool while still delivering {@code handleBlock} calls to the
 * adaptor one at a time, in file order, matching what
 * {@link BlockReaderAdapter#skipBlock} documents -- so an existing
 * {@link BlockReaderAdapter} (such as a {@code crosby.binary.BinaryParser}
 * subclass) needs no changes to benefit.
 *
 * Note: a block's header must be read -- and {@code skipBlock} called on it
 * -- before its body can be read or skipped and the stream advanced past it,
 * so {@code skipBlock} may be invoked for up to {@code pipelineDepth}
 * non-skipped blocks ahead of the block whose {@code handleBlock} result was
 * most recently delivered; a run of skipped blocks doesn't count against
 * that limit (nothing is queued for them), so {@code skipBlock} can in fact
 * run further ahead than {@code pipelineDepth} when the file contains long
 * stretches of skipped blocks. An adaptor whose skip decision depends on
 * mutable state updated inside {@code handleBlock} (rather than just the
 * block's own type or metadata) will see stale state under this read-ahead;
 * such an adaptor should use {@code pipelineDepth} 1, which restores fully
 * sequential delivery at the cost of the parallelism benefit -- pass it via
 * the 4-argument constructor, since the convenience constructor that takes a
 * thread count derives {@code pipelineDepth} as {@code numThreads * 2} and
 * cannot be used to request a depth of 1.
 *
 * Memory use is bounded by keeping at most {@code pipelineDepth} decompressed
 * blocks in flight at once (queued or in progress); tune it down for memory
 * constrained environments, or up to keep more threads fed on files with many
 * small blocks.
 */
public class ParallelBlockInputStream implements Closeable {
    public ParallelBlockInputStream(InputStream input, BlockReaderAdapter adaptor) {
        this(input, adaptor, Runtime.getRuntime().availableProcessors());
    }

    public ParallelBlockInputStream(InputStream input, BlockReaderAdapter adaptor, int numThreads) {
        this(input, adaptor, Executors.newFixedThreadPool(numThreads), true, numThreads * 2);
    }

    /**
     * Use a caller-supplied executor, e.g. to share a thread pool across
     * several files. The executor is not shut down by {@link #close()}.
     */
    public ParallelBlockInputStream(InputStream input, BlockReaderAdapter adaptor,
            ExecutorService executor, int pipelineDepth) {
        this(input, adaptor, executor, false, pipelineDepth);
    }

    private ParallelBlockInputStream(InputStream input, BlockReaderAdapter adaptor,
            ExecutorService executor, boolean ownsExecutor, int pipelineDepth) {
        this.input = input;
        this.adaptor = adaptor;
        this.executor = executor;
        this.ownsExecutor = ownsExecutor;
        this.pipelineDepth = Math.max(1, pipelineDepth);
    }

    public void process() throws IOException {
        Queue<Future<FileBlock>> inflight = new ArrayDeque<>();
        try {
            while (true) {
                try {
                    FileBlockHead head = FileBlockHead.readHead(input);
                    if (adaptor.skipBlock(head)) {
                        head.skipContents(input);
                        continue;
                    }
                    byte[] buf = new byte[head.getDatasize()];
                    new DataInputStream(input).readFully(buf);
                    inflight.add(submit(head, buf));
                } catch (EOFException e) {
                    // Matches BlockInputStream: a clean end of stream, or a truncated
                    // header/body, both just end the read -- the trailing partial
                    // block (if any) is dropped, and anything already in flight is
                    // still delivered below.
                    break;
                }

                if (inflight.size() >= pipelineDepth) {
                    adaptor.handleBlock(take(inflight));
                }
            }
            while (!inflight.isEmpty()) {
                adaptor.handleBlock(take(inflight));
            }
            adaptor.complete();
        } finally {
            for (Future<FileBlock> future : inflight) {
                future.cancel(true);
            }
            if (ownsExecutor) {
                executor.shutdownNow();
            }
        }
    }

    private Future<FileBlock> submit(final FileBlockHead head, final byte[] buf) {
        return executor.submit(new Callable<FileBlock>() {
            @Override
            public FileBlock call() throws IOException {
                return head.parseData(buf);
            }
        });
    }

    private static FileBlock take(Queue<Future<FileBlock>> inflight) throws IOException {
        Future<FileBlock> future = inflight.poll();
        try {
            return future.get();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IOException(e);
        } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            if (cause instanceof IOException) {
                throw (IOException) cause;
            }
            if (cause instanceof RuntimeException) {
                throw (RuntimeException) cause;
            }
            throw new IOException(cause);
        }
    }

    @Override
    public void close() throws IOException {
        try {
            input.close();
        } finally {
            if (ownsExecutor) {
                executor.shutdownNow();
            }
        }
    }

    private final InputStream input;
    private final BlockReaderAdapter adaptor;
    private final ExecutorService executor;
    private final boolean ownsExecutor;
    private final int pipelineDepth;
}