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