1/*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements.  See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License.  You may obtain a copy of the License at
8 *
9 *      http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17package org.apache.commons.io.output;
18
19import java.io.File;
20import java.io.FileInputStream;
21import java.io.FileOutputStream;
22import java.io.IOException;
23import java.io.OutputStream;
24
25import org.apache.commons.io.IOUtils;
26
27
28/**
29 * An output stream which will retain data in memory until a specified
30 * threshold is reached, and only then commit it to disk. If the stream is
31 * closed before the threshold is reached, the data will not be written to
32 * disk at all.
33 * <p>
34 * This class originated in FileUpload processing. In this use case, you do
35 * not know in advance the size of the file being uploaded. If the file is small
36 * you want to store it in memory (for speed), but if the file is large you want
37 * to store it to file (to avoid memory issues).
38 *
39 * @author <a href="mailto:martinc@apache.org">Martin Cooper</a>
40 * @author gaxzerow
41 *
42 * @version $Id: DeferredFileOutputStream.java 606381 2007-12-22 02:03:16Z ggregory $
43 */
44public class DeferredFileOutputStream
45    extends ThresholdingOutputStream
46{
47
48    // ----------------------------------------------------------- Data members
49
50
51    /**
52     * The output stream to which data will be written prior to the theshold
53     * being reached.
54     */
55    private ByteArrayOutputStream memoryOutputStream;
56
57
58    /**
59     * The output stream to which data will be written at any given time. This
60     * will always be one of <code>memoryOutputStream</code> or
61     * <code>diskOutputStream</code>.
62     */
63    private OutputStream currentOutputStream;
64
65
66    /**
67     * The file to which output will be directed if the threshold is exceeded.
68     */
69    private File outputFile;
70
71    /**
72     * The temporary file prefix.
73     */
74    private String prefix;
75
76    /**
77     * The temporary file suffix.
78     */
79    private String suffix;
80
81    /**
82     * The directory to use for temporary files.
83     */
84    private File directory;
85
86
87    /**
88     * True when close() has been called successfully.
89     */
90    private boolean closed = false;
91
92    // ----------------------------------------------------------- Constructors
93
94
95    /**
96     * Constructs an instance of this class which will trigger an event at the
97     * specified threshold, and save data to a file beyond that point.
98     *
99     * @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        super(threshold);
105        this.outputFile = outputFile;
106
107        memoryOutputStream = new ByteArrayOutputStream();
108        currentOutputStream = memoryOutputStream;
109    }
110
111
112    /**
113     * Constructs an instance of this class which will trigger an event at the
114     * specified threshold, and save data to a temporary file beyond that point.
115     *
116     * @param threshold  The number of bytes at which to trigger an event.
117     * @param prefix Prefix to use for the temporary file.
118     * @param suffix Suffix to use for the temporary file.
119     * @param directory Temporary file directory.
120     *
121     * @since Commons IO 1.4
122     */
123    public DeferredFileOutputStream(int threshold, String prefix, String suffix, File directory)
124    {
125        this(threshold, (File)null);
126        if (prefix == null) {
127            throw new IllegalArgumentException("Temporary file prefix is missing");
128        }
129        this.prefix = prefix;
130        this.suffix = suffix;
131        this.directory = directory;
132    }
133
134
135    // --------------------------------------- ThresholdingOutputStream methods
136
137
138    /**
139     * Returns the current output stream. This may be memory based or disk
140     * based, depending on the current state with respect to the threshold.
141     *
142     * @return The underlying output stream.
143     *
144     * @exception IOException if an error occurs.
145     */
146    protected OutputStream getStream() throws IOException
147    {
148        return currentOutputStream;
149    }
150
151
152    /**
153     * Switches the underlying output stream from a memory based stream to one
154     * that is backed by disk. This is the point at which we realise that too
155     * much data is being written to keep in memory, so we elect to switch to
156     * disk-based storage.
157     *
158     * @exception IOException if an error occurs.
159     */
160    protected void thresholdReached() throws IOException
161    {
162        if (prefix != null) {
163            outputFile = File.createTempFile(prefix, suffix, directory);
164        }
165        FileOutputStream fos = new FileOutputStream(outputFile);
166        memoryOutputStream.writeTo(fos);
167        currentOutputStream = fos;
168        memoryOutputStream = null;
169    }
170
171
172    // --------------------------------------------------------- Public methods
173
174
175    /**
176     * Determines whether or not the data for this output stream has been
177     * retained in memory.
178     *
179     * @return <code>true</code> if the data is available in memory;
180     *         <code>false</code> otherwise.
181     */
182    public boolean isInMemory()
183    {
184        return (!isThresholdExceeded());
185    }
186
187
188    /**
189     * Returns the data for this output stream as an array of bytes, assuming
190     * that the data has been retained in memory. If the data was written to
191     * disk, this method returns <code>null</code>.
192     *
193     * @return The data for this output stream, or <code>null</code> if no such
194     *         data is available.
195     */
196    public byte[] getData()
197    {
198        if (memoryOutputStream != null)
199        {
200            return memoryOutputStream.toByteArray();
201        }
202        return null;
203    }
204
205
206    /**
207     * Returns either the output file specified in the constructor or
208     * the temporary file created or null.
209     * <p>
210     * If the constructor specifying the file is used then it returns that
211     * same output file, even when threashold has not been reached.
212     * <p>
213     * If constructor specifying a temporary file prefix/suffix is used
214     * then the temporary file created once the threashold is reached is returned
215     * If the threshold was not reached then <code>null</code> is returned.
216     *
217     * @return The file for this output stream, or <code>null</code> if no such
218     *         file exists.
219     */
220    public File getFile()
221    {
222        return outputFile;
223    }
224
225
226    /**
227     * Closes underlying output stream, and mark this as closed
228     *
229     * @exception IOException if an error occurs.
230     */
231    public void close() throws IOException
232    {
233        super.close();
234        closed = true;
235    }
236
237
238    /**
239     * Writes the data from this output stream to the specified output stream,
240     * after it has been closed.
241     *
242     * @param out output stream to write to.
243     * @exception IOException if this stream is not yet closed or an error occurs.
244     */
245    public void writeTo(OutputStream out) throws IOException
246    {
247        // we may only need to check if this is closed if we are working with a file
248        // but we should force the habit of closing wether we are working with
249        // a file or memory.
250        if (!closed)
251        {
252            throw new IOException("Stream not closed");
253        }
254
255        if(isInMemory())
256        {
257            memoryOutputStream.writeTo(out);
258        }
259        else
260        {
261            FileInputStream fis = new FileInputStream(outputFile);
262            try {
263                IOUtils.copy(fis, out);
264            } finally {
265                IOUtils.closeQuietly(fis);
266            }
267        }
268    }
269}
270