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