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