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 wether 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}