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 static com.google.common.base.Preconditions.checkArgument;
018
019import com.google.common.annotations.Beta;
020import com.google.common.annotations.GwtIncompatible;
021import com.google.common.annotations.VisibleForTesting;
022import java.io.ByteArrayInputStream;
023import java.io.ByteArrayOutputStream;
024import java.io.File;
025import java.io.FileInputStream;
026import java.io.FileOutputStream;
027import java.io.IOException;
028import java.io.InputStream;
029import java.io.OutputStream;
030
031
032/**
033 * An {@link OutputStream} that starts buffering to a byte array, but switches to file buffering
034 * once the data reaches a configurable size.
035 *
036 * <p>This class is thread-safe.
037 *
038 * @author Chris Nokleberg
039 * @since 1.0
040 */
041@Beta
042@GwtIncompatible
043public final class FileBackedOutputStream extends OutputStream {
044
045  private final int fileThreshold;
046  private final boolean resetOnFinalize;
047  private final ByteSource source;
048
049  private OutputStream out;
050  private MemoryOutput memory;
051  private File file;
052
053  /** ByteArrayOutputStream that exposes its internals. */
054  private static class MemoryOutput extends ByteArrayOutputStream {
055    byte[] getBuffer() {
056      return buf;
057    }
058
059    int getCount() {
060      return count;
061    }
062  }
063
064  /** Returns the file holding the data (possibly null). */
065  @VisibleForTesting
066  synchronized File getFile() {
067    return file;
068  }
069
070  /**
071   * Creates a new instance that uses the given file threshold, and does not reset the data when the
072   * {@link ByteSource} returned by {@link #asByteSource} is finalized.
073   *
074   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
075   */
076  public FileBackedOutputStream(int fileThreshold) {
077    this(fileThreshold, false);
078  }
079
080  /**
081   * Creates a new instance that uses the given file threshold, and optionally resets the data when
082   * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
083   *
084   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
085   * @param resetOnFinalize if true, the {@link #reset} method will be called when the {@link
086   *     ByteSource} returned by {@link #asByteSource} is finalized
087   * @throws IllegalArgumentException if {@code fileThreshold} is negative
088   */
089  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
090    checkArgument(
091        fileThreshold >= 0, "fileThreshold must be non-negative, but was %s", fileThreshold);
092    this.fileThreshold = fileThreshold;
093    this.resetOnFinalize = resetOnFinalize;
094    memory = new MemoryOutput();
095    out = memory;
096
097    if (resetOnFinalize) {
098      source =
099          new ByteSource() {
100            @Override
101            public InputStream openStream() throws IOException {
102              return openInputStream();
103            }
104
105            @Override
106            protected void finalize() {
107              try {
108                reset();
109              } catch (Throwable t) {
110                t.printStackTrace(System.err);
111              }
112            }
113          };
114    } else {
115      source =
116          new ByteSource() {
117            @Override
118            public InputStream openStream() throws IOException {
119              return openInputStream();
120            }
121          };
122    }
123  }
124
125  /**
126   * Returns a readable {@link ByteSource} view of the data that has been written to this stream.
127   *
128   * @since 15.0
129   */
130  public ByteSource asByteSource() {
131    return source;
132  }
133
134  private synchronized InputStream openInputStream() throws IOException {
135    if (file != null) {
136      return new FileInputStream(file);
137    } else {
138      return new ByteArrayInputStream(memory.getBuffer(), 0, memory.getCount());
139    }
140  }
141
142  /**
143   * Calls {@link #close} if not already closed, and then resets this object back to its initial
144   * state, for reuse. If data was buffered to a file, it will be deleted.
145   *
146   * @throws IOException if an I/O error occurred while deleting the file buffer
147   */
148  public synchronized void reset() throws IOException {
149    try {
150      close();
151    } finally {
152      if (memory == null) {
153        memory = new MemoryOutput();
154      } else {
155        memory.reset();
156      }
157      out = memory;
158      if (file != null) {
159        File deleteMe = file;
160        file = null;
161        if (!deleteMe.delete()) {
162          throw new IOException("Could not delete: " + deleteMe);
163        }
164      }
165    }
166  }
167
168  @Override
169  public synchronized void write(int b) throws IOException {
170    update(1);
171    out.write(b);
172  }
173
174  @Override
175  public synchronized void write(byte[] b) throws IOException {
176    write(b, 0, b.length);
177  }
178
179  @Override
180  public synchronized void write(byte[] b, int off, int len) throws IOException {
181    update(len);
182    out.write(b, off, len);
183  }
184
185  @Override
186  public synchronized void close() throws IOException {
187    out.close();
188  }
189
190  @Override
191  public synchronized void flush() throws IOException {
192    out.flush();
193  }
194
195  /**
196   * Checks if writing {@code len} bytes would go over threshold, and switches to file buffering if
197   * so.
198   */
199  private void update(int len) throws IOException {
200    if (file == null && (memory.getCount() + len > fileThreshold)) {
201      File temp = TempFileCreator.INSTANCE.createTempFile("FileBackedOutputStream");
202      if (resetOnFinalize) {
203        // Finalizers are not guaranteed to be called on system shutdown;
204        // this is insurance.
205        temp.deleteOnExit();
206      }
207      FileOutputStream transfer = new FileOutputStream(temp);
208      transfer.write(memory.getBuffer(), 0, memory.getCount());
209      transfer.flush();
210
211      // We've successfully transferred the data; switch to writing to file
212      out = transfer;
213      file = temp;
214      memory = null;
215    }
216  }
217}