    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.FileInputStream;
    021import java.io.FileOutputStream;
    022import java.io.IOException;
    023import java.io.InputStream;
    024import java.io.OutputStream;
    025import java.nio.file.Files;
    026
    027import org.apache.commons.io.FileUtils;
    028import org.apache.commons.io.IOUtils;
    029
    030/**
    031 * An output stream which will retain data in memory until a specified threshold is reached, and only then commit it to
    032 * disk. If the stream is closed before the threshold is reached, the data will not be written to disk at all.
    033 * <p>
    034 * This class originated in FileUpload processing. In this use case, you do not know in advance the size of the file
    035 * being uploaded. If the file is small you want to store it in memory (for speed), but if the file is large you want to
    036 * store it to file (to avoid memory issues).
    037 * </p>
    038 */
    039public class DeferredFileOutputStream extends ThresholdingOutputStream {
    040
    041    /**
    042     * The output stream to which data will be written prior to the threshold being reached.
    043     */
    044    private ByteArrayOutputStream memoryOutputStream;
    045
    046    /**
    047     * The output stream to which data will be written at any given time. This will always be one of
    048     * {@code memoryOutputStream} or {@code diskOutputStream}.
    049     */
    050    private OutputStream currentOutputStream;
    051
    052    /**
    053     * The file to which output will be directed if the threshold is exceeded.
    054     */
    055    private File outputFile;
    056
    057    /**
    058     * The temporary file prefix.
    059     */
    060    private final String prefix;
    061
    062    /**
    063     * The temporary file suffix.
    064     */
    065    private final String suffix;
    066
    067    /**
    068     * The directory to use for temporary files.
    069     */
    070    private final File directory;
    071
    072    /**
    073     * True when close() has been called successfully.
    074     */
    075    private boolean closed;
    076
    077    /**
    078     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a
    079     * file beyond that point. The initial buffer size will default to
    080     * {@value AbstractByteArrayOutputStream#DEFAULT_SIZE} bytes which is ByteArrayOutputStream's default buffer size.
    081     *
    082     * @param threshold The number of bytes at which to trigger an event.
    083     * @param outputFile The file to which data is saved beyond the threshold.
    084     */
    085    public DeferredFileOutputStream(final int threshold, final File outputFile) {
    086        this(threshold, outputFile, null, null, null, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    087    }
    088
    089    /**
    090     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data either
    091     * to a file beyond that point.
    092     *
    093     * @param threshold The number of bytes at which to trigger an event.
    094     * @param outputFile The file to which data is saved beyond the threshold.
    095     * @param prefix Prefix to use for the temporary file.
    096     * @param suffix Suffix to use for the temporary file.
    097     * @param directory Temporary file directory.
    098     * @param initialBufferSize The initial size of the in memory buffer.
    099     */
    100    private DeferredFileOutputStream(final int threshold, final File outputFile, final String prefix,
    101        final String suffix, final File directory, final int initialBufferSize) {
    102        super(threshold);
    103        this.outputFile = outputFile;
    104        this.prefix = prefix;
    105        this.suffix = suffix;
    106        this.directory = directory;
    107
    108        memoryOutputStream = new ByteArrayOutputStream(initialBufferSize);
    109        currentOutputStream = memoryOutputStream;
    110    }
    111
    112    /**
    113     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a
    114     * file beyond that point.
    115     *
    116     * @param threshold The number of bytes at which to trigger an event.
    117     * @param initialBufferSize The initial size of the in memory buffer.
    118     * @param outputFile The file to which data is saved beyond the threshold.
    119     *
    120     * @since 2.5
    121     */
    122    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final File outputFile) {
    123        this(threshold, outputFile, null, null, null, initialBufferSize);
    124        if (initialBufferSize < 0) {
    125            throw new IllegalArgumentException("Initial buffer size must be atleast 0.");
    126        }
    127    }
    128
    129    /**
    130     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a
    131     * temporary file beyond that point.
    132     *
    133     * @param threshold The number of bytes at which to trigger an event.
    134     * @param initialBufferSize The initial size of the in memory buffer.
    135     * @param prefix Prefix to use for the temporary file.
    136     * @param suffix Suffix to use for the temporary file.
    137     * @param directory Temporary file directory.
    138     *
    139     * @since 2.5
    140     */
    141    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final String prefix,
    142        final String suffix, final File directory) {
    143        this(threshold, null, prefix, suffix, directory, initialBufferSize);
    144        if (prefix == null) {
    145            throw new IllegalArgumentException("Temporary file prefix is missing");
    146        }
    147        if (initialBufferSize < 0) {
    148            throw new IllegalArgumentException("Initial buffer size must be atleast 0.");
    149        }
    150    }
    151
    152    /**
    153     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a
    154     * temporary file beyond that point. The initial buffer size will default to 32 bytes which is
    155     * ByteArrayOutputStream's default buffer size.
    156     *
    157     * @param threshold The number of bytes at which to trigger an event.
    158     * @param prefix Prefix to use for the temporary file.
    159     * @param suffix Suffix to use for the temporary file.
    160     * @param directory Temporary file directory.
    161     *
    162     * @since 1.4
    163     */
    164    public DeferredFileOutputStream(final int threshold, final String prefix, final String suffix,
    165        final File directory) {
    166        this(threshold, null, prefix, suffix, directory, AbstractByteArrayOutputStream.DEFAULT_SIZE);
    167        if (prefix == null) {
    168            throw new IllegalArgumentException("Temporary file prefix is missing");
    169        }
    170    }
    171
    172    /**
    173     * Closes underlying output stream, and mark this as closed
    174     *
    175     * @throws IOException if an error occurs.
    176     */
    177    @Override
    178    public void close() throws IOException {
    179        super.close();
    180        closed = true;
    181    }
    182
    183    /**
    184     * Returns the data for this output stream as an array of bytes, assuming that the data has been retained in memory.
    185     * If the data was written to disk, this method returns {@code null}.
    186     *
    187     * @return The data for this output stream, or {@code null} if no such data is available.
    188     */
    189    public byte[] getData() {
    190        return memoryOutputStream != null ? memoryOutputStream.toByteArray() : null;
    191    }
    192
    193    /**
    194     * Returns either the output file specified in the constructor or the temporary file created or null.
    195     * <p>
    196     * If the constructor specifying the file is used then it returns that same output file, even when threshold has not
    197     * been reached.
    198     * <p>
    199     * If constructor specifying a temporary file prefix/suffix is used then the temporary file created once the
    200     * threshold is reached is returned If the threshold was not reached then {@code null} is returned.
    201     *
    202     * @return The file for this output stream, or {@code null} if no such file exists.
    203     */
    204    public File getFile() {
    205        return outputFile;
    206    }
    207
    208    /**
    209     * Returns the current output stream. This may be memory based or disk based, depending on the current state with
    210     * respect to the threshold.
    211     *
    212     * @return The underlying output stream.
    213     *
    214     * @throws IOException if an error occurs.
    215     */
    216    @Override
    217    protected OutputStream getStream() throws IOException {
    218        return currentOutputStream;
    219    }
    220
    221    /**
    222     * Determines whether or not the data for this output stream has been retained in memory.
    223     *
    224     * @return {@code true} if the data is available in memory; {@code false} otherwise.
    225     */
    226    public boolean isInMemory() {
    227        return !isThresholdExceeded();
    228    }
    229
    230    /**
    231     * Switches the underlying output stream from a memory based stream to one that is backed by disk. This is the point
    232     * at which we realize that too much data is being written to keep in memory, so we elect to switch to disk-based
    233     * storage.
    234     *
    235     * @throws IOException if an error occurs.
    236     */
    237    @Override
    238    protected void thresholdReached() throws IOException {
    239        if (prefix != null) {
    240            outputFile = File.createTempFile(prefix, suffix, directory);
    241        }
    242        FileUtils.forceMkdirParent(outputFile);
    243        final FileOutputStream fos = new FileOutputStream(outputFile);
    244        try {
    245            memoryOutputStream.writeTo(fos);
    246        } catch (final IOException e) {
    247            fos.close();
    248            throw e;
    249        }
    250        currentOutputStream = fos;
    251        memoryOutputStream = null;
    252    }
    253
    254    /**
    255     * Gets the current contents of this byte stream as an {@link InputStream}.
    256     * If the data for this output stream has been retained in memory, the
    257     * returned stream is backed by buffers of {@code this} stream,
    258     * avoiding memory allocation and copy, thus saving space and time.<br>
    259     * Otherwise, the returned stream will be one that is created from the data
    260     * that has been committed to disk.
    261     *
    262     * @return the current contents of this output stream.
    263     * @throws IOException if this stream is not yet closed or an error occurs.
    264     * @see org.apache.commons.io.output.ByteArrayOutputStream#toInputStream()
    265     *
    266     * @since 2.9.0
    267     */
    268    public InputStream toInputStream() throws IOException {
    269        // we may only need to check if this is closed if we are working with a file
    270        // but we should force the habit of closing whether we are working with
    271        // a file or memory.
    272        if (!closed) {
    273            throw new IOException("Stream not closed");
    274        }
    275
    276        if (isInMemory()) {
    277            return memoryOutputStream.toInputStream();
    278        }
    279        return Files.newInputStream(outputFile.toPath());
    280    }
    281
    282    /**
    283     * Writes the data from this output stream to the specified output stream, after it has been closed.
    284     *
    285     * @param outputStream output stream to write to.
    286     * @throws NullPointerException if the OutputStream is {@code null}.
    287     * @throws IOException if this stream is not yet closed or an error occurs.
    288     */
    289    public void writeTo(final OutputStream outputStream) throws IOException {
    290        // we may only need to check if this is closed if we are working with a file
    291        // but we should force the habit of closing whether we are working with
    292        // a file or memory.
    293        if (!closed) {
    294            throw new IOException("Stream not closed");
    295        }
    296
    297        if (isInMemory()) {
    298            memoryOutputStream.writeTo(outputStream);
    299        } else {
    300            try (FileInputStream fis = new FileInputStream(outputFile)) {
    301                IOUtils.copy(fis, outputStream);
    302            }
    303        }
    304    }
    305}