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