001/*
002 * Licensed under the Apache License, Version 2.0 (the "License");
003 * you may not use this file except in compliance with the License.
004 * You may obtain a copy of the License at
005 *
006 *     http://www.apache.org/licenses/LICENSE-2.0
007 *
008 * Unless required by applicable law or agreed to in writing, software
009 * distributed under the License is distributed on an "AS IS" BASIS,
010 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
011 * See the License for the specific language governing permissions and
012 * limitations under the License.
013 */
014package org.apache.commons.io.input;
015
016import static org.apache.commons.io.IOUtils.EOF;
017
018// import javax.annotation.concurrent.GuardedBy;
019import java.io.EOFException;
020import java.io.FilterInputStream;
021import java.io.IOException;
022import java.io.InputStream;
023import java.io.InterruptedIOException;
024import java.nio.ByteBuffer;
025import java.util.Objects;
026import java.util.concurrent.ExecutorService;
027import java.util.concurrent.Executors;
028import java.util.concurrent.TimeUnit;
029import java.util.concurrent.atomic.AtomicBoolean;
030import java.util.concurrent.locks.Condition;
031import java.util.concurrent.locks.ReentrantLock;
032
033import org.apache.commons.io.build.AbstractStreamBuilder;
034
035/**
036 * Implements {@link InputStream} to asynchronously read ahead from an underlying input stream when a specified amount of data has been read from the current
037 * buffer. It does so by maintaining two buffers: an active buffer and a read ahead buffer. The active buffer contains data which should be returned when a
038 * read() call is issued. The read ahead buffer is used to asynchronously read from the underlying input stream. When the current active buffer is exhausted, we
039 * flip the two buffers so that we can start reading from the read ahead buffer without being blocked by disk I/O.
040 * <p>
041 * To build an instance, use {@link Builder}.
042 * </p>
043 * <p>
044 * This class was ported and adapted from Apache Spark commit 933dc6cb7b3de1d8ccaf73d124d6eb95b947ed19.
045 * </p>
046 *
047 * @see Builder
048 * @since 2.9.0
049 */
050public class ReadAheadInputStream extends FilterInputStream {
051
052    // @formatter:off
053    /**
054     * Builds a new {@link ReadAheadInputStream}.
055     *
056     * <p>
057     * For example:
058     * </p>
059     * <pre>{@code
060     * ReadAheadInputStream s = ReadAheadInputStream.builder()
061     *   .setPath(path)
062     *   .setExecutorService(Executors.newSingleThreadExecutor(ReadAheadInputStream::newThread))
063     *   .get();}
064     * </pre>
065     *
066     * @see #get()
067     * @since 2.12.0
068     */
069    // @formatter:on
070    public static class Builder extends AbstractStreamBuilder<ReadAheadInputStream, Builder> {
071
072        private ExecutorService executorService;
073
074        /**
075         * Builds a new {@link ReadAheadInputStream}.
076         * <p>
077         * You must set input that supports {@link #getInputStream()}, otherwise, this method throws an exception.
078         * </p>
079         * <p>
080         * This builder use the following aspects:
081         * </p>
082         * <ul>
083         * <li>{@link #getInputStream()}</li>
084         * <li>{@link #getBufferSize()}</li>
085         * <li>{@link ExecutorService}</li>
086         * </ul>
087         *
088         * @return a new instance.
089         * @throws IllegalStateException         if the {@code origin} is {@code null}.
090         * @throws UnsupportedOperationException if the origin cannot be converted to an {@link InputStream}.
091         * @throws IOException                   if an I/O error occurs.
092         * @see #getInputStream()
093         * @see #getBufferSize()
094         */
095        @SuppressWarnings("resource")
096        @Override
097        public ReadAheadInputStream get() throws IOException {
098            return new ReadAheadInputStream(getInputStream(), getBufferSize(), executorService != null ? executorService : newExecutorService(),
099                    executorService == null);
100        }
101
102        /**
103         * Sets the executor service for the read-ahead thread.
104         *
105         * @param executorService the executor service for the read-ahead thread.
106         * @return {@code this} instance.
107         */
108        public Builder setExecutorService(final ExecutorService executorService) {
109            this.executorService = executorService;
110            return this;
111        }
112
113    }
114
115    private static final ThreadLocal<byte[]> BYTE_ARRAY_1 = ThreadLocal.withInitial(() -> new byte[1]);
116
117    /**
118     * Constructs a new {@link Builder}.
119     *
120     * @return a new {@link Builder}.
121     * @since 2.12.0
122     */
123    public static Builder builder() {
124        return new Builder();
125    }
126
127    /**
128     * Constructs a new daemon thread.
129     *
130     * @param r the thread's runnable.
131     * @return a new daemon thread.
132     */
133    private static Thread newDaemonThread(final Runnable r) {
134        final Thread thread = new Thread(r, "commons-io-read-ahead");
135        thread.setDaemon(true);
136        return thread;
137    }
138
139    /**
140     * Constructs a new daemon executor service.
141     *
142     * @return a new daemon executor service.
143     */
144    private static ExecutorService newExecutorService() {
145        return Executors.newSingleThreadExecutor(ReadAheadInputStream::newDaemonThread);
146    }
147
148    private final ReentrantLock stateChangeLock = new ReentrantLock();
149
150    // @GuardedBy("stateChangeLock")
151    private ByteBuffer activeBuffer;
152
153    // @GuardedBy("stateChangeLock")
154    private ByteBuffer readAheadBuffer;
155
156    // @GuardedBy("stateChangeLock")
157    private boolean endOfStream;
158
159    // @GuardedBy("stateChangeLock")
160    // true if async read is in progress
161    private boolean readInProgress;
162
163    // @GuardedBy("stateChangeLock")
164    // true if read is aborted due to an exception in reading from underlying input stream.
165    private boolean readAborted;
166
167    // @GuardedBy("stateChangeLock")
168    private Throwable readException;
169
170    // @GuardedBy("stateChangeLock")
171    // whether the close method is called.
172    private boolean isClosed;
173
174    // @GuardedBy("stateChangeLock")
175    // true when the close method will close the underlying input stream. This is valid only if
176    // `isClosed` is true.
177    private boolean isUnderlyingInputStreamBeingClosed;
178
179    // @GuardedBy("stateChangeLock")
180    // whether there is a read ahead task running,
181    private boolean isReading;
182
183    // Whether there is a reader waiting for data.
184    private final AtomicBoolean isWaiting = new AtomicBoolean();
185
186    private final ExecutorService executorService;
187
188    private final boolean shutdownExecutorService;
189
190    private final Condition asyncReadComplete = stateChangeLock.newCondition();
191
192    /**
193     * Constructs an instance with the specified buffer size and read-ahead threshold
194     *
195     * @param inputStream       The underlying input stream.
196     * @param bufferSizeInBytes The buffer size.
197     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
198     */
199    @Deprecated
200    public ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes) {
201        this(inputStream, bufferSizeInBytes, newExecutorService(), true);
202    }
203
204    /**
205     * Constructs an instance with the specified buffer size and read-ahead threshold
206     *
207     * @param inputStream       The underlying input stream.
208     * @param bufferSizeInBytes The buffer size.
209     * @param executorService   An executor service for the read-ahead thread.
210     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
211     */
212    @Deprecated
213    public ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes, final ExecutorService executorService) {
214        this(inputStream, bufferSizeInBytes, executorService, false);
215    }
216
217    /**
218     * Constructs an instance with the specified buffer size and read-ahead threshold
219     *
220     * @param inputStream             The underlying input stream.
221     * @param bufferSizeInBytes       The buffer size.
222     * @param executorService         An executor service for the read-ahead thread.
223     * @param shutdownExecutorService Whether or not to shut down the given ExecutorService on close.
224     */
225    private ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes, final ExecutorService executorService,
226            final boolean shutdownExecutorService) {
227        super(Objects.requireNonNull(inputStream, "inputStream"));
228        if (bufferSizeInBytes <= 0) {
229            throw new IllegalArgumentException("bufferSizeInBytes should be greater than 0, but the value is " + bufferSizeInBytes);
230        }
231        this.executorService = Objects.requireNonNull(executorService, "executorService");
232        this.shutdownExecutorService = shutdownExecutorService;
233        this.activeBuffer = ByteBuffer.allocate(bufferSizeInBytes);
234        this.readAheadBuffer = ByteBuffer.allocate(bufferSizeInBytes);
235        this.activeBuffer.flip();
236        this.readAheadBuffer.flip();
237    }
238
239    @Override
240    public int available() throws IOException {
241        stateChangeLock.lock();
242        // Make sure we have no integer overflow.
243        try {
244            return (int) Math.min(Integer.MAX_VALUE, (long) activeBuffer.remaining() + readAheadBuffer.remaining());
245        } finally {
246            stateChangeLock.unlock();
247        }
248    }
249
250    private void checkReadException() throws IOException {
251        if (readAborted) {
252            if (readException instanceof IOException) {
253                throw (IOException) readException;
254            }
255            throw new IOException(readException);
256        }
257    }
258
259    @Override
260    public void close() throws IOException {
261        boolean isSafeToCloseUnderlyingInputStream = false;
262        stateChangeLock.lock();
263        try {
264            if (isClosed) {
265                return;
266            }
267            isClosed = true;
268            if (!isReading) {
269                // Nobody is reading, so we can close the underlying input stream in this method.
270                isSafeToCloseUnderlyingInputStream = true;
271                // Flip this to make sure the read ahead task will not close the underlying input stream.
272                isUnderlyingInputStreamBeingClosed = true;
273            }
274        } finally {
275            stateChangeLock.unlock();
276        }
277
278        if (shutdownExecutorService) {
279            try {
280                executorService.shutdownNow();
281                executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
282            } catch (final InterruptedException e) {
283                final InterruptedIOException iio = new InterruptedIOException(e.getMessage());
284                iio.initCause(e);
285                throw iio;
286            } finally {
287                if (isSafeToCloseUnderlyingInputStream) {
288                    super.close();
289                }
290            }
291        }
292    }
293
294    private void closeUnderlyingInputStreamIfNecessary() {
295        boolean needToCloseUnderlyingInputStream = false;
296        stateChangeLock.lock();
297        try {
298            isReading = false;
299            if (isClosed && !isUnderlyingInputStreamBeingClosed) {
300                // close method cannot close underlyingInputStream because we were reading.
301                needToCloseUnderlyingInputStream = true;
302            }
303        } finally {
304            stateChangeLock.unlock();
305        }
306        if (needToCloseUnderlyingInputStream) {
307            try {
308                super.close();
309            } catch (final IOException ignored) {
310                // TODO Rethrow as UncheckedIOException?
311            }
312        }
313    }
314
315    private boolean isEndOfStream() {
316        return !activeBuffer.hasRemaining() && !readAheadBuffer.hasRemaining() && endOfStream;
317    }
318
319    @Override
320    public int read() throws IOException {
321        if (activeBuffer.hasRemaining()) {
322            // short path - just get one byte.
323            return activeBuffer.get() & 0xFF;
324        }
325        final byte[] oneByteArray = BYTE_ARRAY_1.get();
326        oneByteArray[0] = 0;
327        return read(oneByteArray, 0, 1) == EOF ? EOF : oneByteArray[0] & 0xFF;
328    }
329
330    @Override
331    public int read(final byte[] b, final int offset, int len) throws IOException {
332        if (offset < 0 || len < 0 || len > b.length - offset) {
333            throw new IndexOutOfBoundsException();
334        }
335        if (len == 0) {
336            return 0;
337        }
338
339        if (!activeBuffer.hasRemaining()) {
340            // No remaining in active buffer - lock and switch to write ahead buffer.
341            stateChangeLock.lock();
342            try {
343                waitForAsyncReadComplete();
344                if (!readAheadBuffer.hasRemaining()) {
345                    // The first read.
346                    readAsync();
347                    waitForAsyncReadComplete();
348                    if (isEndOfStream()) {
349                        return EOF;
350                    }
351                }
352                // Swap the newly read ahead buffer in place of empty active buffer.
353                swapBuffers();
354                // After swapping buffers, trigger another async read for read ahead buffer.
355                readAsync();
356            } finally {
357                stateChangeLock.unlock();
358            }
359        }
360        len = Math.min(len, activeBuffer.remaining());
361        activeBuffer.get(b, offset, len);
362
363        return len;
364    }
365
366    /**
367     * Read data from underlyingInputStream to readAheadBuffer asynchronously.
368     *
369     * @throws IOException if an I/O error occurs.
370     */
371    private void readAsync() throws IOException {
372        stateChangeLock.lock();
373        final byte[] arr;
374        try {
375            arr = readAheadBuffer.array();
376            if (endOfStream || readInProgress) {
377                return;
378            }
379            checkReadException();
380            readAheadBuffer.position(0);
381            readAheadBuffer.flip();
382            readInProgress = true;
383        } finally {
384            stateChangeLock.unlock();
385        }
386        executorService.execute(() -> {
387            stateChangeLock.lock();
388            try {
389                if (isClosed) {
390                    readInProgress = false;
391                    return;
392                }
393                // Flip this so that the close method will not close the underlying input stream when we
394                // are reading.
395                isReading = true;
396            } finally {
397                stateChangeLock.unlock();
398            }
399
400            // Please note that it is safe to release the lock and read into the read ahead buffer
401            // because either of following two conditions will hold:
402            //
403            // 1. The active buffer has data available to read so the reader will not read from the read ahead buffer.
404            //
405            // 2. This is the first time read is called or the active buffer is exhausted, in that case the reader waits
406            // for this async read to complete.
407            //
408            // So there is no race condition in both the situations.
409            int read = 0;
410            int off = 0, len = arr.length;
411            Throwable exception = null;
412            try {
413                // try to fill the read ahead buffer.
414                // if a reader is waiting, possibly return early.
415                do {
416                    read = in.read(arr, off, len);
417                    if (read <= 0) {
418                        break;
419                    }
420                    off += read;
421                    len -= read;
422                } while (len > 0 && !isWaiting.get());
423            } catch (final Throwable ex) {
424                exception = ex;
425                if (ex instanceof Error) {
426                    // `readException` may not be reported to the user. Rethrow Error to make sure at least
427                    // The user can see Error in UncaughtExceptionHandler.
428                    throw (Error) ex;
429                }
430            } finally {
431                stateChangeLock.lock();
432                try {
433                    readAheadBuffer.limit(off);
434                    if (read < 0 || exception instanceof EOFException) {
435                        endOfStream = true;
436                    } else if (exception != null) {
437                        readAborted = true;
438                        readException = exception;
439                    }
440                    readInProgress = false;
441                    signalAsyncReadComplete();
442                } finally {
443                    stateChangeLock.unlock();
444                }
445                closeUnderlyingInputStreamIfNecessary();
446            }
447        });
448    }
449
450    private void signalAsyncReadComplete() {
451        stateChangeLock.lock();
452        try {
453            asyncReadComplete.signalAll();
454        } finally {
455            stateChangeLock.unlock();
456        }
457    }
458
459    @Override
460    public long skip(final long n) throws IOException {
461        if (n <= 0L) {
462            return 0L;
463        }
464        if (n <= activeBuffer.remaining()) {
465            // Only skipping from active buffer is sufficient
466            activeBuffer.position((int) n + activeBuffer.position());
467            return n;
468        }
469        stateChangeLock.lock();
470        final long skipped;
471        try {
472            skipped = skipInternal(n);
473        } finally {
474            stateChangeLock.unlock();
475        }
476        return skipped;
477    }
478
479    /**
480     * Internal skip function which should be called only from skip(). The assumption is that the stateChangeLock is already acquired in the caller before
481     * calling this function.
482     *
483     * @param n the number of bytes to be skipped.
484     * @return the actual number of bytes skipped.
485     * @throws IOException if an I/O error occurs.
486     */
487    private long skipInternal(final long n) throws IOException {
488        if (!stateChangeLock.isLocked()) {
489            throw new IllegalStateException("Expected stateChangeLock to be locked");
490        }
491        waitForAsyncReadComplete();
492        if (isEndOfStream()) {
493            return 0;
494        }
495        if (available() >= n) {
496            // we can skip from the internal buffers
497            int toSkip = (int) n;
498            // We need to skip from both active buffer and read ahead buffer
499            toSkip -= activeBuffer.remaining();
500            if (toSkip <= 0) { // skipping from activeBuffer already handled.
501                throw new IllegalStateException("Expected toSkip > 0, actual: " + toSkip);
502            }
503            activeBuffer.position(0);
504            activeBuffer.flip();
505            readAheadBuffer.position(toSkip + readAheadBuffer.position());
506            swapBuffers();
507            // Trigger async read to emptied read ahead buffer.
508            readAsync();
509            return n;
510        }
511        final int skippedBytes = available();
512        final long toSkip = n - skippedBytes;
513        activeBuffer.position(0);
514        activeBuffer.flip();
515        readAheadBuffer.position(0);
516        readAheadBuffer.flip();
517        final long skippedFromInputStream = in.skip(toSkip);
518        readAsync();
519        return skippedBytes + skippedFromInputStream;
520    }
521
522    /**
523     * Flips the active and read ahead buffer
524     */
525    private void swapBuffers() {
526        final ByteBuffer temp = activeBuffer;
527        activeBuffer = readAheadBuffer;
528        readAheadBuffer = temp;
529    }
530
531    private void waitForAsyncReadComplete() throws IOException {
532        stateChangeLock.lock();
533        try {
534            isWaiting.set(true);
535            // There is only one reader, and one writer, so the writer should signal only once,
536            // but a while loop checking the wake-up condition is still needed to avoid spurious wakeups.
537            while (readInProgress) {
538                asyncReadComplete.await();
539            }
540        } catch (final InterruptedException e) {
541            final InterruptedIOException iio = new InterruptedIOException(e.getMessage());
542            iio.initCause(e);
543            throw iio;
544        } finally {
545            try {
546                isWaiting.set(false);
547            } finally {
548                stateChangeLock.unlock();
549            }
550        }
551        checkReadException();
552    }
553}