001    /*
002     * Copyright (C) 2008 The Guava Authors
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package com.google.common.io;
018    
019    import com.google.common.annotations.Beta;
020    import com.google.common.annotations.VisibleForTesting;
021    
022    import java.io.ByteArrayInputStream;
023    import java.io.ByteArrayOutputStream;
024    import java.io.File;
025    import java.io.FileInputStream;
026    import java.io.FileOutputStream;
027    import java.io.IOException;
028    import java.io.InputStream;
029    import java.io.OutputStream;
030    
031    /**
032     * An {@link OutputStream} that starts buffering to a byte array, but
033     * switches to file buffering once the data reaches a configurable size.
034     *
035     * <p>This class is thread-safe.
036     *
037     * @author Chris Nokleberg
038     * @since 1
039     */
040    @Beta
041    public final class FileBackedOutputStream extends OutputStream {
042    
043      private final int fileThreshold;
044      private final boolean resetOnFinalize;
045      private final InputSupplier<InputStream> supplier;
046    
047      private OutputStream out;
048      private MemoryOutput memory;
049      private File file;
050    
051      /** ByteArrayOutputStream that exposes its internals. */
052      private static class MemoryOutput extends ByteArrayOutputStream {
053        byte[] getBuffer() {
054          return buf;
055        }
056    
057        int getCount() {
058          return count;
059        }
060      }
061    
062      /** Returns the file holding the data (possibly null). */
063      @VisibleForTesting synchronized File getFile() {
064        return file;
065      }
066    
067      /**
068       * Creates a new instance that uses the given file threshold.
069       * Equivalent to {@code ThresholdOutputStream(fileThreshold, false)}.
070       *
071       * @param fileThreshold the number of bytes before the stream should
072       *     switch to buffering to a file
073       */
074      public FileBackedOutputStream(int fileThreshold) {
075        this(fileThreshold, false);
076      }
077    
078      /**
079       * Creates a new instance that uses the given file threshold, and
080       * optionally resets the data when the {@link InputSupplier} returned
081       * by {@link #getSupplier} is finalized.
082       *
083       * @param fileThreshold the number of bytes before the stream should
084       *     switch to buffering to a file
085       * @param resetOnFinalize if true, the {@link #reset} method will
086       *     be called when the {@link InputSupplier} returned by {@link
087       *     #getSupplier} is finalized
088       */
089      public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
090        this.fileThreshold = fileThreshold;
091        this.resetOnFinalize = resetOnFinalize;
092        memory = new MemoryOutput();
093        out = memory;
094    
095        if (resetOnFinalize) {
096          supplier = new InputSupplier<InputStream>() {
097            @Override
098            public InputStream getInput() throws IOException {
099              return openStream();
100            }
101    
102            @Override protected void finalize() {
103              try {
104                reset();
105              } catch (Throwable t) {
106                t.printStackTrace(System.err);
107              }
108            }
109          };
110        } else {
111          supplier = new InputSupplier<InputStream>() {
112            @Override
113            public InputStream getInput() throws IOException {
114              return openStream();
115            }
116          };
117        }
118      }
119    
120      /**
121       * Returns a supplier that may be used to retrieve the data buffered
122       * by this stream.
123       */
124      public InputSupplier<InputStream> getSupplier() {
125        return supplier;
126      }
127    
128      private synchronized InputStream openStream() throws IOException {
129        if (file != null) {
130          return new FileInputStream(file);
131        } else {
132          return new ByteArrayInputStream(
133              memory.getBuffer(), 0, memory.getCount());
134        }
135      }
136    
137      /**
138       * Calls {@link #close} if not already closed, and then resets this
139       * object back to its initial state, for reuse. If data was buffered
140       * to a file, it will be deleted.
141       *
142       * @throws IOException if an I/O error occurred while deleting the file buffer
143       */
144      public synchronized void reset() throws IOException {
145        try {
146          close();
147        } finally {
148          if (memory == null) {
149            memory = new MemoryOutput();
150          } else {
151            memory.reset();
152          }
153          out = memory;
154          if (file != null) {
155            File deleteMe = file;
156            file = null;
157            if (!deleteMe.delete()) {
158              throw new IOException("Could not delete: " + deleteMe);
159            }
160          }
161        }
162      }
163    
164      @Override public synchronized void write(int b) throws IOException {
165        update(1);
166        out.write(b);
167      }
168    
169      @Override public synchronized void write(byte[] b) throws IOException {
170        write(b, 0, b.length);
171      }
172    
173      @Override public synchronized void write(byte[] b, int off, int len)
174          throws IOException {
175        update(len);
176        out.write(b, off, len);
177      }
178    
179      @Override public synchronized void close() throws IOException {
180        out.close();
181      }
182    
183      @Override public synchronized void flush() throws IOException {
184        out.flush();
185      }
186    
187      /**
188       * Checks if writing {@code len} bytes would go over threshold, and
189       * switches to file buffering if so.
190       */
191      private void update(int len) throws IOException {
192        if (file == null && (memory.getCount() + len > fileThreshold)) {
193          File temp = File.createTempFile("FileBackedOutputStream", null);
194          if (resetOnFinalize) {
195            // Finalizers are not guaranteed to be called on system shutdown;
196            // this is insurance.
197            temp.deleteOnExit();
198          }
199          FileOutputStream transfer = new FileOutputStream(temp);
200          transfer.write(memory.getBuffer(), 0, memory.getCount());
201          transfer.flush();
202    
203          // We've successfully transferred the data; switch to writing to file
204          out = transfer;
205          file = temp;
206          memory = null;
207        }
208      }
209    }