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