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