    001/*
    002 * Licensed to the Apache Software Foundation (ASF) under one or more
    003 * contributor license agreements.  See the NOTICE file distributed with
    004 * this work for additional information regarding copyright ownership.
    005 * The ASF licenses this file to You under the Apache License, Version 2.0
    006 * (the "License"); you may not use this file except in compliance with
    007 * the License.  You may obtain a copy of the License at
    008 *
    009 *      http://www.apache.org/licenses/LICENSE-2.0
    010 *
    011 * Unless required by applicable law or agreed to in writing, software
    012 * distributed under the License is distributed on an "AS IS" BASIS,
    013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    014 * See the License for the specific language governing permissions and
    015 * limitations under the License.
    016 */
    017package org.apache.commons.io.output;
    018
    019import java.io.File;
    020import java.io.IOException;
    021import java.io.InputStream;
    022import java.io.OutputStream;
    023import java.nio.file.Files;
    024import java.nio.file.Path;
    025import java.util.Objects;
    026import java.util.function.Supplier;
    027
    028import org.apache.commons.io.build.AbstractStreamBuilder;
    029import org.apache.commons.io.file.PathUtils;
    030
    031/**
    032 * An output stream which will retain data in memory until a specified threshold is reached, and only then commit it to disk. If the stream is closed before the
    033 * threshold is reached, the data will not be written to disk at all.
    034 * <p>
    035 * To build an instance, see {@link Builder}.
    036 * </p>
    037 * <p>
    038 * This class originated in FileUpload processing. In this use case, you do not know in advance the size of the file being uploaded. If the file is small you
    039 * want to store it in memory (for speed), but if the file is large you want to store it to file (to avoid memory issues).
    040 * </p>
    041 */
    042public class DeferredFileOutputStream extends ThresholdingOutputStream {
    043
    044    /**
    045     * Builds a new {@link DeferredFileOutputStream} instance.
    046     * <p>
    047     * For example:
    048     * </p>
    049     * <pre>{@code
    050     * DeferredFileOutputStream s = DeferredFileOutputStream.builder()
    051     *   .setPath(path)
    052     *   .setBufferSize(4096)
    053     *   .setDirectory(dir)
    054     *   .setOutputFile(outputFile)
    055     *   .setPrefix(prefix)
    056     *   .setSuffix(suffix)
    057     *   .setThreshold(threshold)
    058     *   .get();}
    059     * </pre>
    060     * @since 2.12.0
    061     */
    062    public static class Builder extends AbstractStreamBuilder<DeferredFileOutputStream, Builder> {
    063
    064        private int threshold;
    065        private File outputFile;
    066        private String prefix;
    067        private String suffix;
    068        private File directory;
    069
    070        public Builder() {
    071            setBufferSizeDefault(AbstractByteArrayOutputStream.DEFAULT_SIZE);
    072            setBufferSize(AbstractByteArrayOutputStream.DEFAULT_SIZE);
    073        }
    074
    075        /**
    076         * Constructs a new instance.
    077         * <p>
    078         * This builder use the aspects threshold, outputFile, prefix, suffix, directory, buffer size.
    079         * </p>
    080         *
    081         * @return a new instance.
    082         */
    083        @Override
    084        public DeferredFileOutputStream get() {
    085            return new DeferredFileOutputStream(threshold, outputFile, prefix, suffix, directory, getBufferSize());
    086        }
    087
    088        /**
    089         * Sets the temporary file directory.
    090         *
    091         * @param directory Temporary file directory.
    092         * @return this
    093         */
    094        public Builder setDirectory(final File directory) {
    095            this.directory = directory;
    096            return this;
    097        }
    098
    099        /**
    100         * Sets the file to which data is saved beyond the threshold.
    101         *
    102         * @param outputFile The file to which data is saved beyond the threshold.
    103         * @return this
    104         */
    105        public Builder setOutputFile(final File outputFile) {
    106            this.outputFile = outputFile;
    107            return this;
    108        }
    109
    110        /**
    111         * Sets the prefix to use for the temporary file.
    112         *
    113         * @param prefix Prefix to use for the temporary file.
    114         * @return this
    115         */
    116        public Builder setPrefix(final String prefix) {
    117            this.prefix = prefix;
    118            return this;
    119        }
    120
    121        /**
    122         * Sets the suffix to use for the temporary file.
    123         *
    124         * @param suffix Suffix to use for the temporary file.
    125         * @return this
    126         */
    127        public Builder setSuffix(final String suffix) {
    128            this.suffix = suffix;
    129            return this;
    130        }
    131
    132        /**
    133         * Sets the number of bytes at which to trigger an event.
    134         *
    135         * @param threshold The number of bytes at which to trigger an event.
    136         * @return this
    137         */
    138        public Builder setThreshold(final int threshold) {
    139            this.threshold = threshold;
    140            return this;
    141        }
    142
    143    }
    144
    145    /**
    146     * Constructs a new {@link Builder}.
    147     *
    148     * @return a new {@link Builder}.
    149     * @since 2.12.0
    150     */
    151    public static Builder builder() {
    152        return new Builder();
    153    }
    154
    155    private static int checkBufferSize(final int initialBufferSize) {
    156        if (initialBufferSize < 0) {
    157            throw new IllegalArgumentException("Initial buffer size must be at least 0.");
    158        }
    159        return initialBufferSize;
    160    }
    161
    162    /**
    163     * The output stream to which data will be written prior to the threshold being reached.
    164     */
    165    private ByteArrayOutputStream memoryOutputStream;
    166
    167    /**
    168     * The output stream to which data will be written at any given time. This will always be one of {@code memoryOutputStream} or {@code diskOutputStream}.
    169     */
    170    private OutputStream currentOutputStream;
    171
    172    /**
    173     * The file to which output will be directed if the threshold is exceeded.
    174     */
    175    private Path outputPath;
    176
    177    /**
    178     * The temporary file prefix.
    179     */
    180    private final String prefix;
    181
    182    /**
    183     * The temporary file suffix.
    184     */
    185    private final String suffix;
    186
    187    /**
    188     * The directory to use for temporary files.
    189     */
    190    private final Path directory;
    191
    192    /**
    193     * True when close() has been called successfully.
    194     */
    195    private boolean closed;
    196
    197    /**
    198     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a file beyond that point. The initial
    199     * buffer size will default to {@value AbstractByteArrayOutputStream#DEFAULT_SIZE} bytes which is ByteArrayOutputStream's default buffer size.
    200     *
    201     * @param threshold  The number of bytes at which to trigger an event.
    202     * @param outputFile The file to which data is saved beyond the threshold.
    203     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    204     */
    205    @Deprecated
    206    public DeferredFileOutputStream(final int threshold, final File outputFile) {
    207        this(threshold, outputFile, null, null, null, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    208    }
    209
    210    /**
    211     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data either to a file beyond that point.
    212     *
    213     * @param threshold         The number of bytes at which to trigger an event.
    214     * @param outputFile        The file to which data is saved beyond the threshold.
    215     * @param prefix            Prefix to use for the temporary file.
    216     * @param suffix            Suffix to use for the temporary file.
    217     * @param directory         Temporary file directory.
    218     * @param initialBufferSize The initial size of the in memory buffer.
    219     * @throws IllegalArgumentException if initialBufferSize &lt; 0.
    220     */
    221    private DeferredFileOutputStream(final int threshold, final File outputFile, final String prefix, final String suffix, final File directory,
    222            final int initialBufferSize) {
    223        super(threshold);
    224        this.outputPath = toPath(outputFile, null);
    225        this.prefix = prefix;
    226        this.suffix = suffix;
    227        this.directory = toPath(directory, PathUtils::getTempDirectory);
    228        this.memoryOutputStream = new ByteArrayOutputStream(checkBufferSize(initialBufferSize));
    229        this.currentOutputStream = memoryOutputStream;
    230    }
    231
    232    /**
    233     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a file beyond that point.
    234     *
    235     * @param threshold         The number of bytes at which to trigger an event.
    236     * @param initialBufferSize The initial size of the in memory buffer.
    237     * @param outputFile        The file to which data is saved beyond the threshold.
    238     * @since 2.5
    239     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    240     */
    241    @Deprecated
    242    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final File outputFile) {
    243        this(threshold, outputFile, null, null, null, initialBufferSize);
    244    }
    245
    246    /**
    247     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a temporary file beyond that point.
    248     *
    249     * @param threshold         The number of bytes at which to trigger an event.
    250     * @param initialBufferSize The initial size of the in memory buffer.
    251     * @param prefix            Prefix to use for the temporary file.
    252     * @param suffix            Suffix to use for the temporary file.
    253     * @param directory         Temporary file directory.
    254     * @since 2.5
    255     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    256     */
    257    @Deprecated
    258    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final String prefix, final String suffix, final File directory) {
    259        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, initialBufferSize);
    260    }
    261
    262    /**
    263     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a temporary file beyond that point. The
    264     * initial buffer size will default to 32 bytes which is ByteArrayOutputStream's default buffer size.
    265     *
    266     * @param threshold The number of bytes at which to trigger an event.
    267     * @param prefix    Prefix to use for the temporary file.
    268     * @param suffix    Suffix to use for the temporary file.
    269     * @param directory Temporary file directory.
    270     * @since 1.4
    271     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    272     */
    273    @Deprecated
    274    public DeferredFileOutputStream(final int threshold, final String prefix, final String suffix, final File directory) {
    275        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    276    }
    277
    278    /**
    279     * Closes underlying output stream, and mark this as closed
    280     *
    281     * @throws IOException if an error occurs.
    282     */
    283    @Override
    284    public void close() throws IOException {
    285        super.close();
    286        closed = true;
    287    }
    288
    289    /**
    290     * Gets the data for this output stream as an array of bytes, assuming that the data has been retained in memory. If the data was written to disk, this
    291     * method returns {@code null}.
    292     *
    293     * @return The data for this output stream, or {@code null} if no such data is available.
    294     */
    295    public byte[] getData() {
    296        return memoryOutputStream != null ? memoryOutputStream.toByteArray() : null;
    297    }
    298
    299    /**
    300     * Gets either the output file specified in the constructor or the temporary file created or null.
    301     * <p>
    302     * If the constructor specifying the file is used then it returns that same output file, even when threshold has not been reached.
    303     * <p>
    304     * If constructor specifying a temporary file prefix/suffix is used then the temporary file created once the threshold is reached is returned If the
    305     * threshold was not reached then {@code null} is returned.
    306     *
    307     * @return The file for this output stream, or {@code null} if no such file exists.
    308     */
    309    public File getFile() {
    310        return outputPath != null ? outputPath.toFile() : null;
    311    }
    312
    313    /**
    314     * Gets the current output stream. This may be memory based or disk based, depending on the current state with respect to the threshold.
    315     *
    316     * @return The underlying output stream.
    317     *
    318     * @throws IOException if an error occurs.
    319     */
    320    @Override
    321    protected OutputStream getStream() throws IOException {
    322        return currentOutputStream;
    323    }
    324
    325    /**
    326     * Tests whether or not the data for this output stream has been retained in memory.
    327     *
    328     * @return {@code true} if the data is available in memory; {@code false} otherwise.
    329     */
    330    public boolean isInMemory() {
    331        return !isThresholdExceeded();
    332    }
    333
    334    /**
    335     * Switches the underlying output stream from a memory based stream to one that is backed by disk. This is the point at which we realize that too much data
    336     * is being written to keep in memory, so we elect to switch to disk-based storage.
    337     *
    338     * @throws IOException if an error occurs.
    339     */
    340    @Override
    341    protected void thresholdReached() throws IOException {
    342        if (prefix != null) {
    343            outputPath = Files.createTempFile(directory, prefix, suffix);
    344        }
    345        PathUtils.createParentDirectories(outputPath, null, PathUtils.EMPTY_FILE_ATTRIBUTE_ARRAY);
    346        final OutputStream fos = Files.newOutputStream(outputPath);
    347        try {
    348            memoryOutputStream.writeTo(fos);
    349        } catch (final IOException e) {
    350            fos.close();
    351            throw e;
    352        }
    353        currentOutputStream = fos;
    354        memoryOutputStream = null;
    355    }
    356
    357    /**
    358     * Converts the current contents of this byte stream to an {@link InputStream}. If the data for this output stream has been retained in memory, the returned
    359     * stream is backed by buffers of {@code this} stream, avoiding memory allocation and copy, thus saving space and time.<br>
    360     * Otherwise, the returned stream will be one that is created from the data that has been committed to disk.
    361     *
    362     * @return the current contents of this output stream.
    363     * @throws IOException if this stream is not yet closed or an error occurs.
    364     * @see org.apache.commons.io.output.ByteArrayOutputStream#toInputStream()
    365     *
    366     * @since 2.9.0
    367     */
    368    public InputStream toInputStream() throws IOException {
    369        // we may only need to check if this is closed if we are working with a file
    370        // but we should force the habit of closing whether we are working with
    371        // a file or memory.
    372        if (!closed) {
    373            throw new IOException("Stream not closed");
    374        }
    375
    376        if (isInMemory()) {
    377            return memoryOutputStream.toInputStream();
    378        }
    379        return Files.newInputStream(outputPath);
    380    }
    381
    382    private Path toPath(final File file, final Supplier<Path> defaultPathSupplier) {
    383        return file != null ? file.toPath() : defaultPathSupplier == null ? null : defaultPathSupplier.get();
    384    }
    385
    386    /**
    387     * Writes the data from this output stream to the specified output stream, after it has been closed.
    388     *
    389     * @param outputStream output stream to write to.
    390     * @throws NullPointerException if the OutputStream is {@code null}.
    391     * @throws IOException          if this stream is not yet closed or an error occurs.
    392     */
    393    public void writeTo(final OutputStream outputStream) throws IOException {
    394        // we may only need to check if this is closed if we are working with a file
    395        // but we should force the habit of closing whether we are working with
    396        // a file or memory.
    397        if (!closed) {
    398            throw new IOException("Stream not closed");
    399        }
    400
    401        if (isInMemory()) {
    402            memoryOutputStream.writeTo(outputStream);
    403        } else {
    404            Files.copy(outputPath, outputStream);
    405        }
    406    }
    407}