    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     *   .setBufferSize(4096)
    052     *   .setDirectory(dir)
    053     *   .setOutputFile(outputFile)
    054     *   .setPrefix(prefix)
    055     *   .setSuffix(suffix)
    056     *   .setThreshold(threshold)
    057     *   .get();}
    058     * </pre>
    059     * <p>
    060     * The only super's aspect used us buffer size.
    061     * </p>
    062     *
    063     * @since 2.12.0
    064     */
    065    public static class Builder extends AbstractStreamBuilder<DeferredFileOutputStream, Builder> {
    066
    067        private int threshold;
    068        private Path outputFile;
    069        private String prefix;
    070        private String suffix;
    071        private Path directory;
    072
    073        /**
    074         * Constructs a new builder.
    075         */
    076        public Builder() {
    077            setBufferSizeDefault(AbstractByteArrayOutputStream.DEFAULT_SIZE);
    078            setBufferSize(AbstractByteArrayOutputStream.DEFAULT_SIZE);
    079        }
    080
    081        /**
    082         * Constructs a new instance.
    083         * <p>
    084         * This builder use the aspects threshold, outputFile, prefix, suffix, directory, buffer size.
    085         * </p>
    086         *
    087         * @return a new instance.
    088         */
    089        @Override
    090        public DeferredFileOutputStream get() {
    091            return new DeferredFileOutputStream(threshold, outputFile, prefix, suffix, directory, getBufferSize());
    092        }
    093
    094        /**
    095         * Sets the temporary file directory.
    096         *
    097         * @param directory Temporary file directory.
    098         * @return this
    099         */
    100        public Builder setDirectory(final File directory) {
    101            this.directory = toPath(directory, null);
    102            return this;
    103        }
    104
    105        /**
    106         * Sets the temporary file directory.
    107         *
    108         * @param directory Temporary file directory.
    109         * @return this
    110         * @since 2.14.0
    111         */
    112        public Builder setDirectory(final Path directory) {
    113            this.directory = toPath(directory, null);
    114            return this;
    115        }
    116
    117        /**
    118         * Sets the file to which data is saved beyond the threshold.
    119         *
    120         * @param outputFile The file to which data is saved beyond the threshold.
    121         * @return this
    122         */
    123        public Builder setOutputFile(final File outputFile) {
    124            this.outputFile = toPath(outputFile, null);
    125            return this;
    126        }
    127
    128        /**
    129         * Sets the file to which data is saved beyond the threshold.
    130         *
    131         * @param outputFile The file to which data is saved beyond the threshold.
    132         * @return this
    133         * @since 2.14.0
    134         */
    135        public Builder setOutputFile(final Path outputFile) {
    136            this.outputFile = toPath(outputFile, null);
    137            return this;
    138        }
    139
    140        /**
    141         * Sets the prefix to use for the temporary file.
    142         *
    143         * @param prefix Prefix to use for the temporary file.
    144         * @return this
    145         */
    146        public Builder setPrefix(final String prefix) {
    147            this.prefix = prefix;
    148            return this;
    149        }
    150
    151        /**
    152         * Sets the suffix to use for the temporary file.
    153         *
    154         * @param suffix Suffix to use for the temporary file.
    155         * @return this
    156         */
    157        public Builder setSuffix(final String suffix) {
    158            this.suffix = suffix;
    159            return this;
    160        }
    161
    162        /**
    163         * Sets the number of bytes at which to trigger an event.
    164         *
    165         * @param threshold The number of bytes at which to trigger an event.
    166         * @return this
    167         */
    168        public Builder setThreshold(final int threshold) {
    169            this.threshold = threshold;
    170            return this;
    171        }
    172
    173    }
    174
    175    /**
    176     * Constructs a new {@link Builder}.
    177     *
    178     * @return a new {@link Builder}.
    179     * @since 2.12.0
    180     */
    181    public static Builder builder() {
    182        return new Builder();
    183    }
    184
    185    private static int checkBufferSize(final int initialBufferSize) {
    186        if (initialBufferSize < 0) {
    187            throw new IllegalArgumentException("Initial buffer size must be at least 0.");
    188        }
    189        return initialBufferSize;
    190    }
    191
    192    private static Path toPath(final File file, final Supplier<Path> defaultPathSupplier) {
    193        return file != null ? file.toPath() : defaultPathSupplier == null ? null : defaultPathSupplier.get();
    194    }
    195
    196    private static Path toPath(final Path file, final Supplier<Path> defaultPathSupplier) {
    197        return file != null ? file : defaultPathSupplier == null ? null : defaultPathSupplier.get();
    198    }
    199
    200    /**
    201     * The output stream to which data will be written prior to the threshold being reached.
    202     */
    203    private ByteArrayOutputStream memoryOutputStream;
    204
    205    /**
    206     * The output stream to which data will be written at any given time. This will always be one of {@code memoryOutputStream} or {@code diskOutputStream}.
    207     */
    208    private OutputStream currentOutputStream;
    209
    210    /**
    211     * The file to which output will be directed if the threshold is exceeded.
    212     */
    213    private Path outputPath;
    214
    215    /**
    216     * The temporary file prefix.
    217     */
    218    private final String prefix;
    219
    220    /**
    221     * The temporary file suffix.
    222     */
    223    private final String suffix;
    224
    225    /**
    226     * The directory to use for temporary files.
    227     */
    228    private final Path directory;
    229
    230    /**
    231     * True when close() has been called successfully.
    232     */
    233    private boolean closed;
    234
    235    /**
    236     * 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
    237     * buffer size will default to {@value AbstractByteArrayOutputStream#DEFAULT_SIZE} bytes which is ByteArrayOutputStream's default buffer size.
    238     *
    239     * @param threshold  The number of bytes at which to trigger an event.
    240     * @param outputFile The file to which data is saved beyond the threshold.
    241     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    242     */
    243    @Deprecated
    244    public DeferredFileOutputStream(final int threshold, final File outputFile) {
    245        this(threshold, outputFile, null, null, null, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    246    }
    247
    248    /**
    249     * 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.
    250     *
    251     * @param threshold         The number of bytes at which to trigger an event.
    252     * @param outputFile        The file to which data is saved beyond the threshold.
    253     * @param prefix            Prefix to use for the temporary file.
    254     * @param suffix            Suffix to use for the temporary file.
    255     * @param directory         Temporary file directory.
    256     * @param initialBufferSize The initial size of the in memory buffer.
    257     * @throws IllegalArgumentException if initialBufferSize &lt; 0.
    258     */
    259    private DeferredFileOutputStream(final int threshold, final File outputFile, final String prefix, final String suffix, final File directory,
    260            final int initialBufferSize) {
    261        super(threshold);
    262        this.outputPath = toPath(outputFile, null);
    263        this.prefix = prefix;
    264        this.suffix = suffix;
    265        this.directory = toPath(directory, PathUtils::getTempDirectory);
    266        this.memoryOutputStream = new ByteArrayOutputStream(checkBufferSize(initialBufferSize));
    267        this.currentOutputStream = memoryOutputStream;
    268    }
    269
    270    /**
    271     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a file beyond that point.
    272     *
    273     * @param threshold         The number of bytes at which to trigger an event.
    274     * @param initialBufferSize The initial size of the in memory buffer.
    275     * @param outputFile        The file to which data is saved beyond the threshold.
    276     * @since 2.5
    277     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    278     */
    279    @Deprecated
    280    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final File outputFile) {
    281        this(threshold, outputFile, null, null, null, initialBufferSize);
    282    }
    283
    284    /**
    285     * 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.
    286     *
    287     * @param threshold         The number of bytes at which to trigger an event.
    288     * @param initialBufferSize The initial size of the in memory buffer.
    289     * @param prefix            Prefix to use for the temporary file.
    290     * @param suffix            Suffix to use for the temporary file.
    291     * @param directory         Temporary file directory.
    292     * @since 2.5
    293     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    294     */
    295    @Deprecated
    296    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final String prefix, final String suffix, final File directory) {
    297        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, initialBufferSize);
    298    }
    299
    300    /**
    301     * 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.
    302     *
    303     * @param threshold         The number of bytes at which to trigger an event.
    304     * @param outputFile        The file to which data is saved beyond the threshold.
    305     * @param prefix            Prefix to use for the temporary file.
    306     * @param suffix            Suffix to use for the temporary file.
    307     * @param directory         Temporary file directory.
    308     * @param initialBufferSize The initial size of the in memory buffer.
    309     * @throws IllegalArgumentException if initialBufferSize &lt; 0.
    310     */
    311    private DeferredFileOutputStream(final int threshold, final Path outputFile, final String prefix, final String suffix, final Path directory,
    312            final int initialBufferSize) {
    313        super(threshold);
    314        this.outputPath = toPath(outputFile, null);
    315        this.prefix = prefix;
    316        this.suffix = suffix;
    317        this.directory = toPath(directory, PathUtils::getTempDirectory);
    318        this.memoryOutputStream = new ByteArrayOutputStream(checkBufferSize(initialBufferSize));
    319        this.currentOutputStream = memoryOutputStream;
    320    }
    321
    322    /**
    323     * 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
    324     * initial buffer size will default to 32 bytes which is ByteArrayOutputStream's default buffer size.
    325     *
    326     * @param threshold The number of bytes at which to trigger an event.
    327     * @param prefix    Prefix to use for the temporary file.
    328     * @param suffix    Suffix to use for the temporary file.
    329     * @param directory Temporary file directory.
    330     * @since 1.4
    331     * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
    332     */
    333    @Deprecated
    334    public DeferredFileOutputStream(final int threshold, final String prefix, final String suffix, final File directory) {
    335        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    336    }
    337
    338    /**
    339     * Closes underlying output stream, and mark this as closed
    340     *
    341     * @throws IOException if an error occurs.
    342     */
    343    @Override
    344    public void close() throws IOException {
    345        super.close();
    346        closed = true;
    347    }
    348
    349    /**
    350     * 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
    351     * method returns {@code null}.
    352     *
    353     * @return The data for this output stream, or {@code null} if no such data is available.
    354     */
    355    public byte[] getData() {
    356        return memoryOutputStream != null ? memoryOutputStream.toByteArray() : null;
    357    }
    358
    359    /**
    360     * Gets either the output File specified in the constructor or the temporary File created or null.
    361     * <p>
    362     * If the constructor specifying the File is used then it returns that same output File, even when threshold has not been reached.
    363     * <p>
    364     * If constructor specifying a temporary File prefix/suffix is used then the temporary File created once the threshold is reached is returned if the
    365     * threshold was not reached then {@code null} is returned.
    366     *
    367     * @return The File for this output stream, or {@code null} if no such File exists.
    368     */
    369    public File getFile() {
    370        return outputPath != null ? outputPath.toFile() : null;
    371    }
    372
    373    /**
    374     * Gets either the output Path specified in the constructor or the temporary Path created or null.
    375     * <p>
    376     * If the constructor specifying the file is used then it returns that same output Path, even when threshold has not been reached.
    377     * <p>
    378     * If constructor specifying a temporary Path prefix/suffix is used then the temporary Path created once the threshold is reached is returned if the
    379     * threshold was not reached then {@code null} is returned.
    380     *
    381     * @return The Path for this output stream, or {@code null} if no such Path exists.
    382     * @since 2.14.0
    383     */
    384    public Path getPath() {
    385        return outputPath;
    386    }
    387
    388    /**
    389     * Gets the current output stream. This may be memory based or disk based, depending on the current state with respect to the threshold.
    390     *
    391     * @return The underlying output stream.
    392     *
    393     * @throws IOException if an error occurs.
    394     */
    395    @Override
    396    protected OutputStream getStream() throws IOException {
    397        return currentOutputStream;
    398    }
    399
    400    /**
    401     * Tests whether or not the data for this output stream has been retained in memory.
    402     *
    403     * @return {@code true} if the data is available in memory; {@code false} otherwise.
    404     */
    405    public boolean isInMemory() {
    406        return !isThresholdExceeded();
    407    }
    408
    409    /**
    410     * 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
    411     * is being written to keep in memory, so we elect to switch to disk-based storage.
    412     *
    413     * @throws IOException if an error occurs.
    414     */
    415    @Override
    416    protected void thresholdReached() throws IOException {
    417        if (prefix != null) {
    418            outputPath = Files.createTempFile(directory, prefix, suffix);
    419        }
    420        PathUtils.createParentDirectories(outputPath, null, PathUtils.EMPTY_FILE_ATTRIBUTE_ARRAY);
    421        final OutputStream fos = Files.newOutputStream(outputPath);
    422        try {
    423            memoryOutputStream.writeTo(fos);
    424        } catch (final IOException e) {
    425            fos.close();
    426            throw e;
    427        }
    428        currentOutputStream = fos;
    429        memoryOutputStream = null;
    430    }
    431
    432    /**
    433     * 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
    434     * stream is backed by buffers of {@code this} stream, avoiding memory allocation and copy, thus saving space and time.<br>
    435     * Otherwise, the returned stream will be one that is created from the data that has been committed to disk.
    436     *
    437     * @return the current contents of this output stream.
    438     * @throws IOException if this stream is not yet closed or an error occurs.
    439     * @see org.apache.commons.io.output.ByteArrayOutputStream#toInputStream()
    440     *
    441     * @since 2.9.0
    442     */
    443    public InputStream toInputStream() throws IOException {
    444        // we may only need to check if this is closed if we are working with a file
    445        // but we should force the habit of closing whether we are working with
    446        // a file or memory.
    447        if (!closed) {
    448            throw new IOException("Stream not closed");
    449        }
    450
    451        if (isInMemory()) {
    452            return memoryOutputStream.toInputStream();
    453        }
    454        return Files.newInputStream(outputPath);
    455    }
    456
    457    /**
    458     * Writes the data from this output stream to the specified output stream, after it has been closed.
    459     *
    460     * @param outputStream output stream to write to.
    461     * @throws NullPointerException if the OutputStream is {@code null}.
    462     * @throws IOException          if this stream is not yet closed or an error occurs.
    463     */
    464    public void writeTo(final OutputStream outputStream) throws IOException {
    465        // we may only need to check if this is closed if we are working with a file
    466        // but we should force the habit of closing whether we are working with
    467        // a file or memory.
    468        if (!closed) {
    469            throw new IOException("Stream not closed");
    470        }
    471
    472        if (isInMemory()) {
    473            memoryOutputStream.writeTo(outputStream);
    474        } else {
    475            Files.copy(outputPath, outputStream);
    476        }
    477    }
    478}