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