View Javadoc
1   /*
2    * Copyright (C) 2008 The Guava Authors
3    *
4    * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5    * in compliance with the License. You may obtain a copy of the License at
6    *
7    * http://www.apache.org/licenses/LICENSE-2.0
8    *
9    * Unless required by applicable law or agreed to in writing, software distributed under the License
10   * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11   * or implied. See the License for the specific language governing permissions and limitations under
12   * the License.
13   */
14  
15  package com.google.common.io;
16  
17  import com.google.common.annotations.Beta;
18  import com.google.common.annotations.GwtIncompatible;
19  import com.google.common.annotations.VisibleForTesting;
20  import java.io.ByteArrayInputStream;
21  import java.io.ByteArrayOutputStream;
22  import java.io.File;
23  import java.io.FileInputStream;
24  import java.io.FileOutputStream;
25  import java.io.IOException;
26  import java.io.InputStream;
27  import java.io.OutputStream;
28  
29  /**
30   * An {@link OutputStream} that starts buffering to a byte array, but switches to file buffering
31   * once the data reaches a configurable size.
32   *
33   * <p>This class is thread-safe.
34   *
35   * @author Chris Nokleberg
36   * @since 1.0
37   */
38  @Beta
39  @GwtIncompatible
40  public final class FileBackedOutputStream extends OutputStream {
41  
42    private final int fileThreshold;
43    private final boolean resetOnFinalize;
44    private final ByteSource source;
45  
46    private OutputStream out;
47    private MemoryOutput memory;
48    private File file;
49  
50    /** ByteArrayOutputStream that exposes its internals. */
51    private static class MemoryOutput extends ByteArrayOutputStream {
52      byte[] getBuffer() {
53        return buf;
54      }
55  
56      int getCount() {
57        return count;
58      }
59    }
60  
61    /** Returns the file holding the data (possibly null). */
62    @VisibleForTesting
63    synchronized File getFile() {
64      return file;
65    }
66  
67    /**
68     * Creates a new instance that uses the given file threshold, and does not reset the data when the
69     * {@link ByteSource} returned by {@link #asByteSource} is finalized.
70     *
71     * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
72     */
73    public FileBackedOutputStream(int fileThreshold) {
74      this(fileThreshold, false);
75    }
76  
77    /**
78     * Creates a new instance that uses the given file threshold, and optionally resets the data when
79     * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
80     *
81     * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
82     * @param resetOnFinalize if true, the {@link #reset} method will be called when the
83     *     {@link ByteSource} returned by {@link #asByteSource} is finalized
84     */
85    public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
86      this.fileThreshold = fileThreshold;
87      this.resetOnFinalize = resetOnFinalize;
88      memory = new MemoryOutput();
89      out = memory;
90  
91      if (resetOnFinalize) {
92        source =
93            new ByteSource() {
94              @Override
95              public InputStream openStream() throws IOException {
96                return openInputStream();
97              }
98  
99              @Override
100             protected void finalize() {
101               try {
102                 reset();
103               } catch (Throwable t) {
104                 t.printStackTrace(System.err);
105               }
106             }
107           };
108     } else {
109       source =
110           new ByteSource() {
111             @Override
112             public InputStream openStream() throws IOException {
113               return openInputStream();
114             }
115           };
116     }
117   }
118 
119   /**
120    * Returns a readable {@link ByteSource} view of the data that has been written to this stream.
121    *
122    * @since 15.0
123    */
124   public ByteSource asByteSource() {
125     return source;
126   }
127 
128   private synchronized InputStream openInputStream() throws IOException {
129     if (file != null) {
130       return new FileInputStream(file);
131     } else {
132       return new ByteArrayInputStream(memory.getBuffer(), 0, memory.getCount());
133     }
134   }
135 
136   /**
137    * Calls {@link #close} if not already closed, and then resets this object back to its initial
138    * state, for reuse. If data was buffered to a file, it will be deleted.
139    *
140    * @throws IOException if an I/O error occurred while deleting the file buffer
141    */
142   public synchronized void reset() throws IOException {
143     try {
144       close();
145     } finally {
146       if (memory == null) {
147         memory = new MemoryOutput();
148       } else {
149         memory.reset();
150       }
151       out = memory;
152       if (file != null) {
153         File deleteMe = file;
154         file = null;
155         if (!deleteMe.delete()) {
156           throw new IOException("Could not delete: " + deleteMe);
157         }
158       }
159     }
160   }
161 
162   @Override
163   public synchronized void write(int b) throws IOException {
164     update(1);
165     out.write(b);
166   }
167 
168   @Override
169   public synchronized void write(byte[] b) throws IOException {
170     write(b, 0, b.length);
171   }
172 
173   @Override
174   public synchronized void write(byte[] b, int off, int len) throws IOException {
175     update(len);
176     out.write(b, off, len);
177   }
178 
179   @Override
180   public synchronized void close() throws IOException {
181     out.close();
182   }
183 
184   @Override
185   public synchronized void flush() throws IOException {
186     out.flush();
187   }
188 
189   /**
190    * Checks if writing {@code len} bytes would go over threshold, and switches to file buffering if
191    * so.
192    */
193   private void update(int len) throws IOException {
194     if (file == null && (memory.getCount() + len > fileThreshold)) {
195       File temp = File.createTempFile("FileBackedOutputStream", null);
196       if (resetOnFinalize) {
197         // Finalizers are not guaranteed to be called on system shutdown;
198         // this is insurance.
199         temp.deleteOnExit();
200       }
201       FileOutputStream transfer = new FileOutputStream(temp);
202       transfer.write(memory.getBuffer(), 0, memory.getCount());
203       transfer.flush();
204 
205       // We've successfully transferred the data; switch to writing to file
206       out = transfer;
207       file = temp;
208       memory = null;
209     }
210   }
211 }