Source code

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 1686747 2015-06-21 18:44:49Z krosenvold $
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,
135                                     final String suffix, final File directory) {
136        super(threshold);
137        this.outputFile = outputFile;
138
139        memoryOutputStream = new ByteArrayOutputStream();
140        currentOutputStream = memoryOutputStream;
141        this.prefix = prefix;
142        this.suffix = suffix;
143        this.directory = directory;
144    }
145
146
147    // --------------------------------------- ThresholdingOutputStream methods
148
149
150    /**
151     * Returns the current output stream. This may be memory based or disk
152     * based, depending on the current state with respect to the threshold.
153     *
154     * @return The underlying output stream.
155     *
156     * @exception IOException if an error occurs.
157     */
158    @Override
159    protected OutputStream getStream() throws IOException
160    {
161        return currentOutputStream;
162    }
163
164
165    /**
166     * Switches the underlying output stream from a memory based stream to one
167     * that is backed by disk. This is the point at which we realise that too
168     * much data is being written to keep in memory, so we elect to switch to
169     * disk-based storage.
170     *
171     * @exception IOException if an error occurs.
172     */
173    @Override
174    protected void thresholdReached() throws IOException
175    {
176        if (prefix != null) {
177            outputFile = File.createTempFile(prefix, suffix, directory);
178        }
179        final FileOutputStream fos = new FileOutputStream(outputFile);
180        try {
181            memoryOutputStream.writeTo(fos);
182        } catch (IOException e){
183            fos.close();
184            throw e;
185        }
186        currentOutputStream = fos;
187        memoryOutputStream = null;
188    }
189
190
191    // --------------------------------------------------------- Public methods
192
193
194    /**
195     * Determines whether or not the data for this output stream has been
196     * retained in memory.
197     *
198     * @return {@code true} if the data is available in memory;
199     *         {@code false} otherwise.
200     */
201    public boolean isInMemory()
202    {
203        return !isThresholdExceeded();
204    }
205
206
207    /**
208     * Returns the data for this output stream as an array of bytes, assuming
209     * that the data has been retained in memory. If the data was written to
210     * disk, this method returns {@code null}.
211     *
212     * @return The data for this output stream, or {@code null} if no such
213     *         data is available.
214     */
215    public byte[] getData()
216    {
217        if (memoryOutputStream != null)
218        {
219            return memoryOutputStream.toByteArray();
220        }
221        return null;
222    }
223
224
225    /**
226     * Returns either the output file specified in the constructor or
227     * the temporary file created or null.
228     * <p>
229     * If the constructor specifying the file is used then it returns that
230     * same output file, even when threshold has not been reached.
231     * <p>
232     * If constructor specifying a temporary file prefix/suffix is used
233     * then the temporary file created once the threshold is reached is returned
234     * If the threshold was not reached then {@code null} is returned.
235     *
236     * @return The file for this output stream, or {@code null} if no such
237     *         file exists.
238     */
239    public File getFile()
240    {
241        return outputFile;
242    }
243
244
245    /**
246     * Closes underlying output stream, and mark this as closed
247     *
248     * @exception IOException if an error occurs.
249     */
250    @Override
251    public void close() throws IOException
252    {
253        super.close();
254        closed = true;
255    }
256
257
258    /**
259     * Writes the data from this output stream to the specified output stream,
260     * after it has been closed.
261     *
262     * @param out output stream to write to.
263     * @exception IOException if this stream is not yet closed or an error occurs.
264     */
265    public void writeTo(final OutputStream out) throws IOException
266    {
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 wether we are working with
269        // a file or memory.
270        if (!closed)
271        {
272            throw new IOException("Stream not closed");
273        }
274
275        if(isInMemory())
276        {
277            memoryOutputStream.writeTo(out);
278        }
279        else
280        {
281            final FileInputStream fis = new FileInputStream(outputFile);
282            try {
283                IOUtils.copy(fis, out);
284            } finally {
285                IOUtils.closeQuietly(fis);
286            }
287        }
288    }
289}