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