QueueInputStream reads all but the first byte without waiting. (#748)

maxxedev · ppkarwasz · garydgregory · web-flow · commit fd10fed8ffd8 · 2025-05-22T11:01:15.000-04:00
* QueueInputStream reads all but the first byte without waiting.

Fix so that bulk reads avoid getting stuck if a timeout is set
and at least one byte is available.

* Call BlockingQueue::drainTo for improved performance

* Refactor bulk read

* Add benchmark for `QueueInputStream/QueueOutputStream`

Adds a small benchmark to measure how much time does it take to transfer
1 MiB from a `QueueOutputStream` to a `QueueInputStream`.

* More optimizations on bulk read

* Improve test coverage

* Add missing Javadoc since tag

* Use final

Reduce vertical whitespace

* Use the exact same invariant checks as the JDK superclass

Add missing Javadoc `@param` tag

* Update src/main/java/org/apache/commons/io/input/QueueInputStream.java

Co-authored-by: Piotr P. Karwasz &lt;piotr@github.copernik.eu&gt;

* Update src/main/java/org/apache/commons/io/input/QueueInputStream.java

* improve tests

* improve test

---------

Co-authored-by: Piotr P. Karwasz &lt;pkarwasz-github@apache.org&gt;
Co-authored-by: Gary Gregory &lt;garydgregory@users.noreply.github.com&gt;
Co-authored-by: Piotr P. Karwasz &lt;piotr@github.copernik.eu&gt;
diff --git a/src/changes/changes.xml b/src/changes/changes.xml
@@ -47,6 +47,7 @@ The <action> type attribute can be add,update,fix,remove.
   <body>
     <release version="2.20.0" date="YYYY-MM-DD" description="Version 2.19.1: Java 8 or later is required.">
       <!-- FIX -->
+      <action dev="maxxedev" type="fix"                due-to="maxxedev">QueueInputStream reads all but the first byte without waiting</action>
       <action dev="ggregory" type="fix"                due-to="Jesse Glick">[javadoc] Rename parameter of ProxyOutputStream.write(int) #740.</action>
       <action dev="ggregory" type="fix" issue="IO-875" due-to="Pierre Baumard, Gary Gregory">CopyDirectoryVisitor ignores fileFilter #743.</action>
       <action dev="ggregory" type="fix"                due-to="Gary Gregory">org.apache.commons.io.build.AbstractOrigin.getReader(Charset) now maps a null Charset to the default Charset.</action>
diff --git a/src/main/java/org/apache/commons/io/input/QueueInputStream.java b/src/main/java/org/apache/commons/io/input/QueueInputStream.java
@@ -23,6 +23,8 @@
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -224,4 +226,49 @@ public int read() {
         }
     }
 
+    /**
+     * Reads up to {@code length} bytes of data from the input stream into
+     * an array of bytes.  The first byte is read while honoring the timeout; the rest are read while <i>not</i> honoring
+     * the timeout. The number of bytes actually read is returned as an integer.
+     *
+     * @param b     the buffer into which the data is read.
+     * @param offset   the start offset in array {@code b} at which the data is written.
+     * @param length   the maximum number of bytes to read.
+     * @return     the total number of bytes read into the buffer, or {@code -1} if there is no more data because the
+     *              end of the stream has been reached.
+     * @throws NullPointerException If {@code b} is {@code null}.
+     * @throws IllegalStateException if thread is interrupted while waiting for the first byte.
+     * @throws IndexOutOfBoundsException if {@code offset} is negative, {@code length} is negative, or {@code length} is
+     *             greater than {@code b.length - offset}.
+     * @since 2.20.0
+     */
+    @Override
+    public int read(final byte[] b, final int offset, final int length) {
+        if (b == null) {
+            throw new NullPointerException();
+        } else if (offset < 0 || length < 0 || length > b.length - offset) {
+            throw new IndexOutOfBoundsException(
+                    String.format("Range [%d, %<d + %d) out of bounds for length %d", offset, length, b.length));
+        } else if (length == 0) {
+            return 0;
+        }
+        final List<Integer> drain = new ArrayList<>(Math.min(length, blockingQueue.size()));
+        blockingQueue.drainTo(drain, length);
+        if (drain.isEmpty()) {
+            // no data immediately available. wait for first byte
+            final int value = read();
+            if (value == EOF) {
+                return EOF;
+            }
+            drain.add(value);
+            blockingQueue.drainTo(drain, length - 1);
+        }
+        int i = 0;
+        for (final Integer value : drain) {
+            b[offset + i] = (byte) (0xFF & value);
+            i++;
+        }
+        return i;
+    }
+
 }
diff --git a/src/test/java/org/apache/commons/io/input/QueueInputStreamTest.java b/src/test/java/org/apache/commons/io/input/QueueInputStreamTest.java
@@ -16,19 +16,30 @@
  */
 package org.apache.commons.io.input;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTimeout;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
+import static org.junit.jupiter.api.DynamicTest.dynamicTest;
 
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
+import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.time.Duration;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -40,7 +51,9 @@
 import org.apache.commons.io.output.QueueOutputStreamTest;
 import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.DynamicTest;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestFactory;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -124,13 +137,125 @@ public void testAvailableAfterOpen(final String inputData) throws IOException {
     public void testBufferedReads(final String inputData) throws IOException {
         final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>();
         try (BufferedInputStream inputStream = new BufferedInputStream(new QueueInputStream(queue));
-                QueueOutputStream outputStream = new QueueOutputStream(queue)) {
+             QueueOutputStream outputStream = new QueueOutputStream(queue)) {
             outputStream.write(inputData.getBytes(StandardCharsets.UTF_8));
             final String actualData = IOUtils.toString(inputStream, StandardCharsets.UTF_8);
             assertEquals(inputData, actualData);
         }
     }
 
+    @ParameterizedTest(name = "inputData={0}")
+    @MethodSource("inputData")
+    public void testReadLineByLineQueue(final String inputData) throws IOException {
+        final String[] lines = inputData.split("\n");
+        final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>();
+        try (QueueInputStream inputStream = QueueInputStream.builder()
+                .setBlockingQueue(queue)
+                .setTimeout(Duration.ofHours(1))
+                .get();
+             QueueOutputStream outputStream = inputStream.newQueueOutputStream()) {
+
+            doTestReadLineByLine(inputData, inputStream, outputStream);
+        }
+    }
+
+    @ParameterizedTest(name = "inputData={0}")
+    @MethodSource("inputData")
+    public void testReadLineByLineFile(final String inputData) throws IOException {
+        final Path tempFile = Files.createTempFile(getClass().getSimpleName(), ".txt");
+        try (InputStream inputStream = Files.newInputStream(tempFile);
+             OutputStream outputStream = Files.newOutputStream(tempFile)) {
+
+            doTestReadLineByLine(inputData, inputStream, outputStream);
+        } finally {
+            Files.delete(tempFile);
+        }
+    }
+
+    private void doTestReadLineByLine(final String inputData, final InputStream inputStream, final OutputStream outputStream) throws IOException {
+        final String[] lines = inputData.split("\n");
+        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, UTF_8))) {
+            for (String line : lines) {
+                outputStream.write(line.getBytes(UTF_8));
+                outputStream.write('\n');
+
+                final String actualLine = reader.readLine();
+                assertEquals(line, actualLine);
+            }
+        }
+    }
+
+    @TestFactory
+    public DynamicTest[] bulkReadErrorHandlingTests() {
+        final QueueInputStream queueInputStream = new QueueInputStream();
+        return new DynamicTest[] {
+                dynamicTest("Offset too big", () ->
+                        assertThrows(IndexOutOfBoundsException.class, () ->
+                                queueInputStream.read(EMPTY_BYTE_ARRAY, 1, 0))),
+
+                dynamicTest("Offset negative", () ->
+                        assertThrows(IndexOutOfBoundsException.class, () ->
+                                queueInputStream.read(EMPTY_BYTE_ARRAY, -1, 0))),
+
+                dynamicTest("Length too big", () ->
+                        assertThrows(IndexOutOfBoundsException.class, () ->
+                                queueInputStream.read(EMPTY_BYTE_ARRAY, 0, 1))),
+
+                dynamicTest("Length negative", () ->
+                        assertThrows(IndexOutOfBoundsException.class, () ->
+                                queueInputStream.read(EMPTY_BYTE_ARRAY, 0, -1))),
+        };
+    }
+
+    @Test
+    public void testBulkReadZeroLength() {
+        final QueueInputStream queueInputStream = new QueueInputStream();
+        final int read = queueInputStream.read(EMPTY_BYTE_ARRAY, 0, 0);
+        assertEquals(0, read);
+    }
+
+    @ParameterizedTest(name = "inputData={0}")
+    @MethodSource("inputData")
+    public void testBulkReadWaiting(final String inputData) throws IOException {
+        assumeFalse(inputData.isEmpty());
+
+        final CountDownLatch onPollLatch = new CountDownLatch(1);
+        final CountDownLatch afterWriteLatch = new CountDownLatch(1);
+        final LinkedBlockingQueue<Integer> queue = new LinkedBlockingQueue<Integer>() {
+            @Override
+            public Integer poll(final long timeout, final TimeUnit unit) throws InterruptedException {
+                onPollLatch.countDown();
+                afterWriteLatch.await(1, TimeUnit.HOURS);
+                return super.poll(timeout, unit);
+            }
+        };
+
+        // Simulate scenario where there is not data immediately available when bulk reading and QueueInputStream has to
+        // wait.
+        try (QueueInputStream queueInputStream = QueueInputStream.builder()
+                .setBlockingQueue(queue)
+                .setTimeout(Duration.ofHours(1))
+                .get()) {
+            final QueueOutputStream queueOutputStream = queueInputStream.newQueueOutputStream();
+            final CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
+                try {
+                    onPollLatch.await(1, TimeUnit.HOURS);
+                    queueOutputStream.write(inputData.getBytes(UTF_8));
+                    afterWriteLatch.countDown();
+                } catch (final Exception e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            final byte[] data = new byte[inputData.length()];
+            final int read = queueInputStream.read(data, 0, data.length);
+            assertEquals(inputData.length(), read);
+            final String outputData = new String(data, 0, read, StandardCharsets.UTF_8);
+            assertEquals(inputData, outputData);
+            assertDoesNotThrow(() -> future.get());
+        }
+    }
+
     @ParameterizedTest(name = "inputData={0}")
     @MethodSource("inputData")
     public void testBufferedReadWrite(final String inputData) throws IOException {
diff --git a/src/test/java/org/apache/commons/io/jmh/QueueStreamBenchmark.java b/src/test/java/org/apache/commons/io/jmh/QueueStreamBenchmark.java
@@ -0,0 +1,76 @@
+/*
+ * 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.commons.io.jmh;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.io.input.QueueInputStream;
+import org.apache.commons.io.output.QueueOutputStream;
+import org.apache.commons.lang3.RandomUtils;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Group;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.infra.Blackhole;
+
+/**
+ * Measures the amount of time to push 1 MiB to a {@link QueueOutputStream} and read it using a {@link QueueInputStream}
+ */
+@BenchmarkMode(Mode.SampleTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Group)
+public class QueueStreamBenchmark {
+
+    private static final int CAPACITY = 1024 * 1024;
+    private static final int BUFFER_SIZE = 1024;
+
+    private final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>(CAPACITY);
+    private final QueueInputStream inputStream = QueueInputStream.builder()
+            .setBlockingQueue(queue)
+            .get();
+    private final QueueOutputStream outputStream = inputStream.newQueueOutputStream();
+
+    private final byte[] input = RandomUtils.insecure().randomBytes(CAPACITY);
+    private final byte[] output = new byte[BUFFER_SIZE];
+
+    @Benchmark
+    @Group("streams")
+    public void output() throws Exception {
+        int sent = 0;
+        while (sent < CAPACITY) {
+            final int len = Math.min(CAPACITY - sent, BUFFER_SIZE);
+            outputStream.write(input, sent, len);
+            sent += len;
+        }
+    }
+
+    @Benchmark
+    @Group("streams")
+    public void input(Blackhole bh) throws Exception {
+        int received = 0;
+        while (received < CAPACITY) {
+            final int len = inputStream.read(output, 0, BUFFER_SIZE);
+            bh.consume(output);
+            received += len;
+        }
+    }
+}
