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