View Javadoc

1   package com.bradmcevoy.io;
2   
3   import java.io.BufferedInputStream;
4   import java.io.BufferedOutputStream;
5   import java.io.ByteArrayInputStream;
6   import org.apache.commons.io.output.ByteArrayOutputStream;
7   import java.io.File;
8   import java.io.FileNotFoundException;
9   import java.io.FileOutputStream;
10  import java.io.IOException;
11  import java.io.InputStream;
12  import java.io.OutputStream;
13  import org.apache.commons.io.IOUtils;
14  import org.slf4j.Logger;
15  import org.slf4j.LoggerFactory;
16  
17  /**
18   * An output stream which will buffer data, initially using memory up to
19   * maxMemorySize, and then overflowing to a temporary file.
20   *
21   * To use this class you will write to it, and then close it, and then
22   * call getInputStream to read the data.
23   *
24   * The temporary file, if it was created, will be deleted when the inputstream
25   * is closed.
26   *
27   * @author brad
28   */
29  public class BufferingOutputStream extends OutputStream {
30  
31      private static Logger log = LoggerFactory.getLogger( BufferingOutputStream.class );
32      private ByteArrayOutputStream tempMemoryBuffer = new ByteArrayOutputStream();
33      private int maxMemorySize;
34      private File tempFile;
35      private FileOutputStream fout;
36      private BufferedOutputStream bufOut;
37      private Runnable runnable;
38      private long size;
39      private boolean closed;
40  
41      public BufferingOutputStream( int maxMemorySize ) {
42          this.maxMemorySize = maxMemorySize;
43      }
44  
45      public InputStream getInputStream() {
46          if( !closed )
47              throw new IllegalStateException( "this output stream is not yet closed" );
48          if( tempMemoryBuffer == null ) {
49              FileDeletingInputStream fin;
50              try {
51                  fin = new FileDeletingInputStream( tempFile );
52              } catch( FileNotFoundException ex ) {
53                  throw new RuntimeException( tempFile.getAbsolutePath(), ex );
54              }
55              BufferedInputStream bufIn = new BufferedInputStream( fin );
56              return bufIn;
57          } else {
58              return new ByteArrayInputStream( tempMemoryBuffer.toByteArray() );
59          }
60      }
61  
62      @Override
63      public void write( byte[] b ) throws IOException {
64          size += b.length;
65          if( tempMemoryBuffer != null ) {
66              tempMemoryBuffer.write( b );
67          } else {
68              bufOut.write( b );
69          }
70          checkSize();
71      }
72  
73      @Override
74      public void write( int b ) throws IOException {
75          size++;
76          if( tempMemoryBuffer != null ) {
77              tempMemoryBuffer.write( b );
78          } else {
79              bufOut.write( b );
80          }
81          checkSize();
82      }
83  
84      @Override
85      public void write( byte[] b, int off, int len ) throws IOException {
86          size += len;
87          if( tempMemoryBuffer != null ) {
88              tempMemoryBuffer.write( b, off, len );
89          } else {
90              bufOut.write( b, off, len );
91          }
92          checkSize();
93      }
94  
95      private void checkSize() throws IOException {
96          if( log.isTraceEnabled() ) {
97              log.trace( "checkSize: " + size );
98          }
99          if( tempMemoryBuffer == null ) return;
100 
101         if( tempMemoryBuffer.size() < maxMemorySize ) return;
102 
103         tempFile = File.createTempFile( "" + System.currentTimeMillis(), ".buffer" );
104         fout = new FileOutputStream( tempFile );
105         bufOut = new BufferedOutputStream( fout );
106         bufOut.write( tempMemoryBuffer.toByteArray() );
107         tempMemoryBuffer = null;
108     }
109 
110     @Override
111     public void flush() throws IOException {
112         if( tempMemoryBuffer != null ) {
113             tempMemoryBuffer.flush();
114         } else {
115             bufOut.flush();
116             fout.flush();
117         }
118     }
119 
120     @Override
121     public void close() throws IOException {
122         if( tempMemoryBuffer != null ) {
123             tempMemoryBuffer.close();
124         } else {
125             bufOut.close();
126             fout.close();
127         }
128         closed = true;
129         if( runnable != null ) {
130             runnable.run();
131         }
132     }
133 
134     public long getSize() {
135         return size;
136     }
137 
138     File getTempFile() {
139         return tempFile;
140     }
141 
142     ByteArrayOutputStream getTempMemoryBuffer() {
143         return tempMemoryBuffer;
144     }
145 
146     public void setOnClose( Runnable r ) {
147         this.runnable = r;
148     }
149 
150     /**
151      * returns true if the data is completely held in memory
152      *
153      * @return
154      */
155     public boolean isCompleteInMemory() {
156         return tempFile == null;
157     }
158 
159     /**
160      * Gets the data currently held in memory
161      * 
162      * @return
163      */
164     public byte[] getInMemoryData() {
165         return this.tempMemoryBuffer.toByteArray();
166     }
167 
168     @Override
169     protected void finalize() throws Throwable {
170         deleteTempFileIfExists();
171         super.finalize();
172     }
173 
174     public void deleteTempFileIfExists() {
175         if( bufOut != null ) {
176             IOUtils.closeQuietly(bufOut);
177         }
178         if( fout != null ) {
179             IOUtils.closeQuietly(fout);
180         }
181 
182         if( tempFile != null && tempFile.exists() ) {
183             log.error( "temporary file was not deleted. Was close called on the inputstream? Will attempt to delete" );
184             if( !tempFile.delete() ) {
185                 log.error( "Still couldnt delete temporary file: " + tempFile.getAbsolutePath() );
186             }
187         }
188 
189     }
190 }