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