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;
024import java.nio.file.Path;
025import java.util.Objects;
026import java.util.function.Supplier;
027
028import org.apache.commons.io.build.AbstractStreamBuilder;
029import org.apache.commons.io.file.PathUtils;
030
031/**
032 * An output stream which will retain data in memory until a specified threshold is reached, and only then commit it to disk. If the stream is closed before the
033 * threshold is reached, the data will not be written to disk at all.
034 * <p>
035 * This class originated in FileUpload processing. In this use case, you do not know in advance the size of the file being uploaded. If the file is small you
036 * want to store it in memory (for speed), but if the file is large you want to store it to file (to avoid memory issues).
037 * </p>
038 */
039public class DeferredFileOutputStream extends ThresholdingOutputStream {
040
041    /**
042     * Builds a new {@link DeferredFileOutputStream} instance.
043     * <p>
044     * For example:
045     * </p>
046     * <pre>{@code
047     * DeferredFileOutputStream s = DeferredFileOutputStream.builder()
048     *   .setPath(path)
049     *   .setBufferSize(4096)
050     *   .setDirectory(dir)
051     *   .setOutputFile(outputFile)
052     *   .setPrefix(prefix)
053     *   .setSuffix(suffix)
054     *   .setThreshold(threshold)
055     *   .get()}
056     * </pre>
057     * @since 2.12.0
058     */
059    public static class Builder extends AbstractStreamBuilder<DeferredFileOutputStream, Builder> {
060
061        private int threshold;
062        private File outputFile;
063        private String prefix;
064        private String suffix;
065        private File directory;
066
067        public Builder() {
068            setBufferSizeDefault(AbstractByteArrayOutputStream.DEFAULT_SIZE);
069            setBufferSize(AbstractByteArrayOutputStream.DEFAULT_SIZE);
070        }
071
072        @Override
073        public DeferredFileOutputStream get() {
074            return new DeferredFileOutputStream(threshold, outputFile, prefix, suffix, directory, getBufferSize());
075        }
076
077        /**
078         * Sets the temporary file directory.
079         *
080         * @param directory Temporary file directory.
081         * @return this
082         */
083        public Builder setDirectory(final File directory) {
084            this.directory = directory;
085            return this;
086        }
087
088        /**
089         * Sets the file to which data is saved beyond the threshold.
090         *
091         * @param outputFile The file to which data is saved beyond the threshold.
092         * @return this
093         */
094        public Builder setOutputFile(final File outputFile) {
095            this.outputFile = outputFile;
096            return this;
097        }
098
099        /**
100         * Sets the prefix to use for the temporary file.
101         *
102         * @param prefix Prefix to use for the temporary file.
103         * @return this
104         */
105        public Builder setPrefix(final String prefix) {
106            this.prefix = prefix;
107            return this;
108        }
109
110        /**
111         * Sets the suffix to use for the temporary file.
112         *
113         * @param suffix Suffix to use for the temporary file.
114         * @return this
115         */
116        public Builder setSuffix(final String suffix) {
117            this.suffix = suffix;
118            return this;
119        }
120
121        /**
122         * Sets the number of bytes at which to trigger an event.
123         *
124         * @param threshold The number of bytes at which to trigger an event.
125         * @return this
126         */
127        public Builder setThreshold(final int threshold) {
128            this.threshold = threshold;
129            return this;
130        }
131
132    }
133
134    /**
135     * Constructs a new {@link Builder}.
136     *
137     * @return a new {@link Builder}.
138     * @since 2.12.0
139     */
140    public static Builder builder() {
141        return new Builder();
142    }
143
144    private static int checkBufferSize(final int initialBufferSize) {
145        if (initialBufferSize < 0) {
146            throw new IllegalArgumentException("Initial buffer size must be at least 0.");
147        }
148        return initialBufferSize;
149    }
150
151    /**
152     * The output stream to which data will be written prior to the threshold being reached.
153     */
154    private ByteArrayOutputStream memoryOutputStream;
155
156    /**
157     * The output stream to which data will be written at any given time. This will always be one of {@code memoryOutputStream} or {@code diskOutputStream}.
158     */
159    private OutputStream currentOutputStream;
160
161    /**
162     * The file to which output will be directed if the threshold is exceeded.
163     */
164    private Path outputPath;
165
166    /**
167     * The temporary file prefix.
168     */
169    private final String prefix;
170
171    /**
172     * The temporary file suffix.
173     */
174    private final String suffix;
175
176    /**
177     * The directory to use for temporary files.
178     */
179    private final Path directory;
180
181    /**
182     * True when close() has been called successfully.
183     */
184    private boolean closed;
185
186    /**
187     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a file beyond that point. The initial
188     * buffer size will default to {@value AbstractByteArrayOutputStream#DEFAULT_SIZE} bytes which is ByteArrayOutputStream's default buffer size.
189     *
190     * @param threshold  The number of bytes at which to trigger an event.
191     * @param outputFile The file to which data is saved beyond the threshold.
192     * @deprecated Use {@link #builder()}
193     */
194    @Deprecated
195    public DeferredFileOutputStream(final int threshold, final File outputFile) {
196        this(threshold, outputFile, null, null, null, AbstractByteArrayOutputStream.DEFAULT_SIZE);
197    }
198
199    /**
200     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data either to a file beyond that point.
201     *
202     * @param threshold         The number of bytes at which to trigger an event.
203     * @param outputFile        The file to which data is saved beyond the threshold.
204     * @param prefix            Prefix to use for the temporary file.
205     * @param suffix            Suffix to use for the temporary file.
206     * @param directory         Temporary file directory.
207     * @param initialBufferSize The initial size of the in memory buffer.
208     */
209    private DeferredFileOutputStream(final int threshold, final File outputFile, final String prefix, final String suffix, final File directory,
210            final int initialBufferSize) {
211        super(threshold);
212        this.outputPath = toPath(outputFile, null);
213        this.prefix = prefix;
214        this.suffix = suffix;
215        this.directory = toPath(directory, PathUtils::getTempDirectory);
216        memoryOutputStream = new ByteArrayOutputStream(checkBufferSize(initialBufferSize));
217        currentOutputStream = memoryOutputStream;
218    }
219
220    /**
221     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a file beyond that point.
222     *
223     * @param threshold         The number of bytes at which to trigger an event.
224     * @param initialBufferSize The initial size of the in memory buffer.
225     * @param outputFile        The file to which data is saved beyond the threshold.
226     * @since 2.5
227     * @deprecated Use {@link #builder()}
228     */
229    @Deprecated
230    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final File outputFile) {
231        this(threshold, outputFile, null, null, null, initialBufferSize);
232    }
233
234    /**
235     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a temporary file beyond that point.
236     *
237     * @param threshold         The number of bytes at which to trigger an event.
238     * @param initialBufferSize The initial size of the in memory buffer.
239     * @param prefix            Prefix to use for the temporary file.
240     * @param suffix            Suffix to use for the temporary file.
241     * @param directory         Temporary file directory.
242     * @since 2.5
243     * @deprecated Use {@link #builder()}
244     */
245    @Deprecated
246    public DeferredFileOutputStream(final int threshold, final int initialBufferSize, final String prefix, final String suffix, final File directory) {
247        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, initialBufferSize);
248    }
249
250    /**
251     * Constructs an instance of this class which will trigger an event at the specified threshold, and save data to a temporary file beyond that point. The
252     * initial buffer size will default to 32 bytes which is ByteArrayOutputStream's default buffer size.
253     *
254     * @param threshold The number of bytes at which to trigger an event.
255     * @param prefix    Prefix to use for the temporary file.
256     * @param suffix    Suffix to use for the temporary file.
257     * @param directory Temporary file directory.
258     * @since 1.4
259     * @deprecated Use {@link #builder()}
260     */
261    @Deprecated
262    public DeferredFileOutputStream(final int threshold, final String prefix, final String suffix, final File directory) {
263        this(threshold, null, Objects.requireNonNull(prefix, "prefix"), suffix, directory, AbstractByteArrayOutputStream.DEFAULT_SIZE);
264    }
265
266    /**
267     * Closes underlying output stream, and mark this as closed
268     *
269     * @throws IOException if an error occurs.
270     */
271    @Override
272    public void close() throws IOException {
273        super.close();
274        closed = true;
275    }
276
277    /**
278     * Gets the data for this output stream as an array of bytes, assuming that the data has been retained in memory. If the data was written to disk, this
279     * method returns {@code null}.
280     *
281     * @return The data for this output stream, or {@code null} if no such data is available.
282     */
283    public byte[] getData() {
284        return memoryOutputStream != null ? memoryOutputStream.toByteArray() : null;
285    }
286
287    /**
288     * Gets either the output file specified in the constructor or the temporary file created or null.
289     * <p>
290     * If the constructor specifying the file is used then it returns that same output file, even when threshold has not been reached.
291     * <p>
292     * If constructor specifying a temporary file prefix/suffix is used then the temporary file created once the threshold is reached is returned If the
293     * threshold was not reached then {@code null} is returned.
294     *
295     * @return The file for this output stream, or {@code null} if no such file exists.
296     */
297    public File getFile() {
298        return outputPath != null ? outputPath.toFile() : null;
299    }
300
301    /**
302     * Gets the current output stream. This may be memory based or disk based, depending on the current state with respect to the threshold.
303     *
304     * @return The underlying output stream.
305     *
306     * @throws IOException if an error occurs.
307     */
308    @Override
309    protected OutputStream getStream() throws IOException {
310        return currentOutputStream;
311    }
312
313    /**
314     * Tests whether or not the data for this output stream has been retained in memory.
315     *
316     * @return {@code true} if the data is available in memory; {@code false} otherwise.
317     */
318    public boolean isInMemory() {
319        return !isThresholdExceeded();
320    }
321
322    /**
323     * Switches the underlying output stream from a memory based stream to one that is backed by disk. This is the point at which we realize that too much data
324     * is being written to keep in memory, so we elect to switch to disk-based storage.
325     *
326     * @throws IOException if an error occurs.
327     */
328    @Override
329    protected void thresholdReached() throws IOException {
330        if (prefix != null) {
331            outputPath = Files.createTempFile(directory, prefix, suffix);
332        }
333        PathUtils.createParentDirectories(outputPath);
334        final OutputStream fos = Files.newOutputStream(outputPath);
335        try {
336            memoryOutputStream.writeTo(fos);
337        } catch (final IOException e) {
338            fos.close();
339            throw e;
340        }
341        currentOutputStream = fos;
342        memoryOutputStream = null;
343    }
344
345    /**
346     * Converts the current contents of this byte stream to an {@link InputStream}. If the data for this output stream has been retained in memory, the returned
347     * stream is backed by buffers of {@code this} stream, avoiding memory allocation and copy, thus saving space and time.<br>
348     * Otherwise, the returned stream will be one that is created from the data that has been committed to disk.
349     *
350     * @return the current contents of this output stream.
351     * @throws IOException if this stream is not yet closed or an error occurs.
352     * @see org.apache.commons.io.output.ByteArrayOutputStream#toInputStream()
353     *
354     * @since 2.9.0
355     */
356    public InputStream toInputStream() throws IOException {
357        // we may only need to check if this is closed if we are working with a file
358        // but we should force the habit of closing whether we are working with
359        // a file or memory.
360        if (!closed) {
361            throw new IOException("Stream not closed");
362        }
363
364        if (isInMemory()) {
365            return memoryOutputStream.toInputStream();
366        }
367        return Files.newInputStream(outputPath);
368    }
369
370    private Path toPath(final File file, final Supplier<Path> defaultPathSupplier) {
371        return file != null ? file.toPath() : defaultPathSupplier == null ? null : defaultPathSupplier.get();
372    }
373
374    /**
375     * Writes the data from this output stream to the specified output stream, after it has been closed.
376     *
377     * @param outputStream output stream to write to.
378     * @throws NullPointerException if the OutputStream is {@code null}.
379     * @throws IOException          if this stream is not yet closed or an error occurs.
380     */
381    public void writeTo(final OutputStream outputStream) throws IOException {
382        // we may only need to check if this is closed if we are working with a file
383        // but we should force the habit of closing whether we are working with
384        // a file or memory.
385        if (!closed) {
386            throw new IOException("Stream not closed");
387        }
388
389        if (isInMemory()) {
390            memoryOutputStream.writeTo(outputStream);
391        } else {
392            Files.copy(outputPath, outputStream);
393        }
394    }
395}