001/*
002 * Copyright (C) 2008 The Guava Authors
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
005 * in compliance with the License. You may obtain a copy of the License at
006 *
007 * http://www.apache.org/licenses/LICENSE-2.0
008 *
009 * Unless required by applicable law or agreed to in writing, software distributed under the License
010 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
011 * or implied. See the License for the specific language governing permissions and limitations under
012 * the License.
013 */
014
015package com.google.common.io;
016
017import com.google.common.annotations.Beta;
018import com.google.common.annotations.GwtIncompatible;
019import com.google.common.annotations.VisibleForTesting;
020import java.io.ByteArrayInputStream;
021import java.io.ByteArrayOutputStream;
022import java.io.File;
023import java.io.FileInputStream;
024import java.io.FileOutputStream;
025import java.io.IOException;
026import java.io.InputStream;
027import java.io.OutputStream;
028
029/**
030 * An {@link OutputStream} that starts buffering to a byte array, but switches to file buffering
031 * once the data reaches a configurable size.
032 *
033 * <p>This class is thread-safe.
034 *
035 * @author Chris Nokleberg
036 * @since 1.0
037 */
038@Beta
039@GwtIncompatible
040public final class FileBackedOutputStream extends OutputStream {
041
042  private final int fileThreshold;
043  private final boolean resetOnFinalize;
044  private final ByteSource source;
045
046  private OutputStream out;
047  private MemoryOutput memory;
048  private File file;
049
050  /** ByteArrayOutputStream that exposes its internals. */
051  private static class MemoryOutput extends ByteArrayOutputStream {
052    byte[] getBuffer() {
053      return buf;
054    }
055
056    int getCount() {
057      return count;
058    }
059  }
060
061  /** Returns the file holding the data (possibly null). */
062  @VisibleForTesting
063  synchronized File getFile() {
064    return file;
065  }
066
067  /**
068   * Creates a new instance that uses the given file threshold, and does not reset the data when the
069   * {@link ByteSource} returned by {@link #asByteSource} is finalized.
070   *
071   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
072   */
073  public FileBackedOutputStream(int fileThreshold) {
074    this(fileThreshold, false);
075  }
076
077  /**
078   * Creates a new instance that uses the given file threshold, and optionally resets the data when
079   * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
080   *
081   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
082   * @param resetOnFinalize if true, the {@link #reset} method will be called when the
083   *     {@link ByteSource} returned by {@link #asByteSource} is finalized
084   */
085  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
086    this.fileThreshold = fileThreshold;
087    this.resetOnFinalize = resetOnFinalize;
088    memory = new MemoryOutput();
089    out = memory;
090
091    if (resetOnFinalize) {
092      source =
093          new ByteSource() {
094            @Override
095            public InputStream openStream() throws IOException {
096              return openInputStream();
097            }
098
099            @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}