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